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..d16af4f02a82 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] + - 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] 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..1aff3d2b90cc 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,101 @@ 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): + """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 + 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}" + + # 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 ----------------- # # 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"