[TRTLLM-15314][feat] Support native FP8 MoE LoRA weights

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 <typename Config>
+template <typename Config, typename ElementD>
 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 <typename ElementD>
 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,15 +335,15 @@ 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<fp8GroupedGemmConfig::Sm90Config>(
+        fp8CudaGraphGroupedGemmImpl<fp8GroupedGemmConfig::Sm90Config, ElementD>(
             problemSizesPtr, problemCount, ptrAGpu, ptrBGpu, ptrCGpu, ptrDGpu, ldaGpu, ldbGpu, ldcGpu, lddGpu, stream);
         return;
 #endif
     }
     else if (smVersion == fp8GroupedGemmConfig::kSm100)
     {
 #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) && !defined(EXCLUDE_SM_100F)
-        fp8CudaGraphGroupedGemmImpl<fp8GroupedGemmConfig::Sm100Config>(
+        fp8CudaGraphGroupedGemmImpl<fp8GroupedGemmConfig::Sm100Config, ElementD>(
             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<tensorrt_llm::DataType> 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<cutlass::float_e4m3_t>(problemSizesPtr, problemCount, ptrAGpu, ptrBGpu, ptrCGpu,
+                ptrDGpu, ldaGpu, ldbGpu, ldcGpu, lddGpu, stream);
+            break;
+        case tensorrt_llm::DataType::kHALF:
+            fp8CudaGraphGroupedGemm<cutlass::half_t>(problemSizesPtr, problemCount, ptrAGpu, ptrBGpu, ptrCGpu, ptrDGpu,
+                ldaGpu, ldbGpu, ldcGpu, lddGpu, stream);
+            break;
+#ifdef ENABLE_BF16
+        case tensorrt_llm::DataType::kBF16:
+            fp8CudaGraphGroupedGemm<cutlass::bfloat16_t>(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<cutlass::float_e4m3_t>(
             problemSizesPtr, problemCount, ptrAGpu, ptrBGpu, ptrCGpu, ptrDGpu, ldaGpu, ldbGpu, ldcGpu, lddGpu, stream);
         return;
     }

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 <cuda_runtime.h>
+#include <optional>
 
 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<tensorrt_llm::DataType> outputDataType = std::nullopt);
 
 /**
  * @brief CUDA Graph compatible wrapper for split-K grouped GEMM operations.

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_;

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;

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
 

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 <class T>
+__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<int64_t>(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<float>(input[index]) : 0.0f);
+    }
+}
+
+template <class T>
+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<int>((numElements + LORA_KERNELS_THREADS_PER_BLOCK - 1) / LORA_KERNELS_THREADS_PER_BLOCK);
+    convertMoeLoraInputToFp8Kernel<<<blocks, LORA_KERNELS_THREADS_PER_BLOCK, 0, stream>>>(
+        output, input, numValidTokens, hiddenSize, numElements);
+    sync_check_cuda_error(stream);
+}
+#endif
+
 template <class ScaleBiasType, class LoraType, bool IsGated>
 __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<T, WeightType, OutputType, InputType, BackBoneType, Enab
     lora_fc1_result_ = {};
     lora_add_bias_ = {};
     lora_fc2_result_ = {};
+    mLoraFc1ResultBytes = workspaces.at("lora_fc1_result").first;
+    mLoraFc2ResultBytes = workspaces.at("lora_fc2_result").first;
 
     if (use_lora)
     {
@@ -4029,15 +4062,15 @@ inline void runMoeLoraGroupedGemmModule(::tensorrt_llm::kernels::cutlass_kernels
     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::kernels::cutlass_kernels::MoeLoraGroupedGemmRunFn run, 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?).");
     TLLM_CHECK_WITH_INFO(run != nullptr,
         "Grouped-GEMM LoRA GEMM dispatch is unavailable: grouped_gemm.run was not populated (this consumer of "
         "libmoe_gemm_src.a does not link libtorch).");
     run(mod, num_permuted_tokens, in_hidden_size, max_lora_rank, dtype_bytes, splitk_slices, input_base, output_base,
-        data_type, stream);
+        data_type, output_data_type, stream);
 }
 
 // Map the activation/back-bone type to the DataType enum the
@@ -4267,27 +4300,50 @@ auto CutlassMoeFCRunner<T, WeightType, OutputType, InputType, BackBoneType, Enab
     if (lora_params.grouped_gemm.enabled)
     {
         auto const& grouped_gemm = lora_params.grouped_gemm;
-        tensorrt_llm::DataType const data_type = moeLoraDataType<ScaleBiasType>();
+        tensorrt_llm::DataType const data_type = grouped_gemm.data_type;
+        tensorrt_llm::DataType const compute_data_type = moeLoraDataType<ScaleBiasType>();
 
+        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<size_t>(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
         // both the gated and fc1 halves when gated, so one memset covers both.
         size_t const fc1_result_bytes = static_cast<size_t>(expanded_num_rows) * static_cast<size_t>(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<void const*>(input),
-            /*output_base=*/static_cast<void*>(lora_fc1_result), grouped_gemm.run, data_type, stream);
+            /*input_base=*/grouped_input,
+            /*output_base=*/static_cast<void*>(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<void const*>(input),
-                /*output_base=*/static_cast<void*>(lora_gated_out), grouped_gemm.run, data_type, stream);
+                /*input_base=*/grouped_input,
+                /*output_base=*/static_cast<void*>(lora_gated_out), grouped_gemm.run, data_type, compute_data_type,
+                stream);
         }
     }
     else
@@ -4362,18 +4418,41 @@ void CutlassMoeFCRunner<T, WeightType, OutputType, InputType, BackBoneType, Enab
     if (lora_params.grouped_gemm.enabled)
     {
         auto const& grouped_gemm = lora_params.grouped_gemm;
-        tensorrt_llm::DataType const data_type = moeLoraDataType<ScaleBiasType>();
+        tensorrt_llm::DataType const data_type = grouped_gemm.data_type;
+        tensorrt_llm::DataType const compute_data_type = moeLoraDataType<ScaleBiasType>();
 
+        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<size_t>(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<size_t>(num_tokens) * static_cast<size_t>(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<void const*>(input),
-            /*output_base=*/static_cast<void*>(lora_fc2_result_), grouped_gemm.run, data_type, stream);
+            /*input_base=*/grouped_input,
+            /*output_base=*/static_cast<void*>(lora_fc2_result_), grouped_gemm.run, data_type, compute_data_type,
+            stream);
         sync_check_cuda_error(stream);
         return;
     }