From acb294c9faaebff500993cad2304150e5b9d683a Mon Sep 17 00:00:00 2001 From: Aurelien Chartier <2567591+achartier@users.noreply.github.com> Date: Wed, 22 Jul 2026 13:13:07 -0700 Subject: [PATCH 1/8] [None][feat] Support dense FP8 LoRA end to end Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com> --- .../batch_manager/peftCacheManager.h | 11 +- cpp/include/tensorrt_llm/runtime/loraCache.h | 14 +- .../batch_manager/peftCacheManager.cpp | 32 +- .../kernels/cuda_graph_grouped_gemm.cu | 264 +++++++++++++++- cpp/tensorrt_llm/kernels/groupGemm.cu | 286 +++++++++++++++++- cpp/tensorrt_llm/kernels/groupGemm.h | 8 +- cpp/tensorrt_llm/kernels/lora/lora.cpp | 23 +- cpp/tensorrt_llm/kernels/splitkGroupGemm.cu | 252 ++++++++++++++- cpp/tensorrt_llm/kernels/splitkGroupGemm.h | 2 +- .../nanobind/batch_manager/kvCacheManager.cpp | 1 + cpp/tensorrt_llm/runtime/loraCache.cpp | 29 +- cpp/tensorrt_llm/runtime/loraUtils.cpp | 10 +- cpp/tensorrt_llm/runtime/loraUtils.h | 6 +- cpp/tensorrt_llm/thop/loraOp.cpp | 21 +- .../batch_manager/peftCacheManagerTest.cpp | 26 +- tensorrt_llm/_torch/modules/attention.py | 8 +- tensorrt_llm/_torch/modules/gated_mlp.py | 8 +- tensorrt_llm/_torch/modules/linear.py | 5 +- tensorrt_llm/_torch/modules/mlp.py | 5 +- .../peft/lora/cuda_graph_lora_manager.py | 16 + tensorrt_llm/_torch/peft/lora/layer.py | 102 ++++++- .../_torch/pyexecutor/model_engine.py | 5 +- .../_torch/pyexecutor/resource_manager.py | 7 +- tensorrt_llm/lora_manager.py | 21 +- .../test_fp8_lora_grouped_gemm_regressions.py | 174 +++++++++++ .../tests_lora_modules/test_qwen3_sanity.py | 6 +- 26 files changed, 1280 insertions(+), 62 deletions(-) create mode 100644 tests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.py diff --git a/cpp/include/tensorrt_llm/batch_manager/peftCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/peftCacheManager.h index ed928e96d811..81d4ff622d46 100644 --- a/cpp/include/tensorrt_llm/batch_manager/peftCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/peftCacheManager.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. @@ -29,6 +29,8 @@ #include #include +#include +#include #include #include #include @@ -125,6 +127,8 @@ class PeftCacheManager : public BasePeftCacheManager [[nodiscard]] bool isTaskCachedDevice(uint64_t const taskId) const; + [[nodiscard]] tensorrt_llm::DataType getDataType() const; + void resetDeviceCache() override; void markRequestDone(LlmRequest const& llmReq, bool pause = false) override; @@ -157,6 +161,8 @@ class PeftCacheManager : public BasePeftCacheManager void prefetchLoraWeights(std::string const& modelDir, runtime::BufferManager const& bufferManager); private: + void configureDataType(tensorrt_llm::DataType dataType); + std::unique_ptr mHostLoraCache; std::unique_ptr mDeviceLoraCache; @@ -175,6 +181,9 @@ class PeftCacheManager : public BasePeftCacheManager runtime::ModelConfig mModelConfig; runtime::WorldConfig mWorldConfig; + mutable std::mutex mDataTypeMutex; + std::optional mDataType; + int mDevice{-1}; }; diff --git a/cpp/include/tensorrt_llm/runtime/loraCache.h b/cpp/include/tensorrt_llm/runtime/loraCache.h index eb4ef57494ea..d6f8ca964539 100644 --- a/cpp/include/tensorrt_llm/runtime/loraCache.h +++ b/cpp/include/tensorrt_llm/runtime/loraCache.h @@ -1,5 +1,5 @@ -/*loraCac - * Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved. +/* + * Copyright (c) 2022-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. @@ -186,6 +186,16 @@ class LoraCache LoraCache(LoraCachePageManagerConfig const& pageManagerConfig, ModelConfig const& modelConfig, WorldConfig const& worldConfig, BufferManager const& bufferManager); + /** + * \brief Reinitialize an empty cache to store the given data type. + * + * The page geometry remains unchanged. This is intended for selecting a + * homogeneous LoRA cache dtype before the first task is inserted. + */ + void setDataType(tensorrt_llm::DataType dataType); + + [[nodiscard]] tensorrt_llm::DataType getDataType() const; + /** * \brief put a task in the cache, and claim pages for it, and optionally load task weights. * diff --git a/cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp index 89cc475b82ee..bbbe7895e795 100644 --- a/cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -224,6 +224,7 @@ void PeftCacheManager::prefetchLoraWeights(std::string const& modelDir, runtime: TensorPtr weights = runtime::utils::loadNpy(bufferManager, weightsFn, runtime::MemoryType::kCPU); TensorPtr config = runtime::utils::loadNpy(bufferManager, configFn, runtime::MemoryType::kCPU); TLLM_LOG_DEBUG("prefetch lora task %s", tasks[taskId].c_str()); + configureDataType(weights->getDataType()); mHostLoraCache->put(std::stoi(tasks[taskId]), weights, config, true); } TLLM_LOG_DEBUG("%s stop", __PRETTY_FUNCTION__); @@ -261,9 +262,14 @@ void PeftCacheManager::addRequestPeft(std::shared_ptr llmRequest, bo auto optTaskId = llmRequest->getLoraTaskId(); auto optLoraWeights = llmRequest->getLoraWeights(); auto optLoraConfig = llmRequest->getLoraConfig(); + if (optLoraWeights) + { + configureDataType(optLoraWeights.value()->getDataType()); + } if (optTaskId || optLoraWeights || optLoraConfig) { - runtime::lora::loraValidateRequestTensors(optTaskId, optLoraWeights, optLoraConfig, mModelConfig, mWorldConfig); + runtime::lora::loraValidateRequestTensors( + optTaskId, optLoraWeights, optLoraConfig, mModelConfig, mWorldConfig, getDataType()); } else { @@ -374,6 +380,28 @@ void PeftCacheManager::addRequestPeft(std::shared_ptr llmRequest, bo TLLM_LOG_DEBUG("%s stop", __PRETTY_FUNCTION__); } +void PeftCacheManager::configureDataType(tensorrt_llm::DataType dataType) +{ + std::lock_guard lock(mDataTypeMutex); + if (mDataType) + { + TLLM_CHECK_WITH_INFO(mDataType.value() == dataType, + "PEFT cache supports one homogeneous LoRA dtype; cache is %s but received %s", + runtime::IBuffer::getDataTypeName(mDataType.value()), runtime::IBuffer::getDataTypeName(dataType)); + return; + } + + mHostLoraCache->setDataType(dataType); + mDeviceLoraCache->setDataType(dataType); + mDataType = dataType; +} + +tensorrt_llm::DataType PeftCacheManager::getDataType() const +{ + std::lock_guard lock(mDataTypeMutex); + return mDataType.value_or(mModelConfig.getDataType()); +} + std::tuple>, BasePeftCacheManager::TaskIdToReqIds> PeftCacheManager::getTaskMaps(RequestVector const& contextRequests, RequestVector const& generationRequests) { diff --git a/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu b/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu index 8af43b2b4914..d826056969da 100644 --- a/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu +++ b/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu @@ -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. @@ -27,10 +27,51 @@ #include "tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/device/splitk_gemm_grouped.h" #include "tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/default_splitk_gemm_grouped.h" +#ifdef ENABLE_FP8 +#include "cute/tensor.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/gemm/group_array_problem_shape.hpp" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/util/packed_stride.hpp" +#include "tensorrt_llm/common/memoryUtils.h" +#endif // ENABLE_FP8 + TRTLLM_NAMESPACE_BEGIN namespace kernels { +namespace +{ + +constexpr int kFp8TmaAlignment = 16; + +void checkFp8CudaGraphAlignment( + cutlass::gemm::GemmCoord const* hostMaxProblemSizesPtr, int problemCount, int minKN, char const* kernelName) +{ + 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. " + "The minimum K/N dimension is %d.", + kernelName, kFp8TmaAlignment, minKN); + + for (int problemIdx = 0; problemIdx < problemCount; ++problemIdx) + { + auto const& problem = hostMaxProblemSizesPtr[problemIdx]; + if (problem.m() == 0 || problem.n() == 0 || problem.k() == 0) + { + continue; + } + + TLLM_CHECK_WITH_INFO(problem.n() % kFp8TmaAlignment == 0 && problem.k() % kFp8TmaAlignment == 0, + "%s requires GEMM N and K dimensions to be multiples of %d elements for 128-bit TMA alignment. " + "Max problem %d has M=%d, N=%d, K=%d.", + kernelName, kFp8TmaAlignment, problemIdx, problem.m(), problem.n(), problem.k()); + } +} + +} // namespace /** * Template for CUDA Graph compatible grouped GEMM that directly uses GPU tensors @@ -140,10 +181,211 @@ void cudaGraphGroupedGemmType(cutlass::gemm::GemmCoord* problemSizesPtr, int pro } } +#ifdef ENABLE_FP8 +#if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) + +// ==================================================================== +// FP8 CUDA-graph-compatible grouped GEMM using CUTLASS 3.x. +// +// The CUDA graph variant receives live problem sizes, pointers, and leading +// dimensions already on the GPU. The 3.x API needs cute stride arrays on +// device, so a small setup kernel converts the existing graph parameters into +// the CUTLASS 3.x argument layout on the same stream. +// ==================================================================== + +template +__global__ void fillFp8CudaGraphGroupedGemmParams(cutlass::gemm::GemmCoord const* problemSizesPtr, int problemCount, + int64_t const* ldaGpu, int64_t const* ldbGpu, int64_t const* ldcGpu, int64_t const* lddGpu, + typename ProblemShape::UnderlyingProblemShape* problemShapes, StrideA* strideA, StrideB* strideB, StrideC* strideC, + StrideD* strideD) +{ + int const problemIdx = blockIdx.x * blockDim.x + threadIdx.x; + if (problemIdx >= problemCount) + { + return; + } + + int const m = problemSizesPtr[problemIdx].m(); + int const n = problemSizesPtr[problemIdx].n(); + int const k = problemSizesPtr[problemIdx].k(); + + problemShapes[problemIdx] = cute::make_shape(m, n, k); + strideA[problemIdx] + = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(m, static_cast(ldaGpu[problemIdx]), 1)); + strideB[problemIdx] + = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(n, static_cast(ldbGpu[problemIdx]), 1)); + strideC[problemIdx] + = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, static_cast(ldcGpu[problemIdx]), 1)); + strideD[problemIdx] + = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(m, static_cast(lddGpu[problemIdx]), 1)); +} + +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) +{ + TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); + + using namespace cute; + + 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 ElementAccumulator = float; + + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::RowMajor; + using LayoutD = cutlass::layout::RowMajor; + + static constexpr int kAlignmentA = 128 / cutlass::sizeof_bits::value; + static constexpr int kAlignmentB = 128 / cutlass::sizeof_bits::value; + 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 CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder( + sizeof(typename CollectiveEpilogue::SharedStorage))>, + KernelSchedule>::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + using StrideA = typename GemmKernel::InternalStrideA; + using StrideB = typename GemmKernel::InternalStrideB; + using StrideC = typename GemmKernel::InternalStrideC; + using StrideD = typename GemmKernel::InternalStrideD; + + // Allocate device memory for problem shapes, pointer arrays, and strides. + auto const tensorOpts = at::TensorOptions().dtype(at::kByte).device(at::kCUDA); + + auto const align16 = [](int64_t bytes) -> int64_t { return (bytes + 15) / 16 * 16; }; + + int64_t const szProblemShapes = align16(problemCount * sizeof(typename ProblemShape::UnderlyingProblemShape)); + int64_t const szPtrA = align16(problemCount * sizeof(ElementA const*)); + int64_t const szPtrB = align16(problemCount * sizeof(ElementB const*)); + int64_t const szPtrC = align16(problemCount * sizeof(ElementC const*)); + int64_t const szPtrD = align16(problemCount * sizeof(ElementD*)); + int64_t const szStrideA = align16(problemCount * sizeof(StrideA)); + int64_t const szStrideB = align16(problemCount * sizeof(StrideB)); + int64_t const szStrideC = align16(problemCount * sizeof(StrideC)); + int64_t const szStrideD = align16(problemCount * sizeof(StrideD)); + + int64_t const totalParamsBytes + = szProblemShapes + szPtrA + szPtrB + szPtrC + szPtrD + szStrideA + szStrideB + szStrideC + szStrideD; + + at::Tensor paramsTensor = at::empty({totalParamsBytes}, tensorOpts); + char* devBase = static_cast(paramsTensor.data_ptr()); + + int64_t devOff = 0; + auto devPtr = [&devBase, &devOff](int64_t sz) -> void* + { + void* p = devBase + devOff; + devOff += (sz + 15) / 16 * 16; + return p; + }; + + auto* devProblemShapes = static_cast(devPtr(szProblemShapes)); + auto* devPtrA = static_cast(devPtr(szPtrA)); + auto* devPtrB = static_cast(devPtr(szPtrB)); + auto* devPtrC = static_cast(devPtr(szPtrC)); + auto* devPtrD = static_cast(devPtr(szPtrD)); + auto* devStrideA = static_cast(devPtr(szStrideA)); + auto* devStrideB = static_cast(devPtr(szStrideB)); + auto* devStrideC = static_cast(devPtr(szStrideC)); + auto* devStrideD = static_cast(devPtr(szStrideD)); + + cudaMemcpyAsync(devPtrA, ptrAGpu, problemCount * sizeof(void*), cudaMemcpyDeviceToDevice, stream); + cudaMemcpyAsync(devPtrB, ptrBGpu, problemCount * sizeof(void*), cudaMemcpyDeviceToDevice, stream); + cudaMemcpyAsync(devPtrC, ptrCGpu, problemCount * sizeof(void*), cudaMemcpyDeviceToDevice, stream); + cudaMemcpyAsync(devPtrD, ptrDGpu, problemCount * sizeof(void*), cudaMemcpyDeviceToDevice, stream); + + int constexpr kThreads = 256; + int const blocks = (problemCount + kThreads - 1) / kThreads; + fillFp8CudaGraphGroupedGemmParams + <<>>(problemSizesPtr, problemCount, ldaGpu, ldbGpu, ldcGpu, lddGpu, + devProblemShapes, devStrideA, devStrideB, devStrideC, devStrideD); + + cutlass::KernelHardwareInfo hwInfo; + hwInfo.device_id = 0; + cudaGetDevice(&hwInfo.device_id); + hwInfo.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hwInfo.device_id); + + typename Gemm::Arguments arguments{cutlass::gemm::GemmUniversalMode::kGrouped, + {problemCount, devProblemShapes, nullptr}, {devPtrA, devStrideA, devPtrB, devStrideB}, + {{1.0f, 0.0f}, devPtrC, devStrideC, devPtrD, devStrideD}, hwInfo}; + + Gemm gemm; + + size_t const requiredWorkspace = Gemm::get_workspace_size(arguments); + at::Tensor workspace; + void* gemmWorkspace = nullptr; + if (requiredWorkspace > 0) + { + workspace = at::empty({static_cast(requiredWorkspace)}, tensorOpts); + gemmWorkspace = workspace.data_ptr(); + } + + cutlass::Status status = gemm.can_implement(arguments); + TLLM_CHECK_WITH_INFO(status == cutlass::Status::kSuccess, "FP8 CUDA graph grouped GEMM can_implement failed: %s", + cutlass::cutlassGetStatusString(status)); + + status = gemm.initialize(arguments, gemmWorkspace, stream); + TLLM_CHECK_WITH_INFO(status == cutlass::Status::kSuccess, "FP8 CUDA graph grouped GEMM initialize failed: %s", + cutlass::cutlassGetStatusString(status)); + + status = gemm.run(stream); + sync_check_cuda_error(stream); + TLLM_CHECK_WITH_INFO(status == cutlass::Status::kSuccess, "FP8 CUDA graph grouped GEMM run failed: %s", + cutlass::cutlassGetStatusString(status)); + + TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); +} + +#endif // CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED +#endif // ENABLE_FP8 + 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) { +#ifdef ENABLE_FP8 +#if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_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); + return; + } +#else + 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+ " + "kernels)."); + } +#endif // CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED +#endif // ENABLE_FP8 + if (isLoraIn) { if (minKN >= 8) @@ -312,6 +554,26 @@ void cudaGraphSplitKGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int p bool isLoraIn, tensorrt_llm::DataType dataType, int splitKSlices, int minKN, cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, int64_t* splitKOffsetsGpu, cudaStream_t stream) { +#ifdef ENABLE_FP8 +#if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) + if (dataType == tensorrt_llm::DataType::kFP8) + { + // 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( + problemSizesPtr, problemCount, ptrAGpu, ptrBGpu, ptrCGpu, ptrDGpu, ldaGpu, ldbGpu, ldcGpu, lddGpu, stream); + return; + } +#else + 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+ " + "kernels)."); + } +#endif // CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED +#endif // ENABLE_FP8 + if (isLoraIn) { if (minKN >= 8) diff --git a/cpp/tensorrt_llm/kernels/groupGemm.cu b/cpp/tensorrt_llm/kernels/groupGemm.cu index b41021ffe6f7..50334a58547b 100644 --- a/cpp/tensorrt_llm/kernels/groupGemm.cu +++ b/cpp/tensorrt_llm/kernels/groupGemm.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & + * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -22,6 +22,17 @@ #include "cutlass/gemm/kernel/default_gemm_grouped.h" #include "cutlass/gemm/kernel/gemm_grouped.h" +#ifdef ENABLE_FP8 +#include "cute/tensor.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/gemm/group_array_problem_shape.hpp" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/util/packed_stride.hpp" +#endif // ENABLE_FP8 + #include "groupGemm.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/config.h" @@ -34,6 +45,29 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { +namespace +{ + +constexpr int kFp8TmaAlignment = 16; + +void checkFp8GroupedGemmAlignment(std::vector const& problemSizes, char const* kernelName) +{ + for (size_t problemIdx = 0; problemIdx < problemSizes.size(); ++problemIdx) + { + auto const& problem = problemSizes[problemIdx]; + if (problem.m() == 0 || problem.n() == 0 || problem.k() == 0) + { + continue; + } + + TLLM_CHECK_WITH_INFO(problem.n() % kFp8TmaAlignment == 0 && problem.k() % kFp8TmaAlignment == 0, + "%s requires GEMM N and K dimensions to be multiples of %d elements for 128-bit TMA alignment. " + "Problem %zu has M=%d, N=%d, K=%d.", + kernelName, kFp8TmaAlignment, problemIdx, problem.m(), problem.n(), problem.k()); + } +} + +} // namespace int64_t inline getGemmCoordSize(int64_t problemCount) { @@ -59,6 +93,22 @@ int64_t getGroupedGemmParamsWorkSpaceSize(int64_t problemCount) return gemm_coord_size + ptr_size + ldd_size; } +int64_t getFp8GroupedGemmParamsWorkSpaceSize(int64_t problemCount) +{ + // The FP8 grouped GEMM (CUTLASS 3.x) stores on device: + // - problem shapes: problemCount * 12 bytes (Shape) + // - 4 pointer arrays: 4 * problemCount * 8 bytes + // - 4 stride arrays: 4 * problemCount * 16 bytes (conservative upper bound for cute strides) + // All 16-byte aligned. + auto align16 = [](int64_t bytes) -> int64_t { return (bytes + 15) / 16 * 16; }; + + int64_t const szProblemShapes = align16(problemCount * 12); + int64_t const szPtrs = 4 * align16(problemCount * 8); + int64_t const szStrides = 4 * align16(problemCount * 16); + + return szProblemShapes + szPtrs + szStrides; +} + template void groupedGemm_(std::vector problem_sizes, std::vector const& ptrA, @@ -201,12 +251,246 @@ void groupedGemmType_(std::vector problem_sizes, std:: #endif } +#ifdef ENABLE_FP8 + +// ==================================================================== +// FP8 grouped GEMM using CUTLASS 3.x collective API (Hopper sm90+). +// +// The legacy CUTLASS 2.x DefaultGemmGrouped does NOT support fp8 element +// types. This implementation uses the CUTLASS 3.x CollectiveBuilder and +// GemmUniversal kernel with PtrArray schedules, following the pattern of +// cutlass/examples/57_hopper_grouped_gemm. +// +// Layout convention: +// A: RowMajor (M x K) +// B: ColumnMajor (K x N, stored column-major) +// D: RowMajor (M x N, output; C is aliased to D for beta=0) +// +// The 3.x grouped API uses pointer-array arguments: arrays of per-group +// pointers and strides on the GPU. We build these on the host and copy +// them to device memory via the provided workspace. +// ==================================================================== + +#if defined(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) +{ + TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); + + using namespace cute; + + 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 ElementAccumulator = float; + + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::RowMajor; + using LayoutD = cutlass::layout::RowMajor; + + // 128-bit TMA alignment: 128 bits / 8 bits per fp8 element = 16 elements + static constexpr int kAlignmentA = 128 / cutlass::sizeof_bits::value; + static constexpr int kAlignmentB = 128 / cutlass::sizeof_bits::value; + 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 CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder( + sizeof(typename CollectiveEpilogue::SharedStorage))>, + KernelSchedule>::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + using StrideA = typename GemmKernel::InternalStrideA; + using StrideB = typename GemmKernel::InternalStrideB; + using StrideC = typename GemmKernel::InternalStrideC; + using StrideD = typename GemmKernel::InternalStrideD; + + int const problemCount = static_cast(problemSizes.size()); + + // Build host-side arrays for problem shapes, pointers, and strides. + std::vector problemShapesHost(problemCount); + std::vector ptrAHost(problemCount); + std::vector ptrBHost(problemCount); + std::vector ptrCHost(problemCount); + std::vector ptrDHost(problemCount); + std::vector strideAHost(problemCount); + std::vector strideBHost(problemCount); + std::vector strideCHost(problemCount); + std::vector strideDHost(problemCount); + + for (int i = 0; i < problemCount; ++i) + { + auto const& coord = problemSizes[i]; + int const m = coord.m(); + int const n = coord.n(); + int const k = coord.k(); + + problemShapesHost[i] = cute::make_shape(m, n, k); + ptrAHost[i] = static_cast(ptrA[i]); + ptrBHost[i] = static_cast(ptrB[i]); + ptrCHost[i] = static_cast(ptrC[i]); + ptrDHost[i] = static_cast(ptrD[i]); + + strideAHost[i] = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(m, k, 1)); + strideBHost[i] = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(n, k, 1)); + strideCHost[i] = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, n, 1)); + strideDHost[i] = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(m, n, 1)); + } + + // Compute workspace layout sizes (all 16-byte aligned). + auto const align16 = [](int64_t bytes) -> int64_t { return (bytes + 15) / 16 * 16; }; + + int64_t const szProblemShapes = align16(problemCount * sizeof(typename ProblemShape::UnderlyingProblemShape)); + int64_t const szPtrA = align16(problemCount * sizeof(ElementA const*)); + int64_t const szPtrB = align16(problemCount * sizeof(ElementB const*)); + int64_t const szPtrC = align16(problemCount * sizeof(ElementC const*)); + int64_t const szPtrD = align16(problemCount * sizeof(ElementD*)); + int64_t const szStrideA = align16(problemCount * sizeof(StrideA)); + int64_t const szStrideB = align16(problemCount * sizeof(StrideB)); + int64_t const szStrideC = align16(problemCount * sizeof(StrideC)); + int64_t const szStrideD = align16(problemCount * sizeof(StrideD)); + + int64_t const totalParamsBytes + = szProblemShapes + szPtrA + szPtrB + szPtrC + szPtrD + szStrideA + szStrideB + szStrideC + szStrideD; + TLLM_CHECK_WITH_INFO(totalParamsBytes <= gemmParamsWorkSpaceSize, + "FP8 grouped GEMM params workspace too small: need %ld, have %ld", totalParamsBytes, gemmParamsWorkSpaceSize); + + // Stage everything into a single host buffer, then memcpy to device. + std::vector hostBuf(totalParamsBytes); + char* cursor = hostBuf.data(); + + auto copyToHostBuf = [&cursor](void const* src, int64_t bytes) + { + std::memcpy(cursor, src, bytes); + char* base = cursor; + cursor += (bytes + 15) / 16 * 16; + return base; + }; + + copyToHostBuf(problemShapesHost.data(), problemCount * sizeof(typename ProblemShape::UnderlyingProblemShape)); + copyToHostBuf(ptrAHost.data(), problemCount * sizeof(ElementA const*)); + copyToHostBuf(ptrBHost.data(), problemCount * sizeof(ElementB const*)); + copyToHostBuf(ptrCHost.data(), problemCount * sizeof(ElementC const*)); + copyToHostBuf(ptrDHost.data(), problemCount * sizeof(ElementD*)); + copyToHostBuf(strideAHost.data(), problemCount * sizeof(StrideA)); + copyToHostBuf(strideBHost.data(), problemCount * sizeof(StrideB)); + copyToHostBuf(strideCHost.data(), problemCount * sizeof(StrideC)); + copyToHostBuf(strideDHost.data(), problemCount * sizeof(StrideD)); + + tensorrt_llm::common::cudaAutoCpy(reinterpret_cast(gemmParamsWorkSpace), + reinterpret_cast(hostBuf.data()), totalParamsBytes, stream); + + // Compute device pointers from the workspace base. + char* devBase = static_cast(gemmParamsWorkSpace); + int64_t devOff = 0; + + auto devPtr = [&devBase, &devOff](int64_t sz) -> void* + { + void* p = devBase + devOff; + devOff += (sz + 15) / 16 * 16; + return p; + }; + + auto* devProblemShapes = static_cast(devPtr(szProblemShapes)); + auto* devPtrA = static_cast(devPtr(szPtrA)); + auto* devPtrB = static_cast(devPtr(szPtrB)); + auto* devPtrC = static_cast(devPtr(szPtrC)); + auto* devPtrD = static_cast(devPtr(szPtrD)); + auto* devStrideA = static_cast(devPtr(szStrideA)); + auto* devStrideB = static_cast(devPtr(szStrideB)); + auto* devStrideC = static_cast(devPtr(szStrideC)); + auto* devStrideD = static_cast(devPtr(szStrideD)); + + // Hardware info for kernel launch. + cutlass::KernelHardwareInfo hwInfo; + hwInfo.device_id = 0; + cudaGetDevice(&hwInfo.device_id); + hwInfo.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hwInfo.device_id); + + // 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}; + + Gemm gemm; + + size_t const workspaceSize = Gemm::get_workspace_size(arguments); + TLLM_CHECK_WITH_INFO(static_cast(workspaceSize) <= gemmWorkspaceSize, + "FP8 grouped GEMM workspace too small: need %lu, have %ld", workspaceSize, gemmWorkspaceSize); + + cutlass::Status status = gemm.can_implement(arguments); + TLLM_CHECK_WITH_INFO(status == cutlass::Status::kSuccess, "FP8 grouped GEMM can_implement failed: %s", + cutlass::cutlassGetStatusString(status)); + + status = gemm.initialize(arguments, gemmWorkSpace, stream); + TLLM_CHECK_WITH_INFO(status == cutlass::Status::kSuccess, "FP8 grouped GEMM initialize failed: %s", + cutlass::cutlassGetStatusString(status)); + + status = gemm.run(stream); + TLLM_CHECK_WITH_INFO(status == cutlass::Status::kSuccess, "FP8 grouped GEMM run failed: %s", + cutlass::cutlassGetStatusString(status)); + + sync_check_cuda_error(stream); + + TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); +} + +#endif // CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED + +#endif // ENABLE_FP8 + 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, bool isLoraIn, tensorrt_llm::DataType dataType, int minKN, cudaStream_t stream) { 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 (dataType == tensorrt_llm::DataType::kFP8) + { + checkFp8GroupedGemmAlignment(problem_sizes, "FP8 grouped GEMM"); + + fp8GroupedGemm(problem_sizes, ptrA, ptrB, ptrC, ptrD, gemmParamsWorkSpace, gemmParamsWorkSpaceSize, + gemmWorkSpace, gemmWorkspaceSize, stream); + return; + } +#else + 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)."); + } +#endif // CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED +#endif // ENABLE_FP8 + if (isLoraIn) { // K >> N, like K = 1024, N = 8 diff --git a/cpp/tensorrt_llm/kernels/groupGemm.h b/cpp/tensorrt_llm/kernels/groupGemm.h index c526e08e986c..5cf3a8d3e6be 100644 --- a/cpp/tensorrt_llm/kernels/groupGemm.h +++ b/cpp/tensorrt_llm/kernels/groupGemm.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2019-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. @@ -26,6 +26,12 @@ namespace kernels int64_t getGroupedGemmParamsWorkSpaceSize(int64_t problem_count); +//! @brief Returns the parameter workspace size for the FP8 grouped GEMM path. +//! +//! The FP8 path (CUTLASS 3.x) stores problem shapes, pointer arrays, and +//! cute stride arrays on the device. This function returns the required bytes. +int64_t getFp8GroupedGemmParamsWorkSpaceSize(int64_t problemCount); + 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/kernels/lora/lora.cpp b/cpp/tensorrt_llm/kernels/lora/lora.cpp index 7a2b7d330afc..14732922fcc5 100644 --- a/cpp/tensorrt_llm/kernels/lora/lora.cpp +++ b/cpp/tensorrt_llm/kernels/lora/lora.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & + * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -99,6 +99,13 @@ void LoraImpl::setGemmConfig() mCublasWrapper->setBF16GemmConfig(); } #endif +#ifdef ENABLE_FP8 + else if (mType == tensorrt_llm::DataType::kFP8) + { + // FP8 uses the grouped GEMM path (CUTLASS 3.x), not cuBLAS. + // No cuBLAS config needed; this is a no-op. + } +#endif } int64_t getLowRankWorkSpaceSize(int64_t numTokens, int64_t maxLoraModuleNum, int64_t maxLowRank, int64_t typeSize) @@ -108,7 +115,10 @@ int64_t getLowRankWorkSpaceSize(int64_t numTokens, int64_t maxLoraModuleNum, int int64_t getGemmParamsWorkSpaceSize(int64_t nbReq) { - return std::max(getSplitkGroupedGemmParamsWorkSpaceSize(nbReq), getGroupedGemmParamsWorkSpaceSize(nbReq)); + int64_t size = std::max(getSplitkGroupedGemmParamsWorkSpaceSize(nbReq), getGroupedGemmParamsWorkSpaceSize(nbReq)); + // The FP8 path (CUTLASS 3.x) needs additional space for cute strides. + size = std::max(size, getFp8GroupedGemmParamsWorkSpaceSize(nbReq)); + return size; } int64_t getSplitkGroupedGemmWorkSpaceSize( @@ -177,6 +187,15 @@ int LoraImpl::run(int64_t numTokens, int64_t numReqs, void const* input, int32_t char* useUnifiedGemmChar = std::getenv("LORA_USE_UNIFIED_GEMM"); bool useUnifiedGemm = (useUnifiedGemmChar == nullptr || std::string(useUnifiedGemmChar) != "OFF"); +#ifdef ENABLE_FP8 + // The cuBLAS wrapper used by the unified GEMM path does not have an fp8 config. + // Force the grouped GEMM path for fp8, which uses the CUTLASS 3.x collective API. + if (mType == tensorrt_llm::DataType::kFP8) + { + useUnifiedGemm = false; + } +#endif + for (int loraModuleIdx = 0; loraModuleIdx < mNumLoraModules; loraModuleIdx++) { auto const loraRankModule = loraRanks[loraModuleIdx * numTokens]; diff --git a/cpp/tensorrt_llm/kernels/splitkGroupGemm.cu b/cpp/tensorrt_llm/kernels/splitkGroupGemm.cu index 1f63189ec657..a4058c052d30 100644 --- a/cpp/tensorrt_llm/kernels/splitkGroupGemm.cu +++ b/cpp/tensorrt_llm/kernels/splitkGroupGemm.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & + * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -31,10 +31,44 @@ #include "tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/default_splitk_gemm_grouped.h" #include "tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/splitk_gemm_grouped.h" +#ifdef ENABLE_FP8 +#include "cute/tensor.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/gemm/group_array_problem_shape.hpp" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/util/packed_stride.hpp" +#endif // ENABLE_FP8 + TRTLLM_NAMESPACE_BEGIN namespace kernels { +namespace +{ + +constexpr int kFp8TmaAlignment = 16; + +void checkFp8GroupedGemmAlignment(std::vector const& problemSizes, char const* kernelName) +{ + for (size_t problemIdx = 0; problemIdx < problemSizes.size(); ++problemIdx) + { + auto const& problem = problemSizes[problemIdx]; + if (problem.m() == 0 || problem.n() == 0 || problem.k() == 0) + { + continue; + } + + TLLM_CHECK_WITH_INFO(problem.n() % kFp8TmaAlignment == 0 && problem.k() % kFp8TmaAlignment == 0, + "%s requires GEMM N and K dimensions to be multiples of %d elements for 128-bit TMA alignment. " + "Problem %zu has M=%d, N=%d, K=%d.", + kernelName, kFp8TmaAlignment, problemIdx, problem.m(), problem.n(), problem.k()); + } +} + +} // namespace int64_t inline getGemmCoordSize(int64_t problemCount) { @@ -226,12 +260,228 @@ void splitkGroupedGemmType_(std::vector const& problem #endif } +#ifdef ENABLE_FP8 + +// ==================================================================== +// FP8 split-K grouped GEMM using CUTLASS 3.x collective API (Hopper). +// +// CUTLASS 3.x grouped GEMM does not have a split-K variant. We use the +// regular cooperative fp8 grouped GEMM instead: the Hopper cooperative +// schedule distributes work across warps efficiently even for large-K +// small-N problems. The output is fp8, matching the non-split-K path's +// final output type convention used by LoRA. +// ==================================================================== + +#if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) + +void fp8SplitkGroupedGemm(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) +{ + TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); + + using namespace cute; + + 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 ElementAccumulator = float; + + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::RowMajor; + using LayoutD = cutlass::layout::RowMajor; + + static constexpr int kAlignmentA = 128 / cutlass::sizeof_bits::value; + static constexpr int kAlignmentB = 128 / cutlass::sizeof_bits::value; + 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 CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder( + sizeof(typename CollectiveEpilogue::SharedStorage))>, + KernelSchedule>::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + using StrideA = typename GemmKernel::InternalStrideA; + using StrideB = typename GemmKernel::InternalStrideB; + using StrideC = typename GemmKernel::InternalStrideC; + using StrideD = typename GemmKernel::InternalStrideD; + + int const problemCount = static_cast(problemSizes.size()); + + std::vector problemShapesHost(problemCount); + std::vector ptrAHost(problemCount); + std::vector ptrBHost(problemCount); + std::vector ptrCHost(problemCount); + std::vector ptrDHost(problemCount); + std::vector strideAHost(problemCount); + std::vector strideBHost(problemCount); + std::vector strideCHost(problemCount); + std::vector strideDHost(problemCount); + + for (int i = 0; i < problemCount; ++i) + { + auto const& coord = problemSizes[i]; + int const m = coord.m(); + int const n = coord.n(); + int const k = coord.k(); + + problemShapesHost[i] = cute::make_shape(m, n, k); + ptrAHost[i] = static_cast(ptrA[i]); + ptrBHost[i] = static_cast(ptrB[i]); + ptrCHost[i] = static_cast(ptrC[i]); + ptrDHost[i] = static_cast(ptrD[i]); + + strideAHost[i] = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(m, k, 1)); + strideBHost[i] = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(n, k, 1)); + strideCHost[i] = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, n, 1)); + strideDHost[i] = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(m, n, 1)); + } + + auto const align16 = [](int64_t bytes) -> int64_t { return (bytes + 15) / 16 * 16; }; + + int64_t const szProblemShapes = align16(problemCount * sizeof(typename ProblemShape::UnderlyingProblemShape)); + int64_t const szPtrA = align16(problemCount * sizeof(ElementA const*)); + int64_t const szPtrB = align16(problemCount * sizeof(ElementB const*)); + int64_t const szPtrC = align16(problemCount * sizeof(ElementC const*)); + int64_t const szPtrD = align16(problemCount * sizeof(ElementD*)); + int64_t const szStrideA = align16(problemCount * sizeof(StrideA)); + int64_t const szStrideB = align16(problemCount * sizeof(StrideB)); + int64_t const szStrideC = align16(problemCount * sizeof(StrideC)); + int64_t const szStrideD = align16(problemCount * sizeof(StrideD)); + + int64_t const totalParamsBytes + = szProblemShapes + szPtrA + szPtrB + szPtrC + szPtrD + szStrideA + szStrideB + szStrideC + szStrideD; + TLLM_CHECK_WITH_INFO(totalParamsBytes <= gemmParamsWorkSpaceSize, + "FP8 split-K grouped GEMM params workspace too small: need %ld, have %ld", totalParamsBytes, + gemmParamsWorkSpaceSize); + + std::vector hostBuf(totalParamsBytes); + char* cursor = hostBuf.data(); + + auto copyToHostBuf = [&cursor](void const* src, int64_t bytes) + { + std::memcpy(cursor, src, bytes); + cursor += (bytes + 15) / 16 * 16; + }; + + copyToHostBuf(problemShapesHost.data(), problemCount * sizeof(typename ProblemShape::UnderlyingProblemShape)); + copyToHostBuf(ptrAHost.data(), problemCount * sizeof(ElementA const*)); + copyToHostBuf(ptrBHost.data(), problemCount * sizeof(ElementB const*)); + copyToHostBuf(ptrCHost.data(), problemCount * sizeof(ElementC const*)); + copyToHostBuf(ptrDHost.data(), problemCount * sizeof(ElementD*)); + copyToHostBuf(strideAHost.data(), problemCount * sizeof(StrideA)); + copyToHostBuf(strideBHost.data(), problemCount * sizeof(StrideB)); + copyToHostBuf(strideCHost.data(), problemCount * sizeof(StrideC)); + copyToHostBuf(strideDHost.data(), problemCount * sizeof(StrideD)); + + tensorrt_llm::common::cudaAutoCpy(reinterpret_cast(gemmParamsWorkSpace), + reinterpret_cast(hostBuf.data()), totalParamsBytes, stream); + + char* devBase = static_cast(gemmParamsWorkSpace); + int64_t devOff = 0; + + auto devPtr = [&devBase, &devOff](int64_t sz) -> void* + { + void* p = devBase + devOff; + devOff += (sz + 15) / 16 * 16; + return p; + }; + + auto* devProblemShapes = static_cast(devPtr(szProblemShapes)); + auto* devPtrA = static_cast(devPtr(szPtrA)); + auto* devPtrB = static_cast(devPtr(szPtrB)); + auto* devPtrC = static_cast(devPtr(szPtrC)); + auto* devPtrD = static_cast(devPtr(szPtrD)); + auto* devStrideA = static_cast(devPtr(szStrideA)); + auto* devStrideB = static_cast(devPtr(szStrideB)); + auto* devStrideC = static_cast(devPtr(szStrideC)); + auto* devStrideD = static_cast(devPtr(szStrideD)); + + cutlass::KernelHardwareInfo hwInfo; + hwInfo.device_id = 0; + cudaGetDevice(&hwInfo.device_id); + hwInfo.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hwInfo.device_id); + + typename Gemm::Arguments arguments{cutlass::gemm::GemmUniversalMode::kGrouped, + {problemCount, devProblemShapes, nullptr}, {devPtrA, devStrideA, devPtrB, devStrideB}, + {{1.0f, 0.0f}, devPtrC, devStrideC, devPtrD, devStrideD}, hwInfo}; + + Gemm gemm; + + size_t const workspaceSize = Gemm::get_workspace_size(arguments); + TLLM_CHECK_WITH_INFO(static_cast(workspaceSize) <= gemmWorkSpaceSize, + "FP8 split-K grouped GEMM workspace too small: need %lu, have %ld", workspaceSize, gemmWorkSpaceSize); + + cutlass::Status status = gemm.can_implement(arguments); + TLLM_CHECK_WITH_INFO(status == cutlass::Status::kSuccess, "FP8 split-K grouped GEMM can_implement failed: %s", + cutlass::cutlassGetStatusString(status)); + + status = gemm.initialize(arguments, gemmWorkSpace, stream); + TLLM_CHECK_WITH_INFO(status == cutlass::Status::kSuccess, "FP8 split-K grouped GEMM initialize failed: %s", + cutlass::cutlassGetStatusString(status)); + + status = gemm.run(stream); + TLLM_CHECK_WITH_INFO(status == cutlass::Status::kSuccess, "FP8 split-K grouped GEMM run failed: %s", + cutlass::cutlassGetStatusString(status)); + + sync_check_cuda_error(stream); + + TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); +} + +#endif // CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED + +#endif // ENABLE_FP8 + void splitkGroupedGemm(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, bool isLoraIn, tensorrt_llm::DataType dataType, int splitKSlices, int minKN, cudaStream_t stream) { 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 (dataType == tensorrt_llm::DataType::kFP8) + { + checkFp8GroupedGemmAlignment(problemSizes, "FP8 split-K grouped GEMM"); + + fp8SplitkGroupedGemm(problemSizes, ptrA, ptrB, ptrC, ptrD, gemmParamsWorkSpace, gemmParamsWorkSpaceSize, + gemmWorkSpace, gemmWorkSpaceSize, stream); + return; + } +#else + if (dataType == tensorrt_llm::DataType::kFP8) + { + TLLM_CHECK_WITH_INFO(false, + "FP8 split-K grouped GEMM requires CUTLASS modifiable TMA support (CUDA 12.3+ and Hopper sm90+ kernels)."); + } +#endif // CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED +#endif // ENABLE_FP8 + if (isLoraIn) { // K >> N, like K = 1024, N = 8 diff --git a/cpp/tensorrt_llm/kernels/splitkGroupGemm.h b/cpp/tensorrt_llm/kernels/splitkGroupGemm.h index bcde457db7d2..9f78215fc5bd 100644 --- a/cpp/tensorrt_llm/kernels/splitkGroupGemm.h +++ b/cpp/tensorrt_llm/kernels/splitkGroupGemm.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2019-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. diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp index 559934e1ef72..b2112650163d 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp @@ -748,6 +748,7 @@ void tb::BasePeftCacheManagerBindings::initBindings(nb::module_& m) nb::call_guard()) .def("is_task_cached_device", &tb::PeftCacheManager::isTaskCachedDevice, nb::arg("taskId"), nb::call_guard()) // ; + .def_prop_ro("data_type", &tb::PeftCacheManager::getDataType) .def("ensure_batch_map_task_id", &tb::PeftCacheManager::ensureBatchMapTaskId, nb::arg("context_requests"), nb::arg("generation_requests"), nb::arg("reset_gpu_cache") = false, nb::call_guard()); diff --git a/cpp/tensorrt_llm/runtime/loraCache.cpp b/cpp/tensorrt_llm/runtime/loraCache.cpp index 36fb0363816f..53cd23793b06 100644 --- a/cpp/tensorrt_llm/runtime/loraCache.cpp +++ b/cpp/tensorrt_llm/runtime/loraCache.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-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. @@ -511,6 +511,28 @@ LoraCache::LoraCache(LoraCachePageManagerConfig const& pageManagerConfig, ModelC } } +void LoraCache::setDataType(tensorrt_llm::DataType dataType) +{ + std::scoped_lock lock(mPagesMutex, mCacheMutex); + TLLM_CHECK_WITH_INFO(mCacheMap.empty(), "Cannot change LoRA cache dtype after a task has been inserted"); + if (mPageManagerConfig.getDataType() == dataType) + { + return; + } + + // Release the old allocation before creating the replacement to avoid a + // temporary peak equal to both cache allocations. + mCachePageManager.reset(); + mPageManagerConfig.setDataType(dataType); + mCachePageManager = std::make_unique(mPageManagerConfig, *mBufferManager); +} + +tensorrt_llm::DataType LoraCache::getDataType() const +{ + std::lock_guard lock(mPagesMutex); + return mPageManagerConfig.getDataType(); +} + template void LoraCache::splitTransposeCpuInner(ITensor& output, ITensor const& input, SizeType32 tpSize, SizeType32 tpRank) { @@ -562,6 +584,9 @@ std::vector LoraCache::copyToPages(TensorPtr s TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__); TLLM_CHECK_WITH_INFO(!pages.empty(), "empty pages"); + TLLM_CHECK_WITH_INFO(sourceWeights->getDataType() == pages.front()->getDataType(), + "Expected LoRA weights dtype %s to match cache dtype %s", sourceWeights->getDataTypeName(), + pages.front()->getDataTypeName()); TensorPtr weights = sourceWeights->getShape().nbDims == 2 ? sourceWeights @@ -775,6 +800,8 @@ void LoraCache::copyTask(TaskIdType taskId, LoraCache& deviceCache, bool markDon TLLM_CHECK_WITH_INFO(deviceCache.mPageManagerConfig.getMemoryType() == runtime::MemoryType::kGPU && !deviceCache.mDeviceBufferManagers.empty(), "The deviceCache must hold GPU memory and have at least one bufferManager / copy stream"); + TLLM_CHECK_WITH_INFO(mPageManagerConfig.getDataType() == deviceCache.mPageManagerConfig.getDataType(), + "LoRA host and device cache dtypes must match"); // First get the taskValue from this cache // TaskValue& taskValue = copyTaskGetThisTaskValue(taskId); diff --git a/cpp/tensorrt_llm/runtime/loraUtils.cpp b/cpp/tensorrt_llm/runtime/loraUtils.cpp index da7f1475ddec..7b78a34d5ef1 100644 --- a/cpp/tensorrt_llm/runtime/loraUtils.cpp +++ b/cpp/tensorrt_llm/runtime/loraUtils.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -72,7 +72,7 @@ void loraValidateRequestTensorDims(std::optional const& optR void loraValidateRequestTensors(std::optional const& optTaskId, std::optional const& optReqLoraWeights, std::optional const& optReqLoraConfig, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig) + runtime::WorldConfig const& worldConfig, std::optional loraDataType) { TLLM_CHECK_WITH_INFO(optTaskId.has_value(), "lora_task_id must be set for LoRA inference"); if (optReqLoraWeights.has_value() || optReqLoraConfig.has_value()) @@ -86,8 +86,10 @@ void loraValidateRequestTensors(std::optional const& optTaskId, : ITensor::view(config, ITensor::makeShape({config->getShape().d[1], config->getShape().d[2]})); SizeType32 nbModelLayers = modelConfig.getNbLoraLayers(); - TLLM_CHECK_WITH_INFO(weights->getDataType() == modelConfig.getDataType(), - "Expected lora weights to be the same data type as base model"); + auto const expectedDataType = loraDataType.value_or(modelConfig.getDataType()); + TLLM_CHECK_WITH_INFO(weights->getDataType() == expectedDataType, + "Expected LoRA weights dtype %s to match homogeneous cache dtype %s", weights->getDataTypeName(), + IBuffer::getDataTypeName(expectedDataType)); auto loraModules = modelConfig.getLoraModules(); auto maxAdapterSize = modelConfig.getMaxLoraRank(); diff --git a/cpp/tensorrt_llm/runtime/loraUtils.h b/cpp/tensorrt_llm/runtime/loraUtils.h index a3e2edf303aa..21401e332926 100644 --- a/cpp/tensorrt_llm/runtime/loraUtils.h +++ b/cpp/tensorrt_llm/runtime/loraUtils.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -22,6 +22,8 @@ #include "tensorrt_llm/runtime/modelConfig.h" #include "tensorrt_llm/runtime/worldConfig.h" +#include + namespace tensorrt_llm::runtime::lora { @@ -57,5 +59,5 @@ void loraValidateRequestTensorDims(std::optional const& optR void loraValidateRequestTensors(std::optional const& optTaskId, std::optional const& optReqLoraWeights, std::optional const& optReqLoraConfig, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig); + runtime::WorldConfig const& worldConfig, std::optional loraDataType = std::nullopt); } // namespace tensorrt_llm::runtime::lora diff --git a/cpp/tensorrt_llm/thop/loraOp.cpp b/cpp/tensorrt_llm/thop/loraOp.cpp index 6957987be6b8..604e2f2cc80a 100644 --- a/cpp/tensorrt_llm/thop/loraOp.cpp +++ b/cpp/tensorrt_llm/thop/loraOp.cpp @@ -1,6 +1,6 @@ /* - * Copyright (c) 2022-2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-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. @@ -157,7 +157,10 @@ std::vector lora_grouped_gemm(th::Tensor const& input, th::Tensor co { case torch::kFloat16: loraRuntimeDataType = tensorrt_llm::DataType::kHALF; break; case torch::kBFloat16: loraRuntimeDataType = tensorrt_llm::DataType::kBF16; break; - default: throw std::invalid_argument("Invalid dtype, only supports float16, bfloat16"); +#ifdef ENABLE_FP8 + case torch::kFloat8_e4m3fn: loraRuntimeDataType = tensorrt_llm::DataType::kFP8; break; +#endif + default: throw std::invalid_argument("Invalid dtype, only supports float16, bfloat16, float8_e4m3fn"); } auto mLoraImpl = std::make_shared( @@ -227,7 +230,12 @@ void lora_grouped_gemm_cuda_graph(th::Tensor const& lora_in_sizes, // [layer_mod { case torch::kFloat16: loraRuntimeDataType = tensorrt_llm::DataType::kHALF; break; case torch::kBFloat16: loraRuntimeDataType = tensorrt_llm::DataType::kBF16; break; - default: TORCH_CHECK(false, "Invalid dtype, only supports float16, bfloat16, got %s", c10::toString(dtype)); +#ifdef ENABLE_FP8 + case torch::kFloat8_e4m3fn: loraRuntimeDataType = tensorrt_llm::DataType::kFP8; break; +#endif + default: + TORCH_CHECK( + false, "Invalid dtype, only supports float16, bfloat16, float8_e4m3fn, got %s", c10::toString(dtype)); } int const minKnInt = std::max(1, static_cast(minKN)); @@ -307,7 +315,12 @@ void lora_group_gemm_param_fill_row_reorder_fusion(th::Tensor const& in_sizes, / { case torch::kFloat16: loraRuntimeDataType = tensorrt_llm::DataType::kHALF; break; case torch::kBFloat16: loraRuntimeDataType = tensorrt_llm::DataType::kBF16; break; - default: TORCH_CHECK(false, "Invalid dtype, only supports float16, bfloat16, got %s", c10::toString(dtype)); +#ifdef ENABLE_FP8 + case torch::kFloat8_e4m3fn: loraRuntimeDataType = tensorrt_llm::DataType::kFP8; break; +#endif + default: + TORCH_CHECK( + false, "Invalid dtype, only supports float16, bfloat16, float8_e4m3fn, got %s", c10::toString(dtype)); } int64_t const dtype_element_size = input.element_size(); diff --git a/cpp/tests/unit_tests/batch_manager/peftCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/peftCacheManagerTest.cpp index ef513894bfb6..a52f7f461179 100644 --- a/cpp/tests/unit_tests/batch_manager/peftCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/peftCacheManagerTest.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -143,6 +143,30 @@ TEST_F(PeftCacheManagerTest, addRequestPeftMissingTask) mPeftManager->addRequestPeft(llmRequest); } +#if defined(ENABLE_FP8) +TEST_F(PeftCacheManagerTest, adapterSelectsHomogeneousCacheDataType) +{ + SamplingConfig samplingConfig; + auto tokens = std::make_shared>(std::initializer_list{1, 2, 3, 4}); + TensorPtr fp8Weights = mManager->cpu(loraWeightsTp2->getShape(), tensorrt_llm::DataType::kFP8); + auto fp8Request = std::make_shared(0, 4, tokens, samplingConfig, false); + fp8Request->setLoraTaskId(1234); + fp8Request->setLoraWeights(fp8Weights); + fp8Request->setLoraConfig(loraConfigTp2); + + mPeftManager->addRequestPeft(fp8Request); + EXPECT_EQ(mPeftManager->getDataType(), tensorrt_llm::DataType::kFP8); + + auto floatRequest = std::make_shared(1, 4, tokens, samplingConfig, false); + floatRequest->setLoraTaskId(5678); + floatRequest->setLoraWeights(loraWeightsTp2); + floatRequest->setLoraConfig(loraConfigTp2); + EXPECT_THAT([&]() { mPeftManager->addRequestPeft(floatRequest); }, + testing::Throws(testing::Property( + &std::runtime_error::what, testing::HasSubstr("PEFT cache supports one homogeneous LoRA dtype")))); +} +#endif + TEST_F(PeftCacheManagerTest, putGet) { SamplingConfig sampleConfig; diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index 2579a8143451..284eee0b0574 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -22,7 +22,7 @@ from ..distributed import (AllReduceParams, HelixAllToAllNative, alltoall_helix, cp_allgather, reducescatter) from ..model_config import ModelConfig -from ..peft.lora.layer import LoraLayer, LoraModuleType +from ..peft.lora.layer import LoraLayer, LoraModuleType, add_lora_result from ..utils import (Fp4QuantizedTensor, get_model_extra_attrs, is_torch_compiling) from .linear import (Linear, TensorParallelMode, WeightMode, @@ -1053,13 +1053,11 @@ def forward( if bool(lora_params): qkv_lora = self.splitted_qkv_lora(hidden_states, lora_params, self.layer_idx) - if qkv_lora is not None: - qkv = qkv + qkv_lora + qkv = add_lora_result(qkv, qkv_lora) qkv_lora = self.fused_qkv_lora(hidden_states, lora_params, self.layer_idx) - if qkv_lora is not None: - qkv = qkv + qkv_lora + qkv = add_lora_result(qkv, qkv_lora) # For dynamic tree spec decoding with Python RoPE, adjust position_ids # to use tree offsets (same as C++ kernel: past_seq_len + offset). diff --git a/tensorrt_llm/_torch/modules/gated_mlp.py b/tensorrt_llm/_torch/modules/gated_mlp.py index a9655ce45f2b..8c49a116bf4c 100644 --- a/tensorrt_llm/_torch/modules/gated_mlp.py +++ b/tensorrt_llm/_torch/modules/gated_mlp.py @@ -10,7 +10,7 @@ from ..distributed import AllReduceParams from ..model_config import ModelConfig -from ..peft.lora.layer import LoraLayer, LoraModuleType +from ..peft.lora.layer import LoraLayer, LoraModuleType, add_lora_result from ..utils import Fp4QuantizedTensor from .linear import (Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig, is_static_nvfp4_input_eligible) @@ -357,12 +357,10 @@ def forward_lora( h1_lora = self.splitted_gate_up_lora(x, lora_params, self.layer_idx) - if h1_lora is not None: - h1 = h1 + h1_lora + h1 = add_lora_result(h1, h1_lora) h1_lora = self.fused_gate_up_lora(x, lora_params, self.layer_idx) - if h1_lora is not None: - h1 = h1 + h1_lora + h1 = add_lora_result(h1, h1_lora) h2 = self._apply_activation(h1, has_lora=True) output = self.down_proj(h2, diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 9524c4db2851..7378b332fab3 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -14,7 +14,7 @@ import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils from tensorrt_llm._torch.custom_ops.torch_custom_ops import BufferKind -from tensorrt_llm._torch.peft.lora.layer import LoraLayer +from tensorrt_llm._torch.peft.lora.layer import LoraLayer, add_lora_result from tensorrt_llm._utils import is_device_integrated, mpi_disabled from tensorrt_llm.bindings import ipc_nvls_supported from tensorrt_llm.functional import (AllReduceFusionOp, AllReduceParams, @@ -3785,8 +3785,7 @@ def apply_linear(self, output = self.quant_method.apply(self, input, bias) if self.lora is not None and bool(lora_params): lora_result = self.lora(input, lora_params, layer_idx) - if lora_result is not None: - output = output + lora_result + output = add_lora_result(output, lora_result) return output def apply_linear_allreduce(self, diff --git a/tensorrt_llm/_torch/modules/mlp.py b/tensorrt_llm/_torch/modules/mlp.py index 42aae9956219..bad1078027b6 100644 --- a/tensorrt_llm/_torch/modules/mlp.py +++ b/tensorrt_llm/_torch/modules/mlp.py @@ -8,7 +8,7 @@ from tensorrt_llm.mapping import Mapping from ..model_config import ModelConfig -from ..peft.lora.layer import LoraLayer, LoraModuleType +from ..peft.lora.layer import LoraLayer, LoraModuleType, add_lora_result from ..utils import Fp4QuantizedTensor, gelu_tanh, relu2 from .linear import (Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig, is_static_nvfp4_input_eligible) @@ -272,8 +272,7 @@ def forward_lora( assert self.layer_idx is not None, "layer_idx is required for lora" x_up_lora = self.up_lora(x, lora_params, self.layer_idx) - if x_up_lora is not None: - x_up = x_up + x_up_lora + x_up = add_lora_result(x_up, x_up_lora) x_act = self.activation(x_up) x_down = self.down_proj(x_act, diff --git a/tensorrt_llm/_torch/peft/lora/cuda_graph_lora_manager.py b/tensorrt_llm/_torch/peft/lora/cuda_graph_lora_manager.py index 9746e6a2c578..f86042ea8468 100644 --- a/tensorrt_llm/_torch/peft/lora/cuda_graph_lora_manager.py +++ b/tensorrt_llm/_torch/peft/lora/cuda_graph_lora_manager.py @@ -1,3 +1,18 @@ +# 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. + from typing import Dict, Optional import torch @@ -209,6 +224,7 @@ def prepare_cuda_graph_lora_params( "prompt_lens_cpu": attn_metadata.prompt_lens_cpu, "num_seqs": attn_metadata.num_seqs, "use_cuda_graph_mode": True, # Flag to indicate new mode + "data_type": peft_cache_manager.data_type, } return lora_params diff --git a/tensorrt_llm/_torch/peft/lora/layer.py b/tensorrt_llm/_torch/peft/lora/layer.py index a70292c16bc5..e3d3be42c3e7 100644 --- a/tensorrt_llm/_torch/peft/lora/layer.py +++ b/tensorrt_llm/_torch/peft/lora/layer.py @@ -1,3 +1,18 @@ +# 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. + from dataclasses import dataclass from enum import IntEnum from typing import Dict, List, Optional @@ -6,6 +21,45 @@ from .cuda_graph_lora_params import CudaGraphLoraParams +_FP8_LORA_TMA_ALIGNMENT = 16 + + +def _validate_fp8_lora_cuda_graph_alignment(slot_ranks_host: torch.Tensor, + hidden_size: int, + output_hidden_sizes: List[int], + max_rank: int) -> int: + if max_rank < _FP8_LORA_TMA_ALIGNMENT or max_rank % _FP8_LORA_TMA_ALIGNMENT != 0: + raise ValueError( + f"FP8 LoRA CUDA graph mode requires max LoRA rank to be a " + f"multiple of {_FP8_LORA_TMA_ALIGNMENT} and at least " + f"{_FP8_LORA_TMA_ALIGNMENT}. Got max rank {max_rank}.") + + active_ranks = slot_ranks_host[slot_ranks_host > 0] + if active_ranks.numel() > 0: + min_active_rank = int(active_ranks.min().item()) + has_misaligned_rank = bool( + torch.any(active_ranks % _FP8_LORA_TMA_ALIGNMENT != 0).item()) + if min_active_rank < _FP8_LORA_TMA_ALIGNMENT or has_misaligned_rank: + raise ValueError( + f"FP8 LoRA CUDA graph mode requires active LoRA ranks " + f"to be multiples of {_FP8_LORA_TMA_ALIGNMENT} and at " + f"least {_FP8_LORA_TMA_ALIGNMENT}. Got active ranks " + f"{active_ranks.tolist()}.") + else: + min_active_rank = max_rank + + fp8_dims = [hidden_size, *output_hidden_sizes] + misaligned_dims = [ + dim for dim in fp8_dims if dim % _FP8_LORA_TMA_ALIGNMENT != 0 + ] + if misaligned_dims: + raise ValueError( + f"FP8 LoRA CUDA graph mode requires hidden and output sizes " + f"to be multiples of {_FP8_LORA_TMA_ALIGNMENT}. Got " + f"{fp8_dims}.") + + return min(hidden_size, min_active_rank) + @dataclass class GroupedGemmParamsOutput: @@ -157,6 +211,13 @@ def is_mamba(self) -> bool: } +def add_lora_result(output: torch.Tensor, + lora_result: Optional[torch.Tensor]) -> torch.Tensor: + if lora_result is None: + return output + return output + lora_result.to(output.dtype) + + class LoraLayer(torch.nn.Module): def __init__(self, lora_module_types: List[LoraModuleType], @@ -174,16 +235,30 @@ def forward( layer_idx: int, ) -> Optional[torch.Tensor]: - if bool(lora_params): - # Check if we're using CUDA Graph mode - use_cuda_graph_mode = lora_params.get('use_cuda_graph_mode', False) + if not bool(lora_params): + return None - if use_cuda_graph_mode: - return self._forward_cuda_graph_mode(x, lora_params, layer_idx) + input_dtype = x.dtype + data_type = lora_params.get("data_type") + if data_type is not None and input_dtype != data_type: + if data_type == torch.float8_e4m3fn and input_dtype in ( + torch.float16, torch.bfloat16, torch.float32): + fp8_max = torch.finfo(data_type).max + x = x.clamp(min=-fp8_max, max=fp8_max).to(data_type) else: - return self._forward_eager_mode(x, lora_params, layer_idx) + raise TypeError( + f"LoRA input dtype {input_dtype} must match PEFT cache dtype " + f"{data_type}.") + + use_cuda_graph_mode = lora_params.get("use_cuda_graph_mode", False) + if use_cuda_graph_mode: + result = self._forward_cuda_graph_mode(x, lora_params, layer_idx) else: - return None + result = self._forward_eager_mode(x, lora_params, layer_idx) + + if isinstance(result, torch.Tensor) and result.dtype != input_dtype: + result = result.to(input_dtype) + return result def prepare_grouped_gemm_buffers(self, input: GroupedGemmParamsInput): device = input.x.device @@ -429,15 +504,20 @@ def _forward_cuda_graph_mode( # Skip layers that don't have LoRA modules if layer_params is None: - return 0 # Pass-through for layers without LoRA modules + return None # Pass-through for layers without LoRA modules batch_size, hidden_size = x.shape[0], x.shape[-1] num_layer_modules = len(self.lora_module_types) max_rank = cuda_graph_params.max_rank total_output_size = sum(self.output_hidden_sizes) - min_kn = min( - hidden_size, 8, max_rank - ) # TODO: hardcode to 8 for now, for alignments in kernels, might have alignment error if rank is less than 8! + if x.dtype == torch.float8_e4m3fn: + min_kn = _validate_fp8_lora_cuda_graph_alignment( + cuda_graph_params.slot_ranks_host, hidden_size, + self.output_hidden_sizes, max_rank) + else: + min_kn = min( + hidden_size, 8, max_rank + ) # TODO: hardcode to 8 for now, for alignments in kernels, might have alignment error if rank is less than 8! output_buffer = torch.empty(batch_size, total_output_size, diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 6bfa53c107a5..b8f375bda8bb 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -6401,8 +6401,11 @@ def _get_lora_params_from_requests( peft_cache_manager) peft_table = peft_cache_manager.get_and_reset_batch_peft_table( ) if peft_cache_manager is not None else None - return peft_table and self._get_eager_lora_params_from_requests( + lora_params = peft_table and self._get_eager_lora_params_from_requests( scheduled_requests, attn_metadata, peft_table) + if lora_params: + lora_params["data_type"] = peft_cache_manager.data_type + return lora_params def _get_eager_lora_params_from_requests( self, scheduled_requests: ScheduledRequests, diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 153fe93e82ae..a46863d76d7d 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -44,7 +44,8 @@ # isort: on from tensorrt_llm.sampling_params import SamplingParams -from ..._utils import binding_to_str_dtype, mpi_rank, nvtx_range +from ..._utils import (binding_to_str_dtype, binding_to_torch_dtype, mpi_rank, + nvtx_range) from ...logger import logger from ...mapping import CpType, Mapping from .connectors.kv_cache_connector import KvCacheConnectorManager @@ -2753,6 +2754,10 @@ def __init__(self, def get_lora_manager(self) -> LoraManager: return self._lora_manager + @property + def data_type(self) -> torch.dtype: + return binding_to_torch_dtype(self.impl.data_type) + def add_request_peft(self, request: LlmRequest): if request.lora_task_id is not None: is_task_cached = self.impl.is_task_cached(request.lora_task_id) diff --git a/tensorrt_llm/lora_manager.py b/tensorrt_llm/lora_manager.py index f703a85769d5..5537213fb0c1 100644 --- a/tensorrt_llm/lora_manager.py +++ b/tensorrt_llm/lora_manager.py @@ -1060,14 +1060,19 @@ def load_from_model_dir(uid, model_dir, hf_config): else: scale = float(hf_config["lora_alpha"]) / effective_rank - # Cast to model dtype before scaling - # fp8 tensors don't support scalar multiply in PyTorch - model_dtype = str_dtype_to_torch(model_config.dtype) - t_in = t_in.to(model_dtype) - t_out = t_out.to(model_dtype) - t_out = t_out * scale - if is_dora and t_mag is not None: - t_mag = t_mag.to(model_dtype) + is_fp8 = t_out.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) + if is_fp8 and torch.cuda.get_device_capability()[0] >= 9: + # Keep weights in FP8 for the native Hopper kernel. + # FP8 has no scalar multiply, so scale through BF16. + t_out = (t_out.to(torch.bfloat16) * scale).to(t_out.dtype) + else: + # Pre-Hopper kernels require the model compute dtype. + model_dtype = str_dtype_to_torch(model_config.dtype) + t_in = t_in.to(model_dtype) + t_out = t_out.to(model_dtype) + t_out = t_out * scale + if is_dora and t_mag is not None: + t_mag = t_mag.to(model_dtype) self._lora_uid_to_low_ranks[uid][layer_idx][lora_module] = effective_rank if self._retain_device_tensors: diff --git a/tests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.py b/tests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.py new file mode 100644 index 000000000000..2619946736d2 --- /dev/null +++ b/tests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.py @@ -0,0 +1,174 @@ +# 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. + +from pathlib import Path + +import pytest +import torch + +from tensorrt_llm._torch.peft.lora.layer import ( + LoraLayer, + _validate_fp8_lora_cuda_graph_alignment, + add_lora_result, +) + +_REPO_ROOT = Path(__file__).resolve().parents[4] + + +def _kernel_source(filename: str) -> str: + return (_REPO_ROOT / "cpp" / "tensorrt_llm" / "kernels" / filename).read_text() + + +def _function_block(source: str, start: str, end: str) -> str: + start_index = source.index(start) + end_index = source.index(end, start_index) + return source[start_index:end_index] + + +def test_fp8_cuda_graph_alignment_accepts_valid_ranks_and_dims(): + min_kn = _validate_fp8_lora_cuda_graph_alignment( + torch.tensor([0, 16, 32], dtype=torch.int32), 64, [128, 256], 32 + ) + + assert min_kn == 16 + + +@pytest.mark.parametrize("use_cuda_graph_mode", [False, True]) +def test_lora_layer_converts_fp8_cache_input_and_restores_output_dtype( + monkeypatch, use_cuda_graph_mode +): + layer = LoraLayer([], []) + forwarded = {} + + def fake_forward(x, _lora_params, _layer_idx): + forwarded["input"] = x + return x + + method_name = "_forward_cuda_graph_mode" if use_cuda_graph_mode else "_forward_eager_mode" + monkeypatch.setattr(layer, method_name, fake_forward) + + x = torch.tensor([[-500.0, -1.0, 1.0, 500.0]], dtype=torch.bfloat16) + result = layer( + x, + { + "data_type": torch.float8_e4m3fn, + "use_cuda_graph_mode": use_cuda_graph_mode, + }, + layer_idx=0, + ) + + fp8_max = torch.finfo(torch.float8_e4m3fn).max + expected = x.clamp(min=-fp8_max, max=fp8_max).to(torch.float8_e4m3fn) + assert forwarded["input"].dtype == torch.float8_e4m3fn + torch.testing.assert_close(forwarded["input"], expected) + assert result.dtype == torch.bfloat16 + torch.testing.assert_close(result, expected.to(torch.bfloat16)) + + +def test_add_lora_result_casts_fp8_delta_to_base_output_dtype(): + output = torch.zeros((1, 4), dtype=torch.bfloat16) + lora_result = torch.ones((1, 4), dtype=torch.float8_e4m3fn) + + result = add_lora_result(output, lora_result) + + assert result.dtype == torch.bfloat16 + torch.testing.assert_close(result, torch.ones_like(output)) + assert add_lora_result(output, None) is output + + +def test_lora_layer_rejects_non_fp8_activation_cache_dtype_mismatch(): + layer = LoraLayer([], []) + + with pytest.raises(TypeError, match="must match PEFT cache dtype"): + layer( + torch.empty((1, 16), dtype=torch.bfloat16), + {"data_type": torch.float16}, + layer_idx=0, + ) + + +@pytest.mark.parametrize( + "slot_ranks,max_rank,match", + [ + ([16, 24], 32, "active LoRA ranks"), + ([16], 24, "max LoRA rank"), + ], +) +def test_fp8_cuda_graph_alignment_rejects_misaligned_ranks(slot_ranks, max_rank, match): + with pytest.raises(ValueError, match=match): + _validate_fp8_lora_cuda_graph_alignment( + torch.tensor(slot_ranks, dtype=torch.int32), 64, [128], max_rank + ) + + +@pytest.mark.parametrize("hidden_size,output_hidden_sizes", [(24, [128]), (64, [128, 72])]) +def test_fp8_cuda_graph_alignment_rejects_misaligned_hidden_dims(hidden_size, output_hidden_sizes): + with pytest.raises(ValueError, match="hidden and output sizes"): + _validate_fp8_lora_cuda_graph_alignment( + torch.tensor([16], dtype=torch.int32), hidden_size, output_hidden_sizes, 16 + ) + + +def test_fp8_cuda_graph_grouped_gemm_uses_live_device_problem_metadata(): + source = _kernel_source("cuda_graph_grouped_gemm.cu") + fp8_graph_body = _function_block( + source, "void fp8CudaGraphGroupedGemm(", "\nvoid cudaGraphGroupedGemm(" + ) + + assert "hostMaxProblemSizesPtr" not in fp8_graph_body + assert "cudaMemcpyHostToDevice" not in fp8_graph_body + assert "fillFp8CudaGraphGroupedGemmParams" in fp8_graph_body + assert "problemSizesPtr, problemCount, ldaGpu, ldbGpu, ldcGpu, lddGpu" in fp8_graph_body + + +@pytest.mark.parametrize( + "filename,messages", + [ + ("groupGemm.cu", ["FP8 grouped GEMM requires CUTLASS modifiable TMA support"]), + ( + "splitkGroupGemm.cu", + ["FP8 split-K grouped GEMM requires CUTLASS modifiable TMA support"], + ), + ( + "cuda_graph_grouped_gemm.cu", + [ + "FP8 CUDA graph grouped GEMM requires CUTLASS modifiable TMA support", + "FP8 CUDA graph split-K grouped GEMM requires CUTLASS modifiable TMA support", + ], + ), + ], +) +def test_fp8_grouped_gemm_dispatch_has_explicit_unsupported_cutlass_guard(filename, messages): + source = _kernel_source(filename) + + assert "#else" in source + for message in messages: + assert message in source + + +@pytest.mark.parametrize("filename", ["groupGemm.cu", "splitkGroupGemm.cu"]) +def test_fp8_grouped_gemm_alignment_checks_require_multiples_of_16(filename): + source = _kernel_source(filename) + + assert "problem.n() % kFp8TmaAlignment == 0" in source + assert "problem.k() % kFp8TmaAlignment == 0" in source + + +def test_fp8_cuda_graph_alignment_check_requires_rank_multiple_of_16(): + source = _kernel_source("cuda_graph_grouped_gemm.cu") + + assert "minKN >= kFp8TmaAlignment && minKN % kFp8TmaAlignment == 0" in source + assert "problem.n() % kFp8TmaAlignment == 0" in source + assert "problem.k() % kFp8TmaAlignment == 0" in source diff --git a/tests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py b/tests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py index a780bafa8a64..9f95da61a2e0 100644 --- a/tests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py +++ b/tests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py @@ -153,7 +153,11 @@ def _run_lora_test(model_path, target_modules, trtllm_modules, dtype=torch.bfloa """End-to-end helper: create adapter, run inference, assert output differs.""" with tempfile.TemporaryDirectory() as tmpdir: lora_dir = _create_lora_adapter( - os.path.join(tmpdir, "lora"), model_path, target_modules, dtype=dtype + os.path.join(tmpdir, "lora"), + model_path, + target_modules, + lora_rank=16 if dtype == torch.float8_e4m3fn else 8, + dtype=dtype, ) lora_config = LoraConfig( lora_dir=[lora_dir], From 0b9722592792cc040018ddad3d0b9e043d05bb15 Mon Sep 17 00:00:00 2001 From: Aurelien Chartier <2567591+achartier@users.noreply.github.com> Date: Mon, 27 Jul 2026 12:16:22 -0700 Subject: [PATCH 2/8] [None][fix] Address dense FP8 LoRA review findings Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com> --- .../batch_manager/peftCacheManager.cpp | 24 +- .../kernels/cuda_graph_grouped_gemm.cu | 11 +- cpp/tensorrt_llm/kernels/groupGemm.cu | 10 +- cpp/tensorrt_llm/kernels/groupGemm.h | 2 + cpp/tensorrt_llm/kernels/splitkGroupGemm.cu | 248 +---------------- cpp/tensorrt_llm/runtime/loraCache.cpp | 7 +- cpp/tensorrt_llm/thop/loraOp.cpp | 6 +- .../batch_manager/peftCacheManagerTest.cpp | 43 +++ tensorrt_llm/lora_manager.py | 63 ++++- .../test_fp8_lora_grouped_gemm_regressions.py | 34 ++- tests/unittest/others/test_lora_manager.py | 261 +++++++++++++++++- 11 files changed, 436 insertions(+), 273 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp index bbbe7895e795..3e51ddc4c518 100644 --- a/cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp @@ -262,14 +262,17 @@ void PeftCacheManager::addRequestPeft(std::shared_ptr llmRequest, bo auto optTaskId = llmRequest->getLoraTaskId(); auto optLoraWeights = llmRequest->getLoraWeights(); auto optLoraConfig = llmRequest->getLoraConfig(); - if (optLoraWeights) - { - configureDataType(optLoraWeights.value()->getDataType()); - } if (optTaskId || optLoraWeights || optLoraConfig) { + auto const requestDataType = optLoraWeights + ? std::optional{optLoraWeights.value()->getDataType()} + : std::nullopt; runtime::lora::loraValidateRequestTensors( - optTaskId, optLoraWeights, optLoraConfig, mModelConfig, mWorldConfig, getDataType()); + optTaskId, optLoraWeights, optLoraConfig, mModelConfig, mWorldConfig, requestDataType); + if (optLoraWeights) + { + configureDataType(optLoraWeights.value()->getDataType()); + } } else { @@ -382,6 +385,17 @@ void PeftCacheManager::addRequestPeft(std::shared_ptr llmRequest, bo void PeftCacheManager::configureDataType(tensorrt_llm::DataType dataType) { + auto const modelDataType = mModelConfig.getDataType(); +#ifdef ENABLE_FP8 + TLLM_CHECK_WITH_INFO(dataType == modelDataType || dataType == tensorrt_llm::DataType::kFP8, + "Unsupported LoRA weights dtype %s for PEFT cache; expected model dtype %s or FP8", + runtime::IBuffer::getDataTypeName(dataType), runtime::IBuffer::getDataTypeName(modelDataType)); +#else + TLLM_CHECK_WITH_INFO(dataType == modelDataType, + "Unsupported LoRA weights dtype %s for PEFT cache; expected model dtype %s", + runtime::IBuffer::getDataTypeName(dataType), runtime::IBuffer::getDataTypeName(modelDataType)); +#endif + std::lock_guard lock(mDataTypeMutex); if (mDataType) { diff --git a/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu b/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu index d826056969da..28f090648795 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 "groupGemm.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" @@ -43,14 +44,18 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { +#ifdef ENABLE_FP8 +#if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) namespace { -constexpr int kFp8TmaAlignment = 16; - void checkFp8CudaGraphAlignment( cutlass::gemm::GemmCoord const* hostMaxProblemSizesPtr, int problemCount, int minKN, char const* kernelName) { + static int const smVersion = tensorrt_llm::common::getSMVersion(); + TLLM_CHECK_WITH_INFO( + smVersion >= 90, "%s requires Hopper (SM90) or newer, 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. " "The minimum K/N dimension is %d.", @@ -72,6 +77,8 @@ void checkFp8CudaGraphAlignment( } } // namespace +#endif // CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED +#endif // ENABLE_FP8 /** * Template for CUDA Graph compatible grouped GEMM that directly uses GPU tensors diff --git a/cpp/tensorrt_llm/kernels/groupGemm.cu b/cpp/tensorrt_llm/kernels/groupGemm.cu index 50334a58547b..9131dc51494e 100644 --- a/cpp/tensorrt_llm/kernels/groupGemm.cu +++ b/cpp/tensorrt_llm/kernels/groupGemm.cu @@ -45,13 +45,17 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { +#ifdef ENABLE_FP8 +#if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) namespace { -constexpr int kFp8TmaAlignment = 16; - void checkFp8GroupedGemmAlignment(std::vector const& problemSizes, char const* kernelName) { + int const smVersion = tensorrt_llm::common::getSMVersion(); + TLLM_CHECK_WITH_INFO( + smVersion >= 90, "%s requires Hopper (SM90) or newer, but the current device is SM%d", kernelName, smVersion); + for (size_t problemIdx = 0; problemIdx < problemSizes.size(); ++problemIdx) { auto const& problem = problemSizes[problemIdx]; @@ -68,6 +72,8 @@ void checkFp8GroupedGemmAlignment(std::vector const& p } } // namespace +#endif // CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED +#endif // ENABLE_FP8 int64_t inline getGemmCoordSize(int64_t problemCount) { diff --git a/cpp/tensorrt_llm/kernels/groupGemm.h b/cpp/tensorrt_llm/kernels/groupGemm.h index 5cf3a8d3e6be..48b38eb82d65 100644 --- a/cpp/tensorrt_llm/kernels/groupGemm.h +++ b/cpp/tensorrt_llm/kernels/groupGemm.h @@ -24,6 +24,8 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { +inline constexpr int kFp8TmaAlignment = 16; + int64_t getGroupedGemmParamsWorkSpaceSize(int64_t problem_count); //! @brief Returns the parameter workspace size for the FP8 grouped GEMM path. diff --git a/cpp/tensorrt_llm/kernels/splitkGroupGemm.cu b/cpp/tensorrt_llm/kernels/splitkGroupGemm.cu index a4058c052d30..d1aa210665af 100644 --- a/cpp/tensorrt_llm/kernels/splitkGroupGemm.cu +++ b/cpp/tensorrt_llm/kernels/splitkGroupGemm.cu @@ -15,6 +15,7 @@ * limitations under the License. */ +#include "groupGemm.h" #include "splitkGroupGemm.h" #include "cutlass/cutlass.h" @@ -31,45 +32,10 @@ #include "tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/default_splitk_gemm_grouped.h" #include "tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/splitk_gemm_grouped.h" -#ifdef ENABLE_FP8 -#include "cute/tensor.hpp" -#include "cutlass/epilogue/collective/collective_builder.hpp" -#include "cutlass/gemm/collective/collective_builder.hpp" -#include "cutlass/gemm/device/gemm_universal_adapter.h" -#include "cutlass/gemm/dispatch_policy.hpp" -#include "cutlass/gemm/group_array_problem_shape.hpp" -#include "cutlass/gemm/kernel/gemm_universal.hpp" -#include "cutlass/util/packed_stride.hpp" -#endif // ENABLE_FP8 - TRTLLM_NAMESPACE_BEGIN namespace kernels { -namespace -{ - -constexpr int kFp8TmaAlignment = 16; - -void checkFp8GroupedGemmAlignment(std::vector const& problemSizes, char const* kernelName) -{ - for (size_t problemIdx = 0; problemIdx < problemSizes.size(); ++problemIdx) - { - auto const& problem = problemSizes[problemIdx]; - if (problem.m() == 0 || problem.n() == 0 || problem.k() == 0) - { - continue; - } - - TLLM_CHECK_WITH_INFO(problem.n() % kFp8TmaAlignment == 0 && problem.k() % kFp8TmaAlignment == 0, - "%s requires GEMM N and K dimensions to be multiples of %d elements for 128-bit TMA alignment. " - "Problem %zu has M=%d, N=%d, K=%d.", - kernelName, kFp8TmaAlignment, problemIdx, problem.m(), problem.n(), problem.k()); - } -} - -} // namespace - int64_t inline getGemmCoordSize(int64_t problemCount) { return (int64_t) (tensorrt_llm::common::divUp(problemCount * sizeof(cutlass::gemm::GemmCoord), 16) * 16); @@ -260,202 +226,6 @@ void splitkGroupedGemmType_(std::vector const& problem #endif } -#ifdef ENABLE_FP8 - -// ==================================================================== -// FP8 split-K grouped GEMM using CUTLASS 3.x collective API (Hopper). -// -// CUTLASS 3.x grouped GEMM does not have a split-K variant. We use the -// regular cooperative fp8 grouped GEMM instead: the Hopper cooperative -// schedule distributes work across warps efficiently even for large-K -// small-N problems. The output is fp8, matching the non-split-K path's -// final output type convention used by LoRA. -// ==================================================================== - -#if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) - -void fp8SplitkGroupedGemm(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) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - using namespace cute; - - 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 ElementAccumulator = float; - - using LayoutA = cutlass::layout::RowMajor; - using LayoutB = cutlass::layout::ColumnMajor; - using LayoutC = cutlass::layout::RowMajor; - using LayoutD = cutlass::layout::RowMajor; - - static constexpr int kAlignmentA = 128 / cutlass::sizeof_bits::value; - static constexpr int kAlignmentB = 128 / cutlass::sizeof_bits::value; - 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 CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder( - sizeof(typename CollectiveEpilogue::SharedStorage))>, - KernelSchedule>::CollectiveOp; - - using GemmKernel = cutlass::gemm::kernel::GemmUniversal; - using Gemm = cutlass::gemm::device::GemmUniversalAdapter; - - using StrideA = typename GemmKernel::InternalStrideA; - using StrideB = typename GemmKernel::InternalStrideB; - using StrideC = typename GemmKernel::InternalStrideC; - using StrideD = typename GemmKernel::InternalStrideD; - - int const problemCount = static_cast(problemSizes.size()); - - std::vector problemShapesHost(problemCount); - std::vector ptrAHost(problemCount); - std::vector ptrBHost(problemCount); - std::vector ptrCHost(problemCount); - std::vector ptrDHost(problemCount); - std::vector strideAHost(problemCount); - std::vector strideBHost(problemCount); - std::vector strideCHost(problemCount); - std::vector strideDHost(problemCount); - - for (int i = 0; i < problemCount; ++i) - { - auto const& coord = problemSizes[i]; - int const m = coord.m(); - int const n = coord.n(); - int const k = coord.k(); - - problemShapesHost[i] = cute::make_shape(m, n, k); - ptrAHost[i] = static_cast(ptrA[i]); - ptrBHost[i] = static_cast(ptrB[i]); - ptrCHost[i] = static_cast(ptrC[i]); - ptrDHost[i] = static_cast(ptrD[i]); - - strideAHost[i] = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(m, k, 1)); - strideBHost[i] = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(n, k, 1)); - strideCHost[i] = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, n, 1)); - strideDHost[i] = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(m, n, 1)); - } - - auto const align16 = [](int64_t bytes) -> int64_t { return (bytes + 15) / 16 * 16; }; - - int64_t const szProblemShapes = align16(problemCount * sizeof(typename ProblemShape::UnderlyingProblemShape)); - int64_t const szPtrA = align16(problemCount * sizeof(ElementA const*)); - int64_t const szPtrB = align16(problemCount * sizeof(ElementB const*)); - int64_t const szPtrC = align16(problemCount * sizeof(ElementC const*)); - int64_t const szPtrD = align16(problemCount * sizeof(ElementD*)); - int64_t const szStrideA = align16(problemCount * sizeof(StrideA)); - int64_t const szStrideB = align16(problemCount * sizeof(StrideB)); - int64_t const szStrideC = align16(problemCount * sizeof(StrideC)); - int64_t const szStrideD = align16(problemCount * sizeof(StrideD)); - - int64_t const totalParamsBytes - = szProblemShapes + szPtrA + szPtrB + szPtrC + szPtrD + szStrideA + szStrideB + szStrideC + szStrideD; - TLLM_CHECK_WITH_INFO(totalParamsBytes <= gemmParamsWorkSpaceSize, - "FP8 split-K grouped GEMM params workspace too small: need %ld, have %ld", totalParamsBytes, - gemmParamsWorkSpaceSize); - - std::vector hostBuf(totalParamsBytes); - char* cursor = hostBuf.data(); - - auto copyToHostBuf = [&cursor](void const* src, int64_t bytes) - { - std::memcpy(cursor, src, bytes); - cursor += (bytes + 15) / 16 * 16; - }; - - copyToHostBuf(problemShapesHost.data(), problemCount * sizeof(typename ProblemShape::UnderlyingProblemShape)); - copyToHostBuf(ptrAHost.data(), problemCount * sizeof(ElementA const*)); - copyToHostBuf(ptrBHost.data(), problemCount * sizeof(ElementB const*)); - copyToHostBuf(ptrCHost.data(), problemCount * sizeof(ElementC const*)); - copyToHostBuf(ptrDHost.data(), problemCount * sizeof(ElementD*)); - copyToHostBuf(strideAHost.data(), problemCount * sizeof(StrideA)); - copyToHostBuf(strideBHost.data(), problemCount * sizeof(StrideB)); - copyToHostBuf(strideCHost.data(), problemCount * sizeof(StrideC)); - copyToHostBuf(strideDHost.data(), problemCount * sizeof(StrideD)); - - tensorrt_llm::common::cudaAutoCpy(reinterpret_cast(gemmParamsWorkSpace), - reinterpret_cast(hostBuf.data()), totalParamsBytes, stream); - - char* devBase = static_cast(gemmParamsWorkSpace); - int64_t devOff = 0; - - auto devPtr = [&devBase, &devOff](int64_t sz) -> void* - { - void* p = devBase + devOff; - devOff += (sz + 15) / 16 * 16; - return p; - }; - - auto* devProblemShapes = static_cast(devPtr(szProblemShapes)); - auto* devPtrA = static_cast(devPtr(szPtrA)); - auto* devPtrB = static_cast(devPtr(szPtrB)); - auto* devPtrC = static_cast(devPtr(szPtrC)); - auto* devPtrD = static_cast(devPtr(szPtrD)); - auto* devStrideA = static_cast(devPtr(szStrideA)); - auto* devStrideB = static_cast(devPtr(szStrideB)); - auto* devStrideC = static_cast(devPtr(szStrideC)); - auto* devStrideD = static_cast(devPtr(szStrideD)); - - cutlass::KernelHardwareInfo hwInfo; - hwInfo.device_id = 0; - cudaGetDevice(&hwInfo.device_id); - hwInfo.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hwInfo.device_id); - - typename Gemm::Arguments arguments{cutlass::gemm::GemmUniversalMode::kGrouped, - {problemCount, devProblemShapes, nullptr}, {devPtrA, devStrideA, devPtrB, devStrideB}, - {{1.0f, 0.0f}, devPtrC, devStrideC, devPtrD, devStrideD}, hwInfo}; - - Gemm gemm; - - size_t const workspaceSize = Gemm::get_workspace_size(arguments); - TLLM_CHECK_WITH_INFO(static_cast(workspaceSize) <= gemmWorkSpaceSize, - "FP8 split-K grouped GEMM workspace too small: need %lu, have %ld", workspaceSize, gemmWorkSpaceSize); - - cutlass::Status status = gemm.can_implement(arguments); - TLLM_CHECK_WITH_INFO(status == cutlass::Status::kSuccess, "FP8 split-K grouped GEMM can_implement failed: %s", - cutlass::cutlassGetStatusString(status)); - - status = gemm.initialize(arguments, gemmWorkSpace, stream); - TLLM_CHECK_WITH_INFO(status == cutlass::Status::kSuccess, "FP8 split-K grouped GEMM initialize failed: %s", - cutlass::cutlassGetStatusString(status)); - - status = gemm.run(stream); - TLLM_CHECK_WITH_INFO(status == cutlass::Status::kSuccess, "FP8 split-K grouped GEMM run failed: %s", - cutlass::cutlassGetStatusString(status)); - - sync_check_cuda_error(stream); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -#endif // CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED - -#endif // ENABLE_FP8 - void splitkGroupedGemm(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, @@ -464,22 +234,14 @@ void splitkGroupedGemm(std::vector const& problemSizes 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 (dataType == tensorrt_llm::DataType::kFP8) { - checkFp8GroupedGemmAlignment(problemSizes, "FP8 split-K grouped GEMM"); - - fp8SplitkGroupedGemm(problemSizes, ptrA, ptrB, ptrC, ptrD, gemmParamsWorkSpace, gemmParamsWorkSpaceSize, - gemmWorkSpace, gemmWorkSpaceSize, stream); + // CUTLASS 3.x has no split-K grouped FP8 kernel. Reuse the regular + // cooperative implementation, which is also the intended split-K fallback. + groupedGemm(problemSizes, ptrA, ptrB, ptrC, ptrD, gemmParamsWorkSpace, gemmParamsWorkSpaceSize, gemmWorkSpace, + gemmWorkSpaceSize, isLoraIn, dataType, minKN, stream); return; } -#else - if (dataType == tensorrt_llm::DataType::kFP8) - { - TLLM_CHECK_WITH_INFO(false, - "FP8 split-K grouped GEMM requires CUTLASS modifiable TMA support (CUDA 12.3+ and Hopper sm90+ kernels)."); - } -#endif // CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED #endif // ENABLE_FP8 if (isLoraIn) diff --git a/cpp/tensorrt_llm/runtime/loraCache.cpp b/cpp/tensorrt_llm/runtime/loraCache.cpp index 53cd23793b06..6749377445b2 100644 --- a/cpp/tensorrt_llm/runtime/loraCache.cpp +++ b/cpp/tensorrt_llm/runtime/loraCache.cpp @@ -800,9 +800,6 @@ void LoraCache::copyTask(TaskIdType taskId, LoraCache& deviceCache, bool markDon TLLM_CHECK_WITH_INFO(deviceCache.mPageManagerConfig.getMemoryType() == runtime::MemoryType::kGPU && !deviceCache.mDeviceBufferManagers.empty(), "The deviceCache must hold GPU memory and have at least one bufferManager / copy stream"); - TLLM_CHECK_WITH_INFO(mPageManagerConfig.getDataType() == deviceCache.mPageManagerConfig.getDataType(), - "LoRA host and device cache dtypes must match"); - // First get the taskValue from this cache // TaskValue& taskValue = copyTaskGetThisTaskValue(taskId); TaskValuePtr taskValue = [&]() -> TaskValuePtr @@ -832,6 +829,10 @@ void LoraCache::copyTask(TaskIdType taskId, LoraCache& deviceCache, bool markDon std::optional optOtherTaskValuePtr = [&]() -> std::optional { std::lock_guard deviceCacheLock(deviceCache.mCacheMutex); + // setDataType also holds mCacheMutex, so the check and target task + // insertion are atomic with respect to cache reconfiguration. + TLLM_CHECK_WITH_INFO(mPageManagerConfig.getDataType() == deviceCache.mPageManagerConfig.getDataType(), + "LoRA host and device cache dtypes must match"); auto otherStatus = deviceCache.getStatus(taskId); if (kVALUE_STATUS_MISSING != otherStatus) { diff --git a/cpp/tensorrt_llm/thop/loraOp.cpp b/cpp/tensorrt_llm/thop/loraOp.cpp index 604e2f2cc80a..51575a5873ed 100644 --- a/cpp/tensorrt_llm/thop/loraOp.cpp +++ b/cpp/tensorrt_llm/thop/loraOp.cpp @@ -234,8 +234,7 @@ void lora_grouped_gemm_cuda_graph(th::Tensor const& lora_in_sizes, // [layer_mod case torch::kFloat8_e4m3fn: loraRuntimeDataType = tensorrt_llm::DataType::kFP8; break; #endif default: - TORCH_CHECK( - false, "Invalid dtype, only supports float16, bfloat16, float8_e4m3fn, got %s", c10::toString(dtype)); + TORCH_CHECK(false, "Invalid dtype, only supports float16, bfloat16, float8_e4m3fn, got ", c10::toString(dtype)); } int const minKnInt = std::max(1, static_cast(minKN)); @@ -319,8 +318,7 @@ void lora_group_gemm_param_fill_row_reorder_fusion(th::Tensor const& in_sizes, / case torch::kFloat8_e4m3fn: loraRuntimeDataType = tensorrt_llm::DataType::kFP8; break; #endif default: - TORCH_CHECK( - false, "Invalid dtype, only supports float16, bfloat16, float8_e4m3fn, got %s", c10::toString(dtype)); + TORCH_CHECK(false, "Invalid dtype, only supports float16, bfloat16, float8_e4m3fn, got ", c10::toString(dtype)); } int64_t const dtype_element_size = input.element_size(); diff --git a/cpp/tests/unit_tests/batch_manager/peftCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/peftCacheManagerTest.cpp index a52f7f461179..6e4ded72ae26 100644 --- a/cpp/tests/unit_tests/batch_manager/peftCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/peftCacheManagerTest.cpp @@ -143,7 +143,50 @@ TEST_F(PeftCacheManagerTest, addRequestPeftMissingTask) mPeftManager->addRequestPeft(llmRequest); } +TEST_F(PeftCacheManagerTest, unsupportedAdapterDataTypeDoesNotConfigureCache) +{ + SamplingConfig samplingConfig; + auto tokens = std::make_shared>(std::initializer_list{1, 2, 3, 4}); + TensorPtr int8Weights = mManager->cpu(loraWeightsTp2->getShape(), tensorrt_llm::DataType::kINT8); + auto int8Request = std::make_shared(0, 4, tokens, samplingConfig, false); + int8Request->setLoraTaskId(1234); + int8Request->setLoraWeights(int8Weights); + int8Request->setLoraConfig(loraConfigTp2); + + EXPECT_THAT([&]() { mPeftManager->addRequestPeft(int8Request); }, + testing::Throws(testing::Property( + &std::runtime_error::what, testing::HasSubstr("Unsupported LoRA weights dtype int8 for PEFT cache")))); + + auto floatRequest = std::make_shared(1, 4, tokens, samplingConfig, false); + floatRequest->setLoraTaskId(5678); + floatRequest->setLoraWeights(loraWeightsTp2); + floatRequest->setLoraConfig(loraConfigTp2); + EXPECT_NO_THROW(mPeftManager->addRequestPeft(floatRequest)); + EXPECT_EQ(mPeftManager->getDataType(), tensorrt_llm::DataType::kFLOAT); +} + #if defined(ENABLE_FP8) +TEST_F(PeftCacheManagerTest, invalidAdapterDoesNotConfigureCacheDataType) +{ + SamplingConfig samplingConfig; + auto tokens = std::make_shared>(std::initializer_list{1, 2, 3, 4}); + TensorPtr fp8Weights = mManager->cpu(loraWeightsTp2->getShape(), tensorrt_llm::DataType::kFP8); + auto invalidRequest = std::make_shared(0, 4, tokens, samplingConfig, false); + invalidRequest->setLoraTaskId(1234); + invalidRequest->setLoraWeights(fp8Weights); + + EXPECT_THAT([&]() { mPeftManager->addRequestPeft(invalidRequest); }, + testing::Throws(testing::Property( + &std::runtime_error::what, testing::HasSubstr("must have both lora_weights and lora_keys")))); + + auto floatRequest = std::make_shared(1, 4, tokens, samplingConfig, false); + floatRequest->setLoraTaskId(5678); + floatRequest->setLoraWeights(loraWeightsTp2); + floatRequest->setLoraConfig(loraConfigTp2); + EXPECT_NO_THROW(mPeftManager->addRequestPeft(floatRequest)); + EXPECT_EQ(mPeftManager->getDataType(), tensorrt_llm::DataType::kFLOAT); +} + TEST_F(PeftCacheManagerTest, adapterSelectsHomogeneousCacheDataType) { SamplingConfig samplingConfig; diff --git a/tensorrt_llm/lora_manager.py b/tensorrt_llm/lora_manager.py index 5537213fb0c1..dbaa892fd6ba 100644 --- a/tensorrt_llm/lora_manager.py +++ b/tensorrt_llm/lora_manager.py @@ -30,6 +30,7 @@ from .runtime import ModelConfig NEMO_SUPPORTED_LORA_MODULES = {"attn_qkv"} +_FP8_LORA_TMA_ALIGNMENT = 16 logger = logging.getLogger(__name__) @@ -64,6 +65,34 @@ def _is_moe_module_weights(module_weights: Dict) -> bool: ) +def _validate_fp8_lora_alignment( + *, + rank: int, + input_size: int, + output_size: int, + layer_idx: int, + lora_module: str, +) -> None: + dimensions = { + "rank": rank, + "input size": input_size, + "output size": output_size, + } + misaligned_dimensions = { + name: size for name, size in dimensions.items() if size % _FP8_LORA_TMA_ALIGNMENT != 0 + } + if misaligned_dimensions: + formatted_dimensions = ", ".join( + 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"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." + ) + + def get_all_nemo_lora_weights( lora_weights: Dict[str, torch.Tensor], ) -> Dict[int, Dict[str, torch.Tensor]]: @@ -1049,6 +1078,30 @@ def load_from_model_dir(uid, model_dir, hf_config): t_out = prepare_fused_lora_modules_for_tp(lora_module, t_out, rank_dim) effective_rank = t_in.shape[rank_dim] + is_fp8 = t_out.dtype == torch.float8_e4m3fn + use_fp8_kernel = ( + is_fp8 + and not has_expert_indices + and torch.cuda.get_device_capability()[0] >= 9 + ) + if use_fp8_kernel: + if t_in.dtype != t_out.dtype: + raise ValueError( + "FP8 LoRA input and output weights must have the same dtype; " + f"got {t_in.dtype} and {t_out.dtype} for layer {layer_idx} " + f"module {lora_module}" + ) + _validate_fp8_lora_alignment( + rank=effective_rank, + input_size=t_in.shape[-1], + output_size=t_out.shape[-2], + layer_idx=layer_idx, + lora_module=lora_module, + ) + if is_dora: + raise NotImplementedError( + "DoRA is not supported with FP8 LoRA weights on Hopper" + ) t_in = t_in.cuda().contiguous() t_out = t_out.cuda().contiguous() @@ -1060,11 +1113,15 @@ def load_from_model_dir(uid, model_dir, hf_config): else: scale = float(hf_config["lora_alpha"]) / effective_rank - is_fp8 = t_out.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) - if is_fp8 and torch.cuda.get_device_capability()[0] >= 9: + if use_fp8_kernel: # Keep weights in FP8 for the native Hopper kernel. # FP8 has no scalar multiply, so scale through BF16. - t_out = (t_out.to(torch.bfloat16) * scale).to(t_out.dtype) + fp8_max = torch.finfo(t_out.dtype).max + t_out = ( + (t_out.to(torch.bfloat16) * scale) + .clamp(-fp8_max, fp8_max) + .to(t_out.dtype) + ) else: # Pre-Hopper kernels require the model compute dtype. model_dtype = str_dtype_to_torch(model_config.dtype) diff --git a/tests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.py b/tests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.py index 2619946736d2..4eb8eccd1a99 100644 --- a/tests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.py +++ b/tests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.py @@ -137,10 +137,6 @@ def test_fp8_cuda_graph_grouped_gemm_uses_live_device_problem_metadata(): "filename,messages", [ ("groupGemm.cu", ["FP8 grouped GEMM requires CUTLASS modifiable TMA support"]), - ( - "splitkGroupGemm.cu", - ["FP8 split-K grouped GEMM requires CUTLASS modifiable TMA support"], - ), ( "cuda_graph_grouped_gemm.cu", [ @@ -158,14 +154,40 @@ def test_fp8_grouped_gemm_dispatch_has_explicit_unsupported_cutlass_guard(filena assert message in source -@pytest.mark.parametrize("filename", ["groupGemm.cu", "splitkGroupGemm.cu"]) -def test_fp8_grouped_gemm_alignment_checks_require_multiples_of_16(filename): +@pytest.mark.parametrize("filename", ["groupGemm.cu", "cuda_graph_grouped_gemm.cu"]) +def test_fp8_grouped_gemm_dispatch_requires_sm90_or_newer(filename): source = _kernel_source(filename) + assert "getSMVersion()" in source + assert "smVersion >= 90" in source + assert "requires Hopper (SM90) or newer" in source + + +def test_fp8_grouped_gemm_alignment_checks_require_multiples_of_16(): + source = _kernel_source("groupGemm.cu") + assert "problem.n() % kFp8TmaAlignment == 0" in source assert "problem.k() % kFp8TmaAlignment == 0" in source +def test_fp8_splitk_grouped_gemm_delegates_to_regular_grouped_gemm(): + source = _kernel_source("splitkGroupGemm.cu") + + assert "fp8SplitkGroupedGemm" not in source + assert "groupedGemm(problemSizes" in source + + +def test_fp8_tma_alignment_has_one_cpp_definition(): + kernels_dir = _REPO_ROOT / "cpp" / "tensorrt_llm" / "kernels" + definitions = sorted( + path + for path in kernels_dir.iterdir() + if path.suffix in {".cu", ".h"} and "kFp8TmaAlignment = 16" in path.read_text() + ) + + assert definitions == [kernels_dir / "groupGemm.h"] + + def test_fp8_cuda_graph_alignment_check_requires_rank_multiple_of_16(): source = _kernel_source("cuda_graph_grouped_gemm.cu") diff --git a/tests/unittest/others/test_lora_manager.py b/tests/unittest/others/test_lora_manager.py index ff36aeebda86..57304c47d120 100644 --- a/tests/unittest/others/test_lora_manager.py +++ b/tests/unittest/others/test_lora_manager.py @@ -24,7 +24,7 @@ import unittest from dataclasses import dataclass, field from pathlib import Path -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import torch from safetensors.torch import save_file @@ -51,26 +51,81 @@ class MockModelConfig: def _create_dummy_hf_lora_adapter( - adapter_dir: Path, hidden_size: int = 64, rank: int = 8, num_layers: int = 2 + adapter_dir: Path, + hidden_size: int = 64, + output_size: int | None = None, + rank: int = 8, + num_layers: int = 2, + dtype: torch.dtype = torch.float16, + use_dora: bool = False, + input_dtype: torch.dtype | None = None, + lora_alpha: float | None = None, + weight_value: float | None = None, ): """Create a minimal HF-format LoRA adapter on disk.""" + output_size = hidden_size if output_size is None else output_size config = { "r": rank, - "lora_alpha": rank, + "lora_alpha": rank if lora_alpha is None else lora_alpha, "target_modules": ["q_proj", "k_proj", "v_proj"], "bias": "none", "peft_type": "LORA", "task_type": "CAUSAL_LM", + "use_dora": use_dora, } with open(adapter_dir / "adapter_config.json", "w") as f: json.dump(config, f) weights = {} + input_dtype = dtype if input_dtype is None else input_dtype + + def make_weight(shape, target_dtype): + if weight_value is None: + return torch.randn(*shape, dtype=torch.float16).to(target_dtype) + return torch.full(shape, weight_value, dtype=torch.float16).to(target_dtype) + for layer_idx in range(num_layers): for module in ["q_proj", "k_proj", "v_proj"]: prefix = f"base_model.model.model.layers.{layer_idx}.self_attn.{module}" - weights[f"{prefix}.lora_A.weight"] = torch.randn(rank, hidden_size, dtype=torch.float16) - weights[f"{prefix}.lora_B.weight"] = torch.randn(hidden_size, rank, dtype=torch.float16) + weights[f"{prefix}.lora_A.weight"] = make_weight((rank, hidden_size), input_dtype) + weights[f"{prefix}.lora_B.weight"] = make_weight((output_size, rank), dtype) + if use_dora: + weights[f"{prefix}.lora_magnitude_vector"] = torch.ones( + output_size, dtype=torch.bfloat16 + ) + + save_file(weights, str(adapter_dir / "adapter_model.safetensors")) + + +def _create_dummy_hf_moe_lora_adapter( + adapter_dir: Path, + hidden_size: int = 64, + output_size: int = 128, + rank: int = 16, + num_experts: int = 2, + dtype: torch.dtype = torch.float8_e4m3fn, +): + """Create a minimal expert-indexed HF-format LoRA adapter on disk.""" + config = { + "r": rank, + "lora_alpha": rank, + "target_modules": ["gate_proj"], + "bias": "none", + "peft_type": "LORA", + "task_type": "CAUSAL_LM", + } + with open(adapter_dir / "adapter_config.json", "w") as f: + json.dump(config, f) + + weights = {} + for expert_idx in range(num_experts): + prefix = f"base_model.model.model.layers.0.mlp.experts.{expert_idx}.gate_proj" + weights[f"{prefix}.lora_A.weight"] = torch.randn(rank, hidden_size, dtype=torch.float16).to( + dtype + ) + weights[f"{prefix}.lora_B.weight"] = torch.randn(output_size, rank, dtype=torch.float16).to( + dtype + ) save_file(weights, str(adapter_dir / "adapter_model.safetensors")) @@ -162,5 +217,201 @@ def test_many_adapters_no_gpu_accumulation(self): self.assertEqual(len(manager._cpp_lora_weights), num_adapters) +@unittest.skipUnless( + torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 9, + "FP8 LoRA requires Hopper or newer", +) +class TestLoraManagerFp8(unittest.TestCase): + def test_fp8_dora_is_rejected(self): + model_config = MockModelConfig() + 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, + use_dora=True, + ) + + with self.assertRaisesRegex( + NotImplementedError, + "DoRA is not supported with FP8 LoRA weights on Hopper", + ): + manager.load_from_hf( + model_dirs=[str(adapter_dir)], + model_config=model_config, + uids=["fp8-dora"], + ) + + def test_fp8_moe_weights_are_converted_to_model_dtype(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) / "adapter" + adapter_dir.mkdir() + _create_dummy_hf_moe_lora_adapter(adapter_dir) + + manager.load_from_hf( + model_dirs=[str(adapter_dir)], + model_config=model_config, + uids=["fp8-moe"], + ) + + self.assertEqual(manager.cpp_lora_weights["fp8-moe"].dtype, torch.bfloat16) + + def test_fp8_e5m2_weights_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=8, + num_layers=1, + dtype=torch.float8_e5m2, + ) + + manager.load_from_hf( + model_dirs=[str(adapter_dir)], + model_config=model_config, + uids=["fp8-e5m2"], + ) + + self.assertEqual(manager.cpp_lora_weights["fp8-e5m2"].dtype, torch.bfloat16) + + def test_fp8_e4m3_input_output_dtype_mismatch_is_rejected(self): + model_config = MockModelConfig() + 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, + input_dtype=torch.float8_e5m2, + ) + + with self.assertRaisesRegex( + ValueError, "FP8 LoRA input and output weights must have the same dtype" + ): + manager.load_from_hf( + model_dirs=[str(adapter_dir)], + model_config=model_config, + uids=["fp8-mismatched-dtype"], + ) + + def test_scaled_fp8_e4m3_weights_are_clamped_before_cast(self): + model_config = MockModelConfig() + 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, + lora_alpha=32, + weight_value=torch.finfo(torch.float8_e4m3fn).max, + ) + + manager.load_from_hf( + model_dirs=[str(adapter_dir)], + model_config=model_config, + uids=["scaled-fp8"], + ) + + weights = manager.cpp_lora_weights["scaled-fp8"].float() + self.assertTrue(torch.isfinite(weights).all()) + self.assertEqual(weights.abs().max(), torch.finfo(torch.float8_e4m3fn).max) + + +class TestLoraManagerFp8Alignment(unittest.TestCase): + @staticmethod + def _create_manager(model_config): + return LoraManager( + mapping=Mapping(world_size=1, rank=0, tp_size=1), + model_config=model_config, + cpp_peft_cache_manager=MagicMock(), + ) + + def test_misaligned_fp8_adapter_is_rejected_before_cuda_transfer(self): + cases = [ + {"rank": 8, "hidden_size": 64, "output_size": 64, "match": "rank=8"}, + { + "rank": 16, + "hidden_size": 72, + "output_size": 64, + "match": "input size=72", + }, + { + "rank": 16, + "hidden_size": 64, + "output_size": 72, + "match": "output size=72", + }, + ] + + for case in cases: + with self.subTest(case=case), tempfile.TemporaryDirectory() as tmpdir: + adapter_dir = Path(tmpdir) / "adapter" + adapter_dir.mkdir() + _create_dummy_hf_lora_adapter( + adapter_dir, + hidden_size=case["hidden_size"], + output_size=case["output_size"], + rank=case["rank"], + num_layers=1, + dtype=torch.float8_e4m3fn, + ) + model_config = MockModelConfig(hidden_size=case["hidden_size"]) + manager = self._create_manager(model_config) + + with ( + patch("torch.cuda.get_device_capability", return_value=(9, 0)), + self.assertRaisesRegex(ValueError, case["match"]), + ): + manager.load_from_hf( + model_dirs=[str(adapter_dir)], + model_config=model_config, + uids=["misaligned-fp8"], + ) + + if __name__ == "__main__": unittest.main() From ac5e0d1cc418bc65f1e81e9e67b4807aed0463cf Mon Sep 17 00:00:00 2001 From: Aurelien Chartier <2567591+achartier@users.noreply.github.com> Date: Wed, 29 Jul 2026 13:23:03 -0700 Subject: [PATCH 3/8] [None][fix] Restrict dense FP8 LoRA to SM90 Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com> --- .../kernels/cuda_graph_grouped_gemm.cu | 7 +++-- cpp/tensorrt_llm/kernels/groupGemm.cu | 7 +++-- tensorrt_llm/lora_manager.py | 3 +- .../test_fp8_lora_grouped_gemm_regressions.py | 7 +++-- tests/unittest/others/test_lora_manager.py | 31 +++++++++++++++++-- 5 files changed, 43 insertions(+), 12 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu b/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu index 28f090648795..043b0a980c8f 100644 --- a/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu +++ b/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu @@ -53,8 +53,9 @@ void checkFp8CudaGraphAlignment( cutlass::gemm::GemmCoord const* hostMaxProblemSizesPtr, int problemCount, int minKN, char const* kernelName) { 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) or newer, but the current device is SM%d", kernelName, smVersion); + smVersion == 90, "%s requires Hopper (SM90), 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. " @@ -387,7 +388,7 @@ 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 Hopper SM90 " "kernels)."); } #endif // CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED @@ -575,7 +576,7 @@ 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 Hopper SM90 " "kernels)."); } #endif // CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED diff --git a/cpp/tensorrt_llm/kernels/groupGemm.cu b/cpp/tensorrt_llm/kernels/groupGemm.cu index 9131dc51494e..87e5b94c6645 100644 --- a/cpp/tensorrt_llm/kernels/groupGemm.cu +++ b/cpp/tensorrt_llm/kernels/groupGemm.cu @@ -53,8 +53,9 @@ namespace void checkFp8GroupedGemmAlignment(std::vector const& problemSizes, char const* kernelName) { 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) or newer, but the current device is SM%d", kernelName, smVersion); + smVersion == 90, "%s requires Hopper (SM90), but the current device is SM%d", kernelName, smVersion); for (size_t problemIdx = 0; problemIdx < problemSizes.size(); ++problemIdx) { @@ -260,7 +261,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 CUTLASS 3.x collective API (Hopper SM90). // // The legacy CUTLASS 2.x DefaultGemmGrouped does NOT support fp8 element // types. This implementation uses the CUTLASS 3.x CollectiveBuilder and @@ -492,7 +493,7 @@ 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 Hopper SM90 kernels)."); } #endif // CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED #endif // ENABLE_FP8 diff --git a/tensorrt_llm/lora_manager.py b/tensorrt_llm/lora_manager.py index dbaa892fd6ba..b6fd4338c64a 100644 --- a/tensorrt_llm/lora_manager.py +++ b/tensorrt_llm/lora_manager.py @@ -1079,10 +1079,11 @@ def load_from_model_dir(uid, model_dir, hf_config): effective_rank = t_in.shape[rank_dim] is_fp8 = t_out.dtype == torch.float8_e4m3fn + # TODO: Enable SM120/SM121 after validating the native FP8 LoRA kernel there. use_fp8_kernel = ( is_fp8 and not has_expert_indices - and torch.cuda.get_device_capability()[0] >= 9 + and torch.cuda.get_device_capability() == (9, 0) ) if use_fp8_kernel: if t_in.dtype != t_out.dtype: diff --git a/tests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.py b/tests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.py index 4eb8eccd1a99..045bf0fb43c4 100644 --- a/tests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.py +++ b/tests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.py @@ -155,12 +155,13 @@ 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_or_newer(filename): +def test_fp8_grouped_gemm_dispatch_requires_sm90(filename): source = _kernel_source(filename) assert "getSMVersion()" in source - assert "smVersion >= 90" in source - assert "requires Hopper (SM90) or newer" in source + assert "smVersion == 90" in source + assert "requires Hopper (SM90)" in source + assert "SM120/SM121" in source def test_fp8_grouped_gemm_alignment_checks_require_multiples_of_16(): diff --git a/tests/unittest/others/test_lora_manager.py b/tests/unittest/others/test_lora_manager.py index 57304c47d120..7c2d8020bf56 100644 --- a/tests/unittest/others/test_lora_manager.py +++ b/tests/unittest/others/test_lora_manager.py @@ -218,8 +218,8 @@ def test_many_adapters_no_gpu_accumulation(self): @unittest.skipUnless( - torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 9, - "FP8 LoRA requires Hopper or newer", + torch.cuda.is_available() and torch.cuda.get_device_capability() == (9, 0), + "Native FP8 LoRA requires Hopper (SM90)", ) class TestLoraManagerFp8(unittest.TestCase): def test_fp8_dora_is_rejected(self): @@ -302,6 +302,33 @@ 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): + 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)): + 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) + def test_fp8_e4m3_input_output_dtype_mismatch_is_rejected(self): model_config = MockModelConfig() manager = LoraManager( From f32085baad0ec52ea5eb6812993ea8dddc0f3aad Mon Sep 17 00:00:00 2001 From: Aurelien Chartier <2567591+achartier@users.noreply.github.com> Date: Wed, 29 Jul 2026 17:27:29 -0700 Subject: [PATCH 4/8] perf: reuse FP8 CUDA graph GEMM metadata Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com> --- .../kernels/cuda_graph_grouped_gemm.cu | 111 +++++------------- .../test_fp8_lora_grouped_gemm_regressions.py | 10 +- 2 files changed, 36 insertions(+), 85 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu b/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu index 043b0a980c8f..7fd7979ee4bd 100644 --- a/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu +++ b/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu @@ -21,6 +21,8 @@ #include +#include + #include "cutlass/cutlass.h" #include "cutlass/gemm/device/gemm_grouped.h" #include "cutlass/gemm/kernel/default_gemm_grouped.h" @@ -37,7 +39,6 @@ #include "cutlass/gemm/group_array_problem_shape.hpp" #include "cutlass/gemm/kernel/gemm_universal.hpp" #include "cutlass/util/packed_stride.hpp" -#include "tensorrt_llm/common/memoryUtils.h" #endif // ENABLE_FP8 TRTLLM_NAMESPACE_BEGIN @@ -196,38 +197,9 @@ void cudaGraphGroupedGemmType(cutlass::gemm::GemmCoord* problemSizesPtr, int pro // FP8 CUDA-graph-compatible grouped GEMM using CUTLASS 3.x. // // The CUDA graph variant receives live problem sizes, pointers, and leading -// dimensions already on the GPU. The 3.x API needs cute stride arrays on -// device, so a small setup kernel converts the existing graph parameters into -// the CUTLASS 3.x argument layout on the same stream. +// dimensions already on the GPU and reuses their compatible storage directly. // ==================================================================== -template -__global__ void fillFp8CudaGraphGroupedGemmParams(cutlass::gemm::GemmCoord const* problemSizesPtr, int problemCount, - int64_t const* ldaGpu, int64_t const* ldbGpu, int64_t const* ldcGpu, int64_t const* lddGpu, - typename ProblemShape::UnderlyingProblemShape* problemShapes, StrideA* strideA, StrideB* strideB, StrideC* strideC, - StrideD* strideD) -{ - int const problemIdx = blockIdx.x * blockDim.x + threadIdx.x; - if (problemIdx >= problemCount) - { - return; - } - - int const m = problemSizesPtr[problemIdx].m(); - int const n = problemSizesPtr[problemIdx].n(); - int const k = problemSizesPtr[problemIdx].k(); - - problemShapes[problemIdx] = cute::make_shape(m, n, k); - strideA[problemIdx] - = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(m, static_cast(ldaGpu[problemIdx]), 1)); - strideB[problemIdx] - = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(n, static_cast(ldbGpu[problemIdx]), 1)); - strideC[problemIdx] - = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, static_cast(ldcGpu[problemIdx]), 1)); - strideD[problemIdx] - = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(m, static_cast(lddGpu[problemIdx]), 1)); -} - 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) @@ -282,55 +254,29 @@ void fp8CudaGraphGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int prob using StrideC = typename GemmKernel::InternalStrideC; using StrideD = typename GemmKernel::InternalStrideD; - // Allocate device memory for problem shapes, pointer arrays, and strides. - auto const tensorOpts = at::TensorOptions().dtype(at::kByte).device(at::kCUDA); - - auto const align16 = [](int64_t bytes) -> int64_t { return (bytes + 15) / 16 * 16; }; - - int64_t const szProblemShapes = align16(problemCount * sizeof(typename ProblemShape::UnderlyingProblemShape)); - int64_t const szPtrA = align16(problemCount * sizeof(ElementA const*)); - int64_t const szPtrB = align16(problemCount * sizeof(ElementB const*)); - int64_t const szPtrC = align16(problemCount * sizeof(ElementC const*)); - int64_t const szPtrD = align16(problemCount * sizeof(ElementD*)); - int64_t const szStrideA = align16(problemCount * sizeof(StrideA)); - int64_t const szStrideB = align16(problemCount * sizeof(StrideB)); - int64_t const szStrideC = align16(problemCount * sizeof(StrideC)); - int64_t const szStrideD = align16(problemCount * sizeof(StrideD)); - - int64_t const totalParamsBytes - = szProblemShapes + szPtrA + szPtrB + szPtrC + szPtrD + szStrideA + szStrideB + szStrideC + szStrideD; - - at::Tensor paramsTensor = at::empty({totalParamsBytes}, tensorOpts); - char* devBase = static_cast(paramsTensor.data_ptr()); - - int64_t devOff = 0; - auto devPtr = [&devBase, &devOff](int64_t sz) -> void* - { - void* p = devBase + devOff; - devOff += (sz + 15) / 16 * 16; - return p; - }; - - auto* devProblemShapes = static_cast(devPtr(szProblemShapes)); - auto* devPtrA = static_cast(devPtr(szPtrA)); - auto* devPtrB = static_cast(devPtr(szPtrB)); - auto* devPtrC = static_cast(devPtr(szPtrC)); - auto* devPtrD = static_cast(devPtr(szPtrD)); - auto* devStrideA = static_cast(devPtr(szStrideA)); - auto* devStrideB = static_cast(devPtr(szStrideB)); - auto* devStrideC = static_cast(devPtr(szStrideC)); - auto* devStrideD = static_cast(devPtr(szStrideD)); - - cudaMemcpyAsync(devPtrA, ptrAGpu, problemCount * sizeof(void*), cudaMemcpyDeviceToDevice, stream); - cudaMemcpyAsync(devPtrB, ptrBGpu, problemCount * sizeof(void*), cudaMemcpyDeviceToDevice, stream); - cudaMemcpyAsync(devPtrC, ptrCGpu, problemCount * sizeof(void*), cudaMemcpyDeviceToDevice, stream); - cudaMemcpyAsync(devPtrD, ptrDGpu, problemCount * sizeof(void*), cudaMemcpyDeviceToDevice, stream); - - int constexpr kThreads = 256; - int const blocks = (problemCount + kThreads - 1) / kThreads; - fillFp8CudaGraphGroupedGemmParams - <<>>(problemSizesPtr, problemCount, ldaGpu, ldbGpu, ldcGpu, lddGpu, - devProblemShapes, devStrideA, devStrideB, devStrideC, devStrideD); + using UnderlyingProblemShape = typename ProblemShape::UnderlyingProblemShape; + using PackedStride = cute::Stride; + static_assert(std::is_same_v>); + static_assert(sizeof(UnderlyingProblemShape) == sizeof(cutlass::gemm::GemmCoord)); + static_assert(alignof(UnderlyingProblemShape) == alignof(cutlass::gemm::GemmCoord)); + static_assert(std::is_same_v); + static_assert(std::is_same_v); + static_assert(std::is_same_v); + static_assert(std::is_same_v); + static_assert(sizeof(PackedStride) == sizeof(int64_t)); + static_assert(alignof(PackedStride) == alignof(int64_t)); + + // The fused LoRA parameter kernel already stores problem shapes as contiguous M/N/K int32 values and each + // packed stride's single dynamic component as an int64 leading dimension. + auto* problemShapes = reinterpret_cast(problemSizesPtr); + auto* ptrA = const_cast(reinterpret_cast(ptrAGpu)); + auto* ptrB = const_cast(reinterpret_cast(ptrBGpu)); + auto* ptrC = const_cast(reinterpret_cast(ptrCGpu)); + auto* ptrD = reinterpret_cast(ptrDGpu); + auto* strideA = reinterpret_cast(ldaGpu); + auto* strideB = reinterpret_cast(ldbGpu); + auto* strideC = reinterpret_cast(ldcGpu); + auto* strideD = reinterpret_cast(lddGpu); cutlass::KernelHardwareInfo hwInfo; hwInfo.device_id = 0; @@ -338,8 +284,8 @@ void fp8CudaGraphGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int prob hwInfo.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hwInfo.device_id); typename Gemm::Arguments arguments{cutlass::gemm::GemmUniversalMode::kGrouped, - {problemCount, devProblemShapes, nullptr}, {devPtrA, devStrideA, devPtrB, devStrideB}, - {{1.0f, 0.0f}, devPtrC, devStrideC, devPtrD, devStrideD}, hwInfo}; + {problemCount, problemShapes, nullptr}, {ptrA, strideA, ptrB, strideB}, + {{1.0f, 0.0f}, ptrC, strideC, ptrD, strideD}, hwInfo}; Gemm gemm; @@ -348,6 +294,7 @@ void fp8CudaGraphGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int prob void* gemmWorkspace = nullptr; if (requiredWorkspace > 0) { + auto const tensorOpts = at::TensorOptions().dtype(at::kByte).device(at::kCUDA); workspace = at::empty({static_cast(requiredWorkspace)}, tensorOpts); gemmWorkspace = workspace.data_ptr(); } diff --git a/tests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.py b/tests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.py index 045bf0fb43c4..7d97df60402e 100644 --- a/tests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.py +++ b/tests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.py @@ -121,7 +121,7 @@ def test_fp8_cuda_graph_alignment_rejects_misaligned_hidden_dims(hidden_size, ou ) -def test_fp8_cuda_graph_grouped_gemm_uses_live_device_problem_metadata(): +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(" @@ -129,8 +129,12 @@ def test_fp8_cuda_graph_grouped_gemm_uses_live_device_problem_metadata(): assert "hostMaxProblemSizesPtr" not in fp8_graph_body assert "cudaMemcpyHostToDevice" not in fp8_graph_body - assert "fillFp8CudaGraphGroupedGemmParams" in fp8_graph_body - assert "problemSizesPtr, problemCount, ldaGpu, ldbGpu, ldcGpu, lddGpu" in fp8_graph_body + assert "fillFp8CudaGraphGroupedGemmParams" not in source + assert "cudaMemcpyDeviceToDevice" not in fp8_graph_body + assert "reinterpret_cast(problemSizesPtr)" in fp8_graph_body + assert "reinterpret_cast(ptrAGpu)" in fp8_graph_body + assert "reinterpret_cast(ldaGpu)" in fp8_graph_body + assert "std::is_same_v" in fp8_graph_body @pytest.mark.parametrize( From 88e8305e188ab354226d3fc3aec79ff867f94e6f Mon Sep 17 00:00:00 2001 From: Aurelien Chartier <2567591+achartier@users.noreply.github.com> Date: Sun, 2 Aug 2026 13:20:09 -0700 Subject: [PATCH 5/8] [None][fix] Address FP8 LoRA cache and graph regressions Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com> --- tensorrt_llm/_torch/peft/lora/layer.py | 4 + tensorrt_llm/lora_manager.py | 78 +++++++++++---- .../test_fp8_lora_grouped_gemm_regressions.py | 97 +++++++++++++++++++ tests/unittest/others/test_lora_manager.py | 78 ++++++++++++++- 4 files changed, 233 insertions(+), 24 deletions(-) diff --git a/tensorrt_llm/_torch/peft/lora/layer.py b/tensorrt_llm/_torch/peft/lora/layer.py index e3d3be42c3e7..cf580bd8d956 100644 --- a/tensorrt_llm/_torch/peft/lora/layer.py +++ b/tensorrt_llm/_torch/peft/lora/layer.py @@ -562,6 +562,10 @@ def _forward_cuda_graph_mode( host_max_out_sizes, grouped_gemm_params.splitk_offsets, grouped_gemm_params.reordered_input.dtype, min_kn) + # PyTorch does not implement index_copy_ for FP8 tensors. + if output_buffer.dtype == torch.float8_e4m3fn: + output_buffer = output_buffer.to(torch.bfloat16) + # TODO: move to kernel restored_output = torch.zeros_like(output_buffer) restored_output.index_copy_(0, diff --git a/tensorrt_llm/lora_manager.py b/tensorrt_llm/lora_manager.py index b6fd4338c64a..972b0a4cb176 100644 --- a/tensorrt_llm/lora_manager.py +++ b/tensorrt_llm/lora_manager.py @@ -998,11 +998,6 @@ def prepare_fused_lora_modules_for_tp( return t_out def load_from_model_dir(uid, model_dir, hf_config): - if uid not in self._cpp_lora_weights: - self._cpp_lora_weights[uid] = [] # Will be converted to tensor later - if uid not in self._cpp_lora_config: - self._cpp_lora_config[uid] = [] # Will be converted to tensor later - lora_model = load_state_dict(get_model_path(model_dir, "adapter_model")) if lora_model is None: raise ValueError(f"Failed to load adapter_model from {model_dir}") @@ -1010,23 +1005,70 @@ def load_from_model_dir(uid, model_dir, hf_config): all_weights = get_all_hf_lora_weights(lora_model, hf_modules, component) 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) + + def get_output_dtype(module_weights): + if _is_moe_module_weights(module_weights): + output_dtypes = { + weights["out"].dtype + for weights in module_weights.values() + if "out" in weights + } + return next(iter(output_dtypes)) if len(output_dtypes) == 1 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 + ) - self._lora_uid_to_low_ranks[uid] = {} - self._lora_weights_pointers_list[uid] = {} - for layer_idx in sorted(all_weights.keys()): - layer_weights = all_weights[layer_idx] - self._lora_uid_to_low_ranks[uid][layer_idx] = {} - self._lora_weights_pointers_list[uid][layer_idx] = {} - + for layer_weights in all_weights.values(): + placeholder_dtype = next( + ( + weights["out"].dtype + for weights in layer_weights.values() + if not _is_moe_module_weights(weights) and "out" in weights + ), + model_dtype, + ) for lora_module in self.missing_qkv_modules: hf_module = model_config.trtllm_modules_to_hf_modules[lora_module] if isinstance(hf_module, list): hf_module = hf_module[0] layer_weights[hf_module] = { - "in": torch.zeros(rank, model_config.hidden_size), - "out": torch.zeros(model_config.hidden_size, rank), + "in": torch.zeros(rank, model_config.hidden_size, dtype=placeholder_dtype), + "out": torch.zeros(model_config.hidden_size, rank, dtype=placeholder_dtype), } + cache_dtypes = { + torch.float8_e4m3fn if uses_native_fp8(module_weights) else model_dtype + for layer_weights in all_weights.values() + for hf_module, module_weights in layer_weights.items() + if hf_modules_to_trtllm_modules[hf_module] in self.lora_target_modules + } + if len(cache_dtypes) > 1: + 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." + ) + + if uid not in self._cpp_lora_weights: + self._cpp_lora_weights[uid] = [] # Will be converted to tensor later + if uid not in self._cpp_lora_config: + self._cpp_lora_config[uid] = [] # Will be converted to tensor later + + self._lora_uid_to_low_ranks[uid] = {} + self._lora_weights_pointers_list[uid] = {} + for layer_idx in sorted(all_weights.keys()): + layer_weights = all_weights[layer_idx] + self._lora_uid_to_low_ranks[uid][layer_idx] = {} + self._lora_weights_pointers_list[uid][layer_idx] = {} + for hf_module, module_weights in layer_weights.items(): lora_module = hf_modules_to_trtllm_modules[hf_module] if lora_module not in self.lora_target_modules: @@ -1078,13 +1120,8 @@ def load_from_model_dir(uid, model_dir, hf_config): t_out = prepare_fused_lora_modules_for_tp(lora_module, t_out, rank_dim) effective_rank = t_in.shape[rank_dim] - is_fp8 = t_out.dtype == torch.float8_e4m3fn # TODO: Enable SM120/SM121 after validating the native FP8 LoRA kernel there. - use_fp8_kernel = ( - is_fp8 - and not has_expert_indices - and torch.cuda.get_device_capability() == (9, 0) - ) + use_fp8_kernel = uses_native_fp8(module_weights) if use_fp8_kernel: if t_in.dtype != t_out.dtype: raise ValueError( @@ -1125,7 +1162,6 @@ def load_from_model_dir(uid, model_dir, hf_config): ) else: # Pre-Hopper kernels require the model compute dtype. - model_dtype = str_dtype_to_torch(model_config.dtype) t_in = t_in.to(model_dtype) t_out = t_out.to(model_dtype) t_out = t_out * scale diff --git a/tests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.py b/tests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.py index 7d97df60402e..d610fe7e41da 100644 --- a/tests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.py +++ b/tests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.py @@ -18,6 +18,7 @@ import pytest import torch +from tensorrt_llm._torch.peft.lora.cuda_graph_lora_params import CudaGraphLoraParams from tensorrt_llm._torch.peft.lora.layer import ( LoraLayer, _validate_fp8_lora_cuda_graph_alignment, @@ -25,6 +26,7 @@ ) _REPO_ROOT = Path(__file__).resolve().parents[4] +_NATIVE_FP8_AVAILABLE = torch.cuda.is_available() and torch.cuda.get_device_capability() == (9, 0) def _kernel_source(filename: str) -> str: @@ -37,6 +39,101 @@ def _function_block(source: str, start: str, end: str) -> str: return source[start_index:end_index] +def _make_fp8_lora_problem(batch_size=16, hidden_size=32, rank=16, output_size=32): + torch.manual_seed(0) + x = (torch.randn(batch_size, hidden_size, device="cuda") * 0.25).to(torch.float8_e4m3fn) + lora_in = (torch.randn(rank, hidden_size, device="cuda") * 0.25).to(torch.float8_e4m3fn) + lora_out = (torch.randn(output_size, rank, device="cuda") * 0.25).to(torch.float8_e4m3fn) + reference = x.float() @ lora_in.float().T @ lora_out.float().T + return x, lora_in.contiguous(), lora_out.contiguous(), reference + + +def _assert_fp8_gemm_matches_reference(actual, reference): + torch.testing.assert_close( + actual.float(), + reference.to(torch.float8_e4m3fn).float(), + atol=0.25, + rtol=0.25, + ) + + +@pytest.mark.skipif(not _NATIVE_FP8_AVAILABLE, reason="Native FP8 LoRA requires SM90") +def test_fp8_eager_grouped_gemm_matches_reference(): + x, lora_in, lora_out, reference = _make_fp8_lora_problem() + rank = lora_in.shape[0] + weight_pointers = torch.tensor( + [[lora_in.data_ptr(), lora_out.data_ptr(), 0]], dtype=torch.int64 + ) + + actual = torch.ops.trtllm.lora_grouped_gemm( + x, + torch.zeros(1, dtype=torch.int32), + [torch.tensor([rank], dtype=torch.int32)], + [weight_pointers], + torch.tensor([x.shape[0]], dtype=torch.int32), + [lora_out.shape[0]], + False, + True, + rank, + 0, + True, + )[0] + + _assert_fp8_gemm_matches_reference(actual, reference) + + +@pytest.mark.skipif(not _NATIVE_FP8_AVAILABLE, reason="Native FP8 LoRA requires SM90") +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) + rank = lora_in.shape[0] + module_id = 0 + layer_key = CudaGraphLoraParams.LoraLayerKey(layer_idx=0, module_ids=(module_id,)) + cuda_graph_params = CudaGraphLoraParams( + max_batch_size=batch_size, + max_lora_size=1, + max_rank=rank, + layer_info={ + layer_key: CudaGraphLoraParams.LoraLayerInfo( + module_num=1, output_sizes=[lora_out.shape[0]] + ) + }, + ) + cuda_graph_params.update_sorted_indices([0] * batch_size) + cuda_graph_params.update_slots_params([0] * batch_size) + cuda_graph_params.slot_ranks_host[0] = rank + cuda_graph_params.slot_ranks.copy_(cuda_graph_params.slot_ranks_host) + layer_params = cuda_graph_params.get_layer_params(layer_key) + layer_params.h_b_ptrs[0, 0] = lora_in.data_ptr() + layer_params.h_b_prime_ptrs[0, 0] = lora_out.data_ptr() + layer_params.d_b_ptrs.copy_(layer_params.h_b_ptrs) + layer_params.d_b_prime_ptrs.copy_(layer_params.h_b_prime_ptrs) + + layer = LoraLayer([module_id], [lora_out.shape[0]]) + lora_params = { + "cuda_graph_params": cuda_graph_params, + "data_type": torch.float8_e4m3fn, + "use_cuda_graph_mode": True, + } + static_x = x.clone() + for _ in range(3): + layer(static_x, lora_params, layer_idx=0) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output = layer(static_x, lora_params, layer_idx=0) + + replay_x, _, _, reference = _make_fp8_lora_problem(batch_size=batch_size) + replay_x = torch.roll(replay_x.float(), shifts=1, dims=0).to(torch.float8_e4m3fn) + reference = replay_x.float() @ lora_in.float().T @ lora_out.float().T + static_x.copy_(replay_x) + graph.replay() + torch.cuda.synchronize() + + _assert_fp8_gemm_matches_reference(graph_output, reference) + + def test_fp8_cuda_graph_alignment_accepts_valid_ranks_and_dims(): min_kn = _validate_fp8_lora_cuda_graph_alignment( torch.tensor([0, 16, 32], dtype=torch.int32), 64, [128, 256], 32 diff --git a/tests/unittest/others/test_lora_manager.py b/tests/unittest/others/test_lora_manager.py index 7c2d8020bf56..a91a995252c7 100644 --- a/tests/unittest/others/test_lora_manager.py +++ b/tests/unittest/others/test_lora_manager.py @@ -61,13 +61,15 @@ def _create_dummy_hf_lora_adapter( input_dtype: torch.dtype | None = None, lora_alpha: float | None = None, weight_value: float | None = None, + modules: list[str] | None = None, ): """Create a minimal HF-format LoRA adapter on disk.""" output_size = hidden_size if output_size is None else output_size + modules = ["q_proj", "k_proj", "v_proj"] if modules is None else modules config = { "r": rank, "lora_alpha": rank if lora_alpha is None else lora_alpha, - "target_modules": ["q_proj", "k_proj", "v_proj"], + "target_modules": modules, "bias": "none", "peft_type": "LORA", "task_type": "CAUSAL_LM", @@ -85,7 +87,7 @@ def make_weight(shape, target_dtype): return torch.full(shape, weight_value, dtype=torch.float16).to(target_dtype) for layer_idx in range(num_layers): - for module in ["q_proj", "k_proj", "v_proj"]: + for module in modules: prefix = f"base_model.model.model.layers.{layer_idx}.self_attn.{module}" weights[f"{prefix}.lora_A.weight"] = make_weight((rank, hidden_size), input_dtype) weights[f"{prefix}.lora_B.weight"] = make_weight((output_size, rank), dtype) @@ -104,12 +106,13 @@ def _create_dummy_hf_moe_lora_adapter( rank: int = 16, num_experts: int = 2, dtype: torch.dtype = torch.float8_e4m3fn, + include_dense: bool = False, ): """Create a minimal expert-indexed HF-format LoRA adapter on disk.""" config = { "r": rank, "lora_alpha": rank, - "target_modules": ["gate_proj"], + "target_modules": ["gate_proj", "q_proj"] if include_dense else ["gate_proj"], "bias": "none", "peft_type": "LORA", "task_type": "CAUSAL_LM", @@ -126,6 +129,14 @@ def _create_dummy_hf_moe_lora_adapter( weights[f"{prefix}.lora_B.weight"] = torch.randn(output_size, rank, dtype=torch.float16).to( dtype ) + if include_dense: + prefix = "base_model.model.model.layers.0.self_attn.q_proj" + weights[f"{prefix}.lora_A.weight"] = torch.randn(rank, hidden_size, dtype=torch.float16).to( + dtype + ) + weights[f"{prefix}.lora_B.weight"] = torch.randn(hidden_size, rank, dtype=torch.float16).to( + dtype + ) save_file(weights, str(adapter_dir / "adapter_model.safetensors")) @@ -276,6 +287,67 @@ def test_fp8_moe_weights_are_converted_to_model_dtype(self): self.assertEqual(manager.cpp_lora_weights["fp8-moe"].dtype, torch.bfloat16) + def test_mixed_fp8_dense_and_moe_modules_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) / "adapter" + adapter_dir.mkdir() + _create_dummy_hf_moe_lora_adapter(adapter_dir, include_dense=True) + + with self.assertRaisesRegex(ValueError, "one PEFT cache dtype"): + manager.load_from_hf( + model_dirs=[str(adapter_dir)], + model_config=model_config, + uids=["mixed-fp8"], + ) + + self.assertNotIn("mixed-fp8", manager.cpp_lora_weights) + + def test_partial_qkv_fp8_adapter_uses_fp8_placeholders(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, + modules=["q_proj", "v_proj"], + ) + + manager.load_from_hf( + model_dirs=[str(adapter_dir)], + model_config=model_config, + uids=["partial-qkv-fp8"], + ) + + self.assertEqual( + manager.cpp_lora_weights["partial-qkv-fp8"].dtype, + torch.float8_e4m3fn, + ) + def test_fp8_e5m2_weights_are_converted_to_model_dtype(self): model_config = MockModelConfig(dtype="bfloat16") manager = LoraManager( From 9f38cc1f35435595473c82a2233447aa5cb7fd53 Mon Sep 17 00:00:00 2001 From: Aurelien Chartier <2567591+achartier@users.noreply.github.com> Date: Tue, 4 Aug 2026 11:09:25 -0700 Subject: [PATCH 6/8] [None][fix] Address final FP8 LoRA review findings Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com> --- .../_torch/pyexecutor/cuda_graph_runner.py | 14 +++-- .../_torch/pyexecutor/model_engine.py | 7 +++ tensorrt_llm/lora_manager.py | 46 ++++++++------- .../executor/test_pytorch_model_engine.py | 58 ++++++++++++++++++- tests/unittest/others/test_lora_manager.py | 23 ++++++++ 5 files changed, 122 insertions(+), 26 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 4019c037f513..9f626eba6b8a 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -48,6 +48,7 @@ class KeyType(NamedTuple): num_contexts: int = 0 context_query_len: int = 0 num_encoder_tokens: int = 0 + peft_cache_data_type: Optional[torch.dtype] = None def _save_spec_decode_capture_state( @@ -313,7 +314,8 @@ def get_graph_key( new_tensors_device: Optional[SampleStateTensors] = None, spec_resource_manager: Optional[BaseResourceManager] = None, spec_metadata: Optional[SpecMetadata] = None, - promoted_context_request_ids: frozenset[int] = frozenset() + promoted_context_request_ids: frozenset[int] = frozenset(), + peft_cache_data_type: Optional[torch.dtype] = None, ) -> Optional[KeyType]: batch_size = batch.batch_size @@ -339,7 +341,8 @@ def get_graph_key( draft_len=draft_len, is_first_draft=spec_resource_manager.is_first_draft, short_seq_len_mode=short_seq_len_mode, - is_all_greedy_sample=is_all_greedy_sample) + is_all_greedy_sample=is_all_greedy_sample, + peft_cache_data_type=peft_cache_data_type) else: # With dynamic spec decode, the draft length may be zero even when enable_spec_decode is True, # so we need to get the draft length from the batch instead of using enable_spec_decode. @@ -368,7 +371,8 @@ def get_graph_key( is_all_greedy_sample=is_all_greedy_sample, num_contexts=num_contexts, context_query_len=context_query_len, - num_encoder_tokens=num_encoder_tokens) + num_encoder_tokens=num_encoder_tokens, + peft_cache_data_type=peft_cache_data_type) return key def _get_compatible_mixed_encoder_decoder_key(self, @@ -429,6 +433,7 @@ def maybe_get_cuda_graph( new_tensors_device: Optional[SampleStateTensors] = None, spec_resource_manager: Optional[BaseResourceManager] = None, promoted_context_request_ids: frozenset[int] = frozenset(), + peft_cache_data_type: Optional[torch.dtype] = None, ) -> Tuple[Optional[Any], Optional[Any], Optional[KeyType]]: """ Determines if the current batch can be run with a CUDA graph. @@ -474,7 +479,8 @@ def maybe_get_cuda_graph( # callers pass the empty default and retain generation-only behavior. key = self.get_graph_key(batch, new_tensors_device, spec_resource_manager, spec_metadata, - promoted_context_request_ids) + promoted_context_request_ids, + peft_cache_data_type) if key is None: return None, None, None if is_mixed_encoder_decoder: diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index b8f375bda8bb..299124008b05 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -7105,6 +7105,12 @@ def forward(self, padded_graph_requests.all_requests()) self._sync_group_all_greedy_sample(spec_metadata) + peft_cache_data_type = None + if getattr(self, "cuda_graph_lora_manager", None) is not None: + peft_cache_manager = resource_manager.get_resource_manager( + ResourceManagerType.PEFT_CACHE_MANAGER) + peft_cache_data_type = peft_cache_manager.data_type + maybe_attn_metadata, maybe_spec_metadata, key = self.cuda_graph_runner.maybe_get_cuda_graph( padded_graph_requests, enable_spec_decode=self.enable_spec_decode, @@ -7115,6 +7121,7 @@ def forward(self, new_tensors_device=new_tensors_device, spec_resource_manager=spec_resource_manager, promoted_context_request_ids=promoted_context_request_ids, + peft_cache_data_type=peft_cache_data_type, ) can_run_graph = key is not None diff --git a/tensorrt_llm/lora_manager.py b/tensorrt_llm/lora_manager.py index 972b0a4cb176..000fdeff7e74 100644 --- a/tensorrt_llm/lora_manager.py +++ b/tensorrt_llm/lora_manager.py @@ -1057,17 +1057,15 @@ def uses_native_fp8(module_weights): "with compute-dtype modules such as routed-expert MoE is not supported." ) - if uid not in self._cpp_lora_weights: - self._cpp_lora_weights[uid] = [] # Will be converted to tensor later - if uid not in self._cpp_lora_config: - self._cpp_lora_config[uid] = [] # Will be converted to tensor later - - self._lora_uid_to_low_ranks[uid] = {} - self._lora_weights_pointers_list[uid] = {} + cpp_lora_weights = [] + cpp_lora_config = [] + uid_to_low_ranks = {} + lora_weights_pointers = {} + retained_lora_weights = [] for layer_idx in sorted(all_weights.keys()): layer_weights = all_weights[layer_idx] - self._lora_uid_to_low_ranks[uid][layer_idx] = {} - self._lora_weights_pointers_list[uid][layer_idx] = {} + uid_to_low_ranks[layer_idx] = {} + lora_weights_pointers[layer_idx] = {} for hf_module, module_weights in layer_weights.items(): lora_module = hf_modules_to_trtllm_modules[hf_module] @@ -1075,7 +1073,7 @@ def uses_native_fp8(module_weights): warnings.warn( f"LoRA module '{lora_module}' not in target modules {self.lora_target_modules}, skipping." ) - self._lora_uid_to_low_ranks[uid][layer_idx][lora_module] = 0 + uid_to_low_ranks[layer_idx][lora_module] = 0 continue has_expert_indices = _is_moe_module_weights(module_weights) @@ -1168,17 +1166,17 @@ def uses_native_fp8(module_weights): if is_dora and t_mag is not None: t_mag = t_mag.to(model_dtype) - self._lora_uid_to_low_ranks[uid][layer_idx][lora_module] = effective_rank + uid_to_low_ranks[layer_idx][lora_module] = effective_rank if self._retain_device_tensors: - self._lora_weights_pointers_list[uid][layer_idx][lora_module] = [ + lora_weights_pointers[layer_idx][lora_module] = [ t_in.data_ptr(), t_out.data_ptr(), t_mag.data_ptr() if (is_dora and t_mag is not None) else 0, ] - self._lora_weights.append(t_in) - self._lora_weights.append(t_out) + retained_lora_weights.append(t_in) + retained_lora_weights.append(t_out) if is_dora and t_mag is not None: - self._lora_weights.append(t_mag) + retained_lora_weights.append(t_mag) t_in_cpu = t_in.flatten().cpu() t_out_cpu = t_out.flatten().cpu() @@ -1188,22 +1186,28 @@ def uses_native_fp8(module_weights): t_mag_cpu = t_mag.flatten().cpu() weights_to_concat.append(t_mag_cpu) - self._cpp_lora_weights[uid].append(torch.cat(weights_to_concat)) - self._cpp_lora_config[uid].append( + cpp_lora_weights.append(torch.cat(weights_to_concat)) + cpp_lora_config.append( torch.tensor( [self.LORA_MODULE_IDS[lora_module], layer_idx, effective_rank, is_dora], dtype=torch.int32, ) ) - max_weight_size = max(w.size(0) for w in self._cpp_lora_weights[uid]) - self._cpp_lora_weights[uid] = torch.stack( + max_weight_size = max(w.size(0) for w in cpp_lora_weights) + packed_lora_weights = torch.stack( [ torch.nn.functional.pad(w, (0, max_weight_size - w.size(0))) - for w in self._cpp_lora_weights[uid] + for w in cpp_lora_weights ] ) - self._cpp_lora_config[uid] = torch.stack([c for c in self._cpp_lora_config[uid]]) + packed_lora_config = torch.stack(cpp_lora_config) + + self._cpp_lora_weights[uid] = packed_lora_weights + self._cpp_lora_config[uid] = packed_lora_config + self._lora_uid_to_low_ranks[uid] = uid_to_low_ranks + self._lora_weights_pointers_list[uid] = lora_weights_pointers + self._lora_weights.extend(retained_lora_weights) for uid, model_dir, hf_config in zip(new_uids, new_model_dirs, lora_hf_configs): load_from_model_dir(uid, model_dir, hf_config) diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 03d76b437123..ed0f3bd4d00b 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -694,6 +694,62 @@ def test_graph_key_rounds_encoder_tokens_up_to_captured_extent( self.assertEqual(actual_key, compatible_key) + def test_graph_key_includes_peft_cache_dtype(self) -> None: + runner = Mock() + runner.config = SimpleNamespace(is_draft_model=False) + runner._get_seq_len_mode.return_value = False + request = _make_request_stub(7) + batch = ScheduledRequests() + batch.generation_requests = [request] + + model_dtype_key = CUDAGraphRunner.get_graph_key( + runner, batch, peft_cache_data_type=torch.bfloat16) + fp8_key = CUDAGraphRunner.get_graph_key( + runner, batch, peft_cache_data_type=torch.float8_e4m3fn) + + self.assertNotEqual(model_dtype_key, fp8_key) + self.assertEqual( + model_dtype_key._replace(peft_cache_data_type=None), + fp8_key._replace(peft_cache_data_type=None), + ) + + def test_graph_dtype_change_falls_back_to_eager(self) -> None: + runner = Mock() + runner.enabled = True + runner.config = SimpleNamespace( + enable_attention_dp=False, + use_mrope=False, + ) + model_dtype_key = KeyType(batch_size=1, + draft_len=0, + is_first_draft=False, + peft_cache_data_type=torch.bfloat16) + fp8_key = KeyType(batch_size=1, + draft_len=0, + is_first_draft=False, + peft_cache_data_type=torch.float8_e4m3fn) + runner.get_graph_key.return_value = fp8_key + runner.graph_metadata = {model_dtype_key: object()} + runner._capture_allowed = False + runner._is_mixed_encoder_decoder_batch.return_value = False + runner._can_run_cuda_graph_batch.return_value = True + request = _make_request_stub(7) + batch = ScheduledRequests() + batch.generation_requests = [request] + + with patch( + "tensorrt_llm._torch.pyexecutor.cuda_graph_runner.ExpertStatistic.should_record", + return_value=False): + result = CUDAGraphRunner.maybe_get_cuda_graph( + runner, + batch, + enable_spec_decode=False, + attn_metadata=object(), + peft_cache_data_type=torch.float8_e4m3fn, + ) + + self.assertEqual(result, (None, None, None)) + def test_graph_lookup_forwards_promoted_context_ids(self) -> None: runner = Mock() runner.enabled = True @@ -734,7 +790,7 @@ def test_graph_lookup_forwards_promoted_context_ids(self) -> None: ) runner.get_graph_key.assert_called_once_with(batch, None, None, None, - promoted_ids) + promoted_ids, None) self.assertEqual(result, (graph_attn_metadata, graph_spec_metadata, key)) diff --git a/tests/unittest/others/test_lora_manager.py b/tests/unittest/others/test_lora_manager.py index a91a995252c7..ba5809156a5a 100644 --- a/tests/unittest/others/test_lora_manager.py +++ b/tests/unittest/others/test_lora_manager.py @@ -429,6 +429,29 @@ def test_fp8_e4m3_input_output_dtype_mismatch_is_rejected(self): uids=["fp8-mismatched-dtype"], ) + uid = "fp8-mismatched-dtype" + self.assertNotIn(uid, manager._cpp_lora_weights) + self.assertNotIn(uid, manager._cpp_lora_config) + self.assertNotIn(uid, manager._lora_uid_to_low_ranks) + self.assertNotIn(uid, manager._lora_weights_pointers_list) + + _create_dummy_hf_lora_adapter( + adapter_dir, + rank=16, + num_layers=1, + dtype=torch.float8_e4m3fn, + ) + manager.load_from_hf( + model_dirs=[str(adapter_dir)], + model_config=model_config, + uids=[uid], + ) + + self.assertIn(uid, manager._cpp_lora_weights) + self.assertIn(uid, manager._cpp_lora_config) + self.assertIn(uid, manager._lora_uid_to_low_ranks) + self.assertIn(uid, manager._lora_weights_pointers_list) + def test_scaled_fp8_e4m3_weights_are_clamped_before_cast(self): model_config = MockModelConfig() manager = LoraManager( From ff0c51ae4b8b19b26d25d3802029c1ff03d488b0 Mon Sep 17 00:00:00 2001 From: Aurelien Chartier <2567591+achartier@users.noreply.github.com> Date: Wed, 5 Aug 2026 20:02:51 -0700 Subject: [PATCH 7/8] [None][fix] Configure FP8 LoRA cache before graph capture Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com> --- .../batch_manager/peftCacheManager.h | 5 +-- .../nanobind/batch_manager/kvCacheManager.cpp | 2 ++ tensorrt_llm/_torch/pyexecutor/_util.py | 7 ++++- .../_torch/pyexecutor/resource_manager.py | 9 ++++-- tensorrt_llm/lora_manager.py | 31 ++++++++++++++++--- .../_torch/executor/test_resource_manager.py | 13 ++++++++ tests/unittest/others/test_lora_manager.py | 28 ++++++++++++++++- 7 files changed, 84 insertions(+), 11 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/peftCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/peftCacheManager.h index 81d4ff622d46..0438687c7ff9 100644 --- a/cpp/include/tensorrt_llm/batch_manager/peftCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/peftCacheManager.h @@ -129,6 +129,9 @@ class PeftCacheManager : public BasePeftCacheManager [[nodiscard]] tensorrt_llm::DataType getDataType() const; + //! Configure the homogeneous LoRA weight dtype before inserting an adapter. + void configureDataType(tensorrt_llm::DataType dataType); + void resetDeviceCache() override; void markRequestDone(LlmRequest const& llmReq, bool pause = false) override; @@ -161,8 +164,6 @@ class PeftCacheManager : public BasePeftCacheManager void prefetchLoraWeights(std::string const& modelDir, runtime::BufferManager const& bufferManager); private: - void configureDataType(tensorrt_llm::DataType dataType); - std::unique_ptr mHostLoraCache; std::unique_ptr mDeviceLoraCache; diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp index b2112650163d..222f2e528cc3 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp @@ -748,6 +748,8 @@ void tb::BasePeftCacheManagerBindings::initBindings(nb::module_& m) nb::call_guard()) .def("is_task_cached_device", &tb::PeftCacheManager::isTaskCachedDevice, nb::arg("taskId"), nb::call_guard()) // ; + .def("configure_data_type", &tb::PeftCacheManager::configureDataType, nb::arg("data_type"), + nb::call_guard()) .def_prop_ro("data_type", &tb::PeftCacheManager::getDataType) .def("ensure_batch_map_task_id", &tb::PeftCacheManager::ensureBatchMapTaskId, nb::arg("context_requests"), nb::arg("generation_requests"), nb::arg("reset_gpu_cache") = false, diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index c4fb115111fe..ea017d144707 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2717,9 +2717,13 @@ def create_py_executor_instance( # dataclass to avoid ad-hoc getattr + TP-division blocks per model type. from tensorrt_llm.bindings import LoraModule + initial_lora_data_type = None if len(lora_config.lora_dir) == 1: # Route to appropriate loader based on checkpoint source - load_torch_lora(lora_config) + 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 else: assert len(lora_config.lora_target_modules ) >= 1, "Expecting at least one lora target module" @@ -2851,6 +2855,7 @@ def create_py_executor_instance( world_config=world_config, execution_stream=execution_stream, lora_target_modules=target_modules, + initial_data_type=initial_lora_data_type, ) resources[ResourceManagerType.PEFT_CACHE_MANAGER] = peft_cache_manager model_engine.set_lora_model_config( diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index a46863d76d7d..760410d92b8c 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -29,7 +29,8 @@ import tensorrt_llm.bindings from tensorrt_llm._torch.distributed.communicator import Distributed, ReduceOp from tensorrt_llm._utils import (get_size_in_bytes, mpi_comm, mpi_disabled, - prefer_pinned, torch_comm) + prefer_pinned, torch_comm, + torch_dtype_to_binding) from tensorrt_llm.bindings.internal.batch_manager import ( LinearAttentionMetadata, LinearCacheType) from tensorrt_llm.bindings.internal.runtime import TaskLayerModuleConfig @@ -2694,7 +2695,8 @@ def __init__(self, model_config: ModelConfigCpp, world_config: WorldConfig | None = None, execution_stream: Optional[torch.cuda.Stream] = None, - lora_target_modules: Optional[List[str]] = None): + lora_target_modules: Optional[List[str]] = None, + initial_data_type: Optional[torch.dtype] = None): import tensorrt_llm.bindings as _tb peft_cache_config = peft_cache_config._to_pybind() @@ -2729,6 +2731,9 @@ def __init__(self, model_config=model_config, world_config=world_config, buffer_manager=buffer_manager) + if initial_data_type is not None: + self.impl.configure_data_type( + torch_dtype_to_binding(initial_data_type)) self._lora_config = lora_config self._lora_model_config = LoraModelConfig( lora_target_modules if lora_target_modules is not None else diff --git a/tensorrt_llm/lora_manager.py b/tensorrt_llm/lora_manager.py index 000fdeff7e74..128034aa5744 100644 --- a/tensorrt_llm/lora_manager.py +++ b/tensorrt_llm/lora_manager.py @@ -346,6 +346,15 @@ 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.""" + 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"): + lora_dtypes.add(weight.dtype) + return next(iter(lora_dtypes)) if len(lora_dtypes) == 1 else None + @lru_cache(maxsize=128) def _find_nemo_files_single_path(lora_path: str) -> List[str]: @@ -454,12 +463,16 @@ def get_target_modules(self): return self.lora_target_modules -def load_torch_hf_lora(lora_config: LoraConfig): +def load_torch_hf_lora(lora_config: LoraConfig) -> Optional[torch.dtype]: """Load an HF LoRA checkpoint for the PyTorch workflow. Populates lora_config (trtllm_modules_to_hf_modules and inferred lora_target_modules) from the HF adapter directory. The actual weights are 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. """ if not lora_config.trtllm_modules_to_hf_modules: lora_config.trtllm_modules_to_hf_modules = get_default_trtllm_modules_to_hf_modules() @@ -480,9 +493,10 @@ def load_torch_hf_lora(lora_config: LoraConfig): 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() -def load_torch_nemo_lora(lora_config: LoraConfig): +def load_torch_nemo_lora(lora_config: LoraConfig) -> Optional[torch.dtype]: """Load NeMo LoRA checkpoint for PyTorch workflow. This is a PyTorch-specific loader for NeMo LoRA checkpoints, similar to @@ -497,6 +511,9 @@ def load_torch_nemo_lora(lora_config: LoraConfig): Args: lora_config: LoRA configuration with lora_ckpt_source="nemo" + Returns: + ``None`` because NeMo LoRA weights use the model compute dtype. + Raises: ValueError: If NeMo LoRA directory is invalid or unsupported modules are specified """ @@ -524,9 +541,10 @@ def load_torch_nemo_lora(lora_config: LoraConfig): f"but got unsupported modules: {unsupported_modules}. " f"NeMo LoRA does not support embedding, lm_head, or MLP adapters." ) + return None -def load_torch_lora(lora_config: LoraConfig): +def load_torch_lora(lora_config: LoraConfig) -> Optional[torch.dtype]: """Load LoRA checkpoint for PyTorch workflow. This function routes to the appropriate loader based on lora_ckpt_source. @@ -534,13 +552,16 @@ def load_torch_lora(lora_config: LoraConfig): Args: lora_config: LoRA configuration with lora_ckpt_source set to "hf" or "nemo" + Returns: + The configured adapter's homogeneous dense weight dtype, if available. + Raises: ValueError: If lora_ckpt_source is not supported """ if lora_config.lora_ckpt_source == "nemo": - load_torch_nemo_lora(lora_config) + return load_torch_nemo_lora(lora_config) elif lora_config.lora_ckpt_source == "hf": - load_torch_hf_lora(lora_config) + return load_torch_hf_lora(lora_config) else: raise ValueError( f"Unsupported lora_ckpt_source: {lora_config.lora_ckpt_source}. " diff --git a/tests/unittest/_torch/executor/test_resource_manager.py b/tests/unittest/_torch/executor/test_resource_manager.py index f4ab82db1d54..e501ff1e5634 100644 --- a/tests/unittest/_torch/executor/test_resource_manager.py +++ b/tests/unittest/_torch/executor/test_resource_manager.py @@ -423,6 +423,19 @@ def test_successful_mocked_peft_cache_manager_initialization(self): self.assertGreaterEqual(peft_cache_manager.impl.max_host_pages, 1) self.assertGreaterEqual(peft_cache_manager.impl.max_device_pages, 1) + def test_initial_fp8_data_type_configures_cache(self): + if (not torch.cuda.is_available() + or torch.cuda.get_device_capability() != (9, 0)): + self.skipTest("Requires SM90 FP8 support") + peft_cache_manager = PeftCacheManager( + peft_cache_config=self.create_peft_cache_config(), + lora_config=LoraConfig(), + model_config=self.model_config, + initial_data_type=torch.float8_e4m3fn, + ) + + self.assertEqual(peft_cache_manager.data_type, torch.float8_e4m3fn) + def test_add_request_peft_empty_weights_config(self): """Test adding a request with empty LoRA task.""" peft_cache_config = self.create_peft_cache_config() diff --git a/tests/unittest/others/test_lora_manager.py b/tests/unittest/others/test_lora_manager.py index ba5809156a5a..d102108069c8 100644 --- a/tests/unittest/others/test_lora_manager.py +++ b/tests/unittest/others/test_lora_manager.py @@ -29,7 +29,7 @@ import torch from safetensors.torch import save_file -from tensorrt_llm.lora_manager import LoraManager +from tensorrt_llm.lora_manager import HfLoraLoader, LoraManager from tensorrt_llm.mapping import Mapping @@ -233,6 +233,31 @@ def test_many_adapters_no_gpu_accumulation(self): "Native FP8 LoRA requires Hopper (SM90)", ) class TestLoraManagerFp8(unittest.TestCase): + def test_hf_loader_reports_dense_fp8_dtype(self): + 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, + ) + + loader = HfLoraLoader([str(adapter_dir)]) + + self.assertEqual(loader.get_dense_lora_dtype(), torch.float8_e4m3fn) + + def test_hf_loader_ignores_expert_only_fp8_dtype(self): + with tempfile.TemporaryDirectory() as tmpdir: + adapter_dir = Path(tmpdir) / "adapter" + adapter_dir.mkdir() + _create_dummy_hf_moe_lora_adapter(adapter_dir) + + loader = HfLoraLoader([str(adapter_dir)]) + + self.assertIsNone(loader.get_dense_lora_dtype()) + def test_fp8_dora_is_rejected(self): model_config = MockModelConfig() manager = LoraManager( @@ -419,6 +444,7 @@ 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, "FP8 LoRA input and output weights must have the same dtype" From f8cbd4128fcc92b6b60b007eb02470b5f60db4f0 Mon Sep 17 00:00:00 2001 From: Aurelien Chartier <2567591+achartier@users.noreply.github.com> Date: Thu, 6 Aug 2026 08:22:28 -0700 Subject: [PATCH 8/8] [None][fix] Update CUDA graph execution probe Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com> --- .../defs/kv_cache/test_final_single_token_context_cuda_graph.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/integration/defs/kv_cache/test_final_single_token_context_cuda_graph.py b/tests/integration/defs/kv_cache/test_final_single_token_context_cuda_graph.py index 9ef8cf87232b..73522a81991a 100644 --- a/tests/integration/defs/kv_cache/test_final_single_token_context_cuda_graph.py +++ b/tests/integration/defs/kv_cache/test_final_single_token_context_cuda_graph.py @@ -92,6 +92,7 @@ def maybe_get_cuda_graph( new_tensors_device: SampleStateTensors | None = None, spec_resource_manager: BaseResourceManager | None = None, promoted_context_request_ids: frozenset[int] = frozenset(), + peft_cache_data_type: torch.dtype | None = None, ) -> tuple[Any | None, Any | None, KeyType | None]: # A new decision means the preceding one reached eager execution if it # did not call replay. Keep that earlier observation unchanged. @@ -105,6 +106,7 @@ def maybe_get_cuda_graph( new_tensors_device, spec_resource_manager, promoted_context_request_ids=promoted_context_request_ids, + peft_cache_data_type=peft_cache_data_type, ) if promoted_context_request_ids: execution = _PromotedContextGraphExecution(