From 5a7568883d6528e5d4935524282c5e27a646e421 Mon Sep 17 00:00:00 2001 From: Aurelien Chartier <2567591+achartier@users.noreply.github.com> Date: Tue, 1 Sep 2026 12:25:39 -0700 Subject: [PATCH 1/3] [TRTLLM-15314][feat] Support native FP8 MoE LoRA weights Load native FP8 routed-expert adapters and carry their cache dtype through MoE grouped GEMMs. Reuse existing LoRA workspace with capacity checks and store final deltas directly in the activation compute dtype. Validate adapter dtypes, alignment and rank bounds; preserve scratch capacity across dtype changes. Add FP32-reference, mixed-dtype, stride, and CUDA graph regression coverage plus a small synthetic integration fixture. Document FP8 rounding and adapter storage savings. Pre-commit and CPU fixture checks passed. CUDA build and GPU validation remain pending because Computelab reservations were unavailable. Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com> --- .../kernels/cuda_graph_grouped_gemm.cu | 38 ++- .../kernels/cuda_graph_grouped_gemm.h | 7 +- .../kernels/moe/cutlass/include/moe_kernels.h | 2 + .../cutlass/include/moe_lora_grouped_gemm.h | 10 +- .../include/moe_lora_problem_builder.h | 5 +- .../kernels/moe/cutlass/moe_kernels.cu | 103 +++++++- .../moe/cutlass/moe_lora_problem_builder.cu | 11 +- cpp/tensorrt_llm/thop/moe/cutlass/moeOp.cpp | 134 +++++++--- .../kernels/moe/moeLoraProblemBuilderTest.cu | 36 ++- docs/source/features/lora.md | 29 ++- .../_torch/custom_ops/torch_custom_ops.py | 9 +- .../_torch/moe/fused_moe/cutlass/lora.py | 2 + tensorrt_llm/_torch/peft/lora/loaders.py | 11 +- tensorrt_llm/_torch/peft/lora/manager.py | 21 +- .../llmapi/test_llm_api_pytorch_moe_lora.py | 243 +++++++++++++++++- .../test_lists/test-db/l0_h100.yml | 2 + .../peft/test_moe_lora_cuda_graph_params.py | 2 + .../_torch/peft/test_moe_lora_extract.py | 24 ++ .../_torch/peft/test_moe_lora_grouped_gemm.py | 199 ++++++++++++++ .../unittest/_torch/peft/test_moe_lora_op.py | 72 ++++++ tests/unittest/others/test_lora_manager.py | 109 +++++++- 21 files changed, 935 insertions(+), 134 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu b/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu index b7e31f582193..5965a43f3dbd 100644 --- a/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu +++ b/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu @@ -201,7 +201,7 @@ void cudaGraphGroupedGemmType(cutlass::gemm::GemmCoord* problemSizesPtr, int pro // dimensions already on the GPU and reuses their compatible storage directly. // ==================================================================== -template +template void fp8CudaGraphGroupedGemmImpl(cutlass::gemm::GemmCoord* problemSizesPtr, int problemCount, void** ptrAGpu, void** ptrBGpu, void** ptrCGpu, void** ptrDGpu, int64_t* ldaGpu, int64_t* ldbGpu, int64_t* ldcGpu, int64_t* lddGpu, cudaStream_t stream) @@ -218,8 +218,7 @@ void fp8CudaGraphGroupedGemmImpl(cutlass::gemm::GemmCoord* problemSizesPtr, int using ElementA = cutlass::float_e4m3_t; using ElementB = cutlass::float_e4m3_t; - using ElementC = cutlass::float_e4m3_t; - using ElementD = cutlass::float_e4m3_t; + using ElementC = ElementD; using ElementAccumulator = float; using LayoutA = cutlass::layout::RowMajor; @@ -327,6 +326,7 @@ void fp8CudaGraphGroupedGemmImpl(cutlass::gemm::GemmCoord* problemSizesPtr, int TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); } +template void fp8CudaGraphGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int problemCount, void** ptrAGpu, void** ptrBGpu, void** ptrCGpu, void** ptrDGpu, int64_t* ldaGpu, int64_t* ldbGpu, int64_t* ldcGpu, int64_t* lddGpu, cudaStream_t stream) @@ -335,7 +335,7 @@ void fp8CudaGraphGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int prob if (smVersion == fp8GroupedGemmConfig::kSm90) { #if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) && !defined(EXCLUDE_SM_90) - fp8CudaGraphGroupedGemmImpl( + fp8CudaGraphGroupedGemmImpl( problemSizesPtr, problemCount, ptrAGpu, ptrBGpu, ptrCGpu, ptrDGpu, ldaGpu, ldbGpu, ldcGpu, lddGpu, stream); return; #endif @@ -343,7 +343,7 @@ void fp8CudaGraphGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int prob else if (smVersion == fp8GroupedGemmConfig::kSm100) { #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) && !defined(EXCLUDE_SM_100F) - fp8CudaGraphGroupedGemmImpl( + fp8CudaGraphGroupedGemmImpl( problemSizesPtr, problemCount, ptrAGpu, ptrBGpu, ptrCGpu, ptrDGpu, ldaGpu, ldbGpu, ldcGpu, lddGpu, stream); return; #endif @@ -356,15 +356,35 @@ void fp8CudaGraphGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int prob void cudaGraphGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int problemCount, void** ptrAGpu, void** ptrBGpu, void** ptrCGpu, void** ptrDGpu, int64_t* ldaGpu, int64_t* ldbGpu, int64_t* ldcGpu, int64_t* lddGpu, bool isLoraIn, - tensorrt_llm::DataType dataType, int minKN, cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, cudaStream_t stream) + tensorrt_llm::DataType dataType, int minKN, cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, cudaStream_t stream, + std::optional outputDataType) { + auto const outputType = outputDataType.value_or(dataType); + TLLM_CHECK_WITH_INFO(dataType == tensorrt_llm::DataType::kFP8 || outputType == dataType, + "Mixed output dtype is supported only for FP8 grouped GEMM inputs."); #ifdef ENABLE_FP8 #if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) if (dataType == tensorrt_llm::DataType::kFP8) { checkFp8CudaGraphAlignment(hostMaxProblemSizesPtr, problemCount, minKN, "FP8 CUDA graph grouped GEMM"); - fp8CudaGraphGroupedGemm( - problemSizesPtr, problemCount, ptrAGpu, ptrBGpu, ptrCGpu, ptrDGpu, ldaGpu, ldbGpu, ldcGpu, lddGpu, stream); + switch (outputType) + { + case tensorrt_llm::DataType::kFP8: + fp8CudaGraphGroupedGemm(problemSizesPtr, problemCount, ptrAGpu, ptrBGpu, ptrCGpu, + ptrDGpu, ldaGpu, ldbGpu, ldcGpu, lddGpu, stream); + break; + case tensorrt_llm::DataType::kHALF: + fp8CudaGraphGroupedGemm(problemSizesPtr, problemCount, ptrAGpu, ptrBGpu, ptrCGpu, ptrDGpu, + ldaGpu, ldbGpu, ldcGpu, lddGpu, stream); + break; +#ifdef ENABLE_BF16 + case tensorrt_llm::DataType::kBF16: + fp8CudaGraphGroupedGemm(problemSizesPtr, problemCount, ptrAGpu, ptrBGpu, ptrCGpu, + ptrDGpu, ldaGpu, ldbGpu, ldcGpu, lddGpu, stream); + break; +#endif + default: TLLM_THROW("FP8 grouped GEMM output must be fp8, fp16, or bf16."); + } return; } #else @@ -551,7 +571,7 @@ void cudaGraphSplitKGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int p { // Reuse the non-split-K fp8 path; CUTLASS 3.x cooperative schedule handles large-K efficiently. checkFp8CudaGraphAlignment(hostMaxProblemSizesPtr, problemCount, minKN, "FP8 CUDA graph split-K grouped GEMM"); - fp8CudaGraphGroupedGemm( + fp8CudaGraphGroupedGemm( problemSizesPtr, problemCount, ptrAGpu, ptrBGpu, ptrCGpu, ptrDGpu, ldaGpu, ldbGpu, ldcGpu, lddGpu, stream); return; } diff --git a/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.h b/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.h index b447bba3a785..9ded12536407 100644 --- a/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.h +++ b/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -20,6 +20,7 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/tllmDataType.h" #include +#include TRTLLM_NAMESPACE_BEGIN @@ -41,11 +42,13 @@ namespace kernels * @param isLoraIn Whether this is for LoRA input transformation * @param dataType Data type of the matrices * @param minKN Minimum K*N value for kernel selection + * @param outputDataType Optional output dtype; FP8 inputs may emit fp16/bf16. * @param stream CUDA stream */ void cudaGraphGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int problemCount, void** ptrAGpu, void** ptrBGpu, void** ptrCGpu, void** ptrDGpu, int64_t* ldaGpu, int64_t* ldbGpu, int64_t* ldcGpu, int64_t* lddGpu, bool isLoraIn, - tensorrt_llm::DataType dataType, int minKN, cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, cudaStream_t stream); + tensorrt_llm::DataType dataType, int minKN, cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, cudaStream_t stream, + std::optional outputDataType = std::nullopt); /** * @brief CUDA Graph compatible wrapper for split-K grouped GEMM operations. diff --git a/cpp/tensorrt_llm/kernels/moe/cutlass/include/moe_kernels.h b/cpp/tensorrt_llm/kernels/moe/cutlass/include/moe_kernels.h index 5cc977447290..a87f457a35dc 100644 --- a/cpp/tensorrt_llm/kernels/moe/cutlass/include/moe_kernels.h +++ b/cpp/tensorrt_llm/kernels/moe/cutlass/include/moe_kernels.h @@ -1015,6 +1015,8 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface ScaleBiasType* lora_fc1_result_{}; ScaleBiasType* lora_add_bias_{}; ScaleBiasType* lora_fc2_result_{}; + size_t mLoraFc1ResultBytes{}; + size_t mLoraFc2ResultBytes{}; void* smoothed_act_{}; TmaWarpSpecializedGroupedGemmInput tma_ws_grouped_gemm1_input_; diff --git a/cpp/tensorrt_llm/kernels/moe/cutlass/include/moe_lora_grouped_gemm.h b/cpp/tensorrt_llm/kernels/moe/cutlass/include/moe_lora_grouped_gemm.h index f3661e569e89..819a14cfa1f0 100644 --- a/cpp/tensorrt_llm/kernels/moe/cutlass/include/moe_lora_grouped_gemm.h +++ b/cpp/tensorrt_llm/kernels/moe/cutlass/include/moe_lora_grouped_gemm.h @@ -52,11 +52,12 @@ struct MoeLoraGroupedGemmModule; // output_base: base of the module's output buffer; the implementation // accumulates into it, so callers must initialize it. // data_type: scalar dtype expected by the GEMM wrappers -// (fp16/bf16/fp32). +// (fp8/fp16/bf16/fp32). +// output_data_type: scalar dtype of output_base. // stream: CUDA stream to launch onto. using MoeLoraGroupedGemmRunFn = void (*)(MoeLoraGroupedGemmModule const& mod, int64_t num_permuted_tokens, int64_t in_hidden_size, int64_t max_lora_rank, int64_t dtype_bytes, int64_t splitk_slices, void const* input_base, - void* output_base, tensorrt_llm::DataType data_type, cudaStream_t stream); + void* output_base, tensorrt_llm::DataType data_type, tensorrt_llm::DataType output_data_type, cudaStream_t stream); // Per-module device-resident scratch for the MoE LoRA capture-safe path. // Pointers refer to device memory unless noted. @@ -140,11 +141,12 @@ struct MoeLoraGroupedGemm { bool enabled = false; - // Scalars common to all three modules. Fixed for the lifetime of the - // FusedMoeRunner once the scratch is allocated. + // Scalars common to all three modules, populated for each invocation. + // The runner retains scratch capacity across changes in adapter dtype. int64_t in_hidden_size = 0; int64_t max_lora_rank = 0; int64_t dtype_bytes = 0; + tensorrt_llm::DataType data_type = tensorrt_llm::DataType::kBF16; int64_t splitk_slices = 0; bool has_gated = false; diff --git a/cpp/tensorrt_llm/kernels/moe/cutlass/include/moe_lora_problem_builder.h b/cpp/tensorrt_llm/kernels/moe/cutlass/include/moe_lora_problem_builder.h index a71e89b87864..8424424f5aec 100644 --- a/cpp/tensorrt_llm/kernels/moe/cutlass/include/moe_lora_problem_builder.h +++ b/cpp/tensorrt_llm/kernels/moe/cutlass/include/moe_lora_problem_builder.h @@ -102,7 +102,8 @@ struct MoeLoraGemmGroupArrays // per-token rank so the GEMM lands at a known offset. // The out-GEMM's ldb_out[i] is the per-token rank (adapter B // is stored [out_hidden_size, rank]), not out_hidden_size. -// dtype_bytes: scalar size in bytes (2 for bf16/fp16, 4 for fp32). +// dtype_bytes: input/weight scalar size (1 for fp8, 2 for bf16/fp16, 4 for fp32). +// output_dtype_bytes: output scalar size in bytes (must be positive). // splitk_slices: split-K factor for the in-GEMM; drives the per-problem // split-K scratch stride. // @@ -111,7 +112,7 @@ struct MoeLoraGemmGroupArrays void launchMoeLoraProblemBuilder(int32_t const* ranks_dev, int64_t const* ptrs_dev, void const* input_base, void* lowrank_workspace, void* output_base, int64_t num_permuted_tokens, int64_t in_hidden_size, int64_t out_hidden_size, int64_t max_lora_rank, int64_t dtype_bytes, int64_t splitk_slices, - MoeLoraGemmGroupArrays const& out, cudaStream_t stream); + MoeLoraGemmGroupArrays const& out, cudaStream_t stream, int64_t output_dtype_bytes); } // namespace kernels::cutlass_kernels diff --git a/cpp/tensorrt_llm/kernels/moe/cutlass/moe_kernels.cu b/cpp/tensorrt_llm/kernels/moe/cutlass/moe_kernels.cu index 1b4959485862..ca6b8998a6a4 100644 --- a/cpp/tensorrt_llm/kernels/moe/cutlass/moe_kernels.cu +++ b/cpp/tensorrt_llm/kernels/moe/cutlass/moe_kernels.cu @@ -3027,6 +3027,37 @@ void doActivationDynamic(T* output, GemmOutputType const* gemm_result, float con // ============================== Lora Add Bias ================================= constexpr static int LORA_KERNELS_THREADS_PER_BLOCK = 256; +#if defined(ENABLE_FP8) +template +__global__ void convertMoeLoraInputToFp8Kernel( + __nv_fp8_e4m3* output, T const* input, int64_t const* numValidTokens, int64_t hiddenSize, int64_t numElements) +{ + int64_t const index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + int64_t const validElements = numValidTokens ? *numValidTokens * hiddenSize : numElements; + if (index < numElements) + { + // Padded rows can contain uninitialized values and must never be read. + output[index] = __nv_fp8_e4m3(index < validElements ? static_cast(input[index]) : 0.0f); + } +} + +template +void convertMoeLoraInputToFp8(__nv_fp8_e4m3* output, T const* input, int64_t const* numValidTokens, int64_t hiddenSize, + int64_t numRows, cudaStream_t stream) +{ + int64_t const numElements = numRows * hiddenSize; + if (numElements == 0) + { + return; + } + int const blocks + = static_cast((numElements + LORA_KERNELS_THREADS_PER_BLOCK - 1) / LORA_KERNELS_THREADS_PER_BLOCK); + convertMoeLoraInputToFp8Kernel<<>>( + output, input, numValidTokens, hiddenSize, numElements); + sync_check_cuda_error(stream); +} +#endif + template __global__ void loraAddBiasKernel(ScaleBiasType* output, LoraType const* lora_result, ScaleBiasType const* bias, int64_t const* num_valid_tokens_ptr, int* permuted_token_selected_experts, int64_t inter_size) @@ -3534,6 +3565,8 @@ void CutlassMoeFCRunner(); + tensorrt_llm::DataType const data_type = grouped_gemm.data_type; + tensorrt_llm::DataType const compute_data_type = moeLoraDataType(); + void const* grouped_input = input; + if (data_type == tensorrt_llm::DataType::kFP8) + { +#if defined(ENABLE_FP8) + // FC2's result storage must remain unused until FC1's grouped GEMMs finish on this stream. + // It must stay separate from overlapped_gemm1_gemm2_{inputs,outputs} and FC1's LoRA output. + size_t const fp8InputBytes = static_cast(expanded_num_rows) * hidden_size * sizeof(__nv_fp8_e4m3); + TLLM_CHECK_WITH_INFO( + fp8InputBytes <= mLoraFc2ResultBytes, "FC2 LoRA result buffer is too small for the FP8 FC1 input."); + auto* fp8_input = reinterpret_cast<__nv_fp8_e4m3*>(lora_fc2_result_); + convertMoeLoraInputToFp8(fp8_input, input, num_valid_tokens_ptr, hidden_size, expanded_num_rows, stream); + grouped_input = fp8_input; +#else + TLLM_THROW("FP8 MoE LoRA requires an ENABLE_FP8 build."); +#endif + } + else + { + TLLM_CHECK_WITH_INFO(data_type == compute_data_type, + "MoE LoRA cache dtype must match the activation compute dtype unless the cache is FP8."); + } // The grouped-GEMM GEMM skips rank-0 rows, but the bias/reorder paths // read lora_fc1_result_ for every valid row. Zero the buffer first so // skipped rows are a deterministic no-op. It is contiguous and holds @@ -4276,18 +4331,19 @@ auto CutlassMoeFCRunner(expanded_num_rows) * static_cast(inter_size) * (is_gated_activation ? 2u : 1u) * sizeof(ScaleBiasType); TLLM_CUDA_CHECK(cudaMemsetAsync(lora_fc1_result_, 0, fc1_result_bytes, stream)); - runMoeLoraGroupedGemmModule(grouped_gemm.fc1, expanded_num_rows, /*in_hidden_size=*/hidden_size, grouped_gemm.max_lora_rank, grouped_gemm.dtype_bytes, grouped_gemm.splitk_slices, - /*input_base=*/static_cast(input), - /*output_base=*/static_cast(lora_fc1_result), grouped_gemm.run, data_type, stream); + /*input_base=*/grouped_input, + /*output_base=*/static_cast(lora_fc1_result), grouped_gemm.run, data_type, compute_data_type, + stream); if (is_gated_activation) { runMoeLoraGroupedGemmModule(grouped_gemm.gated, expanded_num_rows, /*in_hidden_size=*/hidden_size, grouped_gemm.max_lora_rank, grouped_gemm.dtype_bytes, grouped_gemm.splitk_slices, - /*input_base=*/static_cast(input), - /*output_base=*/static_cast(lora_gated_out), grouped_gemm.run, data_type, stream); + /*input_base=*/grouped_input, + /*output_base=*/static_cast(lora_gated_out), grouped_gemm.run, data_type, compute_data_type, + stream); } } else @@ -4362,18 +4418,41 @@ void CutlassMoeFCRunner(); + tensorrt_llm::DataType const data_type = grouped_gemm.data_type; + tensorrt_llm::DataType const compute_data_type = moeLoraDataType(); + void const* grouped_input = input; + if (data_type == tensorrt_llm::DataType::kFP8) + { +#if defined(ENABLE_FP8) + // GEMM1 and its activation must finish consuming FC1's LoRA delta on this stream before this cast. + // Its storage must stay separate from overlapped_gemm1_gemm2_{inputs,outputs} and FC2's LoRA output, + // and must not be reused again until FC2's grouped GEMMs finish. + size_t const fp8InputBytes = static_cast(num_tokens) * inter_size * sizeof(__nv_fp8_e4m3); + TLLM_CHECK_WITH_INFO( + fp8InputBytes <= mLoraFc1ResultBytes, "FC1 LoRA result buffer is too small for the FP8 FC2 input."); + auto* fp8_input = reinterpret_cast<__nv_fp8_e4m3*>(lora_fc1_result_); + convertMoeLoraInputToFp8(fp8_input, input, num_valid_tokens_ptr, inter_size, num_tokens, stream); + grouped_input = fp8_input; +#else + TLLM_THROW("FP8 MoE LoRA requires an ENABLE_FP8 build."); +#endif + } + else + { + TLLM_CHECK_WITH_INFO(data_type == compute_data_type, + "MoE LoRA cache dtype must match the activation compute dtype unless the cache is FP8."); + } // As in loraFC1, zero the output so rank-0 rows the GEMM skips do not // feed stale data into the downstream add. size_t const fc2_result_bytes = static_cast(num_tokens) * static_cast(hidden_size) * sizeof(ScaleBiasType); TLLM_CUDA_CHECK(cudaMemsetAsync(lora_fc2_result_, 0, fc2_result_bytes, stream)); - runMoeLoraGroupedGemmModule(grouped_gemm.fc2, num_tokens, /*in_hidden_size=*/inter_size, grouped_gemm.max_lora_rank, grouped_gemm.dtype_bytes, grouped_gemm.splitk_slices, - /*input_base=*/static_cast(input), - /*output_base=*/static_cast(lora_fc2_result_), grouped_gemm.run, data_type, stream); + /*input_base=*/grouped_input, + /*output_base=*/static_cast(lora_fc2_result_), grouped_gemm.run, data_type, compute_data_type, + stream); sync_check_cuda_error(stream); return; } diff --git a/cpp/tensorrt_llm/kernels/moe/cutlass/moe_lora_problem_builder.cu b/cpp/tensorrt_llm/kernels/moe/cutlass/moe_lora_problem_builder.cu index 84bc0c99a1b0..023d037a2e24 100644 --- a/cpp/tensorrt_llm/kernels/moe/cutlass/moe_lora_problem_builder.cu +++ b/cpp/tensorrt_llm/kernels/moe/cutlass/moe_lora_problem_builder.cu @@ -43,7 +43,7 @@ __global__ void moeLoraProblemBuilderKernel(int32_t const* __restrict__ ranks, i void** __restrict__ a_ptrs_in, void** __restrict__ b_ptrs_in, void** __restrict__ d_ptrs_in, void** __restrict__ b_ptrs_out, void** __restrict__ d_ptrs_out, int64_t* __restrict__ lda_in, int64_t* __restrict__ ldb_in, int64_t* __restrict__ ldd_in, int64_t* __restrict__ ldb_out, - int64_t* __restrict__ ldd_out, int64_t* __restrict__ splitk_offsets) + int64_t* __restrict__ ldd_out, int64_t* __restrict__ splitk_offsets, int64_t output_dtype_bytes) { int64_t const i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (i >= num_permuted_tokens) @@ -80,10 +80,10 @@ __global__ void moeLoraProblemBuilderKernel(int32_t const* __restrict__ ranks, i problem_sizes_out[i] = cutlass::gemm::GemmCoord(1, out_n, rank); // Pointer rows. dtype_bytes scales the per-row stride so the same - // builder serves bf16/fp16/fp32 adapters without templating. + // builder serves fp8/bf16/fp16/fp32 adapters with a separate output dtype. int64_t const in_row_stride = in_hidden_size * dtype_bytes; int64_t const work_row_stride = max_lora_rank * dtype_bytes; - int64_t const out_row_stride = out_hidden_size * dtype_bytes; + int64_t const out_row_stride = out_hidden_size * output_dtype_bytes; a_ptrs_in[i] = reinterpret_cast(input_base + i * in_row_stride); b_ptrs_in[i] = reinterpret_cast(a_ptr_bits); @@ -121,7 +121,7 @@ __global__ void moeLoraProblemBuilderKernel(int32_t const* __restrict__ ranks, i void launchMoeLoraProblemBuilder(int32_t const* ranks_dev, int64_t const* ptrs_dev, void const* input_base, void* lowrank_workspace, void* output_base, int64_t num_permuted_tokens, int64_t in_hidden_size, int64_t out_hidden_size, int64_t max_lora_rank, int64_t dtype_bytes, int64_t splitk_slices, - MoeLoraGemmGroupArrays const& out, cudaStream_t stream) + MoeLoraGemmGroupArrays const& out, cudaStream_t stream, int64_t output_dtype_bytes) { if (num_permuted_tokens <= 0) { @@ -136,6 +136,7 @@ void launchMoeLoraProblemBuilder(int32_t const* ranks_dev, int64_t const* ptrs_d TLLM_CHECK_WITH_INFO( out.lda_in && out.ldb_in && out.ldd_in && out.ldb_out && out.ldd_out, "All ld* arrays must be non-null"); TLLM_CHECK_WITH_INFO(dtype_bytes > 0, "dtype_bytes must be positive"); + TLLM_CHECK_WITH_INFO(output_dtype_bytes > 0, "output_dtype_bytes must be positive"); TLLM_CHECK_WITH_INFO(max_lora_rank > 0, "max_lora_rank must be positive"); TLLM_CHECK_WITH_INFO(in_hidden_size > 0 && out_hidden_size > 0, "hidden sizes must be positive"); TLLM_CHECK_WITH_INFO(splitk_slices > 0, "splitk_slices must be positive"); @@ -151,7 +152,7 @@ void launchMoeLoraProblemBuilder(int32_t const* ranks_dev, int64_t const* ptrs_d reinterpret_cast(output_base), num_permuted_tokens, in_hidden_size, out_hidden_size, max_lora_rank, dtype_bytes, splitk_slices, out.problem_sizes_in, out.problem_sizes_out, out.a_ptrs_in, out.b_ptrs_in, out.d_ptrs_in, out.b_ptrs_out, out.d_ptrs_out, out.lda_in, out.ldb_in, out.ldd_in, out.ldb_out, out.ldd_out, - out.splitk_offsets); + out.splitk_offsets, output_dtype_bytes); sync_check_cuda_error(stream); } diff --git a/cpp/tensorrt_llm/thop/moe/cutlass/moeOp.cpp b/cpp/tensorrt_llm/thop/moe/cutlass/moeOp.cpp index 16c261163f5a..26271d57a8b4 100644 --- a/cpp/tensorrt_llm/thop/moe/cutlass/moeOp.cpp +++ b/cpp/tensorrt_llm/thop/moe/cutlass/moeOp.cpp @@ -37,6 +37,7 @@ #include "tensorrt_llm/kernels/cuda_graph_grouped_gemm.h" #include "tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h" #include "tensorrt_llm/kernels/cutlass_kernels/include/cutlass_kernel_selector.h" +#include "tensorrt_llm/kernels/groupGemm.h" #include "tensorrt_llm/runtime/torchUtils.h" #include "tensorrt_llm/thop/thUtils.h" @@ -95,7 +96,7 @@ enum class MoeLoraRequestType : int32_t inline void moeLoraGroupedGemmRunImpl(::tensorrt_llm::kernels::cutlass_kernels::MoeLoraGroupedGemmModule const& mod, int64_t num_permuted_tokens, int64_t in_hidden_size, int64_t max_lora_rank, int64_t dtype_bytes, int64_t splitk_slices, void const* input_base, void* output_base, tensorrt_llm::DataType data_type, - cudaStream_t stream) + tensorrt_llm::DataType output_data_type, cudaStream_t stream) { TLLM_CHECK_WITH_INFO(mod.permuted_ranks_dev != nullptr, "Grouped-GEMM LoRA module is missing permuted ranks buffer (forgot to populate grouped_gemm?)."); @@ -120,7 +121,8 @@ inline void moeLoraGroupedGemmRunImpl(::tensorrt_llm::kernels::cutlass_kernels:: ::tensorrt_llm::kernels::cutlass_kernels::launchMoeLoraProblemBuilder(mod.permuted_ranks_dev, mod.permuted_ptrs_dev, input_base, mod.lowrank_workspace_dev, output_base, num_permuted_tokens, in_hidden_size, mod.out_hidden_size, - max_lora_rank, dtype_bytes, splitk_slices, arrays, stream); + max_lora_rank, dtype_bytes, splitk_slices, arrays, stream, + static_cast(common::getDTypeSize(output_data_type))); sync_check_cuda_error(stream); // The cuda_graph_(split_k_)grouped_gemm wrappers accept ldc == ldd when C @@ -146,7 +148,7 @@ inline void moeLoraGroupedGemmRunImpl(::tensorrt_llm::kernels::cutlass_kernels:: ::tensorrt_llm::kernels::cudaGraphGroupedGemm(arrays.problem_sizes_out, static_cast(num_permuted_tokens), arrays.d_ptrs_in /*== a_ptrs_out*/, arrays.b_ptrs_out, arrays.d_ptrs_out, arrays.d_ptrs_out, arrays.ldd_in, arrays.ldb_out, arrays.ldd_out, arrays.ldd_out, - /*isLoraIn=*/false, data_type, kMinKN, host_max_out, stream); + /*isLoraIn=*/false, data_type, kMinKN, host_max_out, stream, output_data_type); sync_check_cuda_error(stream); } @@ -545,6 +547,7 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner torch::optional const& gated_lora_weight_ptrs = torch::nullopt, torch::optional const& host_request_types = torch::nullopt, torch::optional const& host_context_lengths = torch::nullopt, int64_t lora_max_low_rank = 0, + torch::optional const& lora_dtype = torch::nullopt, // Slot-indexed CUDA-graph LoRA inputs (mutually exclusive with the per-request // schema above). When fc1_slot_lora_ranks is provided, the per-token expansion // is performed inside the op via token_to_slot[t] indexed into the slot tables. @@ -797,8 +800,33 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner "MXFP8, and integer quant are not supported."); TORCH_CHECK( mWeightDtype == c10::ScalarType::Half || mWeightDtype == c10::ScalarType::BFloat16 || is_per_tensor_fp8, - "MoE LoRA supports unquantized fp16/bf16 or per-tensor FP8 (qdq) base expert weights only " - "(LoRA adapters are always fp16/bf16)."); + "MoE LoRA supports unquantized fp16/bf16 or per-tensor FP8 (qdq) base expert weights; " + "adapter weights must match the compute dtype or use fp8."); + if (lora_dtype.has_value()) + { +#ifndef ENABLE_FP8 + TORCH_CHECK(*lora_dtype != c10::ScalarType::Float8_e4m3fn, + "Native FP8 MoE LoRA requires an ENABLE_FP8 build; requested fp8 with max rank ", + lora_max_low_rank); +#endif + auto const cache_type = loraTypeFromCacheDtype(*lora_dtype); + auto const compute_type = loraTypeFromActDtype(mActivationDtype); + TORCH_CHECK(cache_type == compute_type || cache_type == tensorrt_llm::DataType::kFP8, + "MoE LoRA cache dtype must match the activation compute dtype or be fp8; got ", *lora_dtype); + if (*lora_dtype == c10::ScalarType::Float8_e4m3fn) + { + int const sm = common::getSMVersion(); + TORCH_CHECK(::tensorrt_llm::kernels::supportsFp8GroupedGemm(sm), + "Native FP8 MoE LoRA requires compiled SM90/SM100 kernels; got SM", sm, " with max rank ", + lora_max_low_rank); + constexpr int64_t kFp8Alignment = 16; + TORCH_CHECK(lora_max_low_rank > 0 && lora_max_low_rank % kFp8Alignment == 0 + && hidden_size % kFp8Alignment == 0 && inter_size % kFp8Alignment == 0, + "Native FP8 MoE LoRA requires rank and hidden dimensions divisible by 16; got max rank ", + lora_max_low_rank, ", hidden size ", hidden_size, ", intermediate size ", inter_size, ", SM", + sm); + } + } // The grouped-GEMM core runs entirely on the stream, so CUDA-graph // capture is supported, but only with the slot-indexed schema. The // per-request schema expands adapters on the host into pinned @@ -815,8 +843,8 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner fc2_lora_weight_ptrs, gated_lora_ranks, gated_lora_weight_ptrs, host_request_types, host_context_lengths, fc1_slot_lora_ranks, fc1_slot_lora_weight_ptrs, fc2_slot_lora_ranks, fc2_slot_lora_weight_ptrs, gated_slot_lora_ranks, gated_slot_lora_weight_ptrs, token_to_slot, - /*num_tokens=*/num_rows, hidden_size, inter_size, mActivationDtype, lora_max_low_rank, is_gated_act, stream, - static_cast(experts_per_token)); + /*num_tokens=*/num_rows, hidden_size, inter_size, mActivationDtype, lora_dtype, lora_max_low_rank, + is_gated_act, stream, static_cast(experts_per_token)); WorkspaceInfo const& workspace_info = getWorkspaceInfo(num_rows, hidden_size, inter_size, num_experts_total, static_cast(experts_per_token), base_activation_type, parallelism_config, min_latency_mode, stream, @@ -1470,10 +1498,8 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner // ===== LoRA helpers ===== - // Map a torch dtype to the TRT-LLM tensorrt_llm::DataType used to size the - // grouped-GEMM low-rank scratch. Kept as a const member (not static) so the - // FP8 case can read mOutputDtype to pick the fp16/bf16 LoRA compute dtype. - tensorrt_llm::DataType loraTypeFromActDtype(c10::ScalarType dtype) const + // Map the PEFT cache scalar type to the grouped-GEMM datatype. + tensorrt_llm::DataType loraTypeFromCacheDtype(c10::ScalarType dtype) const { switch (dtype) { @@ -1483,13 +1509,25 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner case c10::ScalarType::BFloat16: return tensorrt_llm::DataType::kBF16; #endif #ifdef ENABLE_FP8 - case c10::ScalarType::Float8_e4m3fn: + case c10::ScalarType::Float8_e4m3fn: return tensorrt_llm::DataType::kFP8; +#endif + default: C10_THROW_ERROR_FORMATTED(Error, "MoE LoRA only supports fp16/bf16/fp32/fp8 cache dtype."); + } + } + + // FP8 base activations are dequantized to the output dtype for the legacy + // FP16/BF16 LoRA path. An explicit cache dtype bypasses this mapping. + tensorrt_llm::DataType loraTypeFromActDtype(c10::ScalarType dtype) const + { +#ifdef ENABLE_FP8 + if (dtype == c10::ScalarType::Float8_e4m3fn) + { TORCH_CHECK(mOutputDtype != c10::ScalarType::Float8_e4m3fn, "MoE LoRA with FP8 base activations requires an fp16/bf16 output (LoRA compute) dtype."); - return loraTypeFromActDtype(mOutputDtype); -#endif - default: C10_THROW_ERROR_FORMATTED(Error, "MoE LoRA only supports fp16/bf16/fp32 activation dtype."); + return loraTypeFromCacheDtype(mOutputDtype); } +#endif + return loraTypeFromCacheDtype(dtype); } // Expand per-request LoRA ranks and weight-pointer pairs into per-token arrays. @@ -1632,6 +1670,8 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner // realloc guard rejects. TORCH_CHECK(max_lora_rank > 0, "max_lora_rank must be positive to reserve MoE-LoRA device scratch; got ", max_lora_rank); + // The cache must match the compute dtype or be FP8, so the compute + // element size bounds scratch for every supported adapter dtype. int64_t const dtype_bytes = static_cast(common::getDTypeSize(loraTypeFromActDtype(mActivationDtype))); int64_t const capacity = max_num_tokens * experts_per_token; ensureLoraDeviceScratch(capacity, max_lora_rank, dtype_bytes, kGroupedGemmSplitKSlices, has_gated); @@ -1744,8 +1784,9 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner TORCH_CHECK(dtype_bytes > 0, "grouped-GEMM dtype_bytes must be positive; got ", dtype_bytes); TORCH_CHECK(splitk_slices > 0, "grouped-GEMM splitk_slices must be positive; got ", splitk_slices); + // Narrower cache dtypes reuse the byte capacity already reserved for wider ones. bool const need_resize = capacity > mLoraDeviceScratchCapacity || max_lora_rank > mLoraDeviceScratchMaxLoraRank - || dtype_bytes != mLoraDeviceScratchDtypeBytes || splitk_slices != mLoraDeviceScratchSplitKSlices + || dtype_bytes > mLoraDeviceScratchDtypeBytes || splitk_slices != mLoraDeviceScratchSplitKSlices || (has_gated && !mLoraDeviceScratchHasGated); if (!need_resize) { @@ -1758,17 +1799,17 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner // dtype/rank/splitk combo so subsequent calls can early-exit. int64_t const new_capacity = std::max(capacity, mLoraDeviceScratchCapacity); int64_t const new_max_lora_rank = std::max(max_lora_rank, mLoraDeviceScratchMaxLoraRank); + int64_t const new_dtype_bytes = std::max(dtype_bytes, mLoraDeviceScratchDtypeBytes); bool const new_has_gated = mLoraDeviceScratchHasGated || has_gated; // c10::ScalarType for the lowrank workspace. The kernel treats the // buffer opaquely (per-byte stride is dtype_bytes), so we pick a // dtype with matching element size to keep at::Tensor accounting // sensible; consumers cast via .data_ptr(). - c10::ScalarType const dtype_scalar = (dtype_bytes == 2) ? at::kBFloat16 - : (dtype_bytes == 4) ? at::kFloat - : at::kByte; - // Callers should pass bf16/fp16 (2 bytes). Other sizes still work at the - // byte level, but this assertion catches accidental misuse. + c10::ScalarType const dtype_scalar = (new_dtype_bytes == 2) ? at::kBFloat16 + : (new_dtype_bytes == 4) ? at::kFloat + : at::kByte; + // FP8, FP16/BF16 and FP32 use one, two and four bytes respectively. TORCH_CHECK(dtype_bytes == 1 || dtype_bytes == 2 || dtype_bytes == 4, "grouped-GEMM lowrank workspace dtype_bytes must be 1/2/4; got ", dtype_bytes); @@ -1819,7 +1860,7 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner mLoraDeviceScratchCapacity = new_capacity; mLoraDeviceScratchMaxLoraRank = new_max_lora_rank; - mLoraDeviceScratchDtypeBytes = dtype_bytes; + mLoraDeviceScratchDtypeBytes = new_dtype_bytes; mLoraDeviceScratchSplitKSlices = splitk_slices; mLoraDeviceScratchHasGated = new_has_gated; mLoraHostMaxProblemShape.reset(); @@ -1879,8 +1920,8 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner torch::optional const& gated_slot_lora_ranks, torch::optional const& gated_slot_lora_weight_ptrs, torch::optional const& token_to_slot, int64_t num_tokens, int64_t hidden_size, - int64_t inter_size, c10::ScalarType act_dtype, int64_t lora_max_low_rank, bool is_gated_activation, - cudaStream_t stream, int experts_per_token) + int64_t inter_size, c10::ScalarType act_dtype, torch::optional const& lora_dtype, + int64_t lora_max_low_rank, bool is_gated_activation, cudaStream_t stream, int experts_per_token) { #if defined(USING_OSS_CUTLASS_MOE_GEMM) bool const has_per_request = fc1_lora_ranks.has_value(); @@ -1901,6 +1942,26 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner int64_t num_seqs = 0; bool has_gated = false; + // Every rank must fit within lora_max_low_rank, which sizes + // both the lowrank workspace and the max-problem hints. A larger rank + // would make the grouped-GEMM core build GEMM problems wider than + // the allocated scratch and write out of bounds, so reject it up front. + // This host scan runs on every eager invocation, but only at capture time for CUDA graphs. + // Replay does not revalidate updated slot tables: slot updates must preserve the rank-capacity + // and alignment bounds. The loader's _validate_fp8_lora_alignment rejects misaligned FP8 adapters. + auto validate_rank_tensor = [&](char const* name, torch::Tensor const& ranks_tensor) + { + CHECK_CPU_INPUT(ranks_tensor, at::ScalarType::Int) + auto const* rank_data = ranks_tensor.data_ptr(); + for (int64_t i = 0; i < ranks_tensor.size(0); ++i) + { + TORCH_CHECK(rank_data[i] >= 0 && rank_data[i] <= lora_max_low_rank, name, "[", i, "]=", rank_data[i], + " is outside [0, ", lora_max_low_rank, "]."); + TORCH_CHECK(lora_dtype != c10::ScalarType::Float8_e4m3fn || rank_data[i] % 16 == 0, + "Native FP8 MoE LoRA requires active ranks divisible by 16; ", name, "[", i, "]=", rank_data[i]); + } + }; + if (has_per_request) { TORCH_CHECK( @@ -1927,20 +1988,6 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner num_seqs = fc1_lora_ranks->size(0); has_gated = is_gated_activation && gated_lora_ranks.has_value(); - // Every per-request rank must fit within lora_max_low_rank, which sizes - // both the lowrank workspace and the max-problem hints. A larger rank - // would make the grouped-GEMM core build GEMM problems wider than - // the allocated scratch and write out of bounds, so reject it up front. - auto validate_rank_tensor = [&](char const* name, torch::Tensor const& ranks_tensor) - { - CHECK_CPU_INPUT(ranks_tensor, at::ScalarType::Int) - auto const* rank_data = ranks_tensor.data_ptr(); - for (int64_t i = 0; i < ranks_tensor.size(0); ++i) - { - TORCH_CHECK(rank_data[i] >= 0 && rank_data[i] <= lora_max_low_rank, name, "[", i, - "]=", rank_data[i], " is outside [0, ", lora_max_low_rank, "]."); - } - }; validate_rank_tensor("fc1_lora_ranks", *fc1_lora_ranks); validate_rank_tensor("fc2_lora_ranks", *fc2_lora_ranks); if (has_gated) @@ -2019,9 +2066,7 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner // the H2D element sizes are correct and the device kernel reads the // expected types. CHECK_CPU_INPUT((*token_to_slot), at::ScalarType::Int) - CHECK_CPU_INPUT((*fc1_slot_lora_ranks), at::ScalarType::Int) CHECK_CPU_INPUT((*fc1_slot_lora_weight_ptrs), at::ScalarType::Long) - CHECK_CPU_INPUT((*fc2_slot_lora_ranks), at::ScalarType::Int) CHECK_CPU_INPUT((*fc2_slot_lora_weight_ptrs), at::ScalarType::Long) // These tensors are copied to device via async H2D. Under CUDA-graph @@ -2035,6 +2080,8 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner " must be a pinned host tensor for captured async H2D copies."); }; check_pinned(*token_to_slot, "token_to_slot"); + validate_rank_tensor("fc1_slot_lora_ranks", *fc1_slot_lora_ranks); + validate_rank_tensor("fc2_slot_lora_ranks", *fc2_slot_lora_ranks); check_pinned(*fc1_slot_lora_ranks, "fc1_slot_lora_ranks"); check_pinned(*fc1_slot_lora_weight_ptrs, "fc1_slot_lora_weight_ptrs"); check_pinned(*fc2_slot_lora_ranks, "fc2_slot_lora_ranks"); @@ -2085,8 +2132,8 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner h2d_slot(*fc2_slot_lora_weight_ptrs, mLoraSlotFC2PtrsDevice, num_slots * 3); if (has_gated) { - CHECK_CPU_INPUT((*gated_slot_lora_ranks), at::ScalarType::Int) CHECK_CPU_INPUT((*gated_slot_lora_weight_ptrs), at::ScalarType::Long) + validate_rank_tensor("gated_slot_lora_ranks", *gated_slot_lora_ranks); check_pinned(*gated_slot_lora_ranks, "gated_slot_lora_ranks"); check_pinned(*gated_slot_lora_weight_ptrs, "gated_slot_lora_weight_ptrs"); TORCH_CHECK(gated_slot_lora_ranks->size(0) == num_slots, @@ -2168,7 +2215,9 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner // (slot-indexed) inputs feed this capture-safe core. The TensorRT MoE // plugin leaves grouped_gemm.enabled false and uses its own cuBLAS path. { - int64_t const dtype_bytes = static_cast(common::getDTypeSize(loraTypeFromActDtype(act_dtype))); + tensorrt_llm::DataType const lora_data_type + = lora_dtype.has_value() ? loraTypeFromCacheDtype(*lora_dtype) : loraTypeFromActDtype(act_dtype); + int64_t const dtype_bytes = static_cast(common::getDTypeSize(lora_data_type)); int64_t const capacity = num_tokens * static_cast(experts_per_token); // Pass stream so a mid-capture resize (which would invalidate // previously captured graphs) is rejected with a clear error @@ -2181,6 +2230,7 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner grouped_gemm.in_hidden_size = hidden_size; grouped_gemm.max_lora_rank = lora_max_low_rank; grouped_gemm.dtype_bytes = dtype_bytes; + grouped_gemm.data_type = lora_data_type; grouped_gemm.splitk_slices = kGroupedGemmSplitKSlices; grouped_gemm.has_gated = has_gated; // Populate the libtorch-bound GEMM dispatch entry point so diff --git a/cpp/tests/unit_tests/kernels/moe/moeLoraProblemBuilderTest.cu b/cpp/tests/unit_tests/kernels/moe/moeLoraProblemBuilderTest.cu index ef78776304d2..a6a161b57bc9 100644 --- a/cpp/tests/unit_tests/kernels/moe/moeLoraProblemBuilderTest.cu +++ b/cpp/tests/unit_tests/kernels/moe/moeLoraProblemBuilderTest.cu @@ -50,7 +50,7 @@ struct RefOutputs RefOutputs cpuReference(std::vector const& ranks, std::vector const& ptrs, int64_t input_base, int64_t lowrank_workspace, int64_t output_base, int64_t in_hidden_size, int64_t out_hidden_size, - int64_t max_lora_rank, int64_t dtype_bytes, int64_t splitk_slices) + int64_t max_lora_rank, int64_t dtype_bytes, int64_t splitk_slices, int64_t output_dtype_bytes) { int64_t const P = static_cast(ranks.size()); RefOutputs r; @@ -80,7 +80,7 @@ RefOutputs cpuReference(std::vector const& ranks, std::vector int64_t const in_row_stride = in_hidden_size * dtype_bytes; int64_t const work_row_stride = max_lora_rank * dtype_bytes; - int64_t const out_row_stride = out_hidden_size * dtype_bytes; + int64_t const out_row_stride = out_hidden_size * output_dtype_bytes; r.a_ptrs_in[i] = input_base + i * in_row_stride; r.b_ptrs_in[i] = ptrs[2 * i + 0]; @@ -173,11 +173,12 @@ protected: // sentinel thread); the splitk_offsets comparison is then skipped. void runAndCompare(std::vector const& ranks, std::vector const& ptrs, int64_t input_base, int64_t lowrank_workspace, int64_t output_base, int64_t in_hidden_size, int64_t out_hidden_size, - int64_t max_lora_rank, int64_t dtype_bytes, int64_t splitk_slices, bool with_splitk = true) + int64_t max_lora_rank, int64_t dtype_bytes, int64_t splitk_slices, int64_t output_dtype_bytes, + bool with_splitk = true) { auto const P = static_cast(ranks.size()); RefOutputs ref = cpuReference(ranks, ptrs, input_base, lowrank_workspace, output_base, in_hidden_size, - out_hidden_size, max_lora_rank, dtype_bytes, splitk_slices); + out_hidden_size, max_lora_rank, dtype_bytes, splitk_slices, output_dtype_bytes); int32_t* ranks_dev = upload(ranks); int64_t* ptrs_dev = upload(ptrs); @@ -201,7 +202,7 @@ protected: launchMoeLoraProblemBuilder(ranks_dev, ptrs_dev, reinterpret_cast(input_base), reinterpret_cast(lowrank_workspace), reinterpret_cast(output_base), P, in_hidden_size, - out_hidden_size, max_lora_rank, dtype_bytes, splitk_slices, out, mStream); + out_hidden_size, max_lora_rank, dtype_bytes, splitk_slices, out, mStream, output_dtype_bytes); TLLM_CUDA_CHECK(cudaStreamSynchronize(mStream)); // Compare device outputs to host reference. @@ -280,7 +281,22 @@ TEST_F(MoeLoraProblemBuilderTest, Bf16Smoke) } runAndCompare(ranks, ptrs, input_base, lowrank_workspace, output_base, in_hidden_size, out_hidden_size, - max_lora_rank, dtype_bytes, splitk_slices); + max_lora_rank, dtype_bytes, splitk_slices, dtype_bytes); +} + +TEST_F(MoeLoraProblemBuilderTest, Fp8InputWithBf16OutputStrides) +{ + std::vector const ranks = {16, 0, 32}; + std::vector ptrs; + for (int32_t i = 0; i < static_cast(ranks.size()); ++i) + { + ptrs.push_back(fakeAdapter(1, i, 0)); + ptrs.push_back(fakeAdapter(1, i, 1)); + } + // Input and low-rank rows are one byte per element; output rows are two. + runAndCompare(ranks, ptrs, 0x1'0000'0000LL, 0x2'0000'0000LL, 0x3'0000'0000LL, + /*in_hidden_size=*/128, /*out_hidden_size=*/256, /*max_lora_rank=*/32, + /*dtype_bytes=*/1, /*splitk_slices=*/1, /*output_dtype_bytes=*/2, /*with_splitk=*/false); } TEST_F(MoeLoraProblemBuilderTest, Fp32StrideBytes) @@ -304,7 +320,7 @@ TEST_F(MoeLoraProblemBuilderTest, Fp32StrideBytes) } runAndCompare(ranks, ptrs, input_base, lowrank_workspace, output_base, in_hidden_size, out_hidden_size, - max_lora_rank, dtype_bytes, splitk_slices); + max_lora_rank, dtype_bytes, splitk_slices, dtype_bytes); } // Cover an empty call (no-op) and a single-token call (smallest live case) @@ -326,7 +342,7 @@ TEST_F(MoeLoraProblemBuilderTest, BoundaryCases) launchMoeLoraProblemBuilder(nullptr, nullptr, reinterpret_cast(input_base), reinterpret_cast(lowrank_workspace), reinterpret_cast(output_base), /*num_permuted_tokens=*/0, in_hidden_size, out_hidden_size, max_lora_rank, dtype_bytes, splitk_slices, - empty, mStream); + empty, mStream, dtype_bytes); TLLM_CUDA_CHECK(cudaStreamSynchronize(mStream)); } @@ -335,7 +351,7 @@ TEST_F(MoeLoraProblemBuilderTest, BoundaryCases) std::vector ranks = {3}; std::vector ptrs = {fakeAdapter(3, 0, 0), fakeAdapter(3, 0, 1)}; runAndCompare(ranks, ptrs, input_base, lowrank_workspace, output_base, in_hidden_size, out_hidden_size, - max_lora_rank, dtype_bytes, splitk_slices); + max_lora_rank, dtype_bytes, splitk_slices, dtype_bytes); } } @@ -363,7 +379,7 @@ TEST_F(MoeLoraProblemBuilderTest, NullSplitkOffsets) } runAndCompare(ranks, ptrs, input_base, lowrank_workspace, output_base, in_hidden_size, out_hidden_size, - max_lora_rank, dtype_bytes, splitk_slices, /*with_splitk=*/false); + max_lora_rank, dtype_bytes, splitk_slices, dtype_bytes, /*with_splitk=*/false); } } // namespace diff --git a/docs/source/features/lora.md b/docs/source/features/lora.md index d8949ef4b8c4..ff1df37bf74e 100644 --- a/docs/source/features/lora.md +++ b/docs/source/features/lora.md @@ -121,7 +121,7 @@ llm = LLM( #### Native FP8 adapter support -The base model's quantization and the LoRA adapter's data type are independent. Dense LoRA modules can keep +The base model's quantization and the LoRA adapter's data type are independent. Dense and routed-expert MoE LoRA modules can keep FP8 E4M3 adapter weights and execute them with native FP8 grouped GEMM kernels on the following architectures: | GPU architecture | Native FP8 adapter support | @@ -132,8 +132,27 @@ FP8 E4M3 adapter weights and execute them with native FP8 grouped GEMM kernels o | Blackwell (SM120/SM121) | No | If native FP8 LoRA kernels are unavailable for the current device or were excluded from the TensorRT-LLM build, -the adapter weights are converted to the model compute data type. Native FP8 adapter weights are not supported -for routed-expert MoE LoRA modules; see [Routed-Expert MoE LoRA](#routed-expert-moe-lora). +the adapter weights are converted to the model compute data type. Routed-expert modules require the CUTLASS +backend; see [Routed-Expert MoE LoRA](#routed-expert-moe-lora). + +For routed-expert adapters, native FP8 is selected from the checkpoint's `float8_e4m3fn` dtype. +All adapter matrices must have one dtype, and the rank, input dimension, and output dimension must be +multiples of 16. DoRA is unsupported. Save the adapter matrices as BF16 or FP16 to use the compute-dtype +path instead. + +Native FP8 reduces adapter tensor storage from two bytes per element (BF16/FP16) to one byte. This +halves the weight payload in the PEFT cache, not total engine memory or necessarily allocated cache capacity. +This storage reduction does not imply a latency or throughput improvement: each routed-expert FC1/FC2 LoRA stage casts its +input to FP8, adding a full input read/write pass. These casts reuse existing LoRA result buffers; +they do not add full-size workspace allocations. The final grouped-GEMM epilogue writes directly to +BF16/FP16, so no result-conversion pass is needed. + +This execution path is lossy. Activations and the low-rank intermediate are rounded to E4M3 with +unit scale and saturation at magnitude 448. The final LoRA delta is accumulated in FP32 and stored +in the model's BF16/FP16 compute dtype; it is not rounded or saturated to FP8. There is no dynamic +activation scaling. Check application accuracy before deploying an FP8 adapter, especially when +activations or low-rank intermediates approach the saturation limit. Accuracy depends on the adapter +and workload; there is no universal error bound. ### NeMo LoRA Format @@ -162,7 +181,7 @@ LoRA can be applied to the routed-expert projections of a Mixture-of-Experts (Mo | Aspect | Supported | |---|---| | MoE backend | `CUTLASS` only (other backends raise an error at construction). | -| Base-weight dtype | bf16 / fp16. Quantized base weights (FP8, NVFP4, INT4, INT8) are not yet supported. | +| Base-weight dtype | bf16 / fp16 or per-tensor FP8 (QDQ). FP8 block-scale, NVFP4, MXFP8, INT4, and INT8 are not supported. | | Adapter modules | `moe_h_to_4h` (gate side of SwiGLU), `moe_gate` (up side), `moe_4h_to_h` (down). `moe_h_to_4h` and `moe_4h_to_h` must both be present; `moe_gate` is optional (gated activations only). | | Adapter layout | Per-expert (stacked `[num_experts, ...]`). | | Multi-LoRA in flight | Yes. Reuses the existing slot manager. | @@ -229,7 +248,7 @@ The per-request schema is not itself CUDA-graph capturable (its host-side adapte #### What is rejected, and where -If you supply MoE LoRA on a non-Cutlass backend or with quantization, `create_moe` raises at construction with a message pointing at the offending setting. At runtime, the fused MoE op also rejects min-latency mode + LoRA, alltoall + LoRA, and per-request (eager-schema) inputs under CUDA-graph capture (use the slot-indexed schema for capture). +If you supply MoE LoRA on a non-Cutlass backend or with unsupported base quantization, `create_moe` raises at construction with a message pointing at the offending setting. At runtime, the fused MoE op also rejects min-latency mode + LoRA, alltoall + LoRA, and per-request (eager-schema) inputs under CUDA-graph capture (use the slot-indexed schema for capture). ### Cache Management diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 37fc0269eca7..06a4b9de32df 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -276,6 +276,7 @@ def fused_moe( host_request_types: Optional[torch.Tensor] = None, host_context_lengths: Optional[torch.Tensor] = None, lora_max_low_rank: int = 0, + lora_dtype: Optional[torch.dtype] = None, # Slot-indexed CUDA-graph LoRA inputs (mutually exclusive with the per-request # tensors above). When fc1_slot_lora_ranks is provided, the per-token expansion # is performed inside the op via token_to_slot indexed into the slot tables. @@ -383,9 +384,10 @@ def fused_moe( use_dynamic_fc2_scale, fc1_lora_ranks, fc1_lora_weight_ptrs, fc2_lora_ranks, fc2_lora_weight_ptrs, gated_lora_ranks, gated_lora_weight_ptrs, host_request_types, host_context_lengths, - lora_max_low_rank, fc1_slot_lora_ranks, fc1_slot_lora_weight_ptrs, - fc2_slot_lora_ranks, fc2_slot_lora_weight_ptrs, - gated_slot_lora_ranks, gated_slot_lora_weight_ptrs, token_to_slot + lora_max_low_rank, lora_dtype, fc1_slot_lora_ranks, + fc1_slot_lora_weight_ptrs, fc2_slot_lora_ranks, + fc2_slot_lora_weight_ptrs, gated_slot_lora_ranks, + gated_slot_lora_weight_ptrs, token_to_slot ] run_moe_args.append(swiglu_clamp_after_silu) try: @@ -455,6 +457,7 @@ def _(input: torch.Tensor, host_request_types: Optional[torch.Tensor] = None, host_context_lengths: Optional[torch.Tensor] = None, lora_max_low_rank: int = 0, + lora_dtype: Optional[torch.dtype] = None, fc1_slot_lora_ranks: Optional[torch.Tensor] = None, fc1_slot_lora_weight_ptrs: Optional[torch.Tensor] = None, fc2_slot_lora_ranks: Optional[torch.Tensor] = None, diff --git a/tensorrt_llm/_torch/moe/fused_moe/cutlass/lora.py b/tensorrt_llm/_torch/moe/fused_moe/cutlass/lora.py index a99cf30bcac5..7620e4bc6deb 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/cutlass/lora.py +++ b/tensorrt_llm/_torch/moe/fused_moe/cutlass/lora.py @@ -317,6 +317,7 @@ def _slice(t: torch.Tensor | None) -> torch.Tensor | None: "host_request_types": _slice(lora_params["host_request_types"]), "host_context_lengths": _slice(lora_params["prompt_lens_cpu"]), "lora_max_low_rank": active_max_rank, + "lora_dtype": lora_params.get("data_type"), } def _extract_moe_lora_tensors_cuda_graph(self, lora_params: dict) -> dict[str, object] | None: @@ -376,4 +377,5 @@ def _source(module_type: LoraModuleType) -> _RanksAndPtrs | None: ), "token_to_slot": token_to_slot, "lora_max_low_rank": max_rank, + "lora_dtype": lora_params.get("data_type"), } diff --git a/tensorrt_llm/_torch/peft/lora/loaders.py b/tensorrt_llm/_torch/peft/lora/loaders.py index 4a4285560df8..0c703cbca633 100644 --- a/tensorrt_llm/_torch/peft/lora/loaders.py +++ b/tensorrt_llm/_torch/peft/lora/loaders.py @@ -396,14 +396,17 @@ def get_target_modules(self, trtllm_modules_to_hf_modules): lora_target_modules.add(trtllm_module) return list(lora_target_modules) - def get_dense_lora_dtype(self) -> Optional[torch.dtype]: - """Return the common input/output dtype for dense LoRA modules.""" + def get_lora_dtype(self) -> Optional[torch.dtype]: + """Return the common input/output dtype across all LoRA modules.""" lora_dtypes = set() for key, weight in self.lora_weight.items(): match = HF_LORA_PATTERN.match(key) - if match is not None and match.group(6) is None and match.group(8) in ("A", "B"): + if match is not None and match.group(8) in ("A", "B"): lora_dtypes.add(weight.dtype) - return next(iter(lora_dtypes)) if len(lora_dtypes) == 1 else None + if len(lora_dtypes) > 1: + dtype_names = ", ".join(sorted(str(dtype) for dtype in lora_dtypes)) + raise ValueError(f"LoRA checkpoint weights must have one dtype; found {dtype_names}.") + return next(iter(lora_dtypes)) if lora_dtypes else None @lru_cache(maxsize=128) diff --git a/tensorrt_llm/_torch/peft/lora/manager.py b/tensorrt_llm/_torch/peft/lora/manager.py index bf11bb5e63bf..9181e94768c2 100644 --- a/tensorrt_llm/_torch/peft/lora/manager.py +++ b/tensorrt_llm/_torch/peft/lora/manager.py @@ -159,8 +159,8 @@ def load_torch_hf_lora( loaded later by LoraManager when requests arrive with LoRA UIDs. Returns: - The common dense LoRA weight dtype, or ``None`` when no homogeneous - dense LoRA weights are present. + The common LoRA weight dtype, or ``None`` when no homogeneous LoRA + weights are present. """ if not lora_config.trtllm_modules_to_hf_modules: lora_config.trtllm_modules_to_hf_modules = get_default_trtllm_modules_to_hf_modules() @@ -182,7 +182,7 @@ def load_torch_hf_lora( missing_qkv_modules = LoraManager.get_missing_qkv_modules(lora_config.lora_target_modules) lora_config.lora_target_modules.extend(missing_qkv_modules) - return lora_loader.get_dense_lora_dtype() + return lora_loader.get_lora_dtype() def load_torch_nemo_lora( @@ -698,14 +698,17 @@ def get_output_dtype(module_weights): for weights in module_weights.values() if "out" in weights } - return next(iter(output_dtypes)) if len(output_dtypes) == 1 else model_dtype + if len(output_dtypes) > 1: + dtype_names = ", ".join(sorted(str(dtype) for dtype in output_dtypes)) + raise ValueError( + f"MoE LoRA experts must have one output dtype; {model_dir} has {dtype_names}." + ) + return next(iter(output_dtypes)) if output_dtypes else model_dtype return module_weights["out"].dtype if "out" in module_weights else model_dtype def uses_native_fp8(module_weights): return ( - get_output_dtype(module_weights) == torch.float8_e4m3fn - and not _is_moe_module_weights(module_weights) - and supports_native_fp8 + get_output_dtype(module_weights) == torch.float8_e4m3fn and supports_native_fp8 ) for layer_weights in all_weights.values(): @@ -736,8 +739,8 @@ def uses_native_fp8(module_weights): dtype_names = ", ".join(sorted(str(dtype) for dtype in cache_dtypes)) raise ValueError( "A LoRA adapter must use one PEFT cache dtype across all modules; " - f"{model_dir} requires {dtype_names}. Mixing native FP8 dense modules " - "with compute-dtype modules such as routed-expert MoE is not supported." + f"{model_dir} requires {dtype_names}. Mixing native FP8 modules " + "with compute-dtype modules is not supported." ) cpp_lora_weights = [] diff --git a/tests/integration/defs/llmapi/test_llm_api_pytorch_moe_lora.py b/tests/integration/defs/llmapi/test_llm_api_pytorch_moe_lora.py index 0fd9e0adf421..e62684321cdb 100644 --- a/tests/integration/defs/llmapi/test_llm_api_pytorch_moe_lora.py +++ b/tests/integration/defs/llmapi/test_llm_api_pytorch_moe_lora.py @@ -14,10 +14,9 @@ # limitations under the License. """Routed-expert (MoE) LoRA integration tests on the PyTorch CUTLASS backend. -Covers unquantized bf16 Qwen3-MoE base weights, with several adapters of varying -rank applied in one batch. The adapters are fabricated on disk in the per-expert -key layout the TRT-LLM loader expects, so the tests do not depend on a real PEFT -export. +Covers BF16 and native FP8 adapters on unquantized BF16 Qwen3-MoE base weights. +The adapters are fabricated on disk in the per-expert key layout the TRT-LLM +loader expects, so the tests do not depend on a real PEFT export. These tests require their model checkpoints under LLM_MODELS_ROOT and fail (not skip) when a checkpoint is missing, so a misconfigured model root surfaces as a @@ -30,10 +29,15 @@ import pytest import torch +from safetensors.torch import save_file +from transformers import Qwen3MoeConfig, Qwen3MoeForCausalLM from tensorrt_llm import LLM from tensorrt_llm._torch.peft.lora.config import LoraConfig +from tensorrt_llm._torch.peft.lora.manager import supports_native_fp8_lora +from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType from tensorrt_llm.executor.request import LoRARequest +from tensorrt_llm.executor.worker import GenerationExecutorWorker from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig, SamplingParams from tensorrt_llm.llmapi.llm_args import MoeConfig, PeftCacheConfig @@ -48,9 +52,17 @@ # later ones prefill a single token, which is a different forward with a # different result. Off, every call measures the same thing. _KV_CACHE_CONFIG = KvCacheConfig(free_gpu_memory_fraction=0.4, enable_block_reuse=False) +_TARGET_MODULES = ["moe_h_to_4h", "moe_gate", "moe_4h_to_h"] +_MODULE_MAP = { + "moe_h_to_4h": "gate_proj", + "moe_gate": "up_proj", + "moe_4h_to_h": "down_proj", +} # Adapters of varying rank; max_lora_rank must cover the largest. _RANKS = [8, 16, 32, 16, 64] +_FP8_RANK = 16 +_FP8_DELTA_RELATIVE_ERROR = 0.10 def _write_routed_expert_lora_adapter( @@ -240,6 +252,154 @@ def max_divergence(a: tuple[float, ...], b: tuple[float, ...]) -> float: llm.shutdown() +def _write_paired_routed_expert_lora_adapters( + save_dir: str, + *, + config: dict, +) -> tuple[str, str]: + """Write BF16 and FP8 adapters with identical E4M3-rounded values.""" + bf16_dir = os.path.join(save_dir, "adapter-bf16") + fp8_dir = os.path.join(save_dir, "adapter-fp8") + os.makedirs(bf16_dir) + os.makedirs(fp8_dir) + + hidden_size = config["hidden_size"] + intermediate_size = config["moe_intermediate_size"] + num_experts = config["num_experts"] + layer_idx = config["num_hidden_layers"] - 1 + projections = ( + ("gate_proj", hidden_size, intermediate_size), + ("up_proj", hidden_size, intermediate_size), + ("down_proj", intermediate_size, hidden_size), + ) + generator = torch.Generator().manual_seed(15314) + bf16_weights = {} + fp8_weights = {} + for expert_idx in range(num_experts): + prefix = f"base_model.model.model.layers.{layer_idx}.mlp.experts.{expert_idx}" + for projection, in_features, out_features in projections: + for suffix, shape in ( + ("A", (_FP8_RANK, in_features)), + ("B", (out_features, _FP8_RANK)), + ): + master = torch.randn(shape, generator=generator, dtype=torch.float32) * 0.05 + fp8_weight = master.to(torch.float8_e4m3fn) + key = f"{prefix}.{projection}.lora_{suffix}.weight" + fp8_weights[key] = fp8_weight + bf16_weights[key] = fp8_weight.to(torch.bfloat16) + + assert all(torch.count_nonzero(weight.float()) for weight in fp8_weights.values()) + save_file(bf16_weights, os.path.join(bf16_dir, "adapter_model.safetensors")) + save_file(fp8_weights, os.path.join(fp8_dir, "adapter_model.safetensors")) + adapter_config = { + "bias": "none", + "inference_mode": True, + "lora_alpha": _FP8_RANK, + "lora_dropout": 0.0, + "peft_type": "LORA", + "r": _FP8_RANK, + "target_modules": ["gate_proj", "up_proj", "down_proj"], + "task_type": "CAUSAL_LM", + "use_rslora": False, + } + for adapter_dir in (bf16_dir, fp8_dir): + with open(os.path.join(adapter_dir, "adapter_config.json"), "w") as config_file: + json.dump(adapter_config, config_file) + return bf16_dir, fp8_dir + + +def _local_worker(llm: LLM) -> GenerationExecutorWorker: + assert isinstance(llm._executor, GenerationExecutorWorker) + return llm._executor + + +def _run_paired_adapter_engine( + model_dir: str, + adapter_dir: str, + adapter_dtype: torch.dtype, + mode: str, + num_hidden_layers: int, +) -> torch.Tensor: + """Run prefill/decode and return the adapter's first-token logit delta.""" + lora_config = LoraConfig( + lora_dir=[adapter_dir], + lora_target_modules=_TARGET_MODULES, + trtllm_modules_to_hf_modules=_MODULE_MAP, + max_lora_rank=_FP8_RANK, + max_loras=1, + max_cpu_loras=1, + cuda_graph_specialize_lora=mode == "cudagraph", + ) + cuda_graph_config = CudaGraphConfig(max_batch_size=1) if mode == "cudagraph" else None + module_capacity = _FP8_RANK * num_hidden_layers * len(_TARGET_MODULES) + llm = LLM( + model=model_dir, + lora_config=lora_config, + moe_config=MoeConfig(backend="CUTLASS"), + kv_cache_config=KvCacheConfig(max_tokens=128, enable_block_reuse=False), + skip_tokenizer_init=True, + peft_cache_config=PeftCacheConfig( + num_host_module_layer=module_capacity, + num_device_module_layer=module_capacity, + ), + cuda_graph_config=cuda_graph_config, + gather_generation_logits=True, + max_batch_size=1, + max_num_tokens=64, + ) + try: + replay_count = [0] + if mode == "cudagraph": + runner = _local_worker(llm).engine.model_engine.cuda_graph_runner + original_replay = runner.replay + + def counted_replay(key, current_inputs, *, position_id_offsets=None): + replay_count[0] += 1 + return original_replay(key, current_inputs, position_id_offsets=position_id_offsets) + + runner.replay = counted_replay + + sampling_params = SamplingParams( + max_tokens=4, + temperature=0.0, + ignore_eos=True, + return_generation_logits=True, + ) + prompt = [1, 2, 3, 4] + base_output = llm.generate([prompt], sampling_params)[0].outputs[0] + replay_before_lora = replay_count[0] + request = LoRARequest(f"moe-lora-{adapter_dtype}", 0, adapter_dir) + lora_output = llm.generate([prompt], sampling_params, lora_request=[request])[0].outputs[0] + + assert len(base_output.token_ids) == sampling_params.max_tokens + assert len(lora_output.token_ids) == sampling_params.max_tokens + if mode == "cudagraph": + assert replay_count[0] > replay_before_lora + + cache_manager = _local_worker(llm).engine.resource_manager.get_resource_manager( + ResourceManagerType.PEFT_CACHE_MANAGER + ) + assert cache_manager is not None + assert cache_manager.data_type == adapter_dtype + assert cache_manager.impl.is_task_cached(request.lora_int_id) + low_ranks = cache_manager.get_lora_manager().uid_to_low_ranks(str(request.lora_int_id)) + nonzero_ranks = [ + rank for layer_ranks in low_ranks.values() for rank in layer_ranks.values() if rank > 0 + ] + assert nonzero_ranks == [_FP8_RANK] * len(_TARGET_MODULES) + + assert base_output.generation_logits is not None + assert lora_output.generation_logits is not None + base_logits = base_output.generation_logits[0].to(dtype=torch.float32, device="cpu") + lora_logits = lora_output.generation_logits[0].to(dtype=torch.float32, device="cpu") + assert torch.isfinite(lora_logits).all() + delta = lora_logits - base_logits + assert torch.count_nonzero(delta) + return delta + finally: + llm.shutdown() + + # Each parametrization holds the 30B weights plus a 6.6GB LoRA device cache (the # rank-64 adapter sits exactly at the PEFT cache floor for this model, so the # cache cannot be sized down). Reusing one MPI worker across both leaves ~8GiB of @@ -255,13 +415,6 @@ def test_qwen_moe_routed_expert_multi_lora_varying_ranks( cuda_graph_config = CudaGraphConfig(max_batch_size=10) if moe_lora_mode == "cudagraph" else None model_dir = f"{llm_models_root()}/Qwen3/Qwen3-30B-A3B" - target_modules = ["moe_h_to_4h", "moe_gate", "moe_4h_to_h"] - trtllm_modules_to_hf_modules = { - "moe_h_to_4h": "gate_proj", - "moe_gate": "up_proj", - "moe_4h_to_h": "down_proj", - } - with open(f"{model_dir}/config.json") as f: cfg = json.load(f) num_experts = cfg["num_experts"] @@ -272,7 +425,7 @@ def test_qwen_moe_routed_expert_multi_lora_varying_ranks( # logits directly, without fabricating adapters for the entire 30B model. moe_layers = [cfg["num_hidden_layers"] - 1] ranks = _RANKS - num_module_layers = max(ranks) * cfg["num_hidden_layers"] * len(target_modules) + num_module_layers = max(ranks) * cfg["num_hidden_layers"] * len(_TARGET_MODULES) peft_cache_config = PeftCacheConfig( num_host_module_layer=num_module_layers, num_device_module_layer=num_module_layers, @@ -298,9 +451,71 @@ def test_qwen_moe_routed_expert_multi_lora_varying_ranks( model_dir, lora_paths, max_rank=max(ranks), - target_modules=target_modules, - trtllm_modules_to_hf_modules=trtllm_modules_to_hf_modules, + target_modules=_TARGET_MODULES, + trtllm_modules_to_hf_modules=_MODULE_MAP, cuda_graph_config=cuda_graph_config, preallocate_all_adapters=False, peft_cache_config=peft_cache_config, ) + + +@pytest.mark.skipif( + not torch.cuda.is_available() + or not supports_native_fp8_lora(torch.cuda.get_device_capability()), + reason="Native FP8 LoRA requires SM90/SM100 kernels.", +) +@pytest.mark.parametrize("moe_lora_mode", ["eager", "cudagraph"]) +def test_qwen_moe_native_fp8_lora_matches_bf16_reference( + moe_lora_mode: str, tmp_path, monkeypatch +) -> None: + """Exercise the full cache/engine path with a small, deterministic Qwen3 MoE.""" + monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") + config = Qwen3MoeConfig( + vocab_size=128, + hidden_size=128, + intermediate_size=256, + moe_intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=32, + num_experts=4, + num_experts_per_tok=2, + max_position_embeddings=128, + bos_token_id=1, + eos_token_id=2, + tie_word_embeddings=False, + ) + model_dir = str(tmp_path / "model") + with torch.random.fork_rng(devices=[]): + torch.manual_seed(15314) + model = Qwen3MoeForCausalLM(config).to(torch.bfloat16) + weights = model.state_dict() + # Export the per-expert layout consumed by the Qwen3-MoE weight loader. + for layer_idx in range(config.num_hidden_layers): + prefix = f"model.layers.{layer_idx}.mlp.experts" + gate_up = weights.pop(f"{prefix}.gate_up_proj") + down = weights.pop(f"{prefix}.down_proj") + for expert_idx in range(config.num_experts): + gate, up = gate_up[expert_idx].chunk(2, dim=0) + for projection, tensor in ( + ("gate_proj", gate), + ("up_proj", up), + ("down_proj", down[expert_idx]), + ): + weights[f"{prefix}.{expert_idx}.{projection}.weight"] = tensor.contiguous() + model.save_pretrained(model_dir, state_dict=weights) + del model + bf16_dir, fp8_dir = _write_paired_routed_expert_lora_adapters( + str(tmp_path), config=config.to_dict() | {"num_experts": config.num_experts} + ) + bf16_delta = _run_paired_adapter_engine( + model_dir, bf16_dir, torch.bfloat16, moe_lora_mode, config.num_hidden_layers + ) + fp8_delta = _run_paired_adapter_engine( + model_dir, fp8_dir, torch.float8_e4m3fn, moe_lora_mode, config.num_hidden_layers + ) + relative_error = torch.linalg.vector_norm(fp8_delta - bf16_delta) / torch.linalg.vector_norm( + bf16_delta + ) + assert relative_error < _FP8_DELTA_RELATIVE_ERROR diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 209c866aa838..29643f5689dc 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -376,6 +376,8 @@ l0_h100: - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_field_completeness - llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_routed_expert_multi_lora_varying_ranks[eager] TIMEOUT (300) - llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_routed_expert_multi_lora_varying_ranks[cudagraph] TIMEOUT (300) + - llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_native_fp8_lora_matches_bf16_reference[eager] TIMEOUT (180) + - llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_native_fp8_lora_matches_bf16_reference[cudagraph] TIMEOUT (180) - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[eagle3_one_model=True-enable_chunked_prefill=True-enable_max_concurrency=False-enable_draft_len_schedule=False] - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[eagle3_one_model=True-enable_chunked_prefill=False-enable_max_concurrency=False-enable_draft_len_schedule=True] - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[eagle3_one_model=True-enable_chunked_prefill=False-enable_max_concurrency=True-enable_draft_len_schedule=False] diff --git a/tests/unittest/_torch/peft/test_moe_lora_cuda_graph_params.py b/tests/unittest/_torch/peft/test_moe_lora_cuda_graph_params.py index 0aa4d7021c9d..a8edaf138067 100644 --- a/tests/unittest/_torch/peft/test_moe_lora_cuda_graph_params.py +++ b/tests/unittest/_torch/peft/test_moe_lora_cuda_graph_params.py @@ -328,12 +328,14 @@ def test_extract_moe_lora_tensors_cuda_graph_wires_slot_tables(): "use_cuda_graph_mode": True, "cuda_graph_params": params, "num_seqs": num_seqs, + "data_type": torch.float8_e4m3fn, } kwargs = CutlassMoELoraMixin._extract_moe_lora_tensors_cuda_graph(_ExtractStub(0), lora_params) assert kwargs is not None # max_rank is the global cap, not a per-step value. assert kwargs["lora_max_low_rank"] == rank + assert kwargs["lora_dtype"] == torch.float8_e4m3fn # All three modules present, mapped to their kernel slots. for kernel, module_id in ( diff --git a/tests/unittest/_torch/peft/test_moe_lora_extract.py b/tests/unittest/_torch/peft/test_moe_lora_extract.py index ead8a71c355e..4ab68a3c05b9 100644 --- a/tests/unittest/_torch/peft/test_moe_lora_extract.py +++ b/tests/unittest/_torch/peft/test_moe_lora_extract.py @@ -73,6 +73,30 @@ def _make_lora_params(layer_idx, modules, num_seqs=1): } +# --------------------------------------------------------------------------- +# Cache dtype propagation. +# --------------------------------------------------------------------------- + + +def test_extract_propagates_fp8_cache_dtype(): + layer_idx = 0 + modules = { + LoraModuleType.MOE_H_TO_4H: (16, 0x1110, 0x1111), + LoraModuleType.MOE_4H_TO_H: (16, 0x2220, 0x2221), + } + lora_params = _make_lora_params(layer_idx, modules) + lora_params["data_type"] = torch.float8_e4m3fn + + result = _extract(layer_idx, lora_params) + + assert result is not None + assert result["lora_dtype"] == torch.float8_e4m3fn + + +def test_extract_allows_fp8_cache_without_active_moe_weights(): + assert _extract(0, {"data_type": torch.float8_e4m3fn}) is None + + # --------------------------------------------------------------------------- # Routing correctness (the swap regression). # --------------------------------------------------------------------------- diff --git a/tests/unittest/_torch/peft/test_moe_lora_grouped_gemm.py b/tests/unittest/_torch/peft/test_moe_lora_grouped_gemm.py index a46cdf97e560..f8b58947c937 100644 --- a/tests/unittest/_torch/peft/test_moe_lora_grouped_gemm.py +++ b/tests/unittest/_torch/peft/test_moe_lora_grouped_gemm.py @@ -32,6 +32,7 @@ from utils.util import skip_pre_ada from tensorrt_llm._torch.moe.fused_moe.cutlass import TrtllmCutlassFp8Impl +from tensorrt_llm._torch.peft.lora.manager import supports_native_fp8_lora from tensorrt_llm._torch.peft.lora.moe_layout import make_per_expert_lora, reference_swiglu_moe_lora from tensorrt_llm._torch.utils import ActivationType @@ -43,6 +44,13 @@ ) +requires_native_fp8 = pytest.mark.skipif( + not torch.cuda.is_available() + or not supports_native_fp8_lora(torch.cuda.get_device_capability()), + reason="Native FP8 LoRA requires SM90/SM100 kernels.", +) + + @pytest.fixture(autouse=True) def _isolate_moe_runner_cache(): """Give every test a fresh cached FusedMoeRunner and release captured graphs @@ -104,6 +112,17 @@ def _scaled(*args, seed): return {"fc1": fc1, "gated": gated, "fc2": fc2} +def _adapter_set_to_fp8(adapters): + fp8_max = torch.finfo(torch.float8_e4m3fn).max + return { + module: { + name: tensor.clamp(min=-fp8_max, max=fp8_max).to(torch.float8_e4m3fn).contiguous() + for name, tensor in weights.items() + } + for module, weights in adapters.items() + } + + def _per_request_kwargs(num_tokens, adapters, rank): """Single-request per-request schema covering all tokens with one adapter.""" fc1, gated, fc2 = adapters["fc1"], adapters["gated"], adapters["fc2"] @@ -337,6 +356,186 @@ def call_layers() -> list[torch.Tensor]: torch.testing.assert_close(replay_output, eager_output, rtol=0, atol=0) +@requires_cuda_and_op +def test_cuda_graph_activates_empty_slot_on_replay(): + """Engine warmup captures empty LoRA slots before adapters enter the PEFT cache.""" + device = torch.device("cuda") + dtype = torch.bfloat16 + num_tokens, hidden_size, inter_size = 16, 128, 256 + num_experts, top_k, rank = 4, 2, 8 + x, w3_w1, w2, topk_ids, topk_scores = _build_base_inputs( + num_tokens, hidden_size, inter_size, num_experts, top_k, dtype, device + ) + adapters = _make_adapter_set( + num_experts, rank, hidden_size, inter_size, dtype, device, base_seed=425 + ) + active_kwargs = _slot_kwargs(torch.zeros(num_tokens, dtype=torch.int32), [adapters], rank) + active_tables = { + key: value.clone() + for key, value in active_kwargs.items() + if isinstance(value, torch.Tensor) and ("ranks" in key or "weight_ptrs" in key) + } + replay_kwargs = active_kwargs + for key, value in replay_kwargs.items(): + if isinstance(value, torch.Tensor) and ("ranks" in key or "weight_ptrs" in key): + value.zero_() + + def call(): + return _call_fused_moe(x, w3_w1, w2, topk_ids, topk_scores, dtype, replay_kwargs) + + graph, captured = _warmup_and_capture(call) + for key, value in active_tables.items(): + replay_kwargs[key].copy_(value) + graph.replay() + torch.cuda.synchronize() + + out_ref = _reference(x, w3_w1, w2, topk_ids, topk_scores, adapters) + torch.testing.assert_close(captured, out_ref, rtol=_RTOL, atol=_ATOL) + + +@requires_cuda_and_op +@requires_native_fp8 +def test_fp8_adapter_cuda_graph_replay_matches_eager(): + device = torch.device("cuda") + dtype = torch.bfloat16 + num_tokens, hidden_size, inter_size = 16, 128, 256 + num_experts, top_k, rank = 4, 2, 16 + + x, w3_w1, w2, topk_ids, topk_scores = _build_base_inputs( + num_tokens, hidden_size, inter_size, num_experts, top_k, dtype, device + ) + adapters = _adapter_set_to_fp8( + _make_adapter_set(num_experts, rank, hidden_size, inter_size, dtype, device, base_seed=450) + ) + slot_kwargs = _slot_kwargs(torch.zeros(num_tokens, dtype=torch.int32), [adapters], rank) + slot_kwargs["lora_dtype"] = torch.float8_e4m3fn + + def call(): + return _call_fused_moe(x, w3_w1, w2, topk_ids, topk_scores, dtype, dict(slot_kwargs)) + + baseline = _call_fused_moe(x, w3_w1, w2, topk_ids, topk_scores, dtype, {}) + out_eager = call().clone() + graph, captured = _warmup_and_capture(call) + graph.replay() + torch.cuda.synchronize() + out_replay = captured.clone() + + assert torch.isfinite(out_replay).all() + assert (out_replay.float() - baseline.float()).abs().mean().item() > 1e-3 + torch.testing.assert_close(out_replay, out_eager, rtol=0, atol=0) + reference = _reference(x.float(), w3_w1, w2, topk_ids, topk_scores, adapters) + zero_adapters = { + module: {name: torch.zeros_like(weight) for name, weight in weights.items()} + for module, weights in adapters.items() + } + reference_base = _reference(x.float(), w3_w1, w2, topk_ids, topk_scores, zero_adapters) + reference_delta = reference - reference_base + actual_delta = out_eager.float() - baseline.float() + relative_error = torch.linalg.vector_norm( + actual_delta - reference_delta + ) / torch.linalg.vector_norm(reference_delta) + assert relative_error < 0.10 + + +@requires_cuda_and_op +@requires_native_fp8 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("schema", ["per_request", "slot_indexed"]) +def test_fp8_moe_delta_matches_fp32_without_output_clipping(dtype, schema): + """Bound delta error while forcing the final GEMM result above E4M3's range.""" + device = torch.device("cuda") + num_tokens, hidden_size, inter_size, num_experts, top_k, rank = 4, 128, 128, 4, 2, 16 + x, w3_w1, w2, topk_ids, topk_scores = _build_base_inputs( + num_tokens, hidden_size, inter_size, num_experts, top_k, dtype, device + ) + x.fill_(1) + w3_w1.zero_() + w3_w1[:, :, 0] = 1 + w2.zero_() + adapters = _make_adapter_set( + num_experts, rank, hidden_size, inter_size, dtype, device, base_seed=475 + ) + for module in ("fc1", "gated"): + adapters[module]["A"].zero_() + adapters[module]["B"].zero_() + adapters["fc2"]["A"].fill_(1 / 16) + adapters["fc2"]["B"].fill_(16) + adapters = _adapter_set_to_fp8(adapters) + if schema == "per_request": + kwargs = _per_request_kwargs(num_tokens, adapters, rank) + else: + kwargs = _slot_kwargs(torch.zeros(num_tokens, dtype=torch.int32), [adapters], rank) + kwargs["lora_dtype"] = torch.float8_e4m3fn + baseline = _call_fused_moe(x, w3_w1, w2, topk_ids, topk_scores, dtype, {}) + actual = _call_fused_moe(x, w3_w1, w2, topk_ids, topk_scores, dtype, kwargs) + # Passing float32 inputs keeps the independent reference's output in FP32. + expected = _reference(x.float(), w3_w1, w2, topk_ids, topk_scores, adapters) + delta = actual.float() - baseline.float() + assert expected.abs().min() > torch.finfo(torch.float8_e4m3fn).max + assert torch.isfinite(delta).all() + torch.testing.assert_close(delta, expected, atol=0, rtol=0.05) + + +@requires_cuda_and_op +@requires_native_fp8 +def test_fp8_scratch_growth_preserves_compute_dtype_capacity(): + """A larger FP8 run must retain scratch width reserved by an earlier BF16 run.""" + device = torch.device("cuda") + dtype = torch.bfloat16 + hidden_size, inter_size, num_experts, top_k, rank = 128, 256, 4, 2, 16 + adapters = _make_adapter_set( + num_experts, rank, hidden_size, inter_size, dtype, device, base_seed=480 + ) + fp8_adapters = _adapter_set_to_fp8(adapters) + small = _build_base_inputs(2, hidden_size, inter_size, num_experts, top_k, dtype, device) + large = _build_base_inputs(16, hidden_size, inter_size, num_experts, top_k, dtype, device) + small_kwargs = _slot_kwargs(torch.zeros(2, dtype=torch.int32), [adapters], rank) + _call_fused_moe(*small, dtype, small_kwargs) + fp8_kwargs = _slot_kwargs(torch.zeros(16, dtype=torch.int32), [fp8_adapters], rank) + fp8_kwargs["lora_dtype"] = torch.float8_e4m3fn + graph, captured = _warmup_and_capture(lambda: _call_fused_moe(*large, dtype, fp8_kwargs)) + graph.replay() + torch.cuda.synchronize() + expected_fp8 = captured.clone() + bf16_kwargs = _slot_kwargs(torch.zeros(16, dtype=torch.int32), [adapters], rank) + actual_bf16 = _call_fused_moe(*large, dtype, bf16_kwargs) + expected_bf16 = _reference(*large, adapters) + torch.testing.assert_close(actual_bf16, expected_bf16, rtol=_RTOL, atol=_ATOL) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(captured, expected_fp8, rtol=0, atol=0) + + +@requires_cuda_and_op +@requires_native_fp8 +@pytest.mark.parametrize("schema", ["per_request", "slot_indexed"]) +@pytest.mark.parametrize("max_rank", [8, 16]) +def test_fp8_moe_rejects_misaligned_rank(schema, max_rank): + device = torch.device("cuda") + dtype = torch.bfloat16 + num_tokens, hidden_size, inter_size, num_experts, top_k, rank = 4, 128, 256, 4, 2, 8 + inputs = _build_base_inputs( + num_tokens, hidden_size, inter_size, num_experts, top_k, dtype, device + ) + adapters = _adapter_set_to_fp8( + _make_adapter_set(num_experts, rank, hidden_size, inter_size, dtype, device, base_seed=490) + ) + kwargs = ( + _per_request_kwargs(num_tokens, adapters, rank) + if schema == "per_request" + else _slot_kwargs(torch.zeros(num_tokens, dtype=torch.int32), [adapters], rank) + ) + kwargs["lora_dtype"] = torch.float8_e4m3fn + kwargs["lora_max_low_rank"] = max_rank + error = ( + "requires rank and hidden dimensions divisible by 16" + if max_rank == 8 + else "requires active ranks divisible by 16" + ) + with pytest.raises(RuntimeError, match=error): + _call_fused_moe(*inputs, dtype, kwargs) + + @requires_cuda_and_op def test_cuda_graph_multi_adapter(): """Two adapter slots routed per-token via token_to_slot, captured + replayed. diff --git a/tests/unittest/_torch/peft/test_moe_lora_op.py b/tests/unittest/_torch/peft/test_moe_lora_op.py index 00c8baaa7105..535fb1844d33 100644 --- a/tests/unittest/_torch/peft/test_moe_lora_op.py +++ b/tests/unittest/_torch/peft/test_moe_lora_op.py @@ -18,6 +18,7 @@ import pytest import torch +from tensorrt_llm._torch.peft.lora.manager import supports_native_fp8_lora from tensorrt_llm._torch.peft.lora.moe_layout import make_per_expert_lora, reference_swiglu_moe_lora _TRTLLM_AVAILABLE = hasattr(torch.ops, "trtllm") and hasattr(torch.ops.trtllm, "fused_moe") @@ -27,6 +28,12 @@ reason="Requires CUDA and built TensorRT-LLM C++ extension (torch.ops.trtllm.fused_moe).", ) +requires_native_fp8 = pytest.mark.skipif( + not torch.cuda.is_available() + or not supports_native_fp8_lora(torch.cuda.get_device_capability()), + reason="Native FP8 LoRA requires SM90/SM100 kernels.", +) + # ---------- shared fixtures ---------------------------------------------------- @@ -1239,6 +1246,71 @@ def test_moe_all_three_lora_with_gated_aliased_to_fc1_eager_matches_pytorch_refe MoERunner.runner_dict.clear() +# -- Native FP8 routed-expert LoRA adapter weights ------------------------- + + +def _adapter_to_fp8(adapter): + fp8_max = torch.finfo(torch.float8_e4m3fn).max + return { + name: tensor.clamp(min=-fp8_max, max=fp8_max).to(torch.float8_e4m3fn).contiguous() + for name, tensor in adapter.items() + } + + +@requires_cuda_and_op +@requires_native_fp8 +@pytest.mark.parametrize("schema", ["per_request", "slot_indexed"]) +def test_moe_fp8_adapter_weights_change_output(schema): + device = torch.device("cuda") + dtype = torch.bfloat16 + num_tokens, hidden_size, inter_size = 16, 128, 256 + num_experts, top_k = 4, 2 + rank = 16 + + x, w3_w1, w2, topk_ids, topk_scores = _build_base_inputs( + num_tokens, hidden_size, inter_size, num_experts, top_k, dtype, device + ) + out_baseline = _call_fused_moe(x, w3_w1, w2, topk_ids, topk_scores, output_dtype=dtype)[0] + + fc1_adapter = _adapter_to_fp8( + _make_per_expert_lora_scaled( + num_experts, rank, hidden_size, inter_size, dtype=dtype, device=device, seed=510 + ) + ) + gated_adapter = _adapter_to_fp8( + _make_per_expert_lora_scaled( + num_experts, rank, hidden_size, inter_size, dtype=dtype, device=device, seed=511 + ) + ) + fc2_adapter = _adapter_to_fp8( + _make_per_expert_lora_scaled( + num_experts, rank, inter_size, hidden_size, dtype=dtype, device=device, seed=512 + ) + ) + + builder = _build_lora_request_buffers if schema == "per_request" else _build_lora_slot_buffers + lora_kwargs = builder( + num_tokens, + fc1_adapter["A"], + fc1_adapter["B"], + fc2_adapter["A"], + fc2_adapter["B"], + rank=rank, + gated_a=gated_adapter["A"], + gated_b=gated_adapter["B"], + ) + lora_kwargs["lora_dtype"] = torch.float8_e4m3fn + + out_lora = _call_fused_moe( + x, w3_w1, w2, topk_ids, topk_scores, output_dtype=dtype, lora_kwargs=lora_kwargs + )[0] + + assert out_lora.dtype == dtype + assert torch.isfinite(out_lora).all() + diff = (out_lora.float() - out_baseline.float()).abs().mean().item() + assert diff > 1e-3, f"Native FP8 MoE LoRA ({schema}) had no observable effect: {diff}" + + # -- Per-tensor FP8 (qdq) base weights + routed-expert LoRA ----------------- # # The CUTLASS MoE LoRA kernel dequantizes per-tensor FP8 activations to the diff --git a/tests/unittest/others/test_lora_manager.py b/tests/unittest/others/test_lora_manager.py index c9288b5a511e..90c612b9efb1 100644 --- a/tests/unittest/others/test_lora_manager.py +++ b/tests/unittest/others/test_lora_manager.py @@ -27,7 +27,7 @@ from unittest.mock import MagicMock, patch import torch -from safetensors.torch import save_file +from safetensors.torch import load_file, save_file from tensorrt_llm._torch.peft.lora.loaders import HfLoraLoader from tensorrt_llm._torch.peft.lora.manager import LoraManager, supports_native_fp8_lora @@ -234,7 +234,7 @@ def test_many_adapters_no_gpu_accumulation(self): "Native FP8 LoRA requires SM90 or SM100", ) class TestLoraManagerFp8(unittest.TestCase): - def test_hf_loader_reports_dense_fp8_dtype(self): + def test_hf_loader_reports_fp8_dtype(self): with tempfile.TemporaryDirectory() as tmpdir: adapter_dir = Path(tmpdir) / "adapter" adapter_dir.mkdir() @@ -247,9 +247,9 @@ def test_hf_loader_reports_dense_fp8_dtype(self): loader = HfLoraLoader([str(adapter_dir)]) - self.assertEqual(loader.get_dense_lora_dtype(), torch.float8_e4m3fn) + self.assertEqual(loader.get_lora_dtype(), torch.float8_e4m3fn) - def test_hf_loader_ignores_expert_only_fp8_dtype(self): + def test_hf_loader_reports_expert_only_fp8_adapter_dtype(self): with tempfile.TemporaryDirectory() as tmpdir: adapter_dir = Path(tmpdir) / "adapter" adapter_dir.mkdir() @@ -257,7 +257,7 @@ def test_hf_loader_ignores_expert_only_fp8_dtype(self): loader = HfLoraLoader([str(adapter_dir)]) - self.assertIsNone(loader.get_dense_lora_dtype()) + self.assertEqual(loader.get_lora_dtype(), torch.float8_e4m3fn) def test_fp8_dora_is_rejected(self): model_config = MockModelConfig() @@ -288,7 +288,7 @@ def test_fp8_dora_is_rejected(self): uids=["fp8-dora"], ) - def test_fp8_moe_weights_are_converted_to_model_dtype(self): + def test_fp8_moe_weights_remain_fp8(self): model_config = MockModelConfig( lora_target_modules=["moe_h_to_4h"], trtllm_modules_to_hf_modules={"moe_h_to_4h": "mlp.experts.gate_proj"}, @@ -311,9 +311,9 @@ def test_fp8_moe_weights_are_converted_to_model_dtype(self): uids=["fp8-moe"], ) - self.assertEqual(manager.cpp_lora_weights["fp8-moe"].dtype, torch.bfloat16) + self.assertEqual(manager.cpp_lora_weights["fp8-moe"].dtype, torch.float8_e4m3fn) - def test_mixed_fp8_dense_and_moe_modules_are_rejected_before_cache_mutation(self): + def test_mixed_fp8_dense_and_moe_modules_use_one_fp8_cache_dtype(self): model_config = MockModelConfig( lora_target_modules=["attn_q", "moe_h_to_4h"], trtllm_modules_to_hf_modules={ @@ -335,14 +335,96 @@ def test_mixed_fp8_dense_and_moe_modules_are_rejected_before_cache_mutation(self adapter_dir.mkdir() _create_dummy_hf_moe_lora_adapter(adapter_dir, include_dense=True) + manager.load_from_hf( + model_dirs=[str(adapter_dir)], + model_config=model_config, + uids=["mixed-fp8"], + ) + + self.assertEqual(manager.cpp_lora_weights["mixed-fp8"].dtype, torch.float8_e4m3fn) + + def test_fp8_moe_misaligned_rank_is_rejected_before_cache_mutation(self): + model_config = MockModelConfig( + lora_target_modules=["moe_h_to_4h"], + trtllm_modules_to_hf_modules={"moe_h_to_4h": "mlp.experts.gate_proj"}, + dtype="bfloat16", + ) + manager = LoraManager( + mapping=Mapping(world_size=1, rank=0, tp_size=1), + model_config=model_config, + cpp_peft_cache_manager=MagicMock(), + ) + with tempfile.TemporaryDirectory() as tmpdir: + adapter_dir = Path(tmpdir) + _create_dummy_hf_moe_lora_adapter(adapter_dir, rank=8) + with self.assertRaisesRegex(ValueError, "multiples of 16.*rank=8"): + manager.load_from_hf( + model_dirs=[str(adapter_dir)], model_config=model_config, uids=["misaligned"] + ) + self.assertNotIn("misaligned", manager.cpp_lora_weights) + + def test_fp8_moe_and_bf16_dense_are_rejected_before_cache_mutation(self): + model_config = MockModelConfig( + lora_target_modules=["attn_q", "moe_h_to_4h"], + trtllm_modules_to_hf_modules={ + "attn_q": "q_proj", + "attn_k": "k_proj", + "attn_v": "v_proj", + "moe_h_to_4h": "mlp.experts.gate_proj", + }, + dtype="bfloat16", + ) + manager = LoraManager( + mapping=Mapping(world_size=1, rank=0, tp_size=1), + model_config=model_config, + cpp_peft_cache_manager=MagicMock(), + ) + with tempfile.TemporaryDirectory() as tmpdir: + adapter_dir = Path(tmpdir) + _create_dummy_hf_moe_lora_adapter(adapter_dir, include_dense=True) + weights_path = str(adapter_dir / "adapter_model.safetensors") + weights = load_file(weights_path) + weights = { + name: tensor.to(torch.bfloat16) if "self_attn" in name else tensor + for name, tensor in weights.items() + } + save_file(weights, weights_path) + with self.assertRaisesRegex(ValueError, "checkpoint weights must have one dtype"): + HfLoraLoader([str(adapter_dir)]).get_lora_dtype() with self.assertRaisesRegex(ValueError, "one PEFT cache dtype"): manager.load_from_hf( - model_dirs=[str(adapter_dir)], - model_config=model_config, - uids=["mixed-fp8"], + model_dirs=[str(adapter_dir)], model_config=model_config, uids=["mixed"] ) + self.assertNotIn("mixed", manager.cpp_lora_weights) - self.assertNotIn("mixed-fp8", manager.cpp_lora_weights) + def test_heterogeneous_expert_dtypes_are_rejected(self): + model_config = MockModelConfig( + lora_target_modules=["moe_h_to_4h"], + trtllm_modules_to_hf_modules={"moe_h_to_4h": "mlp.experts.gate_proj"}, + dtype="bfloat16", + ) + manager = LoraManager( + mapping=Mapping(world_size=1, rank=0, tp_size=1), + model_config=model_config, + cpp_peft_cache_manager=MagicMock(), + ) + with tempfile.TemporaryDirectory() as tmpdir: + adapter_dir = Path(tmpdir) + _create_dummy_hf_moe_lora_adapter(adapter_dir) + weights_path = str(adapter_dir / "adapter_model.safetensors") + weights = load_file(weights_path) + weights = { + name: tensor.to(torch.bfloat16) if ".experts.1." in name else tensor + for name, tensor in weights.items() + } + save_file(weights, weights_path) + with self.assertRaisesRegex(ValueError, "checkpoint weights must have one dtype"): + HfLoraLoader([str(adapter_dir)]).get_lora_dtype() + with self.assertRaisesRegex(ValueError, "MoE LoRA experts must have one output dtype"): + manager.load_from_hf( + model_dirs=[str(adapter_dir)], model_config=model_config, uids=["heterogeneous"] + ) + self.assertNotIn("heterogeneous", manager.cpp_lora_weights) def test_partial_qkv_fp8_adapter_uses_fp8_placeholders(self): model_config = MockModelConfig(dtype="bfloat16") @@ -517,7 +599,8 @@ def test_fp8_e4m3_input_output_dtype_mismatch_is_rejected(self): dtype=torch.float8_e4m3fn, input_dtype=torch.float8_e5m2, ) - self.assertIsNone(HfLoraLoader([str(adapter_dir)]).get_dense_lora_dtype()) + with self.assertRaisesRegex(ValueError, "checkpoint weights must have one dtype"): + HfLoraLoader([str(adapter_dir)]).get_lora_dtype() with self.assertRaisesRegex( ValueError, "FP8 LoRA input and output weights must have the same dtype" From ad73fa6446298cf374418baee5c10b625b890c9c Mon Sep 17 00:00:00 2001 From: Aurelien Chartier <2567591+achartier@users.noreply.github.com> Date: Tue, 20 Oct 2026 08:31:55 -0700 Subject: [PATCH 2/3] [TRTLLM-15314][test] Compare FP8 MoE LoRA with rounded BF16 adapters Run the BF16 path with the same E4M3-rounded adapter values for both LoRA schemas and bound the relative delta error. Keep the finite-output and observable-effect checks. Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com> --- .../unittest/_torch/peft/test_moe_lora_op.py | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/tests/unittest/_torch/peft/test_moe_lora_op.py b/tests/unittest/_torch/peft/test_moe_lora_op.py index 535fb1844d33..1aff3d2b90cc 100644 --- a/tests/unittest/_torch/peft/test_moe_lora_op.py +++ b/tests/unittest/_torch/peft/test_moe_lora_op.py @@ -1261,6 +1261,7 @@ def _adapter_to_fp8(adapter): @requires_native_fp8 @pytest.mark.parametrize("schema", ["per_request", "slot_indexed"]) def test_moe_fp8_adapter_weights_change_output(schema): + """Compare native FP8 LoRA with BF16 execution of the same rounded weights.""" device = torch.device("cuda") dtype = torch.bfloat16 num_tokens, hidden_size, inter_size = 16, 128, 256 @@ -1310,6 +1311,35 @@ def test_moe_fp8_adapter_weights_change_output(schema): diff = (out_lora.float() - out_baseline.float()).abs().mean().item() assert diff > 1e-3, f"Native FP8 MoE LoRA ({schema}) had no observable effect: {diff}" + # Upcasting preserves the E4M3 weight values; only execution precision differs. + fc1_ref, gated_ref, fc2_ref = [ + {name: tensor.to(dtype).contiguous() for name, tensor in adapter.items()} + for adapter in (fc1_adapter, gated_adapter, fc2_adapter) + ] + ref_kwargs = builder( + num_tokens, + fc1_ref["A"], + fc1_ref["B"], + fc2_ref["A"], + fc2_ref["B"], + rank=rank, + gated_a=gated_ref["A"], + gated_b=gated_ref["B"], + ) + ref_kwargs["lora_dtype"] = dtype + out_ref = _call_fused_moe( + x, w3_w1, w2, topk_ids, topk_scores, output_dtype=dtype, lora_kwargs=ref_kwargs + )[0] + + reference_delta = out_ref.float() - out_baseline.float() + actual_delta = out_lora.float() - out_baseline.float() + relative_error = torch.linalg.vector_norm( + actual_delta - reference_delta + ) / torch.linalg.vector_norm(reference_delta) + assert relative_error < 0.10, ( + f"Native FP8 MoE LoRA ({schema}) differs from BF16 reference: {relative_error}" + ) + # -- Per-tensor FP8 (qdq) base weights + routed-expert LoRA ----------------- # From ea7cc812d25b968e373683d64eca3db383654e61 Mon Sep 17 00:00:00 2001 From: Aurelien Chartier <2567591+achartier@users.noreply.github.com> Date: Thu, 22 Oct 2026 04:54:15 -0700 Subject: [PATCH 3/3] [TRTLLM-15314][test] Use default timeout for native FP8 MoE LoRA tests Remove the 180-minute overrides for the tiny synthetic eager and CUDA graph cases. Use the standard 3600-second CI timeout until measured runtime warrants an override. Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com> --- tests/integration/test_lists/test-db/l0_h100.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 29643f5689dc..d16af4f02a82 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -376,8 +376,8 @@ l0_h100: - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_field_completeness - llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_routed_expert_multi_lora_varying_ranks[eager] TIMEOUT (300) - llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_routed_expert_multi_lora_varying_ranks[cudagraph] TIMEOUT (300) - - llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_native_fp8_lora_matches_bf16_reference[eager] TIMEOUT (180) - - llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_native_fp8_lora_matches_bf16_reference[cudagraph] TIMEOUT (180) + - llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_native_fp8_lora_matches_bf16_reference[eager] + - llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_native_fp8_lora_matches_bf16_reference[cudagraph] - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[eagle3_one_model=True-enable_chunked_prefill=True-enable_max_concurrency=False-enable_draft_len_schedule=False] - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[eagle3_one_model=True-enable_chunked_prefill=False-enable_max_concurrency=False-enable_draft_len_schedule=True] - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[eagle3_one_model=True-enable_chunked_prefill=False-enable_max_concurrency=True-enable_draft_len_schedule=False]