From 1630fc378873d4127120137b3b62d393513ede82 Mon Sep 17 00:00:00 2001 From: Aurelien Chartier <2567591+achartier@users.noreply.github.com> Date: Tue, 11 Aug 2026 18:32:51 -0700 Subject: [PATCH 1/5] feat: add FP8 LoRA support for B200 Add native SM100 FP8 grouped GEMM dispatch for eager and CUDA graph LoRA paths, mirror CUTLASS cluster settings, and size persistent grids using active clusters. Share the runtime kernel capability gate between adapter loading and cache initialization, preserve compute-dtype fallback for unsupported builds, and warn when the capability query is unavailable. Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com> --- .../kernels/cuda_graph_grouped_gemm.cu | 89 ++++++++++----- cpp/tensorrt_llm/kernels/groupGemm.cu | 105 +++++++++++++----- cpp/tensorrt_llm/kernels/groupGemm.h | 3 + cpp/tensorrt_llm/thop/loraOp.cpp | 8 ++ tensorrt_llm/_torch/peft/lora/manager.py | 39 ++++++- tensorrt_llm/_torch/pyexecutor/_util.py | 19 +++- .../test_fp8_lora_grouped_gemm_regressions.py | 33 ++++-- .../others/test_lora_device_capability.py | 82 ++++++++++++++ tests/unittest/others/test_lora_manager.py | 86 ++++++++++++-- 9 files changed, 385 insertions(+), 79 deletions(-) create mode 100644 tests/unittest/others/test_lora_device_capability.py diff --git a/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu b/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu index 7fd7979ee4bd..98a703eeefc8 100644 --- a/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu +++ b/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu @@ -46,7 +46,7 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { #ifdef ENABLE_FP8 -#if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) +#if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) namespace { @@ -55,8 +55,8 @@ void checkFp8CudaGraphAlignment( { static int const smVersion = tensorrt_llm::common::getSMVersion(); // CUTLASS also exposes this kernel level on SM120/SM121; enable those after validation. - TLLM_CHECK_WITH_INFO( - smVersion == 90, "%s requires Hopper (SM90), but the current device is SM%d", kernelName, smVersion); + TLLM_CHECK_WITH_INFO(smVersion == 90 || smVersion == 100, + "%s requires Hopper (SM90) or B200 (SM100), but the current device is SM%d", kernelName, smVersion); TLLM_CHECK_WITH_INFO(minKN >= kFp8TmaAlignment && minKN % kFp8TmaAlignment == 0, "%s requires active LoRA ranks to be multiples of %d elements for 128-bit TMA alignment. " @@ -79,7 +79,7 @@ void checkFp8CudaGraphAlignment( } } // namespace -#endif // CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED +#endif // CUTLASS architecture support #endif // ENABLE_FP8 /** @@ -191,7 +191,7 @@ void cudaGraphGroupedGemmType(cutlass::gemm::GemmCoord* problemSizesPtr, int pro } #ifdef ENABLE_FP8 -#if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) +#if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) // ==================================================================== // FP8 CUDA-graph-compatible grouped GEMM using CUTLASS 3.x. @@ -200,7 +200,9 @@ void cudaGraphGroupedGemmType(cutlass::gemm::GemmCoord* problemSizesPtr, int pro // dimensions already on the GPU and reuses their compatible storage directly. // ==================================================================== -void fp8CudaGraphGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int problemCount, void** ptrAGpu, +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) { @@ -224,21 +226,14 @@ void fp8CudaGraphGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int prob static constexpr int kAlignmentC = 128 / cutlass::sizeof_bits::value; static constexpr int kAlignmentD = 128 / cutlass::sizeof_bits::value; - using ArchTag = cutlass::arch::Sm90; using OperatorClass = cutlass::arch::OpClassTensorOp; - using TileShape = Shape<_128, _128, _128>; - using ClusterShape = Shape<_1, _2, _1>; - - using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperativeFP8FastAccum; - using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecializedCooperative; - using ProblemShape = cutlass::gemm::GroupProblemShape>; - using CollectiveEpilogue = - typename cutlass::epilogue::collective::CollectiveBuilder::CollectiveOp; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder>::CollectiveOp; using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder(&cutlass::device_kernel)); + } - typename Gemm::Arguments arguments{cutlass::gemm::GemmUniversalMode::kGrouped, - {problemCount, problemShapes, nullptr}, {ptrA, strideA, ptrB, strideB}, - {{1.0f, 0.0f}, ptrC, strideC, ptrD, strideD}, hwInfo}; + typename Gemm::Arguments arguments; + decltype(arguments.epilogue.thread) fusionArgs{}; + fusionArgs.alpha = 1.0f; + fusionArgs.beta = 0.0f; + arguments = + typename Gemm::Arguments{cutlass::gemm::GemmUniversalMode::kGrouped, {problemCount, problemShapes, nullptr}, + {ptrA, strideA, ptrB, strideB}, {fusionArgs, ptrC, strideC, ptrD, strideD}, hwInfo}; Gemm gemm; @@ -315,7 +321,36 @@ void fp8CudaGraphGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int prob TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); } -#endif // CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED +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) +{ + int const smVersion = tensorrt_llm::common::getSMVersion(); + if (smVersion == 90) + { +#if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) && !defined(EXCLUDE_SM_90) + fp8CudaGraphGroupedGemmImpl, + cute::Shape, + cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperativeFP8FastAccum, + cutlass::epilogue::PtrArrayTmaWarpSpecializedCooperative>( + problemSizesPtr, problemCount, ptrAGpu, ptrBGpu, ptrCGpu, ptrDGpu, ldaGpu, ldbGpu, ldcGpu, lddGpu, stream); + return; +#endif + } + else if (smVersion == 100) + { +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) && !defined(EXCLUDE_SM_100F) + fp8CudaGraphGroupedGemmImpl, + cute::Shape, cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmSm100, + cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm>( + problemSizesPtr, problemCount, ptrAGpu, ptrBGpu, ptrCGpu, ptrDGpu, ldaGpu, ldbGpu, ldcGpu, lddGpu, stream); + return; +#endif + } + TLLM_CHECK_WITH_INFO(false, "FP8 CUDA graph grouped GEMM was not compiled for SM%d", smVersion); +} + +#endif // CUTLASS architecture support #endif // ENABLE_FP8 void cudaGraphGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int problemCount, void** ptrAGpu, void** ptrBGpu, @@ -323,7 +358,7 @@ void cudaGraphGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int problem tensorrt_llm::DataType dataType, int minKN, cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, cudaStream_t stream) { #ifdef ENABLE_FP8 -#if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) +#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"); @@ -335,10 +370,10 @@ void cudaGraphGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int problem if (dataType == tensorrt_llm::DataType::kFP8) { TLLM_CHECK_WITH_INFO(false, - "FP8 CUDA graph grouped GEMM requires CUTLASS modifiable TMA support (CUDA 12.3+ and Hopper SM90 " + "FP8 CUDA graph grouped GEMM requires CUTLASS modifiable TMA support (CUDA 12.3+ and SM90/SM100 " "kernels)."); } -#endif // CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED +#endif // CUTLASS architecture support #endif // ENABLE_FP8 if (isLoraIn) @@ -510,7 +545,7 @@ void cudaGraphSplitKGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int p cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, int64_t* splitKOffsetsGpu, cudaStream_t stream) { #ifdef ENABLE_FP8 -#if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) +#if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) if (dataType == tensorrt_llm::DataType::kFP8) { // Reuse the non-split-K fp8 path; CUTLASS 3.x cooperative schedule handles large-K efficiently. @@ -523,10 +558,10 @@ void cudaGraphSplitKGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int p if (dataType == tensorrt_llm::DataType::kFP8) { TLLM_CHECK_WITH_INFO(false, - "FP8 CUDA graph split-K grouped GEMM requires CUTLASS modifiable TMA support (CUDA 12.3+ and Hopper SM90 " + "FP8 CUDA graph split-K grouped GEMM requires CUTLASS modifiable TMA support (CUDA 12.3+ and SM90/SM100 " "kernels)."); } -#endif // CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED +#endif // CUTLASS architecture support #endif // ENABLE_FP8 if (isLoraIn) diff --git a/cpp/tensorrt_llm/kernels/groupGemm.cu b/cpp/tensorrt_llm/kernels/groupGemm.cu index 87e5b94c6645..0caaea63afe5 100644 --- a/cpp/tensorrt_llm/kernels/groupGemm.cu +++ b/cpp/tensorrt_llm/kernels/groupGemm.cu @@ -45,8 +45,25 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { +bool supportsFp8GroupedGemm(int smVersion) +{ + if (smVersion == 90) + { +#if defined(ENABLE_FP8) && defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) && !defined(EXCLUDE_SM_90) + return true; +#endif + } + else if (smVersion == 100) + { +#if defined(ENABLE_FP8) && defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) && !defined(EXCLUDE_SM_100F) + return true; +#endif + } + return false; +} + #ifdef ENABLE_FP8 -#if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) +#if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) namespace { @@ -54,8 +71,8 @@ void checkFp8GroupedGemmAlignment(std::vector const& p { int const smVersion = tensorrt_llm::common::getSMVersion(); // CUTLASS also exposes this kernel level on SM120/SM121; enable those after validation. - TLLM_CHECK_WITH_INFO( - smVersion == 90, "%s requires Hopper (SM90), but the current device is SM%d", kernelName, smVersion); + TLLM_CHECK_WITH_INFO(smVersion == 90 || smVersion == 100, + "%s requires Hopper (SM90) or B200 (SM100), but the current device is SM%d", kernelName, smVersion); for (size_t problemIdx = 0; problemIdx < problemSizes.size(); ++problemIdx) { @@ -73,7 +90,7 @@ void checkFp8GroupedGemmAlignment(std::vector const& p } } // namespace -#endif // CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED +#endif // CUTLASS architecture support #endif // ENABLE_FP8 int64_t inline getGemmCoordSize(int64_t problemCount) @@ -261,7 +278,7 @@ void groupedGemmType_(std::vector problem_sizes, std:: #ifdef ENABLE_FP8 // ==================================================================== -// FP8 grouped GEMM using CUTLASS 3.x collective API (Hopper SM90). +// FP8 grouped GEMM using architecture-specific CUTLASS 3.x collectives. // // The legacy CUTLASS 2.x DefaultGemmGrouped does NOT support fp8 element // types. This implementation uses the CUTLASS 3.x CollectiveBuilder and @@ -278,9 +295,11 @@ void groupedGemmType_(std::vector problem_sizes, std:: // them to device memory via the provided workspace. // ==================================================================== -#if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) +#if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) -void fp8GroupedGemm(std::vector const& problemSizes, std::vector const& ptrA, +template +void fp8GroupedGemmImpl(std::vector const& problemSizes, std::vector const& ptrA, std::vector const& ptrB, std::vector const& ptrC, std::vector const& ptrD, void* gemmParamsWorkSpace, int64_t gemmParamsWorkSpaceSize, void* gemmWorkSpace, int64_t gemmWorkspaceSize, cudaStream_t stream) @@ -306,23 +325,14 @@ void fp8GroupedGemm(std::vector const& problemSizes, s static constexpr int kAlignmentC = 128 / cutlass::sizeof_bits::value; static constexpr int kAlignmentD = 128 / cutlass::sizeof_bits::value; - using ArchTag = cutlass::arch::Sm90; using OperatorClass = cutlass::arch::OpClassTensorOp; - // Tile and cluster shapes chosen for fp8 on Hopper. - using TileShape = Shape<_128, _128, _128>; - using ClusterShape = Shape<_1, _2, _1>; - - // Kernel and epilogue schedule for fp8 grouped GEMM with PtrArray TMA. - using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperativeFP8FastAccum; - using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecializedCooperative; - using ProblemShape = cutlass::gemm::GroupProblemShape>; - using CollectiveEpilogue = - typename cutlass::epilogue::collective::CollectiveBuilder::CollectiveOp; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder>::CollectiveOp; using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder const& problemSizes, s hwInfo.device_id = 0; cudaGetDevice(&hwInfo.device_id); hwInfo.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hwInfo.device_id); + if constexpr (kIsSm100) + { + hwInfo.cluster_shape = dim3(4, 2, 1); + hwInfo.cluster_shape_fallback = dim3(2, 1, 1); + hwInfo.max_active_clusters = cutlass::KernelHardwareInfo::query_device_max_active_clusters(hwInfo.cluster_shape, + GemmKernel::MaxThreadsPerBlock, reinterpret_cast(&cutlass::device_kernel)); + } // Build CUTLASS 3.x arguments. alpha=1, beta=0 for LoRA GEMM. - typename Gemm::Arguments arguments{cutlass::gemm::GemmUniversalMode::kGrouped, - {problemCount, devProblemShapes, nullptr}, {devPtrA, devStrideA, devPtrB, devStrideB}, - {{1.0f, 0.0f}, devPtrC, devStrideC, devPtrD, devStrideD}, hwInfo}; + typename Gemm::Arguments arguments; + decltype(arguments.epilogue.thread) fusionArgs{}; + fusionArgs.alpha = 1.0f; + fusionArgs.beta = 0.0f; + arguments = + typename Gemm::Arguments{cutlass::gemm::GemmUniversalMode::kGrouped, {problemCount, devProblemShapes, nullptr}, + {devPtrA, devStrideA, devPtrB, devStrideB}, {fusionArgs, devPtrC, devStrideC, devPtrD, devStrideD}, hwInfo}; Gemm gemm; @@ -468,7 +489,37 @@ void fp8GroupedGemm(std::vector const& problemSizes, s TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); } -#endif // CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED +void fp8GroupedGemm(std::vector const& problemSizes, std::vector const& ptrA, + std::vector const& ptrB, std::vector const& ptrC, std::vector const& ptrD, + void* gemmParamsWorkSpace, int64_t gemmParamsWorkSpaceSize, void* gemmWorkSpace, int64_t gemmWorkspaceSize, + cudaStream_t stream) +{ + int const smVersion = tensorrt_llm::common::getSMVersion(); + if (smVersion == 90) + { +#if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) && !defined(EXCLUDE_SM_90) + fp8GroupedGemmImpl, + cute::Shape, + cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperativeFP8FastAccum, + cutlass::epilogue::PtrArrayTmaWarpSpecializedCooperative>(problemSizes, ptrA, ptrB, ptrC, ptrD, + gemmParamsWorkSpace, gemmParamsWorkSpaceSize, gemmWorkSpace, gemmWorkspaceSize, stream); + return; +#endif + } + else if (smVersion == 100) + { +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) && !defined(EXCLUDE_SM_100F) + fp8GroupedGemmImpl, + cute::Shape, cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmSm100, + cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm>(problemSizes, ptrA, ptrB, ptrC, ptrD, gemmParamsWorkSpace, + gemmParamsWorkSpaceSize, gemmWorkSpace, gemmWorkspaceSize, stream); + return; +#endif + } + TLLM_CHECK_WITH_INFO(false, "FP8 grouped GEMM was not compiled for SM%d", smVersion); +} + +#endif // CUTLASS architecture support #endif // ENABLE_FP8 @@ -480,7 +531,7 @@ void groupedGemm(std::vector problem_sizes, std::vecto TLLM_LOG_TRACE("%s start, isLoraIn: %d, minKN = %d", __PRETTY_FUNCTION__, static_cast(isLoraIn), minKN); #ifdef ENABLE_FP8 -#if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) +#if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) if (dataType == tensorrt_llm::DataType::kFP8) { checkFp8GroupedGemmAlignment(problem_sizes, "FP8 grouped GEMM"); @@ -493,9 +544,9 @@ void groupedGemm(std::vector problem_sizes, std::vecto if (dataType == tensorrt_llm::DataType::kFP8) { TLLM_CHECK_WITH_INFO( - false, "FP8 grouped GEMM requires CUTLASS modifiable TMA support (CUDA 12.3+ and Hopper SM90 kernels)."); + false, "FP8 grouped GEMM requires CUTLASS modifiable TMA support (CUDA 12.3+ and SM90/SM100 kernels)."); } -#endif // CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED +#endif // CUTLASS architecture support #endif // ENABLE_FP8 if (isLoraIn) diff --git a/cpp/tensorrt_llm/kernels/groupGemm.h b/cpp/tensorrt_llm/kernels/groupGemm.h index 48b38eb82d65..6020955525d4 100644 --- a/cpp/tensorrt_llm/kernels/groupGemm.h +++ b/cpp/tensorrt_llm/kernels/groupGemm.h @@ -34,6 +34,9 @@ int64_t getGroupedGemmParamsWorkSpaceSize(int64_t problem_count); //! cute stride arrays on the device. This function returns the required bytes. int64_t getFp8GroupedGemmParamsWorkSpaceSize(int64_t problemCount); +//! @brief Returns whether the FP8 grouped GEMM kernel for an SM version is present in this build. +bool supportsFp8GroupedGemm(int smVersion); + void groupedGemm(std::vector problem_sizes, std::vector const& ptrA, std::vector const& ptrB, std::vector const& ptrC, std::vector const& ptrD, void* gemmParamsWorkspace, int64_t gemmParamsWorkSpaceSize, void* gemmWorkSpace, int64_t gemmWorkspaceSize, diff --git a/cpp/tensorrt_llm/thop/loraOp.cpp b/cpp/tensorrt_llm/thop/loraOp.cpp index 51575a5873ed..a22b6e5a5146 100644 --- a/cpp/tensorrt_llm/thop/loraOp.cpp +++ b/cpp/tensorrt_llm/thop/loraOp.cpp @@ -20,6 +20,7 @@ #include "tensorrt_llm/common/opUtils.h" #include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/cuda_graph_grouped_gemm.h" +#include "tensorrt_llm/kernels/groupGemm.h" #include "tensorrt_llm/kernels/lora/lora.h" #include "tensorrt_llm/kernels/lora/loraGroupGEMMParamFillRowReorderFusion.h" #include "tensorrt_llm/kernels/selectiveScan/selectiveScan.h" @@ -41,6 +42,11 @@ enum class RequestType : int32_t kGENERATION = 1 }; +bool loraGroupedGemmSupportsFp8(int64_t smVersion) +{ + return tk::supportsFp8GroupedGemm(static_cast(smVersion)); +} + int64_t getNumTokens(th::Tensor const& input) { int ndim = input.sizes().size(); @@ -360,6 +366,8 @@ TRTLLM_NAMESPACE_END TORCH_LIBRARY_FRAGMENT(trtllm, m) { + m.def("lora_grouped_gemm_supports_fp8", &tensorrt_llm::torch_ext::loraGroupedGemmSupportsFp8); + m.def( "lora_grouped_gemm(Tensor input, " "Tensor host_request_types, " diff --git a/tensorrt_llm/_torch/peft/lora/manager.py b/tensorrt_llm/_torch/peft/lora/manager.py index 369f883fb633..05f889ab22ab 100644 --- a/tensorrt_llm/_torch/peft/lora/manager.py +++ b/tensorrt_llm/_torch/peft/lora/manager.py @@ -19,8 +19,9 @@ import logging import warnings from dataclasses import dataclass +from functools import lru_cache from pathlib import Path -from typing import TYPE_CHECKING, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union import numpy as np import torch @@ -54,6 +55,32 @@ _FP8_LORA_TMA_ALIGNMENT = 16 +_NATIVE_FP8_LORA_DEVICE_CAPABILITIES = ((9, 0), (10, 0)) + + +@lru_cache(maxsize=1) +def _warn_native_fp8_lora_capability_query_unavailable() -> None: + logger.warning( + "Native FP8 LoRA capability query is unavailable; adapter weights " + "will fall back to the model compute dtype. Check that the " + "TensorRT-LLM libraries match the Python package and are loaded." + ) + + +def _native_fp8_lora_kernels_available(device_capability: Tuple[int, int]) -> bool: + kernel_support_query = getattr(torch.ops.trtllm, "lora_grouped_gemm_supports_fp8", None) + if kernel_support_query is None: + _warn_native_fp8_lora_capability_query_unavailable() + return False + major, minor = device_capability + return kernel_support_query(major * 10 + minor) + + +def supports_native_fp8_lora(device_capability: Tuple[int, int]) -> bool: + """Return whether native FP8 LoRA kernels support a CUDA capability.""" + return device_capability in _NATIVE_FP8_LORA_DEVICE_CAPABILITIES and ( + _native_fp8_lora_kernels_available(device_capability) + ) def _check_lora_in_out( @@ -107,7 +134,7 @@ def _validate_fp8_lora_alignment( f"{name}={size}" for name, size in misaligned_dimensions.items() ) raise ValueError( - f"FP8 LoRA weights on Hopper require rank, input size, and output size " + f"FP8 LoRA weights on SM90/SM100 require rank, input size, and output size " f"to be multiples of {_FP8_LORA_TMA_ALIGNMENT} for 128-bit TMA alignment. " f"Layer {layer_idx} module '{lora_module}' has {formatted_dimensions}. " f"Use aligned adapter dimensions or non-FP8 LoRA weights." @@ -647,7 +674,7 @@ def load_from_model_dir(uid, model_dir, hf_config): rank = int(hf_config["r"]) rs_lora = bool(hf_config.get("use_rslora", False)) model_dtype = str_dtype_to_torch(model_config.dtype) - supports_native_fp8 = torch.cuda.get_device_capability() == (9, 0) + supports_native_fp8 = supports_native_fp8_lora(torch.cuda.get_device_capability()) def get_output_dtype(module_weights): if _is_moe_module_weights(module_weights): @@ -777,7 +804,7 @@ def uses_native_fp8(module_weights): ) if is_dora: raise NotImplementedError( - "DoRA is not supported with FP8 LoRA weights on Hopper" + "DoRA is not supported with FP8 LoRA weights on SM90/SM100" ) t_in = t_in.cuda().contiguous() @@ -791,7 +818,7 @@ def uses_native_fp8(module_weights): scale = float(hf_config["lora_alpha"]) / effective_rank if use_fp8_kernel: - # Keep weights in FP8 for the native Hopper kernel. + # Keep weights in FP8 for the native SM90/SM100 kernel. # FP8 has no scalar multiply, so scale through BF16. fp8_max = torch.finfo(t_out.dtype).max t_out = ( @@ -800,7 +827,7 @@ def uses_native_fp8(module_weights): .to(t_out.dtype) ) else: - # Pre-Hopper kernels require the model compute dtype. + # Other architectures require the model compute dtype. t_in = t_in.to(model_dtype) t_out = t_out.to(model_dtype) t_out = t_out * scale diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index eabdc677ae2c..faa36672a820 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -37,7 +37,8 @@ # isort: on from tensorrt_llm._torch.peft.lora.config import ( LoraConfig, get_default_trtllm_modules_to_hf_modules) -from tensorrt_llm._torch.peft.lora.manager import load_torch_lora +from tensorrt_llm._torch.peft.lora.manager import (load_torch_lora, + supports_native_fp8_lora) from tensorrt_llm.logger import logger from tensorrt_llm.mapping import CpType, Mapping @@ -89,6 +90,16 @@ def ceil_div(a: int, b: int) -> int: return (a + b - 1) // b +def _get_initial_lora_data_type( + configured_lora_data_type: Optional[torch.dtype], +) -> Optional[torch.dtype]: + if configured_lora_data_type != torch.float8_e4m3fn: + return None + if supports_native_fp8_lora(torch.cuda.get_device_capability()): + return configured_lora_data_type + return None + + def _non_hybrid_kv_cache_manager_cls(config, kv_cache_config: KvCacheConfig): # Models with per-layer head_dim (e.g., Gemma4 hybrid attention) # require KVCacheManagerV2 for per-layer buffer sizes. @@ -2819,10 +2830,8 @@ def create_py_executor_instance( initial_lora_data_type = None if len(lora_config.lora_dir) == 1: # Route to appropriate loader based on checkpoint source - configured_lora_data_type = load_torch_lora(lora_config) - if (configured_lora_data_type == torch.float8_e4m3fn - and torch.cuda.get_device_capability() == (9, 0)): - initial_lora_data_type = configured_lora_data_type + initial_lora_data_type = _get_initial_lora_data_type( + load_torch_lora(lora_config)) else: assert len(lora_config.lora_target_modules ) >= 1, "Expecting at least one lora target module" diff --git a/tests/unittest/_torch/peft/test_fp8_lora_grouped_gemm_regressions.py b/tests/unittest/_torch/peft/test_fp8_lora_grouped_gemm_regressions.py index d610fe7e41da..2d96d3224eb1 100644 --- a/tests/unittest/_torch/peft/test_fp8_lora_grouped_gemm_regressions.py +++ b/tests/unittest/_torch/peft/test_fp8_lora_grouped_gemm_regressions.py @@ -24,9 +24,12 @@ _validate_fp8_lora_cuda_graph_alignment, add_lora_result, ) +from tensorrt_llm._torch.peft.lora.manager import supports_native_fp8_lora _REPO_ROOT = Path(__file__).resolve().parents[4] -_NATIVE_FP8_AVAILABLE = torch.cuda.is_available() and torch.cuda.get_device_capability() == (9, 0) +_NATIVE_FP8_AVAILABLE = torch.cuda.is_available() and supports_native_fp8_lora( + torch.cuda.get_device_capability() +) def _kernel_source(filename: str) -> str: @@ -57,7 +60,16 @@ def _assert_fp8_gemm_matches_reference(actual, reference): ) -@pytest.mark.skipif(not _NATIVE_FP8_AVAILABLE, reason="Native FP8 LoRA requires SM90") +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ((9, 0), (10, 0)), + reason="Native FP8 LoRA requires SM90 or SM100", +) +def test_fp8_grouped_gemm_kernel_is_reported_available(): + major, minor = torch.cuda.get_device_capability() + assert torch.ops.trtllm.lora_grouped_gemm_supports_fp8(major * 10 + minor) + + +@pytest.mark.skipif(not _NATIVE_FP8_AVAILABLE, reason="Native FP8 LoRA requires SM90 or SM100") def test_fp8_eager_grouped_gemm_matches_reference(): x, lora_in, lora_out, reference = _make_fp8_lora_problem() rank = lora_in.shape[0] @@ -82,7 +94,7 @@ def test_fp8_eager_grouped_gemm_matches_reference(): _assert_fp8_gemm_matches_reference(actual, reference) -@pytest.mark.skipif(not _NATIVE_FP8_AVAILABLE, reason="Native FP8 LoRA requires SM90") +@pytest.mark.skipif(not _NATIVE_FP8_AVAILABLE, reason="Native FP8 LoRA requires SM90 or SM100") def test_fp8_cuda_graph_grouped_gemm_matches_reference_after_replay(): batch_size = 16 x, lora_in, lora_out, _ = _make_fp8_lora_problem(batch_size=batch_size) @@ -221,7 +233,9 @@ def test_fp8_cuda_graph_alignment_rejects_misaligned_hidden_dims(hidden_size, ou def test_fp8_cuda_graph_grouped_gemm_reuses_live_device_metadata(): source = _kernel_source("cuda_graph_grouped_gemm.cu") fp8_graph_body = _function_block( - source, "void fp8CudaGraphGroupedGemm(", "\nvoid cudaGraphGroupedGemm(" + source, + "void fp8CudaGraphGroupedGemmImpl(", + "\nvoid fp8CudaGraphGroupedGemm(", ) assert "hostMaxProblemSizesPtr" not in fp8_graph_body @@ -256,12 +270,17 @@ def test_fp8_grouped_gemm_dispatch_has_explicit_unsupported_cutlass_guard(filena @pytest.mark.parametrize("filename", ["groupGemm.cu", "cuda_graph_grouped_gemm.cu"]) -def test_fp8_grouped_gemm_dispatch_requires_sm90(filename): +def test_fp8_grouped_gemm_dispatch_supports_sm90_and_sm100(filename): source = _kernel_source(filename) assert "getSMVersion()" in source - assert "smVersion == 90" in source - assert "requires Hopper (SM90)" in source + assert "smVersion == 90 || smVersion == 100" in source + assert "requires Hopper (SM90) or B200 (SM100)" in source + assert "cutlass::arch::Sm90" in source + assert "cutlass::arch::Sm100" in source + assert "KernelPtrArrayTmaWarpSpecializedCooperativeFP8FastAccum" in source + assert "KernelPtrArrayTmaWarpSpecialized1SmSm100" in source + assert "PtrArrayTmaWarpSpecialized1Sm" in source assert "SM120/SM121" in source diff --git a/tests/unittest/others/test_lora_device_capability.py b/tests/unittest/others/test_lora_device_capability.py new file mode 100644 index 000000000000..4c47167abbe6 --- /dev/null +++ b/tests/unittest/others/test_lora_device_capability.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +import tensorrt_llm._torch.peft.lora.manager as lora_manager +from tensorrt_llm._torch.pyexecutor._util import _get_initial_lora_data_type +from tensorrt_llm._torch.peft.lora.manager import supports_native_fp8_lora + +pytestmark = pytest.mark.cpu_only + + +@pytest.mark.parametrize( + "device_capability,kernels_available,expected", + [ + ((9, 0), True, True), + ((10, 0), True, True), + ((10, 0), False, False), + ((8, 0), True, False), + ((12, 0), True, False), + ], +) +def test_supports_native_fp8_lora(device_capability, kernels_available, expected, monkeypatch): + monkeypatch.setattr( + lora_manager, + "_native_fp8_lora_kernels_available", + lambda _: kernels_available, + ) + assert supports_native_fp8_lora(device_capability) is expected + + +def test_missing_native_fp8_lora_capability_query_warns_once(caplog, monkeypatch): + monkeypatch.setattr( + torch.ops.trtllm, + "lora_grouped_gemm_supports_fp8", + None, + raising=False, + ) + lora_manager._warn_native_fp8_lora_capability_query_unavailable.cache_clear() + + with caplog.at_level("WARNING", logger=lora_manager.__name__): + assert not lora_manager._native_fp8_lora_kernels_available((9, 0)) + assert not lora_manager._native_fp8_lora_kernels_available((10, 0)) + + assert caplog.messages == [ + "Native FP8 LoRA capability query is unavailable; adapter weights " + "will fall back to the model compute dtype. Check that the " + "TensorRT-LLM libraries match the Python package and are loaded." + ] + + +@pytest.mark.parametrize("device_capability", [(9, 0), (10, 0)]) +def test_native_fp8_lora_initializes_fp8_cache(device_capability, monkeypatch): + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: device_capability) + monkeypatch.setattr(lora_manager, "_native_fp8_lora_kernels_available", lambda _: True) + + assert _get_initial_lora_data_type(torch.float8_e4m3fn) == torch.float8_e4m3fn + + +@pytest.mark.parametrize("device_capability", [(8, 0), (12, 0)]) +def test_unsupported_device_does_not_initialize_fp8_cache(device_capability, monkeypatch): + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: device_capability) + monkeypatch.setattr(lora_manager, "_native_fp8_lora_kernels_available", lambda _: True) + + assert _get_initial_lora_data_type(torch.float8_e4m3fn) is None + + +def test_missing_native_fp8_lora_kernels_do_not_initialize_fp8_cache(monkeypatch): + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (10, 0)) + monkeypatch.setattr(lora_manager, "_native_fp8_lora_kernels_available", lambda _: False) + + assert _get_initial_lora_data_type(torch.float8_e4m3fn) is None + + +def test_non_fp8_lora_does_not_query_device_capability(monkeypatch): + def fail_if_called(): + raise AssertionError("device capability should not be queried for non-FP8 LoRA") + + monkeypatch.setattr(torch.cuda, "get_device_capability", fail_if_called) + + assert _get_initial_lora_data_type(torch.bfloat16) is None diff --git a/tests/unittest/others/test_lora_manager.py b/tests/unittest/others/test_lora_manager.py index ca90ac205d72..2254ea02b432 100644 --- a/tests/unittest/others/test_lora_manager.py +++ b/tests/unittest/others/test_lora_manager.py @@ -30,7 +30,7 @@ from safetensors.torch import save_file from tensorrt_llm._torch.peft.lora.loaders import HfLoraLoader -from tensorrt_llm._torch.peft.lora.manager import LoraManager +from tensorrt_llm._torch.peft.lora.manager import LoraManager, supports_native_fp8_lora from tensorrt_llm.mapping import Mapping @@ -230,8 +230,8 @@ def test_many_adapters_no_gpu_accumulation(self): @unittest.skipUnless( - torch.cuda.is_available() and torch.cuda.get_device_capability() == (9, 0), - "Native FP8 LoRA requires Hopper (SM90)", + torch.cuda.is_available() and supports_native_fp8_lora(torch.cuda.get_device_capability()), + "Native FP8 LoRA requires SM90 or SM100", ) class TestLoraManagerFp8(unittest.TestCase): def test_hf_loader_reports_dense_fp8_dtype(self): @@ -280,7 +280,7 @@ def test_fp8_dora_is_rejected(self): with self.assertRaisesRegex( NotImplementedError, - "DoRA is not supported with FP8 LoRA weights on Hopper", + "DoRA is not supported with FP8 LoRA weights on SM90/SM100", ): manager.load_from_hf( model_dirs=[str(adapter_dir)], @@ -400,7 +400,7 @@ def test_fp8_e5m2_weights_are_converted_to_model_dtype(self): self.assertEqual(manager.cpp_lora_weights["fp8-e5m2"].dtype, torch.bfloat16) - def test_fp8_e4m3_weights_on_sm100_are_converted_to_model_dtype(self): + def test_fp8_e4m3_weights_on_sm100_remain_fp8(self): model_config = MockModelConfig(dtype="bfloat16") manager = LoraManager( mapping=Mapping(world_size=1, rank=0, tp_size=1), @@ -418,14 +418,86 @@ def test_fp8_e4m3_weights_on_sm100_are_converted_to_model_dtype(self): dtype=torch.float8_e4m3fn, ) - with patch("torch.cuda.get_device_capability", return_value=(10, 0)): + with ( + patch("torch.cuda.get_device_capability", return_value=(10, 0)), + patch( + "tensorrt_llm._torch.peft.lora.manager._native_fp8_lora_kernels_available", + return_value=True, + ), + ): manager.load_from_hf( model_dirs=[str(adapter_dir)], model_config=model_config, uids=["fp8-e4m3-sm100"], ) - self.assertEqual(manager.cpp_lora_weights["fp8-e4m3-sm100"].dtype, torch.bfloat16) + self.assertEqual( + manager.cpp_lora_weights["fp8-e4m3-sm100"].dtype, + torch.float8_e4m3fn, + ) + + def test_fp8_e4m3_weights_on_sm100_without_kernels_use_model_dtype(self): + model_config = MockModelConfig(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) / "adapter" + adapter_dir.mkdir() + _create_dummy_hf_lora_adapter( + adapter_dir, + rank=16, + num_layers=1, + dtype=torch.float8_e4m3fn, + ) + + with ( + patch("torch.cuda.get_device_capability", return_value=(10, 0)), + patch( + "tensorrt_llm._torch.peft.lora.manager._native_fp8_lora_kernels_available", + return_value=False, + ), + ): + manager.load_from_hf( + model_dirs=[str(adapter_dir)], + model_config=model_config, + uids=["fp8-e4m3-sm100-fallback"], + ) + + self.assertEqual( + manager.cpp_lora_weights["fp8-e4m3-sm100-fallback"].dtype, + torch.bfloat16, + ) + + def test_fp8_e4m3_weights_on_sm120_are_converted_to_model_dtype(self): + model_config = MockModelConfig(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) / "adapter" + adapter_dir.mkdir() + _create_dummy_hf_lora_adapter( + adapter_dir, + rank=16, + num_layers=1, + dtype=torch.float8_e4m3fn, + ) + + with patch("torch.cuda.get_device_capability", return_value=(12, 0)): + manager.load_from_hf( + model_dirs=[str(adapter_dir)], + model_config=model_config, + uids=["fp8-e4m3-sm120"], + ) + + self.assertEqual(manager.cpp_lora_weights["fp8-e4m3-sm120"].dtype, torch.bfloat16) def test_fp8_e4m3_input_output_dtype_mismatch_is_rejected(self): model_config = MockModelConfig() From 216b6d070392235960a68e1bd132dc95acbf38d7 Mon Sep 17 00:00:00 2001 From: Aurelien Chartier <2567591+achartier@users.noreply.github.com> Date: Wed, 12 Aug 2026 13:55:50 -0700 Subject: [PATCH 2/5] fix: address FP8 LoRA review feedback Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com> --- .../kernels/cuda_graph_grouped_gemm.cu | 25 ++++--- .../kernels/fp8GroupedGemmConfig.h | 73 +++++++++++++++++++ cpp/tensorrt_llm/kernels/groupGemm.cu | 27 +++---- docs/source/features/lora.md | 16 ++++ tensorrt_llm/_torch/peft/lora/manager.py | 7 +- .../test_lists/test-db/l0_b200.yml | 2 + .../test_fp8_lora_grouped_gemm_regressions.py | 7 +- .../others/test_lora_device_capability.py | 29 ++++---- tests/unittest/others/test_lora_manager.py | 5 +- 9 files changed, 142 insertions(+), 49 deletions(-) create mode 100644 cpp/tensorrt_llm/kernels/fp8GroupedGemmConfig.h diff --git a/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu b/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu index 98a703eeefc8..5374bd1fb20c 100644 --- a/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu +++ b/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu @@ -39,6 +39,7 @@ #include "cutlass/gemm/group_array_problem_shape.hpp" #include "cutlass/gemm/kernel/gemm_universal.hpp" #include "cutlass/util/packed_stride.hpp" +#include "fp8GroupedGemmConfig.h" #endif // ENABLE_FP8 TRTLLM_NAMESPACE_BEGIN @@ -200,8 +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) @@ -210,6 +210,12 @@ void fp8CudaGraphGroupedGemmImpl(cutlass::gemm::GemmCoord* problemSizesPtr, int using namespace cute; + using ArchTag = typename Config::ArchTag; + using TileShape = typename Config::TileShape; + using ClusterShape = typename Config::ClusterShape; + using KernelSchedule = typename Config::KernelSchedule; + using EpilogueSchedule = typename Config::EpilogueSchedule; + using ElementA = cutlass::float_e4m3_t; using ElementB = cutlass::float_e4m3_t; using ElementC = cutlass::float_e4m3_t; @@ -277,10 +283,10 @@ void fp8CudaGraphGroupedGemmImpl(cutlass::gemm::GemmCoord* problemSizesPtr, int hwInfo.device_id = 0; cudaGetDevice(&hwInfo.device_id); hwInfo.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hwInfo.device_id); - if constexpr (kIsSm100) + if constexpr (Config::kUsesDynamicClusterShape) { - hwInfo.cluster_shape = dim3(4, 2, 1); - hwInfo.cluster_shape_fallback = dim3(2, 1, 1); + hwInfo.cluster_shape = Config::clusterShape(); + hwInfo.cluster_shape_fallback = Config::clusterShapeFallback(); hwInfo.max_active_clusters = cutlass::KernelHardwareInfo::query_device_max_active_clusters(hwInfo.cluster_shape, GemmKernel::MaxThreadsPerBlock, reinterpret_cast(&cutlass::device_kernel)); } @@ -329,10 +335,7 @@ void fp8CudaGraphGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int prob if (smVersion == 90) { #if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) && !defined(EXCLUDE_SM_90) - fp8CudaGraphGroupedGemmImpl, - cute::Shape, - cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperativeFP8FastAccum, - cutlass::epilogue::PtrArrayTmaWarpSpecializedCooperative>( + fp8CudaGraphGroupedGemmImpl( problemSizesPtr, problemCount, ptrAGpu, ptrBGpu, ptrCGpu, ptrDGpu, ldaGpu, ldbGpu, ldcGpu, lddGpu, stream); return; #endif @@ -340,9 +343,7 @@ void fp8CudaGraphGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int prob else if (smVersion == 100) { #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) && !defined(EXCLUDE_SM_100F) - fp8CudaGraphGroupedGemmImpl, - cute::Shape, cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmSm100, - cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm>( + fp8CudaGraphGroupedGemmImpl( problemSizesPtr, problemCount, ptrAGpu, ptrBGpu, ptrCGpu, ptrDGpu, ldaGpu, ldbGpu, ldcGpu, lddGpu, stream); return; #endif diff --git a/cpp/tensorrt_llm/kernels/fp8GroupedGemmConfig.h b/cpp/tensorrt_llm/kernels/fp8GroupedGemmConfig.h new file mode 100644 index 000000000000..d690b1e39a70 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/fp8GroupedGemmConfig.h @@ -0,0 +1,73 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "cute/tensor.hpp" +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "tensorrt_llm/common/config.h" + +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::fp8_grouped_gemm +{ + +#if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) +struct Sm90Config +{ + using ArchTag = cutlass::arch::Sm90; + using TileShape = cute::Shape; + using ClusterShape = cute::Shape; + using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperativeFP8FastAccum; + using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecializedCooperative; + + static constexpr bool kUsesDynamicClusterShape = false; +}; +#endif + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) +struct Sm100Config +{ + using ArchTag = cutlass::arch::Sm100; + using TileShape = cute::Shape; + using ClusterShape = cute::Shape; + using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmSm100; + using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm; + + static constexpr bool kUsesDynamicClusterShape = true; + + static dim3 clusterShape() + { + return {4, 2, 1}; + } + + static dim3 clusterShapeFallback() + { + return {2, 1, 1}; + } +}; +#endif + +} // namespace kernels::fp8_grouped_gemm + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/groupGemm.cu b/cpp/tensorrt_llm/kernels/groupGemm.cu index 0caaea63afe5..7224af4be499 100644 --- a/cpp/tensorrt_llm/kernels/groupGemm.cu +++ b/cpp/tensorrt_llm/kernels/groupGemm.cu @@ -31,6 +31,7 @@ #include "cutlass/gemm/group_array_problem_shape.hpp" #include "cutlass/gemm/kernel/gemm_universal.hpp" #include "cutlass/util/packed_stride.hpp" +#include "fp8GroupedGemmConfig.h" #endif // ENABLE_FP8 #include "groupGemm.h" @@ -297,8 +298,7 @@ void groupedGemmType_(std::vector problem_sizes, std:: #if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) -template +template void fp8GroupedGemmImpl(std::vector const& problemSizes, std::vector const& ptrA, std::vector const& ptrB, std::vector const& ptrC, std::vector const& ptrD, void* gemmParamsWorkSpace, int64_t gemmParamsWorkSpaceSize, void* gemmWorkSpace, int64_t gemmWorkspaceSize, @@ -308,6 +308,12 @@ void fp8GroupedGemmImpl(std::vector const& problemSize using namespace cute; + using ArchTag = typename Config::ArchTag; + using TileShape = typename Config::TileShape; + using ClusterShape = typename Config::ClusterShape; + using KernelSchedule = typename Config::KernelSchedule; + using EpilogueSchedule = typename Config::EpilogueSchedule; + using ElementA = cutlass::float_e4m3_t; using ElementB = cutlass::float_e4m3_t; using ElementC = cutlass::float_e4m3_t; @@ -449,10 +455,10 @@ void fp8GroupedGemmImpl(std::vector const& problemSize hwInfo.device_id = 0; cudaGetDevice(&hwInfo.device_id); hwInfo.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hwInfo.device_id); - if constexpr (kIsSm100) + if constexpr (Config::kUsesDynamicClusterShape) { - hwInfo.cluster_shape = dim3(4, 2, 1); - hwInfo.cluster_shape_fallback = dim3(2, 1, 1); + hwInfo.cluster_shape = Config::clusterShape(); + hwInfo.cluster_shape_fallback = Config::clusterShapeFallback(); hwInfo.max_active_clusters = cutlass::KernelHardwareInfo::query_device_max_active_clusters(hwInfo.cluster_shape, GemmKernel::MaxThreadsPerBlock, reinterpret_cast(&cutlass::device_kernel)); } @@ -498,20 +504,15 @@ void fp8GroupedGemm(std::vector const& problemSizes, s if (smVersion == 90) { #if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) && !defined(EXCLUDE_SM_90) - fp8GroupedGemmImpl, - cute::Shape, - cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperativeFP8FastAccum, - cutlass::epilogue::PtrArrayTmaWarpSpecializedCooperative>(problemSizes, ptrA, ptrB, ptrC, ptrD, - gemmParamsWorkSpace, gemmParamsWorkSpaceSize, gemmWorkSpace, gemmWorkspaceSize, stream); + fp8GroupedGemmImpl(problemSizes, ptrA, ptrB, ptrC, ptrD, gemmParamsWorkSpace, + gemmParamsWorkSpaceSize, gemmWorkSpace, gemmWorkspaceSize, stream); return; #endif } else if (smVersion == 100) { #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) && !defined(EXCLUDE_SM_100F) - fp8GroupedGemmImpl, - cute::Shape, cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmSm100, - cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm>(problemSizes, ptrA, ptrB, ptrC, ptrD, gemmParamsWorkSpace, + fp8GroupedGemmImpl(problemSizes, ptrA, ptrB, ptrC, ptrD, gemmParamsWorkSpace, gemmParamsWorkSpaceSize, gemmWorkSpace, gemmWorkspaceSize, stream); return; #endif diff --git a/docs/source/features/lora.md b/docs/source/features/lora.md index c5f596bef982..a8f3bd614d65 100644 --- a/docs/source/features/lora.md +++ b/docs/source/features/lora.md @@ -118,6 +118,22 @@ 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 +FP8 E4M3 adapter weights and execute them with native FP8 grouped GEMM kernels on the following architectures: + +| GPU architecture | Native FP8 adapter support | +|---|---| +| Hopper (SM90) | Yes | +| Blackwell B200 (SM100) | Yes | +| Blackwell (SM103/SM107) | No | +| 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). + ### NeMo LoRA Format ```python diff --git a/tensorrt_llm/_torch/peft/lora/manager.py b/tensorrt_llm/_torch/peft/lora/manager.py index 05f889ab22ab..ca47c3812c21 100644 --- a/tensorrt_llm/_torch/peft/lora/manager.py +++ b/tensorrt_llm/_torch/peft/lora/manager.py @@ -55,7 +55,6 @@ _FP8_LORA_TMA_ALIGNMENT = 16 -_NATIVE_FP8_LORA_DEVICE_CAPABILITIES = ((9, 0), (10, 0)) @lru_cache(maxsize=1) @@ -77,10 +76,8 @@ def _native_fp8_lora_kernels_available(device_capability: Tuple[int, int]) -> bo def supports_native_fp8_lora(device_capability: Tuple[int, int]) -> bool: - """Return whether native FP8 LoRA kernels support a CUDA capability.""" - return device_capability in _NATIVE_FP8_LORA_DEVICE_CAPABILITIES and ( - _native_fp8_lora_kernels_available(device_capability) - ) + """Return whether compiled native FP8 LoRA kernels support a CUDA capability.""" + return _native_fp8_lora_kernels_available(device_capability) def _check_lora_in_out( diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index f1b02c22217e..a669a1d631f8 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -16,6 +16,7 @@ l0_b200: tests: # ------------- PyTorch tests --------------- - unittest/others/test_kv_cache_transceiver.py::test_cpp_nixl_sync_transfer_stress + - unittest/others/test_lora_manager.py - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8B::test_nvfp4 - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8B::test_nvfp4_streaming[stream_interval_4] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8B::test_nvfp4_streaming[stream_interval_64] @@ -79,6 +80,7 @@ l0_b200: - unittest/_torch/attention - unittest/_torch/compilation - unittest/_torch/debugger + - unittest/_torch/lora - unittest/_torch/kv_cache_compression - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_2_model_mtp - unittest/disaggregated/test_deepseek_v4_kv_transfer.py diff --git a/tests/unittest/_torch/peft/test_fp8_lora_grouped_gemm_regressions.py b/tests/unittest/_torch/peft/test_fp8_lora_grouped_gemm_regressions.py index 2d96d3224eb1..1220e8cc9e37 100644 --- a/tests/unittest/_torch/peft/test_fp8_lora_grouped_gemm_regressions.py +++ b/tests/unittest/_torch/peft/test_fp8_lora_grouped_gemm_regressions.py @@ -276,11 +276,8 @@ def test_fp8_grouped_gemm_dispatch_supports_sm90_and_sm100(filename): assert "getSMVersion()" in source assert "smVersion == 90 || smVersion == 100" in source assert "requires Hopper (SM90) or B200 (SM100)" in source - assert "cutlass::arch::Sm90" in source - assert "cutlass::arch::Sm100" in source - assert "KernelPtrArrayTmaWarpSpecializedCooperativeFP8FastAccum" in source - assert "KernelPtrArrayTmaWarpSpecialized1SmSm100" in source - assert "PtrArrayTmaWarpSpecialized1Sm" in source + assert "fp8_grouped_gemm::Sm90Config" in source + assert "fp8_grouped_gemm::Sm100Config" in source assert "SM120/SM121" in source diff --git a/tests/unittest/others/test_lora_device_capability.py b/tests/unittest/others/test_lora_device_capability.py index 4c47167abbe6..9c4fad0d638b 100644 --- a/tests/unittest/others/test_lora_device_capability.py +++ b/tests/unittest/others/test_lora_device_capability.py @@ -12,20 +12,21 @@ @pytest.mark.parametrize( - "device_capability,kernels_available,expected", + "device_capability,expected", [ - ((9, 0), True, True), - ((10, 0), True, True), - ((10, 0), False, False), - ((8, 0), True, False), - ((12, 0), True, False), + ((9, 0), True), + ((10, 0), True), + ((10, 3), True), + ((8, 0), False), + ((12, 0), False), ], ) -def test_supports_native_fp8_lora(device_capability, kernels_available, expected, monkeypatch): +def test_supports_native_fp8_lora(device_capability, expected, monkeypatch): monkeypatch.setattr( - lora_manager, - "_native_fp8_lora_kernels_available", - lambda _: kernels_available, + torch.ops.trtllm, + "lora_grouped_gemm_supports_fp8", + lambda sm_version: sm_version in {90, 100, 103}, + raising=False, ) assert supports_native_fp8_lora(device_capability) is expected @@ -50,7 +51,7 @@ def test_missing_native_fp8_lora_capability_query_warns_once(caplog, monkeypatch ] -@pytest.mark.parametrize("device_capability", [(9, 0), (10, 0)]) +@pytest.mark.parametrize("device_capability", [(9, 0), (10, 0), (10, 3)]) def test_native_fp8_lora_initializes_fp8_cache(device_capability, monkeypatch): monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: device_capability) monkeypatch.setattr(lora_manager, "_native_fp8_lora_kernels_available", lambda _: True) @@ -59,9 +60,11 @@ def test_native_fp8_lora_initializes_fp8_cache(device_capability, monkeypatch): @pytest.mark.parametrize("device_capability", [(8, 0), (12, 0)]) -def test_unsupported_device_does_not_initialize_fp8_cache(device_capability, monkeypatch): +def test_device_without_native_kernels_does_not_initialize_fp8_cache( + device_capability, monkeypatch +): monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: device_capability) - monkeypatch.setattr(lora_manager, "_native_fp8_lora_kernels_available", lambda _: True) + monkeypatch.setattr(lora_manager, "_native_fp8_lora_kernels_available", lambda _: False) assert _get_initial_lora_data_type(torch.float8_e4m3fn) is None diff --git a/tests/unittest/others/test_lora_manager.py b/tests/unittest/others/test_lora_manager.py index 2254ea02b432..c9288b5a511e 100644 --- a/tests/unittest/others/test_lora_manager.py +++ b/tests/unittest/others/test_lora_manager.py @@ -624,7 +624,10 @@ def test_misaligned_fp8_adapter_is_rejected_before_cuda_transfer(self): manager = self._create_manager(model_config) with ( - patch("torch.cuda.get_device_capability", return_value=(9, 0)), + patch( + "tensorrt_llm._torch.peft.lora.manager.supports_native_fp8_lora", + return_value=True, + ), self.assertRaisesRegex(ValueError, case["match"]), ): manager.load_from_hf( From 22a758436965158d579f82b7b7325190ecbf5787 Mon Sep 17 00:00:00 2001 From: Aurelien Chartier <2567591+achartier@users.noreply.github.com> Date: Wed, 12 Aug 2026 14:47:39 -0700 Subject: [PATCH 3/5] style: clean up FP8 grouped GEMM constants Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com> --- .../kernels/cuda_graph_grouped_gemm.cu | 12 ++++++------ cpp/tensorrt_llm/kernels/fp8GroupedGemmConfig.h | 7 +++++-- cpp/tensorrt_llm/kernels/groupGemm.cu | 16 ++++++++-------- .../test_fp8_lora_grouped_gemm_regressions.py | 7 ++++--- 4 files changed, 23 insertions(+), 19 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu b/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu index 5374bd1fb20c..b7e31f582193 100644 --- a/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu +++ b/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu @@ -15,6 +15,7 @@ */ #include "cuda_graph_grouped_gemm.h" +#include "fp8GroupedGemmConfig.h" #include "groupGemm.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" @@ -39,7 +40,6 @@ #include "cutlass/gemm/group_array_problem_shape.hpp" #include "cutlass/gemm/kernel/gemm_universal.hpp" #include "cutlass/util/packed_stride.hpp" -#include "fp8GroupedGemmConfig.h" #endif // ENABLE_FP8 TRTLLM_NAMESPACE_BEGIN @@ -56,7 +56,7 @@ void checkFp8CudaGraphAlignment( { static int const smVersion = tensorrt_llm::common::getSMVersion(); // CUTLASS also exposes this kernel level on SM120/SM121; enable those after validation. - TLLM_CHECK_WITH_INFO(smVersion == 90 || smVersion == 100, + TLLM_CHECK_WITH_INFO(smVersion == fp8GroupedGemmConfig::kSm90 || smVersion == fp8GroupedGemmConfig::kSm100, "%s requires Hopper (SM90) or B200 (SM100), but the current device is SM%d", kernelName, smVersion); TLLM_CHECK_WITH_INFO(minKN >= kFp8TmaAlignment && minKN % kFp8TmaAlignment == 0, @@ -332,18 +332,18 @@ void fp8CudaGraphGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int prob cudaStream_t stream) { int const smVersion = tensorrt_llm::common::getSMVersion(); - if (smVersion == 90) + 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 } - else if (smVersion == 100) + 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 diff --git a/cpp/tensorrt_llm/kernels/fp8GroupedGemmConfig.h b/cpp/tensorrt_llm/kernels/fp8GroupedGemmConfig.h index d690b1e39a70..2599296baac5 100644 --- a/cpp/tensorrt_llm/kernels/fp8GroupedGemmConfig.h +++ b/cpp/tensorrt_llm/kernels/fp8GroupedGemmConfig.h @@ -29,9 +29,12 @@ TRTLLM_NAMESPACE_BEGIN -namespace kernels::fp8_grouped_gemm +namespace kernels::fp8GroupedGemmConfig { +inline constexpr int kSm90 = 90; +inline constexpr int kSm100 = 100; + #if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) struct Sm90Config { @@ -68,6 +71,6 @@ struct Sm100Config }; #endif -} // namespace kernels::fp8_grouped_gemm +} // namespace kernels::fp8GroupedGemmConfig TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/groupGemm.cu b/cpp/tensorrt_llm/kernels/groupGemm.cu index 7224af4be499..1450628ed45e 100644 --- a/cpp/tensorrt_llm/kernels/groupGemm.cu +++ b/cpp/tensorrt_llm/kernels/groupGemm.cu @@ -21,6 +21,7 @@ #include "cutlass/gemm/gemm.h" #include "cutlass/gemm/kernel/default_gemm_grouped.h" #include "cutlass/gemm/kernel/gemm_grouped.h" +#include "fp8GroupedGemmConfig.h" #ifdef ENABLE_FP8 #include "cute/tensor.hpp" @@ -31,7 +32,6 @@ #include "cutlass/gemm/group_array_problem_shape.hpp" #include "cutlass/gemm/kernel/gemm_universal.hpp" #include "cutlass/util/packed_stride.hpp" -#include "fp8GroupedGemmConfig.h" #endif // ENABLE_FP8 #include "groupGemm.h" @@ -48,13 +48,13 @@ namespace kernels { bool supportsFp8GroupedGemm(int smVersion) { - if (smVersion == 90) + if (smVersion == fp8GroupedGemmConfig::kSm90) { #if defined(ENABLE_FP8) && defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) && !defined(EXCLUDE_SM_90) return true; #endif } - else if (smVersion == 100) + else if (smVersion == fp8GroupedGemmConfig::kSm100) { #if defined(ENABLE_FP8) && defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) && !defined(EXCLUDE_SM_100F) return true; @@ -72,7 +72,7 @@ void checkFp8GroupedGemmAlignment(std::vector const& p { int const smVersion = tensorrt_llm::common::getSMVersion(); // CUTLASS also exposes this kernel level on SM120/SM121; enable those after validation. - TLLM_CHECK_WITH_INFO(smVersion == 90 || smVersion == 100, + TLLM_CHECK_WITH_INFO(smVersion == fp8GroupedGemmConfig::kSm90 || smVersion == fp8GroupedGemmConfig::kSm100, "%s requires Hopper (SM90) or B200 (SM100), but the current device is SM%d", kernelName, smVersion); for (size_t problemIdx = 0; problemIdx < problemSizes.size(); ++problemIdx) @@ -501,18 +501,18 @@ void fp8GroupedGemm(std::vector const& problemSizes, s cudaStream_t stream) { int const smVersion = tensorrt_llm::common::getSMVersion(); - if (smVersion == 90) + if (smVersion == fp8GroupedGemmConfig::kSm90) { #if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) && !defined(EXCLUDE_SM_90) - fp8GroupedGemmImpl(problemSizes, ptrA, ptrB, ptrC, ptrD, gemmParamsWorkSpace, + fp8GroupedGemmImpl(problemSizes, ptrA, ptrB, ptrC, ptrD, gemmParamsWorkSpace, gemmParamsWorkSpaceSize, gemmWorkSpace, gemmWorkspaceSize, stream); return; #endif } - else if (smVersion == 100) + else if (smVersion == fp8GroupedGemmConfig::kSm100) { #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) && !defined(EXCLUDE_SM_100F) - fp8GroupedGemmImpl(problemSizes, ptrA, ptrB, ptrC, ptrD, gemmParamsWorkSpace, + fp8GroupedGemmImpl(problemSizes, ptrA, ptrB, ptrC, ptrD, gemmParamsWorkSpace, gemmParamsWorkSpaceSize, gemmWorkSpace, gemmWorkspaceSize, stream); return; #endif diff --git a/tests/unittest/_torch/peft/test_fp8_lora_grouped_gemm_regressions.py b/tests/unittest/_torch/peft/test_fp8_lora_grouped_gemm_regressions.py index 1220e8cc9e37..b1d4eb7e9c5c 100644 --- a/tests/unittest/_torch/peft/test_fp8_lora_grouped_gemm_regressions.py +++ b/tests/unittest/_torch/peft/test_fp8_lora_grouped_gemm_regressions.py @@ -274,10 +274,11 @@ def test_fp8_grouped_gemm_dispatch_supports_sm90_and_sm100(filename): source = _kernel_source(filename) assert "getSMVersion()" in source - assert "smVersion == 90 || smVersion == 100" in source + assert "smVersion == fp8GroupedGemmConfig::kSm90" in source + assert "smVersion == fp8GroupedGemmConfig::kSm100" in source assert "requires Hopper (SM90) or B200 (SM100)" in source - assert "fp8_grouped_gemm::Sm90Config" in source - assert "fp8_grouped_gemm::Sm100Config" in source + assert "fp8GroupedGemmConfig::Sm90Config" in source + assert "fp8GroupedGemmConfig::Sm100Config" in source assert "SM120/SM121" in source From c3893c2389d5f5536091908e33bf45111d378b10 Mon Sep 17 00:00:00 2001 From: Aurelien Chartier <2567591+achartier@users.noreply.github.com> Date: Tue, 18 Aug 2026 16:22:42 -0700 Subject: [PATCH 4/5] test: remove redundant FP8 LoRA capability test Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com> --- .../others/test_lora_device_capability.py | 21 ------------------- 1 file changed, 21 deletions(-) diff --git a/tests/unittest/others/test_lora_device_capability.py b/tests/unittest/others/test_lora_device_capability.py index 9c4fad0d638b..b74c0e1119da 100644 --- a/tests/unittest/others/test_lora_device_capability.py +++ b/tests/unittest/others/test_lora_device_capability.py @@ -6,31 +6,10 @@ import tensorrt_llm._torch.peft.lora.manager as lora_manager from tensorrt_llm._torch.pyexecutor._util import _get_initial_lora_data_type -from tensorrt_llm._torch.peft.lora.manager import supports_native_fp8_lora pytestmark = pytest.mark.cpu_only -@pytest.mark.parametrize( - "device_capability,expected", - [ - ((9, 0), True), - ((10, 0), True), - ((10, 3), True), - ((8, 0), False), - ((12, 0), False), - ], -) -def test_supports_native_fp8_lora(device_capability, expected, monkeypatch): - monkeypatch.setattr( - torch.ops.trtllm, - "lora_grouped_gemm_supports_fp8", - lambda sm_version: sm_version in {90, 100, 103}, - raising=False, - ) - assert supports_native_fp8_lora(device_capability) is expected - - def test_missing_native_fp8_lora_capability_query_warns_once(caplog, monkeypatch): monkeypatch.setattr( torch.ops.trtllm, From a4483f2c3f43bbd340f077d3c20a995f7c1e46bb Mon Sep 17 00:00:00 2001 From: Aurelien Chartier <2567591+achartier@users.noreply.github.com> Date: Wed, 19 Aug 2026 07:25:55 -0700 Subject: [PATCH 5/5] test: fix B200 FP8 LoRA unit test path Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com> --- tests/integration/test_lists/test-db/l0_b200.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index a669a1d631f8..ff8d023ae64d 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -80,7 +80,7 @@ l0_b200: - unittest/_torch/attention - unittest/_torch/compilation - unittest/_torch/debugger - - unittest/_torch/lora + - unittest/_torch/peft/test_fp8_lora_grouped_gemm_regressions.py - unittest/_torch/kv_cache_compression - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_2_model_mtp - unittest/disaggregated/test_deepseek_v4_kv_transfer.py