diff --git a/cpp/include/tensorrt_llm/batch_manager/peftCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/peftCacheManager.h index ed928e96d811..0438687c7ff9 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,11 @@ class PeftCacheManager : public BasePeftCacheManager [[nodiscard]] bool isTaskCachedDevice(uint64_t const taskId) const; + [[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; @@ -175,6 +182,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..3e51ddc4c518 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__); @@ -263,7 +264,15 @@ void PeftCacheManager::addRequestPeft(std::shared_ptr llmRequest, bo auto optLoraConfig = llmRequest->getLoraConfig(); if (optTaskId || optLoraWeights || optLoraConfig) { - runtime::lora::loraValidateRequestTensors(optTaskId, optLoraWeights, optLoraConfig, mModelConfig, mWorldConfig); + auto const requestDataType = optLoraWeights + ? std::optional{optLoraWeights.value()->getDataType()} + : std::nullopt; + runtime::lora::loraValidateRequestTensors( + optTaskId, optLoraWeights, optLoraConfig, mModelConfig, mWorldConfig, requestDataType); + if (optLoraWeights) + { + configureDataType(optLoraWeights.value()->getDataType()); + } } else { @@ -374,6 +383,39 @@ void PeftCacheManager::addRequestPeft(std::shared_ptr llmRequest, bo TLLM_LOG_DEBUG("%s stop", __PRETTY_FUNCTION__); } +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) + { + 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..7fd7979ee4bd 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. @@ -15,11 +15,14 @@ */ #include "cuda_graph_grouped_gemm.h" +#include "groupGemm.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" #include +#include + #include "cutlass/cutlass.h" #include "cutlass/gemm/device/gemm_grouped.h" #include "cutlass/gemm/kernel/default_gemm_grouped.h" @@ -27,10 +30,57 @@ #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" +#endif // ENABLE_FP8 + TRTLLM_NAMESPACE_BEGIN namespace kernels { +#ifdef ENABLE_FP8 +#if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) +namespace +{ + +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), 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.", + 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 +#endif // CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED +#endif // ENABLE_FP8 /** * Template for CUDA Graph compatible grouped GEMM that directly uses GPU tensors @@ -140,10 +190,157 @@ 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 and reuses their compatible storage directly. +// ==================================================================== + +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; + + 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; + 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, problemShapes, nullptr}, {ptrA, strideA, ptrB, strideB}, + {{1.0f, 0.0f}, ptrC, strideC, ptrD, strideD}, hwInfo}; + + Gemm gemm; + + size_t const requiredWorkspace = Gemm::get_workspace_size(arguments); + at::Tensor workspace; + 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(); + } + + 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 +509,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..87e5b94c6645 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,36 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { +#ifdef ENABLE_FP8 +#if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED) +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), but the current device is SM%d", kernelName, smVersion); + + 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 +#endif // CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED +#endif // ENABLE_FP8 int64_t inline getGemmCoordSize(int64_t problemCount) { @@ -59,6 +100,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 +258,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..48b38eb82d65 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. @@ -24,8 +24,16 @@ 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. +//! +//! 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..d1aa210665af 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"); @@ -15,6 +15,7 @@ * limitations under the License. */ +#include "groupGemm.h" #include "splitkGroupGemm.h" #include "cutlass/cutlass.h" @@ -35,7 +36,6 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { - int64_t inline getGemmCoordSize(int64_t problemCount) { return (int64_t) (tensorrt_llm::common::divUp(problemCount * sizeof(cutlass::gemm::GemmCoord), 16) * 16); @@ -232,6 +232,18 @@ void splitkGroupedGemm(std::vector const& problemSizes 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 (dataType == tensorrt_llm::DataType::kFP8) + { + // 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; + } +#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..222f2e528cc3 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp @@ -748,6 +748,9 @@ 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, nb::call_guard()); diff --git a/cpp/tensorrt_llm/runtime/loraCache.cpp b/cpp/tensorrt_llm/runtime/loraCache.cpp index 36fb0363816f..6749377445b2 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,7 +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"); - // First get the taskValue from this cache // TaskValue& taskValue = copyTaskGetThisTaskValue(taskId); TaskValuePtr taskValue = [&]() -> TaskValuePtr @@ -805,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/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..51575a5873ed 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,11 @@ 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 ", c10::toString(dtype)); } int const minKnInt = std::max(1, static_cast(minKN)); @@ -307,7 +314,11 @@ 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 ", 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..6e4ded72ae26 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,73 @@ 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; + 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..cf580bd8d956 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, @@ -482,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/_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/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 6bfa53c107a5..299124008b05 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, @@ -7102,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, @@ -7112,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/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 153fe93e82ae..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 @@ -44,7 +45,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 @@ -2693,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() @@ -2728,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 @@ -2753,6 +2759,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..128034aa5744 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]]: @@ -317,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]: @@ -425,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() @@ -451,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 @@ -468,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 """ @@ -495,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. @@ -505,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}. " @@ -969,11 +1019,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}") @@ -981,30 +1026,75 @@ 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." + ) + + 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] + 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] if lora_module not in self.lora_target_modules: 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) @@ -1049,6 +1139,26 @@ 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] + # TODO: Enable SM120/SM121 after validating the native FP8 LoRA kernel there. + use_fp8_kernel = uses_native_fp8(module_weights) + 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,26 +1170,34 @@ 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) + if use_fp8_kernel: + # Keep weights in FP8 for the native Hopper kernel. + # FP8 has no scalar multiply, so scale through BF16. + 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. + 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 + 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() @@ -1089,22 +1207,28 @@ def load_from_model_dir(uid, model_dir, hf_config): 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/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( 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/_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/_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..d610fe7e41da --- /dev/null +++ b/tests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.py @@ -0,0 +1,298 @@ +# 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.cuda_graph_lora_params import CudaGraphLoraParams +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] +_NATIVE_FP8_AVAILABLE = torch.cuda.is_available() and torch.cuda.get_device_capability() == (9, 0) + + +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 _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 + ) + + 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_reuses_live_device_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" 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( + "filename,messages", + [ + ("groupGemm.cu", ["FP8 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", "cuda_graph_grouped_gemm.cu"]) +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)" in source + assert "SM120/SM121" 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") + + 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], diff --git a/tests/unittest/others/test_lora_manager.py b/tests/unittest/others/test_lora_manager.py index ff36aeebda86..d102108069c8 100644 --- a/tests/unittest/others/test_lora_manager.py +++ b/tests/unittest/others/test_lora_manager.py @@ -24,12 +24,12 @@ 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 -from tensorrt_llm.lora_manager import LoraManager +from tensorrt_llm.lora_manager import HfLoraLoader, LoraManager from tensorrt_llm.mapping import Mapping @@ -51,26 +51,92 @@ 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, + 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, - "target_modules": ["q_proj", "k_proj", "v_proj"], + "lora_alpha": rank if lora_alpha is None else lora_alpha, + "target_modules": modules, "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"]: + for module in modules: 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, + 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", "q_proj"] if include_dense else ["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 + ) + 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")) @@ -162,5 +228,338 @@ 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() == (9, 0), + "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( + 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_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( + 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_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( + 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, + ) + 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" + ): + manager.load_from_hf( + model_dirs=[str(adapter_dir)], + model_config=model_config, + 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( + 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()