From 4a706e3814d2f89646f15a65b213bd5a30f97f1c Mon Sep 17 00:00:00 2001 From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> Date: Tue, 12 May 2026 21:02:08 +0000 Subject: [PATCH 01/25] [None][feat] add NVFP4 W4A16 SM120 support Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> --- .../tensorrt_llm/common/quantization.h | 16 +- .../cudaCoreGemmW4A16NVFP4.cu | 303 ++++++++++++++++++ .../cudaCoreGemmW4A16NVFP4.h | 71 ++++ cpp/tensorrt_llm/nanobind/bindings.cpp | 3 +- cpp/tensorrt_llm/thop/CMakeLists.txt | 1 + cpp/tensorrt_llm/thop/w4a16Nvfp4Gemm.cpp | 113 +++++++ .../_torch/custom_ops/cpp_custom_ops.py | 10 + tensorrt_llm/_torch/model_config.py | 21 +- .../hf/nemotron_h_weight_mapper.py | 18 +- .../_torch/models/modeling_nemotron_h.py | 47 ++- tensorrt_llm/_torch/models/modeling_utils.py | 10 +- tensorrt_llm/_torch/modules/embedding.py | 14 +- tensorrt_llm/_torch/modules/linear.py | 71 ++++ tensorrt_llm/models/quant_config_utils.py | 21 +- tensorrt_llm/quantization/mode.py | 19 +- .../test_modeling_nemotron_h_moe_quant.py | 74 +++++ .../hf/test_nemotron_h_weight_mapper.py | 80 +++++ .../_torch/modules/test_w4a16_nvfp4_linear.py | 142 ++++++++ tests/unittest/_torch/test_model_config.py | 84 +++++ 19 files changed, 1084 insertions(+), 34 deletions(-) create mode 100644 cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.cu create mode 100644 cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.h create mode 100644 cpp/tensorrt_llm/thop/w4a16Nvfp4Gemm.cpp create mode 100644 tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py create mode 100644 tests/unittest/_torch/models/checkpoints/hf/test_nemotron_h_weight_mapper.py create mode 100644 tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py diff --git a/cpp/include/tensorrt_llm/common/quantization.h b/cpp/include/tensorrt_llm/common/quantization.h index df13a674d688..18a28c8c552f 100644 --- a/cpp/include/tensorrt_llm/common/quantization.h +++ b/cpp/include/tensorrt_llm/common/quantization.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2023-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. @@ -134,6 +134,11 @@ class QuantMode return QuantMode(BaseType(1u) << 16); } + static constexpr QuantMode w4a16Nvfp4() noexcept + { + return QuantMode(BaseType(1u) << 17); + } + constexpr BaseType value() const noexcept { return mValue; @@ -224,6 +229,11 @@ class QuantMode return isSet(w4a16Mxfp4()); } + constexpr bool hasW4a16Nvfp4() const noexcept + { + return isSet(w4a16Nvfp4()); + } + constexpr bool hasKvCacheQuant() const noexcept { return hasInt8KvCache() || hasFp8KvCache() || hasFp4KvCache(); @@ -419,6 +429,10 @@ class QuantMode quantMode = fromDescription(false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, true); } + else if (quantAlgo == "W4A16_NVFP4") + { + quantMode = w4a16Nvfp4(); + } if (kvCacheQuantAlgo == "INT8") { diff --git a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.cu b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.cu new file mode 100644 index 000000000000..8ed5a5d5a318 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.cu @@ -0,0 +1,303 @@ +/* + * Copyright (c) 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. + * 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. + */ + +#include "cutlass/numeric_conversion.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h" +#include "tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.h" + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace cuda_core_gemm_w4a16_nvfp4 +{ +namespace +{ + +template +__device__ float toFloat(T value) +{ + return static_cast(value); +} + +template <> +__device__ float toFloat(half value) +{ + return __half2float(value); +} + +template <> +__device__ float toFloat<__nv_bfloat16>(__nv_bfloat16 value) +{ + return __bfloat162float(value); +} + +} // namespace + +template +__device__ void cudaCoreGemmImpl(ActType const* __restrict__ act, __nv_fp4_e2m1 const* __restrict__ weight, + ScaleType const* __restrict__ weightScale, float const weightGlobalScale, OutputType* __restrict__ output, + SizeType32 m, SizeType32 n, SizeType32 k) +{ + using VecType = int4; + using ScaleVecType = __nv_fp8x2_e4m3; + using CvtWeightType = + typename tensorrt_llm::kernels::cutlass_kernels::TllmToCutlassTypeAdapter<__nv_fp4_e2m1>::type; + using Converter = cutlass::NumericArrayConverter; + using CvtSrcType = typename Converter::source_type; + using CvtResType = typename Converter::result_type; + + static constexpr SizeType32 kNvfp4ScaleGranularity = 16; + static constexpr SizeType32 kStepK = 32; + static constexpr SizeType32 kStepKScale = kStepK / kNvfp4ScaleGranularity; + static constexpr SizeType32 kTileK = kStepK * kBlockSize; + static constexpr SizeType32 kCvtCount = static_cast(sizeof(VecType) / sizeof(CvtSrcType)); + + static_assert(kStepK % kNvfp4ScaleGranularity == 0); + + auto const tileIdM = static_cast(blockIdx.x * kTileM); + auto const tileIdN = static_cast(blockIdx.y * kTileN); + auto const tid = static_cast(threadIdx.x); + (void) m; + + float tileAct[kStepK]; + float tileWeight[kTileN * kStepK]; + float tileWeightScale[kTileN * kStepKScale]; + float acc[kTileM * kTileN]; + +#pragma unroll + for (SizeType32 i = 0; i < kTileM * kTileN; ++i) + { + acc[i] = 0.0F; + } + + act += tileIdM * k; + weight += tileIdN * k / 2; + output += tileIdM * n + tileIdN; + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaGridDependencySynchronize(); +#endif + + SizeType32 const numColsSf = k / kNvfp4ScaleGranularity; + SizeType32 const numSfTilesK = (numColsSf + 4 - 1) / 4; + + for (SizeType32 idxK = tid * kStepK; idxK < k; idxK += kTileK) + { +#pragma unroll + for (SizeType32 j = 0; j < kTileN; ++j) + { + auto tileWeightQuantized = reinterpret_cast(weight + (j * k + idxK) / 2)[0]; +#pragma unroll + for (SizeType32 cvtIdx = 0; cvtIdx < kCvtCount; ++cvtIdx) + { + reinterpret_cast(tileWeight)[j * kCvtCount + cvtIdx] + = Converter::convert(reinterpret_cast(&tileWeightQuantized)[cvtIdx]); + } + } + +#pragma unroll + for (SizeType32 j = 0; j < kTileN; ++j) + { + SizeType32 const rowIdx = tileIdN + j; + SizeType32 const colIdx = idxK / kNvfp4ScaleGranularity; + SizeType32 const tileOffset = ((rowIdx / 128) * numSfTilesK + colIdx / 4) * 512; + SizeType32 const dstIdx = tileOffset + (rowIdx % 32) * 16 + ((rowIdx % 128) / 32) * 4 + colIdx % 4; + auto const tileWeightScaleFp8x2 = reinterpret_cast(weightScale + dstIdx)[0]; + char2 const tmp = reinterpret_cast(tileWeightScaleFp8x2); + tileWeightScale[j * kStepKScale + 0] = static_cast(reinterpret_cast<__nv_fp8_e4m3 const&>(tmp.x)); + tileWeightScale[j * kStepKScale + 1] = static_cast(reinterpret_cast<__nv_fp8_e4m3 const&>(tmp.y)); + } + +#pragma unroll + for (SizeType32 i = 0; i < kTileM; ++i) + { +#pragma unroll + for (SizeType32 l = 0; l < kStepK; ++l) + { + tileAct[l] = toFloat(act[i * k + idxK + l]); + } + +#pragma unroll + for (SizeType32 j = 0; j < kTileN; ++j) + { +#pragma unroll + for (SizeType32 l = 0; l < kStepK; ++l) + { + float const scaledWeight = tileWeight[j * kStepK + l] + * tileWeightScale[j * kStepKScale + l / kNvfp4ScaleGranularity] * weightGlobalScale; + acc[i * kTileN + j] = fma(tileAct[l], scaledWeight, acc[i * kTileN + j]); + } + } + } + } + + using WarpReduce = cub::WarpReduce; + static constexpr SizeType32 kWarpSize = 32; + static constexpr SizeType32 kWarpNum = kBlockSize / kWarpSize; + SizeType32 const warpId = tid / kWarpSize; + SizeType32 const laneId = tid % kWarpSize; + __shared__ float shmem[kTileM * kTileN * kWarpNum]; + __shared__ typename WarpReduce::TempStorage tempStorage[kWarpNum]; + +#pragma unroll + for (SizeType32 mi = 0; mi < kTileM; ++mi) + { +#pragma unroll + for (SizeType32 ni = 0; ni < kTileN; ++ni) + { + float const val = WarpReduce(tempStorage[warpId]).Sum(acc[mi * kTileN + ni]); + if (laneId == 0) + { + shmem[mi * kTileN + ni + warpId * kTileM * kTileN] = val; + } + } + } + __syncthreads(); + + for (SizeType32 ii = tid; ii < kTileM * kTileN; ii += kBlockSize) + { + SizeType32 const mid = ii / kTileN; + SizeType32 const nid = ii % kTileN; + float val = 0.0F; +#pragma unroll + for (SizeType32 jj = 0; jj < kWarpNum; ++jj) + { + val += shmem[jj * kTileM * kTileN + ii]; + } + output[mid * n + nid] = static_cast(val); + } + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +template +__global__ void cudaCoreGemm(ActType const* __restrict__ act, __nv_fp4_e2m1 const* __restrict__ weight, + ScaleType const* __restrict__ weightScale, float const* weightGlobalScale, OutputType* __restrict__ output, + SizeType32 m, SizeType32 n, SizeType32 k) +{ + cudaCoreGemmImpl( + act, weight, weightScale, weightGlobalScale[0], output, m, n, k); +} + +template +void cudaCoreGemmKernel(Params const& params, cudaStream_t stream) +{ + dim3 const block(kBlockSize); + dim3 const grid(params.m / kTileM, params.n / kTileN); + cudaCoreGemm<<>>( + reinterpret_cast(params.act), reinterpret_cast<__nv_fp4_e2m1 const*>(params.weight), + reinterpret_cast(params.weightScale), params.weightGlobalScale, + reinterpret_cast(params.output), params.m, params.n, params.k); +} + +template +bool cudaCoreGemmTemplateCaller(Params const& params, cudaStream_t stream) +{ + constexpr int kCudaCoreGemmTemplateMaxM = 16; + if (params.m == kTileM) + { + cudaCoreGemmKernel(params, stream); + return true; + } + if constexpr (kTileM < kCudaCoreGemmTemplateMaxM) + { + return cudaCoreGemmTemplateCaller( + params, stream); + } + return false; +} + +template +bool cudaCoreGemmLauncher(Params const& params, cudaStream_t stream) +{ + return cudaCoreGemmTemplateCaller(params, stream); +} + +template +bool dispatchOutputType(Params const& params, cudaStream_t stream) +{ + if (params.outputType == CUDA_R_16F) + { + return cudaCoreGemmLauncher(params, stream); + } + if (params.outputType == CUDA_R_16BF) + { + return cudaCoreGemmLauncher(params, stream); + } + if (params.outputType == CUDA_R_32F) + { + return cudaCoreGemmLauncher(params, stream); + } + return false; +} + +bool cudaCoreGemmDispatcher(Params const& params, cudaStream_t stream) +{ + bool dispatched = true; + int const smVersion = tensorrt_llm::common::getSMVersion(); + if (smVersion != 120 && smVersion != 121) + { + dispatched = false; + } + else if (params.n % 2 != 0 || params.k % 32 != 0) + { + dispatched = false; + } + else if (params.m < 1 || params.m > 16) + { + dispatched = false; + } + else if (params.weightScale == nullptr || params.weightGlobalScale == nullptr) + { + dispatched = false; + } + else if (params.inputType == CUDA_R_16F) + { + dispatched = dispatchOutputType(params, stream); + } + else if (params.inputType == CUDA_R_16BF) + { + dispatched = dispatchOutputType<__nv_bfloat16>(params, stream); + } + else + { + dispatched = false; + } + + if (!dispatched) + { + TLLM_LOG_WARNING( + "tensorrt_llm::kernels::cuda_core_gemm_w4a16_nvfp4::cudaCoreGemmDispatcher [NOT DISPATCHED], " + "inputType=%d, outputType=%d, m=%d, n=%d, k=%d, sm=%d", + params.inputType, params.outputType, params.m, params.n, params.k, smVersion); + } + return dispatched; +} + +} // namespace cuda_core_gemm_w4a16_nvfp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.h b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.h new file mode 100644 index 000000000000..2f525c50c9b7 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.h @@ -0,0 +1,71 @@ +/* + * Copyright (c) 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. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/runtime/common.h" + +#include + +#include +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace cuda_core_gemm_w4a16_nvfp4 +{ +using SizeType32 = tensorrt_llm::runtime::SizeType32; + +struct Params +{ + void const* act; + void const* weight; + void const* weightScale; + float const* weightGlobalScale; + void* output; + SizeType32 m, n, k; + cudaDataType_t inputType; + cudaDataType_t outputType; + + Params(void const* act_, void const* weight_, void const* weightScale_, float const* weightGlobalScale_, + void* output_, SizeType32 m_, SizeType32 n_, SizeType32 k_, cudaDataType_t inputType_, + cudaDataType_t outputType_) + : act(act_) + , weight(weight_) + , weightScale(weightScale_) + , weightGlobalScale(weightGlobalScale_) + , output(output_) + , m(m_) + , n(n_) + , k(k_) + , inputType(inputType_) + , outputType(outputType_) + { + } +}; + +bool cudaCoreGemmDispatcher(Params const& params, cudaStream_t stream); + +} // namespace cuda_core_gemm_w4a16_nvfp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/nanobind/bindings.cpp b/cpp/tensorrt_llm/nanobind/bindings.cpp index db263fa639f9..5d9607a17181 100644 --- a/cpp/tensorrt_llm/nanobind/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/bindings.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -273,6 +273,7 @@ NB_MODULE(TRTLLM_NB_MODULE, m) .def_prop_ro("has_w4a8_mxfp4_mxfp8", &tc::QuantMode::hasW4a8Mxfp4Mxfp8) .def_prop_ro("has_w4a16_mxfp4", &tc::QuantMode::hasW4a16Mxfp4) + .def_prop_ro("has_w4a16_nvfp4", &tc::QuantMode::hasW4a16Nvfp4) .def_prop_ro("has_kv_cache_quant", &tc::QuantMode::hasKvCacheQuant) .def_static("from_description", &tc::QuantMode::fromDescription, nb::arg("quantize_weights"), diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index 6849415992b7..119fad3b9a2d 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -50,6 +50,7 @@ add_library( cublasFp4ScaledMM.cpp cudaNvfp4MM.cpp cudaScaledMM.cpp + w4a16Nvfp4Gemm.cpp dynamicDecodeOp.cpp fmhaPackMaskOp.cpp fp8Op.cpp diff --git a/cpp/tensorrt_llm/thop/w4a16Nvfp4Gemm.cpp b/cpp/tensorrt_llm/thop/w4a16Nvfp4Gemm.cpp new file mode 100644 index 000000000000..797593a4b5db --- /dev/null +++ b/cpp/tensorrt_llm/thop/w4a16Nvfp4Gemm.cpp @@ -0,0 +1,113 @@ +/* + * Copyright (c) 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. + * 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. + */ + +#include "tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include +#include + +using torch::Tensor; + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ +namespace +{ + +void checkActDtype(Tensor const& act) +{ + TORCH_CHECK(act.scalar_type() == torch::kFloat16 || act.scalar_type() == torch::kBFloat16, + "w4a16_nvfp4_gemm only supports FP16/BF16 activations, got ", act.scalar_type()); +} + +void w4a16Nvfp4GemmCaller( + Tensor& out, Tensor const& act, Tensor const& weight, Tensor const& weightScale, Tensor const& weightScale2) +{ + auto const m = static_cast(act.sizes()[0]); + auto const k = static_cast(act.sizes()[1]); + auto const n = static_cast(weight.sizes()[0]); + TORCH_CHECK(weight.sizes()[1] * 2 == k, "weight shape [N, K/2] must match activation shape [M, K]"); + + auto stream = at::cuda::getCurrentCUDAStream(act.get_device()); + + auto* actPtr = static_cast(act.data_ptr()); + auto* weightPtr = static_cast(weight.data_ptr()); + auto* weightScalePtr = static_cast(weightScale.data_ptr()); + auto* weightScale2Ptr = static_cast(weightScale2.data_ptr()); + auto* outPtr = static_cast(out.data_ptr()); + + auto const inputType = convert_torch_dtype(act.scalar_type()); + auto const outType = convert_torch_dtype(out.scalar_type()); + + tensorrt_llm::kernels::cuda_core_gemm_w4a16_nvfp4::Params params( + actPtr, weightPtr, weightScalePtr, weightScale2Ptr, outPtr, m, n, k, inputType, outType); + bool const dispatched = tensorrt_llm::kernels::cuda_core_gemm_w4a16_nvfp4::cudaCoreGemmDispatcher(params, stream); + TORCH_CHECK(dispatched, "Failed to dispatch w4a16_nvfp4_gemm CUDA-core kernel"); +} + +} // namespace + +Tensor& w4a16_nvfp4_gemm_out(Tensor const& act, Tensor const& weight, Tensor const& weightScale, + Tensor const& weightScale2, std::optional outDtype, std::optional const& bias, Tensor& out) +{ + CHECK_TH_CUDA(act); + CHECK_CONTIGUOUS(act); + checkActDtype(act); + CHECK_INPUT(weight, FLOAT4_E2M1X2); + CHECK_INPUT(weightScale, SF_DTYPE); + CHECK_INPUT(weightScale2, torch::kFloat32); + CHECK_TH_CUDA(out); + CHECK_CONTIGUOUS(out); + + TORCH_CHECK(act.dim() == 2 && weight.dim() == 2 && out.dim() == 2); + TORCH_CHECK(act.sizes()[0] == out.sizes()[0]); + TORCH_CHECK(weight.sizes()[0] == out.sizes()[1]); + TORCH_CHECK(weight.sizes()[1] * 2 == act.sizes()[1]); + TORCH_CHECK(weightScale2.numel() == 1, "weight_scale_2 must be a scalar tensor"); + TORCH_CHECK(!bias.has_value(), "w4a16_nvfp4_gemm does not support bias"); + TORCH_CHECK(!outDtype.has_value() || out.scalar_type() == outDtype.value()); + + w4a16Nvfp4GemmCaller(out, act, weight, weightScale, weightScale2); + return out; +} + +Tensor w4a16_nvfp4_gemm(Tensor const& act, Tensor const& weight, Tensor const& weightScale, Tensor const& weightScale2, + std::optional outDtype, std::optional const& bias) +{ + TORCH_CHECK(act.dim() == 2 && weight.dim() == 2); + auto const outDtypeValue = outDtype.value_or(act.scalar_type()); + std::vector outputSize = {act.sizes()[0], weight.sizes()[0]}; + Tensor out = at::empty(outputSize, act.options().dtype(outDtypeValue)); + return w4a16_nvfp4_gemm_out(act, weight, weightScale, weightScale2, outDtype, bias, out); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "w4a16_nvfp4_gemm(Tensor act, Tensor weight, Tensor weight_scale, Tensor weight_scale_2, ScalarType? " + "out_dtype, Tensor? bias=None) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("w4a16_nvfp4_gemm", &tensorrt_llm::torch_ext::w4a16_nvfp4_gemm); +} diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 644155e9a4a7..ac159a009539 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -1044,6 +1044,16 @@ def _(mat_a: torch.Tensor, n = mat_b.shape[0] return mat_a.new_empty((m, n), dtype=out_dtype) + @torch.library.register_fake("trtllm::w4a16_nvfp4_gemm") + def _(act: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + weight_scale_2: torch.Tensor, + out_dtype: Optional[torch.dtype], + bias: Optional[torch.Tensor] = None): + del weight_scale, weight_scale_2, bias + return act.new_empty((act.shape[0], weight.shape[0]), dtype=out_dtype) + @torch.library.register_fake("trtllm::mla_rope_generation") def _( fused_q: torch.Tensor, diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index f2a15151d781..23e0eb4dbae8 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -24,7 +24,7 @@ from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.models.quant_config_utils import \ - update_quant_config_from_compressed_tensors + update_quant_config_from_compressed_tensors, is_w4a16_nvfp4_hf_quant_config from tensorrt_llm.quantization.mode import QuantAlgo from tensorrt_llm.quantization.modelopt_config import ( is_modelopt_quant_config, read_modelopt_quant_config, @@ -314,16 +314,16 @@ def resolve_moe_backend(moe_backend: str, @staticmethod def load_modelopt_quant_config(quant_config_file, checkpoint_dir, - moe_backend): + moe_backend, hf_quant_config=None): with open(quant_config_file) as f: quant_config_dict = json.load(f) return ModelConfig._build_modelopt_quant_config( read_modelopt_quant_config(quant_config_dict), checkpoint_dir, - moe_backend) + moe_backend, hf_quant_config) @staticmethod def _build_modelopt_quant_config(json_quant_configs, checkpoint_dir, - moe_backend): + moe_backend, hf_quant_config=None): """Build (quant_config, layer_quant_config) from a normalized modelopt 'quantization' inner dict. ``json_quant_configs`` should be a dict as produced by @@ -348,6 +348,13 @@ def _build_modelopt_quant_config(json_quant_configs, checkpoint_dir, if 'pre_quant_scale' in json_quant_configs: quant_config.pre_quant_scale = json_quant_configs['pre_quant_scale'] + if (quant_config.quant_algo in (QuantAlgo.NVFP4, "NVFP4") and + is_w4a16_nvfp4_hf_quant_config(hf_quant_config)): + quant_config.quant_algo = QuantAlgo.W4A16_NVFP4 + quant_config.group_size = 16 + quant_config.exclude_modules = hf_quant_config.get( + "ignore", quant_config.exclude_modules) + if quant_config.quant_algo == QuantAlgo.MIXED_PRECISION: json_extended_quant_configs: dict = {} # See tests/unittest/llmapi/test_llm_quant.py @@ -536,6 +543,7 @@ def override_quant_algo(): new_algo = os.environ.get("OVERRIDE_QUANT_ALGO", None) supported_algos = { "W4A16_MXFP4": QuantAlgo.W4A16_MXFP4, + "W4A16_NVFP4": QuantAlgo.W4A16_NVFP4, "W4A8_MXFP4_MXFP8": QuantAlgo.W4A8_MXFP4_MXFP8, "W4A8_MXFP4_FP8": QuantAlgo.W4A8_MXFP4_FP8, } @@ -692,6 +700,9 @@ def _recursive_update_config(config: transformers.PretrainedConfig, moe_backend_hint = cls.resolve_moe_backend(requested_moe_backend, architecture) + hf_quant_config = getattr(pretrained_config, "quantization_config", + None) + # quantized ckpt in modelopt format if quant_config_file := cached_file(checkpoint_dir, 'hf_quant_config.json'): @@ -706,7 +717,7 @@ def _recursive_update_config(config: transformers.PretrainedConfig, source_file="hf_quant_config.json", ) quant_config, layer_quant_config = cls._build_modelopt_quant_config( - normalized, checkpoint_dir, moe_backend_hint) + normalized, checkpoint_dir, moe_backend_hint, hf_quant_config=hf_quant_config) # quantized ckpt in other formats elif getattr(pretrained_config, "quantization_config", None) is not None: diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py index 852d8cb4b5d3..b8788373d88d 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py @@ -5,6 +5,7 @@ HfWeightMapper from tensorrt_llm._torch.models.modeling_utils import register_mapper from tensorrt_llm._torch.utils import split +from tensorrt_llm.quantization.mode import QuantAlgo @register_mapper("HF", "NemotronHPuzzleForCausalLM") @@ -40,11 +41,24 @@ def _split_mamba2_mixer_in_proj(w: torch.Tensor) -> torch.Tensor: w = torch.concat(w).contiguous() return w - is_nvfp4 = self.config.quant_config.quant_algo == "NVFP4" + quant_algo = getattr(self.config.quant_config, "quant_algo", None) + is_nvfp4 = quant_algo in (QuantAlgo.NVFP4, QuantAlgo.W4A16_NVFP4, + "NVFP4", "W4A16_NVFP4") n_groups = config.n_groups d_state = config.ssm_state_size nheads = config.mamba_num_heads + def _canonicalize_quant_key(key: str) -> str: + replacements = { + ".weight_packed": ".weight", + ".weight_global_scale": ".weight_scale_2", + ".input_global_scale": ".input_scale_2", + } + for suffix, replacement in replacements.items(): + if key.endswith(suffix): + return f"{key[:-len(suffix)]}{replacement}" + return key + new_weights = {} for name, _ in weights.items(): key = name @@ -70,6 +84,8 @@ def _split_mamba2_mixer_in_proj(w: torch.Tensor) -> torch.Tensor: if "A_log" in key: key = key.replace("A_log", "A") + key = _canonicalize_quant_key(key) + if ("mixer.in_proj" in key or "mixer.out_proj" in key) and "_scale" in key: # Special handing for nvfp4 Mamba2 mixer in_proj.weight_scale. diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index 574c74d9b190..1309339a54a9 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py @@ -31,6 +31,7 @@ from tensorrt_llm._utils import get_sm_version from tensorrt_llm.logger import logger from tensorrt_llm.lora_helper import LoraConfig +from tensorrt_llm.quantization.mode import QuantAlgo from ..attention_backend import AttentionMetadata from ..distributed import AllReduce, AllReduceFusionOp, AllReduceParams @@ -155,6 +156,28 @@ def forward( **kwargs) +def _get_nemotron_h_moe_model_config( + model_config: ModelConfig[PretrainedConfig], + layer_idx: int) -> ModelConfig[PretrainedConfig]: + # Per-expert mixed precision config is more specific than the global config. + if model_config.quant_config_dict is not None: + experts_prefix = f"model.layers.{layer_idx}.mixer.experts." + for key, cfg in model_config.quant_config_dict.items(): + if key.startswith(experts_prefix): + return replace(model_config, quant_config=cfg) + + quant_config = model_config.quant_config + if (quant_config is not None and quant_config.quant_algo + in (QuantAlgo.W4A16_NVFP4, "W4A16_NVFP4")): + moe_quant_config = quant_config.model_copy(deep=True) + moe_quant_config.quant_algo = QuantAlgo.NVFP4 + moe_quant_config.__dict__.pop("quant_mode", None) + moe_quant_config.__dict__.pop("layer_quant_mode", None) + return replace(model_config, quant_config=moe_quant_config) + + return model_config + + # Ref code: https://huggingface.co/nvidia/Nemotron-Nano-3-30B-A3.5B-dev-1024/blob/main/modeling_nemotron_h.py#L818 class NemotronHMOE(nn.Module): @@ -240,18 +263,8 @@ def _moe(name): moe_backend=model_config.moe_backend, ) - # For MIXED_PRECISION models, the global quant_config has quant_algo=MIXED_PRECISION - # which maps to QuantMode(0) (no quant). This would cause the MoE backend to select - # UnquantizedFusedMoEMethod and allocate BF16 weight buffers, causing a shape mismatch - # when loading NVFP4/W4A8_NVFP4_FP8 quantized expert weights. - # Look up the per-expert quant config from quant_config_dict and use it for create_moe. - moe_model_config = model_config - if model_config.quant_config_dict is not None: - experts_prefix = f"model.layers.{layer_idx}.mixer.experts." - for key, cfg in model_config.quant_config_dict.items(): - if key.startswith(experts_prefix): - moe_model_config = replace(model_config, quant_config=cfg) - break + moe_model_config = _get_nemotron_h_moe_model_config( + model_config, layer_idx) # Setup MoE experts. self.experts = create_moe( @@ -1117,14 +1130,16 @@ def _get_mtp_sublayer_quant_config(self, model_config: NemotronHModelConfig, layer_idx: int): """ Get quantization config for MTP sublayer. - The MTP layer in the nvfp4 checkpoint is unquantized. Because the TRTLLM - moe_backend only supports fp8/fp4 quantization, we need to override - the quant_config for the MTP layer. + The Nano3.5 NVFP4 W4A16 checkpoint stores the MTP body tensors in BF16. + The shared MTP head still receives the checkpoint-backed lm_head, so its + logits path keeps the lm_head precision instead of inheriting this + sublayer override. """ from tensorrt_llm.models.modeling_utils import QuantConfig quant_config = model_config.quant_config - # MTP layers are always unquantized, force quant_algo=None + # This checkpoint's MTP body is unquantized, so force quant_algo=None + # only for the MTP sublayers constructed here. if quant_config is None: return None return QuantConfig( diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 1b79818d4288..0fae7075ed13 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -18,6 +18,7 @@ from ...logger import logger from ...models.modeling_utils import QuantConfig +from ...quantization.mode import QuantAlgo from ..attention_backend import AttentionMetadata from ..distributed.communicator import pp_recv_tensors, pp_send_tensors from ..model_config import ModelConfig, TConfig @@ -368,16 +369,20 @@ def __init__(self, model: TModel, *, config: ModelConfig[TConfig], self.pp_rank = config.mapping.pp_rank self.pp_size = config.mapping.pp_size self.has_custom_lm_head = False + lm_head_quant_config = None + if (config.quant_config is not None and config.quant_config.quant_algo + in (QuantAlgo.W4A16_NVFP4, "W4A16_NVFP4") and not config. + quant_config.is_module_excluded_from_quantization("lm_head")): + lm_head_quant_config = config.quant_config if config.mapping.enable_attention_dp and not config.mapping.enable_lm_head_tp_in_adp: self.lm_head = LMHead( vocab_size, hidden_size, dtype=config.pretrained_config.torch_dtype, + quant_config=lm_head_quant_config, ) else: - # TODO(zhenhuanc): Currently lm_head Linear will not accept QuantConfig - # will considering per layer QuantConfig in the future. if (hasattr(config, 'lora_config') and config.lora_config is not None and len(config.lora_config.lora_dir) == 1): @@ -399,6 +404,7 @@ def __init__(self, model: TModel, *, config: ModelConfig[TConfig], reduce_output=False, use_custom_cublas_mm=getattr(model, 'use_custom_cublas_mm', False), + quant_config=lm_head_quant_config, ) if self.has_custom_lm_head: diff --git a/tensorrt_llm/_torch/modules/embedding.py b/tensorrt_llm/_torch/modules/embedding.py index 900bdf02af8a..0d878dee4cfd 100644 --- a/tensorrt_llm/_torch/modules/embedding.py +++ b/tensorrt_llm/_torch/modules/embedding.py @@ -9,6 +9,7 @@ from tensorrt_llm.mapping import Mapping from tensorrt_llm.math_utils import ceil_div +from ...models.modeling_utils import QuantConfig from ..distributed import allgather from .linear import Linear, TensorParallelMode @@ -22,6 +23,8 @@ class LMHead(Linear): dtype (Optional[torch.dtype]): type of the parameters. mapping (Optional[Mapping]): parallelism configuration. If not provided, the embedding is not parallelized. + quant_config (Optional[QuantConfig]): quantization configuration for + the lm head projection. """ def __init__( @@ -34,6 +37,7 @@ def __init__( gather_output: bool = False, reduce_output: bool = True, use_custom_cublas_mm: bool = False, + quant_config: Optional[QuantConfig] = None, ): local_in_features = embedding_dim local_out_features = num_embeddings @@ -66,6 +70,7 @@ def __init__( gather_output=gather_output, reduce_output=reduce_output, use_custom_cublas_mm=use_custom_cublas_mm, + quant_config=quant_config, ) if tensor_parallel_mode == TensorParallelMode.ROW: @@ -76,9 +81,10 @@ def __init__( self.num_embeddings = num_embeddings self.embedding_dim = embedding_dim - weight_shape = (self.out_features, self.in_features) - self.weight = Parameter(torch.empty(weight_shape, dtype=dtype)) - self.register_parameter("bias", None) + if not self.has_any_quant: + weight_shape = (self.out_features, self.in_features) + self.weight = Parameter(torch.empty(weight_shape, dtype=dtype)) + self.register_parameter("bias", None) @property def vocab_size_padded(self) -> int: @@ -126,7 +132,7 @@ def load_weights(self, weights: List[Dict], allow_partial_loading: bool = False): original_weight = None - if self.tp_mode == TensorParallelMode.COLUMN: + if self.tp_mode == TensorParallelMode.COLUMN and not self.has_any_quant: if self.tp_rank == self.tp_size - 1 and self.padding_size > 0: original_weight = self.weight.data.zero_() self.weight.data = self.weight[:-self.padding_size, :] diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 2d3482fe1b34..043df5ced20d 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -1861,6 +1861,75 @@ def post_load_weights(self, module: Linear): module.rebuild_tensor_metadata) +class W4A16NVFP4LinearMethod(NVFP4LinearMethod): + + def create_weights(self, module: Linear, in_features: int, + out_features: int, bias: bool, dtype: torch.dtype): + module.scaling_vector_size = 16 + assert in_features % module.scaling_vector_size == 0, ( + f"in_features {in_features} must be divisible by scaling_vector_size {module.scaling_vector_size}" + ) + + module.weight = Parameter(torch.empty([out_features, in_features // 2], + dtype=fp4_utils.float4_e2m1x2), + requires_grad=False) + + nrows = fp4_utils.pad_up(out_features, 128) + ncols = fp4_utils.pad_up(in_features // module.scaling_vector_size, 4) + module.weight_scale = Parameter(torch.empty( + [nrows * ncols], dtype=fp4_utils.float4_sf_dtype), + requires_grad=False) + + module.weight_scale_2 = Parameter(torch.empty([1], dtype=torch.float32), + requires_grad=False) + + module.input_scale = None + module.inv_input_scale = None + module.alpha = None + module.pre_quant_scale = None + + module.kv_scales = Parameter(torch.ones(3, dtype=torch.float32), + requires_grad=False) + module.inv_kv_scales = Parameter(torch.ones(3, dtype=torch.float32), + requires_grad=False) + + if bias: + module.bias = Parameter(torch.empty((out_features), dtype=dtype), + requires_grad=False) + else: + module.register_parameter("bias", None) + + def apply(self, module: Linear, input: torch.Tensor, + bias: Optional[torch.Tensor]): + original_shape = None + if input.dim() > 2: + original_shape = input.shape + input = input.reshape(-1, input.shape[-1]) + + if module.pre_quant_scale is not None: + assert input.dtype == module.pre_quant_scale.dtype, "Input dtype and pre_quant_scale dtype must match" + input = input * module.pre_quant_scale + + output = torch.ops.trtllm.w4a16_nvfp4_gemm( + input, + module.weight, + module.weight_scale, + module.weight_scale_2, + module.dtype, + bias=None, + ) + + if output.shape[-1] > module.out_features: + output = output[..., :module.out_features].contiguous() + + if original_shape is not None: + output = output.reshape(*original_shape[:-1], output.shape[-1]) + + if bias is not None: + output = output + bias + return output + + class W4A8NVFP4FP8LinearMethod(LinearMethodBase): def create_weights(self, module: Linear, in_features: int, @@ -2720,6 +2789,8 @@ def get_quant_method(quant_config: Optional[QuantConfig] = None): return NVFP4ARCLinearMethod() else: return NVFP4LinearMethod() + if quant_config.layer_quant_mode.has_w4a16_nvfp4(): + return W4A16NVFP4LinearMethod() if quant_config.layer_quant_mode.has_w4a8_nvfp4_fp8(): return W4A8NVFP4FP8LinearMethod() if quant_config.layer_quant_mode.has_w4a8_mxfp4_fp8(): diff --git a/tensorrt_llm/models/quant_config_utils.py b/tensorrt_llm/models/quant_config_utils.py index ffebe91ec786..e946e5e70d7a 100644 --- a/tensorrt_llm/models/quant_config_utils.py +++ b/tensorrt_llm/models/quant_config_utils.py @@ -18,6 +18,22 @@ from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization.mode import QuantAlgo +def is_w4a16_nvfp4_hf_quant_config(hf_quant_config): + if hf_quant_config is None: + return False + config_groups = hf_quant_config.get("config_groups") + if config_groups is None: + return False + + group_config = config_groups.get("group_0") + if group_config is None: + return False + weights_quant_config = group_config.get("weights", {}) + inputs_quant_config = group_config.get("input_activations") + return (hf_quant_config.get("format") == "nvfp4-pack-quantized" + and weights_quant_config.get("num_bits") == 4 + and weights_quant_config.get("group_size") == 16 + and inputs_quant_config is None) def update_quant_config_from_compressed_tensors( quant_config: QuantConfig, hf_quant_config: Mapping[str, Any] @@ -32,7 +48,10 @@ def update_quant_config_from_compressed_tensors( weights_quant_strategy = weights_quant_config["strategy"] inputs_quant_strategy = inputs_quant_config["strategy"] - if weights_quant_config["num_bits"] == 8: + if is_w4a16_nvfp4_hf_quant_config(hf_quant_config): + quant_config.quant_algo = QuantAlgo.W4A16_NVFP4 + quant_config.group_size = 16 + elif weights_quant_config["num_bits"] == 8: if weights_quant_strategy == "channel": if inputs_quant_strategy != "token": raise ValueError(f"Unsupported inputs_quant_strategy: {inputs_quant_strategy}.") diff --git a/tensorrt_llm/quantization/mode.py b/tensorrt_llm/quantization/mode.py index e4e8fbc89d17..b2945000e2d6 100644 --- a/tensorrt_llm/quantization/mode.py +++ b/tensorrt_llm/quantization/mode.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -44,6 +44,7 @@ class QuantAlgo(StrEnum, metaclass=BaseEnumMeta): W4A8_MXFP4_FP8 = auto() W4A8_MXFP4_MXFP8 = auto() W4A16_MXFP4 = auto() + W4A16_NVFP4 = auto() NVFP4_AWQ = auto() NVFP4_ARC = auto() NO_QUANT = auto() @@ -99,6 +100,7 @@ class QuantMode(IntFlag): W4A8_MXFP4_FP8 = auto() W4A8_MXFP4_MXFP8 = auto() W4A16_MXFP4 = auto() + W4A16_NVFP4 = auto() # The smallest power-of-two that is not used by a flag. Do not call auto() after that line. COUNT = auto() @@ -196,6 +198,9 @@ def has_w4a8_mxfp4_mxfp8(self): def has_w4a16_mxfp4(self): return self._any(self.W4A16_MXFP4) + def has_w4a16_nvfp4(self): + return self._any(self.W4A16_NVFP4) + def has_mxfp4(self): return self._any(self.W4A8_MXFP4_FP8 | self.W4A8_MXFP4_MXFP8 | self.W4A16_MXFP4) @@ -214,6 +219,7 @@ def has_any_quant(self, exclude_kv_cache: bool = False): | self.W4A8_NVFP4_FP8 | self.W4A8_MXFP4_FP8 | self.W4A16_MXFP4 + | self.W4A16_NVFP4 | self.W4A8_MXFP4_MXFP8) if exclude_kv_cache: return has_quant @@ -253,7 +259,8 @@ def from_description(quantize_weights=False, use_w4a8_qserve=False, use_w4a8_mxfp4_fp8=False, use_w4a8_mxfp4_mxfp8=False, - use_w4a16_mxfp4=False): + use_w4a16_mxfp4=False, + use_w4a16_nvfp4=False): def raise_error(): raise ValueError(f"Unsupported combination of QuantMode args: " @@ -272,7 +279,8 @@ def raise_error(): f"{use_w4a8_qserve=}, " f"{use_w4a8_mxfp4_fp8=}, " f"{use_w4a8_mxfp4_mxfp8=}, " - f"{use_w4a16_mxfp4=}") + f"{use_w4a16_mxfp4=}, " + f"{use_w4a16_nvfp4=}") # We must quantize weights when we quantize activations. if quantize_activations and not quantize_weights: @@ -339,6 +347,9 @@ def raise_error(): if use_w4a16_mxfp4: mode = mode | QuantMode.W4A16_MXFP4 + if use_w4a16_nvfp4: + mode = mode | QuantMode.W4A16_NVFP4 + return mode @staticmethod @@ -426,6 +437,8 @@ def from_quant_algo( quant_mode = QuantMode.from_description(use_w4a8_mxfp4_mxfp8=True) elif quant_algo == QuantAlgo.W4A16_MXFP4: quant_mode = QuantMode.from_description(use_w4a16_mxfp4=True) + elif quant_algo == QuantAlgo.W4A16_NVFP4: + quant_mode = QuantMode.from_description(use_w4a16_nvfp4=True) else: quant_mode = QuantMode(0) diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py new file mode 100644 index 000000000000..5c9f5f6a02eb --- /dev/null +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py @@ -0,0 +1,74 @@ +# 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 types import SimpleNamespace +from unittest.mock import patch + +import torch +from torch import nn + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_nemotron_h import NemotronHMOE +from tensorrt_llm._torch.utils import AuxStreamType +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + + +def _make_nemotron_h_moe_config(quant_config: QuantConfig) -> ModelConfig: + return ModelConfig( + pretrained_config=SimpleNamespace( + hidden_size=16, + intermediate_size=32, + mlp_bias=False, + moe_intermediate_size=64, + moe_latent_size=None, + n_group=1, + n_routed_experts=4, + n_shared_experts=0, + num_experts_per_tok=1, + routed_scaling_factor=1.0, + topk_group=1, + torch_dtype=torch.float16, + ), + moe_backend="CUTLASS", + quant_config=quant_config, + ) + + +def test_nemotron_h_moe_uses_w4a4_nvfp4_expert_config_for_w4a16_checkpoint(): + quant_config = QuantConfig( + quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16, exclude_modules=["lm_head"] + ) + quant_config.mamba_ssm_cache_dtype = torch.float32 + model_config = _make_nemotron_h_moe_config(quant_config) + captured = {} + + def fake_create_moe(**kwargs): + captured["model_config"] = kwargs["model_config"] + return nn.Identity() + + with patch( + "tensorrt_llm._torch.models.modeling_nemotron_h.create_moe", side_effect=fake_create_moe + ): + with patch("torch.cuda.Event", side_effect=lambda: object()): + aux_stream_dict = {AuxStreamType.MoeShared: None} + NemotronHMOE(model_config=model_config, layer_idx=1, aux_stream_dict=aux_stream_dict) + + moe_quant_config = captured["model_config"].quant_config + assert moe_quant_config is not quant_config + assert moe_quant_config.quant_algo == QuantAlgo.NVFP4 + assert moe_quant_config.group_size == 16 + assert moe_quant_config.exclude_modules == ["lm_head"] + assert moe_quant_config.mamba_ssm_cache_dtype == torch.float32 + assert model_config.quant_config.quant_algo == QuantAlgo.W4A16_NVFP4 diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_nemotron_h_weight_mapper.py b/tests/unittest/_torch/models/checkpoints/hf/test_nemotron_h_weight_mapper.py new file mode 100644 index 000000000000..02b28bd49920 --- /dev/null +++ b/tests/unittest/_torch/models/checkpoints/hf/test_nemotron_h_weight_mapper.py @@ -0,0 +1,80 @@ +# 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 types import SimpleNamespace + +import torch + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.checkpoints.hf.nemotron_h_weight_mapper import ( + NemotronHHfWeightMapper, +) +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + + +def _make_mapper() -> NemotronHHfWeightMapper: + mapper = NemotronHHfWeightMapper() + mapper._config = ModelConfig( + pretrained_config=SimpleNamespace( + mamba_head_dim=1, + mamba_num_heads=1, + n_groups=1, + num_hidden_layers=52, + ssm_state_size=1, + ), + mapping=Mapping(), + moe_backend="CUTLASS", + quant_config=QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4), + ) + return mapper + + +def test_nemotron_h_mapper_canonicalizes_w4a16_nvfp4_checkpoint_keys(): + mapper = _make_mapper() + weights = { + "lm_head.weight_packed": torch.empty((8, 4), dtype=torch.uint8), + "lm_head.weight_scale": torch.empty((8, 1), dtype=torch.float8_e4m3fn), + "lm_head.weight_global_scale": torch.tensor(0.25, dtype=torch.float32), + } + + mapped = mapper.preprocess_weights(weights) + + assert "lm_head.weight" in mapped + assert "lm_head.weight_scale" in mapped + assert "lm_head.weight_scale_2" in mapped + assert "lm_head.weight_packed" not in mapped + assert "lm_head.weight_global_scale" not in mapped + assert mapped["lm_head.weight"] is weights["lm_head.weight_packed"] + assert mapped["lm_head.weight_scale_2"] is weights["lm_head.weight_global_scale"] + + +def test_nemotron_h_mapper_handles_scalar_w4a16_nvfp4_moe_global_scales(): + mapper = _make_mapper() + prefix = "backbone.layers.1.mixer.experts.0.up_proj" + weights = { + f"{prefix}.weight_packed": torch.empty((8, 4), dtype=torch.uint8), + f"{prefix}.weight_scale": torch.empty((8, 1), dtype=torch.float8_e4m3fn), + f"{prefix}.weight_global_scale": torch.tensor(0.25, dtype=torch.float32), + } + + mapped = mapper.preprocess_weights(weights) + + assert "model.layers.1.mixer.experts.0.w1.weight" in mapped + assert "model.layers.1.mixer.experts.0.w3.weight" in mapped + assert "model.layers.1.mixer.experts.0.w1.weight_scale_2" in mapped + assert "model.layers.1.mixer.experts.0.w3.weight_scale_2" in mapped + assert mapped["model.layers.1.mixer.experts.0.w3.weight"].shape == (0, 4) + assert mapped["model.layers.1.mixer.experts.0.w3.weight_scale_2"].shape == () diff --git a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py new file mode 100644 index 000000000000..911d0d4bf250 --- /dev/null +++ b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py @@ -0,0 +1,142 @@ +# 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 types import SimpleNamespace +from unittest.mock import patch + +import torch + +from tensorrt_llm._torch.modules.embedding import LMHead +from tensorrt_llm._torch.modules.linear import W4A16NVFP4LinearMethod, get_quant_method +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + + +def test_get_quant_method_returns_w4a16_nvfp4_linear_method(): + quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) + + method = get_quant_method(quant_config) + + assert isinstance(method, W4A16NVFP4LinearMethod) + + +def test_w4a16_nvfp4_linear_uses_high_precision_activation_without_fp4_quantize(): + method = W4A16NVFP4LinearMethod() + input_tensor = torch.ones((2, 4), dtype=torch.bfloat16) + bias = torch.tensor([1.0, 2.0, 3.0], dtype=torch.bfloat16) + module = SimpleNamespace( + weight=torch.empty((5, 2), dtype=torch.uint8), + weight_scale=torch.empty((128 * 4,), dtype=torch.uint8), + weight_scale_2=torch.tensor([0.25], dtype=torch.float32), + dtype=torch.bfloat16, + out_features=3, + pre_quant_scale=None, + ) + captured = {} + + def fake_w4a16_nvfp4_gemm( + input_arg, weight, weight_scale, weight_scale_2, out_dtype, bias=None + ): + captured["input"] = input_arg + captured["weight"] = weight + captured["weight_scale"] = weight_scale + captured["weight_scale_2"] = weight_scale_2 + captured["out_dtype"] = out_dtype + captured["bias"] = bias + return torch.ones((input_arg.shape[0], weight.shape[0]), dtype=out_dtype) + + def fail_fp4_quantize(*args, **kwargs): + raise AssertionError("W4A16 NVFP4 must not quantize activations") + + with patch("torch.ops.trtllm.w4a16_nvfp4_gemm", side_effect=fake_w4a16_nvfp4_gemm, create=True): + with patch("torch.ops.trtllm.fp4_quantize", side_effect=fail_fp4_quantize, create=True): + output = method.apply(module, input_tensor, bias) + + assert captured["input"] is input_tensor + assert captured["weight"] is module.weight + assert captured["weight_scale"] is module.weight_scale + assert captured["weight_scale_2"] is module.weight_scale_2 + assert captured["out_dtype"] is torch.bfloat16 + assert captured["bias"] is None + expected = torch.tensor([[2.0, 3.0, 4.0], [2.0, 3.0, 4.0]], dtype=torch.bfloat16) + torch.testing.assert_close(output, expected) + + +def test_w4a16_nvfp4_linear_restores_high_rank_input_shape(): + method = W4A16NVFP4LinearMethod() + input_tensor = torch.ones((2, 3, 4), dtype=torch.float16) + module = SimpleNamespace( + weight=torch.empty((7, 2), dtype=torch.uint8), + weight_scale=torch.empty((128 * 4,), dtype=torch.uint8), + weight_scale_2=torch.tensor([0.5], dtype=torch.float32), + dtype=torch.float16, + out_features=5, + pre_quant_scale=None, + ) + + def fake_w4a16_nvfp4_gemm( + input_arg, weight, weight_scale, weight_scale_2, out_dtype, bias=None + ): + assert input_arg.shape == (6, 4) + return torch.ones((input_arg.shape[0], weight.shape[0]), dtype=out_dtype) + + with patch("torch.ops.trtllm.w4a16_nvfp4_gemm", side_effect=fake_w4a16_nvfp4_gemm, create=True): + output = method.apply(module, input_tensor, bias=None) + + assert output.shape == (2, 3, 5) + + +def test_lm_head_uses_w4a16_nvfp4_quant_method_for_packed_lm_head(): + quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) + + lm_head = LMHead( + num_embeddings=3, embedding_dim=16, dtype=torch.float16, quant_config=quant_config + ) + + assert isinstance(lm_head.quant_method, W4A16NVFP4LinearMethod) + assert lm_head.weight.dtype == torch.uint8 + assert lm_head.weight.shape == (3, 8) + assert lm_head.weight_scale.shape == (128 * 4,) + assert lm_head.weight_scale_2.shape == (1,) + + +def test_lm_head_w4a16_nvfp4_forward_dispatches_to_dense_op(): + quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) + lm_head = LMHead( + num_embeddings=3, embedding_dim=16, dtype=torch.float16, quant_config=quant_config + ) + input_tensor = torch.ones((2, 16), dtype=torch.float16) + captured = {} + + def fake_w4a16_nvfp4_gemm( + input_arg, weight, weight_scale, weight_scale_2, out_dtype, bias=None + ): + captured["input"] = input_arg + captured["weight"] = weight + captured["weight_scale"] = weight_scale + captured["weight_scale_2"] = weight_scale_2 + captured["out_dtype"] = out_dtype + captured["bias"] = bias + return torch.ones((input_arg.shape[0], weight.shape[0]), dtype=out_dtype) + + with patch("torch.ops.trtllm.w4a16_nvfp4_gemm", side_effect=fake_w4a16_nvfp4_gemm, create=True): + output = lm_head(input_tensor) + + assert captured["input"] is input_tensor + assert captured["weight"] is lm_head.weight + assert captured["weight_scale"] is lm_head.weight_scale + assert captured["weight_scale_2"] is lm_head.weight_scale_2 + assert captured["out_dtype"] is torch.float16 + assert captured["bias"] is None + assert output.shape == (2, 3) diff --git a/tests/unittest/_torch/test_model_config.py b/tests/unittest/_torch/test_model_config.py index ba879df12c0d..cdfed7d9d8be 100644 --- a/tests/unittest/_torch/test_model_config.py +++ b/tests/unittest/_torch/test_model_config.py @@ -1,3 +1,4 @@ +import json import types import pytest @@ -7,6 +8,7 @@ from tensorrt_llm._torch.pyexecutor.model_loader import validate_and_set_kv_cache_quant from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig +from tensorrt_llm.quantization.mode import QuantMode def make_pretrained_config( @@ -116,3 +118,85 @@ def test_validate_and_set_kv_cache_quant_rejects_invalid_dtype(): model_config = _make_model_config_with_kv_quant(QuantAlgo.FP8) with pytest.raises(ValueError, match="Accepted types are"): validate_and_set_kv_cache_quant(model_config, "invalid_dtype") + + +def test_quant_mode_distinguishes_w4a16_nvfp4_from_w4a4_nvfp4(): + w4a4_mode = QuantMode.from_quant_algo(QuantAlgo.NVFP4) + assert w4a4_mode.has_nvfp4() + assert not w4a4_mode.has_w4a16_nvfp4() + + w4a16_mode = QuantMode.from_quant_algo(QuantAlgo.W4A16_NVFP4) + assert w4a16_mode.has_w4a16_nvfp4() + assert not w4a16_mode.has_nvfp4() + assert w4a16_mode.has_any_quant(exclude_kv_cache=True) + + +def test_load_hf_quant_config_detects_nvfp4_w4a16_compressed_tensors(): + hf_quant_config = { + "quant_method": "compressed-tensors", + "format": "nvfp4-pack-quantized", + "ignore": ["mtp.layers"], + "config_groups": { + "group_0": { + "targets": ["Linear", "lm_head"], + "weights": { + "type": "float", + "num_bits": 4, + "strategy": "group", + "group_size": 16, + }, + "input_activations": None, + }, + }, + } + + quant_config, layer_quant_config = ModelConfig.load_hf_quant_config( + hf_quant_config, moe_backend="CUTLASS" + ) + + assert quant_config.quant_algo == QuantAlgo.W4A16_NVFP4 + assert quant_config.group_size == 16 + assert quant_config.exclude_modules == ["mtp.layers"] + assert layer_quant_config is None + + +def test_load_modelopt_quant_config_respects_hf_w4a16_metadata(tmp_path): + modelopt_quant_config = { + "quantization": { + "quant_algo": "NVFP4", + "kv_cache_quant_algo": None, + "group_size": 16, + "exclude_modules": ["mtp*"], + } + } + hf_quant_config = { + "quant_method": "compressed-tensors", + "format": "nvfp4-pack-quantized", + "ignore": ["mtp.layers.0.mixer.q_proj"], + "config_groups": { + "group_0": { + "targets": ["Linear", "lm_head"], + "weights": { + "type": "float", + "num_bits": 4, + "strategy": "tensor_group", + "group_size": 16, + }, + "input_activations": None, + }, + }, + } + quant_config_file = tmp_path / "hf_quant_config.json" + quant_config_file.write_text(json.dumps(modelopt_quant_config), encoding="utf-8") + + quant_config, layer_quant_config = ModelConfig.load_modelopt_quant_config( + str(quant_config_file), + str(tmp_path), + moe_backend="CUTLASS", + hf_quant_config=hf_quant_config, + ) + + assert quant_config.quant_algo == QuantAlgo.W4A16_NVFP4 + assert quant_config.group_size == 16 + assert quant_config.exclude_modules == ["mtp.layers.0.mixer.q_proj"] + assert layer_quant_config is None From 97149b51d0d83d2b108ae4d0e3fb0962525302a1 Mon Sep 17 00:00:00 2001 From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> Date: Fri, 15 May 2026 20:11:39 +0000 Subject: [PATCH 02/25] fix cuda core w4a16 Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> --- .../cudaCoreGemmW4A16NVFP4.cu | 14 ++- tensorrt_llm/_torch/model_config.py | 24 +++++ .../_torch/models/modeling_nemotron_h.py | 15 +++- .../modules/fused_moe/fused_moe_cutlass.py | 4 +- tensorrt_llm/_torch/modules/linear.py | 59 ++++++++++++ .../test_modeling_nemotron_h_moe_quant.py | 76 +++++++++++++++- .../hf/test_nemotron_h_weight_mapper.py | 8 +- .../moe/test_cute_dsl_b12x_moe_backend.py | 17 ++++ .../_torch/modules/test_w4a16_nvfp4_linear.py | 90 +++++++++++++++++++ tests/unittest/_torch/test_model_config.py | 24 +++++ 10 files changed, 321 insertions(+), 10 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.cu b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.cu index 8ed5a5d5a318..70bc95fbdac3 100644 --- a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.cu +++ b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.cu @@ -211,6 +211,7 @@ void cudaCoreGemmKernel(Params const& params, cudaStream_t stream) reinterpret_cast(params.act), reinterpret_cast<__nv_fp4_e2m1 const*>(params.weight), reinterpret_cast(params.weightScale), params.weightGlobalScale, reinterpret_cast(params.output), params.m, params.n, params.k); + TLLM_CUDA_CHECK(cudaGetLastError()); } template @@ -233,7 +234,18 @@ bool cudaCoreGemmTemplateCaller(Params const& params, cudaStream_t stream) template bool cudaCoreGemmLauncher(Params const& params, cudaStream_t stream) { - return cudaCoreGemmTemplateCaller(params, stream); + constexpr int kDefaultTileN = 2; + constexpr int kWideTileN = 4; + constexpr int kMaxGridDimY = 65535; + if (params.n / kDefaultTileN <= kMaxGridDimY) + { + return cudaCoreGemmTemplateCaller(params, stream); + } + if (params.n % kWideTileN == 0 && params.n / kWideTileN <= kMaxGridDimY) + { + return cudaCoreGemmTemplateCaller(params, stream); + } + return false; } template diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 23e0eb4dbae8..4af19424fa01 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -312,6 +312,22 @@ def resolve_moe_backend(moe_backend: str, return "CUTLASS" + def resolve_moe_backend_after_quant_config( + moe_backend: str, architecture: str, + quant_config: QuantConfig) -> str: + """Resolve AUTO moe_backend after quantization metadata is known.""" + if moe_backend.upper() != "AUTO": + return moe_backend + + is_nemotron_h = architecture in ("NemotronHForCausalLM", + "NemotronHPuzzleForCausalLM") + is_w4a16_nvfp4 = quant_config.quant_algo in (QuantAlgo.W4A16_NVFP4, + "W4A16_NVFP4") + if is_nemotron_h and is_w4a16_nvfp4 and get_sm_version() in (120, 121): + return "FLASHINFER_NVFP4SM12X" + + return ModelConfig.resolve_moe_backend(moe_backend, architecture) + @staticmethod def load_modelopt_quant_config(quant_config_file, checkpoint_dir, moe_backend, hf_quant_config=None): @@ -692,6 +708,11 @@ def _recursive_update_config(config: transformers.PretrainedConfig, quant_config = QuantConfig() layer_quant_config = None requested_moe_backend = kwargs.get('moe_backend', 'AUTO') +<<<<<<< HEAD +======= + moe_backend = requested_moe_backend + # Resolve AUTO to specific backend based on model architecture +>>>>>>> ff309a699b (fix cuda core w4a16) architecture = pretrained_config.architectures[ 0] if pretrained_config.architectures else "" # Use an architecture-only backend hint for quant config parsing. Some @@ -736,6 +757,9 @@ def _recursive_update_config(config: transformers.PretrainedConfig, quant_config=quant_config, ) + kwargs['moe_backend'] = cls.resolve_moe_backend_after_quant_config( + requested_moe_backend, architecture, quant_config) + model_config = cls(pretrained_config=pretrained_config, quant_config=quant_config, quant_config_dict=layer_quant_config, diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index 1309339a54a9..92a6c88dbd7b 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py @@ -168,7 +168,8 @@ def _get_nemotron_h_moe_model_config( quant_config = model_config.quant_config if (quant_config is not None and quant_config.quant_algo - in (QuantAlgo.W4A16_NVFP4, "W4A16_NVFP4")): + in (QuantAlgo.W4A16_NVFP4, "W4A16_NVFP4") + and model_config.moe_backend.upper() != "FLASHINFER_NVFP4SM12X"): moe_quant_config = quant_config.model_copy(deep=True) moe_quant_config.quant_algo = QuantAlgo.NVFP4 moe_quant_config.__dict__.pop("quant_mode", None) @@ -1103,14 +1104,20 @@ def __init__( sublayer_quant_config = self._get_mtp_sublayer_quant_config( model_config, self.layer_idx) + sublayer_moe_backend = model_config.moe_backend + if (sublayer_quant_config is None + or sublayer_quant_config.quant_algo is None): + sublayer_moe_backend = "CUTLASS" # Create a model_config copy with quant_config overridden and # spec_config cleared. All other fields (use_cuda_graph, - # moe_backend, moe_max_num_tokens, etc.) must be inherited - # so MoE layers are configured correctly for CUDA graph - # capture and communication (e.g., DeepEP). + # moe_max_num_tokens, etc.) must be inherited so MoE layers are + # configured correctly for CUDA graph capture and communication + # (e.g., DeepEP). BF16 MTP body layers cannot use the NVFP4-only + # FlashInfer MoE backend, so route those sublayers to CUTLASS. sublayer_model_config = replace(model_config, quant_config=sublayer_quant_config, + moe_backend=sublayer_moe_backend, spec_config=None) self.layers[str(step_rel_idx)] = NemotronHMTPDecoderLayer( diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py index 5b9574a10065..9890146e6fab 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py @@ -368,6 +368,7 @@ def _check_configs(self): if self.quant_config and self.quant_config.quant_mode.has_any_quant( exclude_kv_cache=True): if not (self.quant_config.quant_mode.has_nvfp4() + | self.quant_config.quant_mode.has_w4a16_nvfp4() | self.quant_config.quant_mode.has_fp8_block_scales() | self.quant_config.quant_mode.has_fp8_qdq() | self.quant_config.quant_mode.is_weight_only() @@ -531,7 +532,8 @@ def _get_quant_method(self): return FP8QDQFusedMoEMethod() elif self.quant_config.layer_quant_mode.has_fp8_block_scales(): return DeepSeekFP8BlockScalesFusedMoEMethod() - elif self.quant_config.layer_quant_mode.has_nvfp4(): + elif (self.quant_config.layer_quant_mode.has_nvfp4() + or self.quant_config.layer_quant_mode.has_w4a16_nvfp4()): return NVFP4CutlassFusedMoEMethod() elif self.quant_config.layer_quant_mode.is_int4_weight_only_per_group( ): diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 043df5ced20d..618a6eaa4fdf 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -1863,6 +1863,8 @@ def post_load_weights(self, module: Linear): class W4A16NVFP4LinearMethod(NVFP4LinearMethod): + CUDA_CORE_MAX_M: ClassVar[int] = 16 + def create_weights(self, module: Linear, in_features: int, out_features: int, bias: bool, dtype: torch.dtype): module.scaling_vector_size = 16 @@ -1899,8 +1901,65 @@ def create_weights(self, module: Linear, in_features: int, else: module.register_parameter("bias", None) + def _process_weights_without_static_activation_scale( + self, module: Linear, process_fn): + original_input_scale = module.input_scale + original_inv_input_scale = module.inv_input_scale + original_alpha = module.alpha + had_scalar_alpha = hasattr(module, "scalar_alpha") + original_scalar_alpha = getattr(module, "scalar_alpha", None) + + device = module.weight_scale_2.device + module.input_scale = Parameter(torch.empty([1], + dtype=torch.float32, + device=device), + requires_grad=False) + module.inv_input_scale = Parameter(torch.empty([1], + dtype=torch.float32, + device=device), + requires_grad=False) + module.alpha = Parameter(torch.empty([1], + dtype=torch.float32, + device=device), + requires_grad=False) + try: + process_fn(module) + finally: + module.input_scale = original_input_scale + module.inv_input_scale = original_inv_input_scale + module.alpha = original_alpha + if had_scalar_alpha: + module.scalar_alpha = original_scalar_alpha + elif hasattr(module, "scalar_alpha"): + delattr(module, "scalar_alpha") + + def process_weights_after_loading_vanilla(self, module: Linear): + self._process_weights_without_static_activation_scale( + module, + super().process_weights_after_loading_vanilla) + + def process_weights_after_loading_fused_qkv_linear(self, module: Linear): + self._process_weights_without_static_activation_scale( + module, + super().process_weights_after_loading_fused_qkv_linear) + + def process_weights_after_loading_fused_gate_up_linear( + self, module: Linear): + self._process_weights_without_static_activation_scale( + module, + super().process_weights_after_loading_fused_gate_up_linear) + def apply(self, module: Linear, input: torch.Tensor, bias: Optional[torch.Tensor]): + if input.dim() == 1: + m = 1 + elif input.dim() > 2: + m = math.prod(input.shape[:-1]) + else: + m = input.shape[0] + if m > self.CUDA_CORE_MAX_M: + return super().apply(module, input, bias) + original_shape = None if input.dim() > 2: original_shape = input.shape diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py index 5c9f5f6a02eb..02882c834813 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py @@ -20,12 +20,14 @@ from torch import nn from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.models.modeling_nemotron_h import NemotronHMOE +from tensorrt_llm._torch.models.modeling_nemotron_h import NemotronHMOE, NemotronHMTP from tensorrt_llm._torch.utils import AuxStreamType from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig -def _make_nemotron_h_moe_config(quant_config: QuantConfig) -> ModelConfig: +def _make_nemotron_h_moe_config( + quant_config: QuantConfig, moe_backend: str = "CUTLASS" +) -> ModelConfig: return ModelConfig( pretrained_config=SimpleNamespace( hidden_size=16, @@ -41,7 +43,7 @@ def _make_nemotron_h_moe_config(quant_config: QuantConfig) -> ModelConfig: topk_group=1, torch_dtype=torch.float16, ), - moe_backend="CUTLASS", + moe_backend=moe_backend, quant_config=quant_config, ) @@ -72,3 +74,71 @@ def fake_create_moe(**kwargs): assert moe_quant_config.exclude_modules == ["lm_head"] assert moe_quant_config.mamba_ssm_cache_dtype == torch.float32 assert model_config.quant_config.quant_algo == QuantAlgo.W4A16_NVFP4 + + +def test_nemotron_h_moe_preserves_w4a16_config_for_flashinfer_sm12x(): + quant_config = QuantConfig( + quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16, exclude_modules=["lm_head"] + ) + model_config = _make_nemotron_h_moe_config(quant_config, moe_backend="FLASHINFER_NVFP4SM12X") + captured = {} + + def fake_create_moe(**kwargs): + captured["model_config"] = kwargs["model_config"] + return nn.Identity() + + with patch( + "tensorrt_llm._torch.models.modeling_nemotron_h.create_moe", side_effect=fake_create_moe + ): + with patch("torch.cuda.Event", side_effect=lambda: object()): + aux_stream_dict = {AuxStreamType.MoeShared: None} + NemotronHMOE(model_config=model_config, layer_idx=1, aux_stream_dict=aux_stream_dict) + + moe_quant_config = captured["model_config"].quant_config + assert moe_quant_config is quant_config + assert moe_quant_config.quant_algo == QuantAlgo.W4A16_NVFP4 + + +def test_nemotron_h_mtp_bf16_body_uses_cutlass_moe_backend(): + quant_config = QuantConfig( + quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16, exclude_modules=["lm_head"] + ) + model_config = ModelConfig( + pretrained_config=SimpleNamespace( + mtp_hybrid_override_pattern="*E", + torch_dtype=torch.bfloat16, + ), + moe_backend="FLASHINFER_NVFP4SM12X", + quant_config=quant_config, + ) + captured = [] + + def fake_decoder_layer(**kwargs): + captured.append(kwargs) + return nn.Identity() + + with patch( + "tensorrt_llm._torch.models.modeling_nemotron_h.NemotronHMTPDecoderLayer", + side_effect=fake_decoder_layer, + ): + with patch( + "tensorrt_llm._torch.models.modeling_nemotron_h.DeepseekV3MTPHead", + side_effect=lambda model_config: nn.Identity(), + ): + with patch( + "tensorrt_llm._torch.models.modeling_nemotron_h.get_sm_version", + return_value=121, + ): + NemotronHMTP( + model_config=model_config, + layer_idx=52, + aux_stream_dict={}, + ) + + assert len(captured) == 2 + for layer_kwargs in captured: + sublayer_model_config = layer_kwargs["model_config"] + assert sublayer_model_config.quant_config.quant_algo is None + assert sublayer_model_config.moe_backend == "CUTLASS" + assert model_config.quant_config.quant_algo == QuantAlgo.W4A16_NVFP4 + assert model_config.moe_backend == "FLASHINFER_NVFP4SM12X" diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_nemotron_h_weight_mapper.py b/tests/unittest/_torch/models/checkpoints/hf/test_nemotron_h_weight_mapper.py index 02b28bd49920..3e14fab36cb4 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_nemotron_h_weight_mapper.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_nemotron_h_weight_mapper.py @@ -55,10 +55,16 @@ def test_nemotron_h_mapper_canonicalizes_w4a16_nvfp4_checkpoint_keys(): assert "lm_head.weight" in mapped assert "lm_head.weight_scale" in mapped assert "lm_head.weight_scale_2" in mapped + assert "lm_head.input_scale" in mapped assert "lm_head.weight_packed" not in mapped assert "lm_head.weight_global_scale" not in mapped assert mapped["lm_head.weight"] is weights["lm_head.weight_packed"] - assert mapped["lm_head.weight_scale_2"] is weights["lm_head.weight_global_scale"] + torch.testing.assert_close( + mapped["lm_head.weight_scale_2"], torch.tensor(4.0, dtype=torch.float32) + ) + torch.testing.assert_close( + mapped["lm_head.input_scale"], torch.tensor([1.0], dtype=torch.float32) + ) def test_nemotron_h_mapper_handles_scalar_w4a16_nvfp4_moe_global_scales(): diff --git a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py index 1feac564b176..2a075f8b7c39 100644 --- a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py @@ -55,6 +55,14 @@ def test_can_implement_accepts_supported_sm_with_nvfp4(sm_version): assert reason is None +@pytest.mark.parametrize("sm_version", sorted(FlashInferNvfp4Sm12xFusedMoE._SUPPORTED_SM_VERSIONS)) +def test_can_implement_accepts_supported_sm_with_w4a16_nvfp4(sm_version): + with patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=sm_version): + ok, reason = FlashInferNvfp4Sm12xFusedMoE.can_implement(QuantAlgo.W4A16_NVFP4) + assert ok + assert reason is None + + @pytest.mark.parametrize( "quant_algo", [ @@ -162,6 +170,15 @@ def _raise_on_flashinfer(name, *args, **kwargs): assert cls is CuteDslFusedMoE +def test_get_moe_cls_returns_flashinfer_for_w4a16_nvfp4_on_supported_sm(): + cfg = ModelConfig() + cfg.moe_backend = "FLASHINFER_NVFP4SM12X" + cfg.quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) + with patch("tensorrt_llm._utils.get_sm_version", return_value=120): + cls = get_moe_cls(cfg) + assert cls is FlashInferNvfp4Sm12xFusedMoE + + # -------------------------------------------------------------------------- # Hybrid CUTLASS-prefill / b12x-decode dispatch predicate tests # diff --git a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py index 911d0d4bf250..fa78c964f9f3 100644 --- a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py +++ b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py @@ -97,6 +97,96 @@ def fake_w4a16_nvfp4_gemm( assert output.shape == (2, 3, 5) +def test_w4a16_nvfp4_linear_falls_back_to_w4a4_nvfp4_for_large_m(): + method = W4A16NVFP4LinearMethod() + input_tensor = torch.ones((17, 16), dtype=torch.bfloat16) + module = SimpleNamespace( + weight=torch.empty((5, 8), dtype=torch.uint8), + weight_scale=torch.empty((128 * 4,), dtype=torch.uint8), + weight_scale_2=torch.tensor([0.5], dtype=torch.float32), + dtype=torch.bfloat16, + out_features=3, + pre_quant_scale=None, + input_scale=None, + alpha=None, + force_dynamic_quantization=False, + scaling_vector_size=16, + nvfp4_allowed_backends=["cuda_core"], + all_reduce=None, + mapping=None, + ) + captured = {} + + def fail_w4a16_gemm(*args, **kwargs): + raise AssertionError("large-M dense fallback must not call w4a16_nvfp4_gemm") + + def fake_fp4_quantize(input_arg, input_scale, scaling_vector_size, is_sf_swizzled): + captured["quant_input"] = input_arg + captured["input_scale"] = input_scale + captured["scaling_vector_size"] = scaling_vector_size + captured["is_sf_swizzled"] = is_sf_swizzled + return torch.empty( + (input_arg.shape[0], input_arg.shape[1] // 2), dtype=torch.uint8 + ), torch.empty((input_arg.shape[0], 4), dtype=torch.uint8) + + def fake_nvfp4_gemm( + act_fp4, + weight, + act_sf, + weight_scale, + alpha, + out_dtype, + output_buffer_kind=0, + allowed_backends="", + group=None, + ): + captured["act_fp4"] = act_fp4 + captured["weight"] = weight + captured["act_sf"] = act_sf + captured["weight_scale"] = weight_scale + captured["alpha"] = alpha + captured["out_dtype"] = out_dtype + captured["allowed_backends"] = allowed_backends + captured["group"] = group + return torch.ones((act_fp4.shape[0], weight.shape[0]), dtype=out_dtype) + + with patch("torch.ops.trtllm.w4a16_nvfp4_gemm", side_effect=fail_w4a16_gemm, create=True): + with patch("torch.ops.trtllm.fp4_quantize", side_effect=fake_fp4_quantize, create=True): + with patch("torch.ops.trtllm.nvfp4_gemm", side_effect=fake_nvfp4_gemm, create=True): + output = method.apply(module, input_tensor, bias=None) + + assert captured["quant_input"] is input_tensor + assert captured["scaling_vector_size"] == 16 + assert captured["is_sf_swizzled"] is False + assert captured["weight"] is module.weight + assert captured["weight_scale"] is module.weight_scale + assert captured["out_dtype"] is torch.bfloat16 + assert captured["allowed_backends"] == "cuda_core" + assert captured["group"] is None + assert output.shape == (17, 3) + + +def test_w4a16_nvfp4_post_load_ignores_checkpoint_activation_scale(): + method = W4A16NVFP4LinearMethod() + module = SimpleNamespace( + input_scale=None, + inv_input_scale=None, + alpha=None, + weight_scale_2=torch.empty([1], dtype=torch.float32), + tmp_nvfp4_input_scales_list=[torch.tensor(1.0, dtype=torch.float32)], + tmp_nvfp4_weight_scale_2_list=[torch.tensor(0.25, dtype=torch.float32)], + ) + + method.process_weights_after_loading_vanilla(module) + + assert module.input_scale is None + assert module.inv_input_scale is None + assert module.alpha is None + torch.testing.assert_close(module.weight_scale_2, torch.tensor([0.25], dtype=torch.float32)) + assert not hasattr(module, "tmp_nvfp4_input_scales_list") + assert not hasattr(module, "tmp_nvfp4_weight_scale_2_list") + + def test_lm_head_uses_w4a16_nvfp4_quant_method_for_packed_lm_head(): quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) diff --git a/tests/unittest/_torch/test_model_config.py b/tests/unittest/_torch/test_model_config.py index cdfed7d9d8be..912fc2a42f90 100644 --- a/tests/unittest/_torch/test_model_config.py +++ b/tests/unittest/_torch/test_model_config.py @@ -200,3 +200,27 @@ def test_load_modelopt_quant_config_respects_hf_w4a16_metadata(tmp_path): assert quant_config.group_size == 16 assert quant_config.exclude_modules == ["mtp.layers.0.mixer.q_proj"] assert layer_quant_config is None + + +def test_auto_moe_backend_selects_flashinfer_for_nemotron_h_w4a16_sm12x(): + quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16) + + with pytest.MonkeyPatch.context() as mp: + mp.setattr("tensorrt_llm._torch.model_config.get_sm_version", lambda: 121) + moe_backend = ModelConfig.resolve_moe_backend_after_quant_config( + "AUTO", "NemotronHForCausalLM", quant_config + ) + + assert moe_backend == "FLASHINFER_NVFP4SM12X" + + +def test_auto_moe_backend_keeps_cutlass_for_nemotron_h_w4a16_on_other_sm(): + quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16) + + with pytest.MonkeyPatch.context() as mp: + mp.setattr("tensorrt_llm._torch.model_config.get_sm_version", lambda: 100) + moe_backend = ModelConfig.resolve_moe_backend_after_quant_config( + "AUTO", "NemotronHForCausalLM", quant_config + ) + + assert moe_backend == "CUTLASS" From 5a2738694434024a5553e77374418fcf7ade0db0 Mon Sep 17 00:00:00 2001 From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> Date: Thu, 21 May 2026 18:58:23 +0000 Subject: [PATCH 03/25] [None][feat] Add W4A16 NVFP4 cuda-core and CUTLASS paths Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> --- .../w4a16_nvfp4_gemm_sm120.cuh | 183 ++++++++++ .../cudaCoreGemmW4A16NVFP4.cu | 16 +- .../cutlassGemmW4A16NVFP4.cu | 59 +++ .../cutlassGemmW4A16NVFP4.h | 35 ++ .../weightOnlyBatchedGemv/nvfp4ScaleLayout.h | 50 +++ cpp/tensorrt_llm/thop/w4a16Nvfp4Gemm.cpp | 108 +++++- tensorrt_llm/_torch/modules/linear.py | 27 +- .../_torch/modules/test_w4a16_nvfp4_linear.py | 335 +++++++++++++++--- 8 files changed, 750 insertions(+), 63 deletions(-) create mode 100644 cpp/tensorrt_llm/kernels/cutlass_kernels/w4a16_nvfp4_gemm/w4a16_nvfp4_gemm_sm120.cuh create mode 100644 cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cutlassGemmW4A16NVFP4.cu create mode 100644 cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cutlassGemmW4A16NVFP4.h create mode 100644 cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/nvfp4ScaleLayout.h diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/w4a16_nvfp4_gemm/w4a16_nvfp4_gemm_sm120.cuh b/cpp/tensorrt_llm/kernels/cutlass_kernels/w4a16_nvfp4_gemm/w4a16_nvfp4_gemm_sm120.cuh new file mode 100644 index 000000000000..0d2a3caeaac6 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/w4a16_nvfp4_gemm/w4a16_nvfp4_gemm_sm120.cuh @@ -0,0 +1,183 @@ +/* + * Copyright (c) 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. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/thread/linear_combination.h" +#include "cutlass/epilogue/thread/scale_type.h" +#include "cutlass/gemm/device/gemm.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/kernels/weightOnlyBatchedGemv/cutlassGemmW4A16NVFP4.h" +#include "tensorrt_llm/kernels/weightOnlyBatchedGemv/nvfp4ScaleLayout.h" + +#include +#include +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace cutlass_gemm_w4a16_nvfp4 +{ +namespace sm120 +{ +namespace +{ + +using SizeType32 = tensorrt_llm::runtime::SizeType32; + +__device__ float loadNvfp4Weight(uint8_t const* weight, __nv_fp8_e4m3 const* weightScale, float const weightGlobalScale, + SizeType32 nIdx, SizeType32 kIdx, SizeType32 k) +{ + size_t const packedOffset = (static_cast(nIdx) * static_cast(k) + static_cast(kIdx)) / 2; + uint8_t const packed = weight[packedOffset]; + __nv_fp4_storage_t const nibble = (kIdx % 2 == 0) ? (packed & 0x0FU) : (packed >> 4U); + half const weightValue = static_cast(__nv_cvt_fp4_to_halfraw(nibble, __NV_E2M1)); + + SizeType32 const scaleIdx = w4a16_nvfp4::getScaleIndex(nIdx, kIdx / w4a16_nvfp4::kScaleGranularity, k); + float const scale = static_cast(weightScale[scaleIdx]); + return __half2float(weightValue) * scale * weightGlobalScale; +} + +__global__ void dequantizeWeightToBf16ColumnMajor(__nv_bfloat16* dequantizedWeight, uint8_t const* weight, + __nv_fp8_e4m3 const* weightScale, float const* weightGlobalScale, SizeType32 n, SizeType32 k) +{ + size_t const total = static_cast(n) * static_cast(k); + float const globalScale = weightGlobalScale[0]; + for (size_t idx + = static_cast(blockIdx.x) * static_cast(blockDim.x) + static_cast(threadIdx.x); + idx < total; idx += static_cast(gridDim.x) * static_cast(blockDim.x)) + { + SizeType32 const nIdx = static_cast(idx / static_cast(k)); + SizeType32 const kIdx = static_cast(idx % static_cast(k)); + size_t const dequantizedOffset = static_cast(kIdx) + static_cast(nIdx) * static_cast(k); + dequantizedWeight[dequantizedOffset] + = __float2bfloat16(loadNvfp4Weight(weight, weightScale, globalScale, nIdx, kIdx, k)); + } +} + +void dequantizeWeight(Params const& params, __nv_bfloat16* dequantizedWeight, cudaStream_t stream) +{ + constexpr SizeType32 kBlockSize = 256; + constexpr SizeType32 kMaxGridSize = 65535; + size_t const total = static_cast(params.n) * static_cast(params.k); + auto const gridSize = static_cast( + std::min((total + static_cast(kBlockSize) - 1) / static_cast(kBlockSize), + static_cast(kMaxGridSize))); + dequantizeWeightToBf16ColumnMajor<<>>(dequantizedWeight, + reinterpret_cast(params.weight), reinterpret_cast<__nv_fp8_e4m3 const*>(params.weightScale), + params.weightGlobalScale, params.n, params.k); + TLLM_CUDA_CHECK(cudaGetLastError()); +} + +bool runCutlassBf16Gemm(Params const& params, __nv_bfloat16 const* dequantizedWeight, cudaStream_t stream) +{ + using ElementA = cutlass::bfloat16_t; + using ElementB = cutlass::bfloat16_t; + using ElementOutput = cutlass::bfloat16_t; + using ElementAccumulator = float; + using ElementCompute = float; + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutOutput = cutlass::layout::RowMajor; + using ThreadblockShape = cutlass::gemm::GemmShape<128, 128, 32>; + using WarpShape = cutlass::gemm::GemmShape<64, 64, 32>; + using InstructionShape = cutlass::gemm::GemmShape<16, 8, 16>; + using EpilogueOp + = cutlass::epilogue::thread::LinearCombination::value, + ElementAccumulator, ElementCompute, cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling>; + using Gemm = cutlass::gemm::device::Gemm, 4, 8, 8>; + + typename Gemm::Arguments arguments({params.m, params.n, params.k}, + {reinterpret_cast(params.act), params.k}, + {reinterpret_cast(dequantizedWeight), params.k}, + {reinterpret_cast(params.output), params.n}, + {reinterpret_cast(params.output), params.n}, {ElementCompute(1), ElementCompute(0)}, 1); + + Gemm gemm; + cutlass::Status status = gemm.can_implement(arguments); + if (status != cutlass::Status::kSuccess) + { + TLLM_LOG_WARNING("w4a16_nvfp4 transient BF16 CUTLASS GEMM cannot implement shape m=%d n=%d k=%d: %s", params.m, + params.n, params.k, cutlass::cutlassGetStatusString(status)); + return false; + } + + size_t const workspaceBytes = Gemm::get_workspace_size(arguments); + void* workspace = nullptr; + if (workspaceBytes > 0) + { + TLLM_CUDA_CHECK(cudaMallocAsync(&workspace, workspaceBytes, stream)); + } + + status = gemm.initialize(arguments, workspace, stream); + if (status == cutlass::Status::kSuccess) + { + status = gemm.run(stream); + } + + if (workspace != nullptr) + { + TLLM_CUDA_CHECK(cudaFreeAsync(workspace, stream)); + } + + if (status != cutlass::Status::kSuccess) + { + TLLM_LOG_WARNING("w4a16_nvfp4 transient BF16 CUTLASS GEMM failed for shape m=%d n=%d k=%d: %s", params.m, + params.n, params.k, cutlass::cutlassGetStatusString(status)); + return false; + } + TLLM_CUDA_CHECK(cudaGetLastError()); + return true; +} + +bool runTransientDequantCutlassGemm(Params const& params, cudaStream_t stream) +{ + void* dequantizedWeight = nullptr; + size_t const dequantizedBytes + = static_cast(params.n) * static_cast(params.k) * sizeof(__nv_bfloat16); + TLLM_CUDA_CHECK(cudaMallocAsync(&dequantizedWeight, dequantizedBytes, stream)); + dequantizeWeight(params, reinterpret_cast<__nv_bfloat16*>(dequantizedWeight), stream); + bool const dispatched + = runCutlassBf16Gemm(params, reinterpret_cast<__nv_bfloat16 const*>(dequantizedWeight), stream); + TLLM_CUDA_CHECK(cudaFreeAsync(dequantizedWeight, stream)); + return dispatched; +} + +} // namespace + +inline bool dispatch(Params const& params, cudaStream_t stream) +{ + if (params.inputType == CUDA_R_16BF && params.outputType == CUDA_R_16BF) + { + return runTransientDequantCutlassGemm(params, stream); + } + return false; +} + +} // namespace sm120 +} // namespace cutlass_gemm_w4a16_nvfp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.cu b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.cu index 70bc95fbdac3..984179f7393b 100644 --- a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.cu +++ b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.cu @@ -19,6 +19,7 @@ #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h" #include "tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.h" +#include "tensorrt_llm/kernels/weightOnlyBatchedGemv/nvfp4ScaleLayout.h" #include @@ -65,13 +66,12 @@ __device__ void cudaCoreGemmImpl(ActType const* __restrict__ act, __nv_fp4_e2m1 using CvtSrcType = typename Converter::source_type; using CvtResType = typename Converter::result_type; - static constexpr SizeType32 kNvfp4ScaleGranularity = 16; static constexpr SizeType32 kStepK = 32; - static constexpr SizeType32 kStepKScale = kStepK / kNvfp4ScaleGranularity; + static constexpr SizeType32 kStepKScale = kStepK / w4a16_nvfp4::kScaleGranularity; static constexpr SizeType32 kTileK = kStepK * kBlockSize; static constexpr SizeType32 kCvtCount = static_cast(sizeof(VecType) / sizeof(CvtSrcType)); - static_assert(kStepK % kNvfp4ScaleGranularity == 0); + static_assert(kStepK % w4a16_nvfp4::kScaleGranularity == 0); auto const tileIdM = static_cast(blockIdx.x * kTileM); auto const tileIdN = static_cast(blockIdx.y * kTileN); @@ -97,9 +97,6 @@ __device__ void cudaCoreGemmImpl(ActType const* __restrict__ act, __nv_fp4_e2m1 cudaGridDependencySynchronize(); #endif - SizeType32 const numColsSf = k / kNvfp4ScaleGranularity; - SizeType32 const numSfTilesK = (numColsSf + 4 - 1) / 4; - for (SizeType32 idxK = tid * kStepK; idxK < k; idxK += kTileK) { #pragma unroll @@ -118,9 +115,8 @@ __device__ void cudaCoreGemmImpl(ActType const* __restrict__ act, __nv_fp4_e2m1 for (SizeType32 j = 0; j < kTileN; ++j) { SizeType32 const rowIdx = tileIdN + j; - SizeType32 const colIdx = idxK / kNvfp4ScaleGranularity; - SizeType32 const tileOffset = ((rowIdx / 128) * numSfTilesK + colIdx / 4) * 512; - SizeType32 const dstIdx = tileOffset + (rowIdx % 32) * 16 + ((rowIdx % 128) / 32) * 4 + colIdx % 4; + SizeType32 const colIdx = idxK / w4a16_nvfp4::kScaleGranularity; + SizeType32 const dstIdx = w4a16_nvfp4::getScaleIndex(rowIdx, colIdx, k); auto const tileWeightScaleFp8x2 = reinterpret_cast(weightScale + dstIdx)[0]; char2 const tmp = reinterpret_cast(tileWeightScaleFp8x2); tileWeightScale[j * kStepKScale + 0] = static_cast(reinterpret_cast<__nv_fp8_e4m3 const&>(tmp.x)); @@ -143,7 +139,7 @@ __device__ void cudaCoreGemmImpl(ActType const* __restrict__ act, __nv_fp4_e2m1 for (SizeType32 l = 0; l < kStepK; ++l) { float const scaledWeight = tileWeight[j * kStepK + l] - * tileWeightScale[j * kStepKScale + l / kNvfp4ScaleGranularity] * weightGlobalScale; + * tileWeightScale[j * kStepKScale + l / w4a16_nvfp4::kScaleGranularity] * weightGlobalScale; acc[i * kTileN + j] = fma(tileAct[l], scaledWeight, acc[i * kTileN + j]); } } diff --git a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cutlassGemmW4A16NVFP4.cu b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cutlassGemmW4A16NVFP4.cu new file mode 100644 index 000000000000..2c42d8cf8c7d --- /dev/null +++ b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cutlassGemmW4A16NVFP4.cu @@ -0,0 +1,59 @@ +/* + * Copyright (c) 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. + * 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. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/kernels/cutlass_kernels/w4a16_nvfp4_gemm/w4a16_nvfp4_gemm_sm120.cuh" +#include "tensorrt_llm/kernels/weightOnlyBatchedGemv/cutlassGemmW4A16NVFP4.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace cutlass_gemm_w4a16_nvfp4 +{ +namespace +{ + +bool isSupported(Params const& params) +{ + int const smVersion = tensorrt_llm::common::getSMVersion(); + return (smVersion == 120 || smVersion == 121) && params.inputType == CUDA_R_16BF && params.outputType == CUDA_R_16BF + && params.m > 16 && params.n > 0 && params.k > 0 && params.n % 32 == 0 && params.k % 32 == 0 + && params.weightScale != nullptr && params.weightGlobalScale != nullptr; +} + +} // namespace + +bool cutlassGemmDispatcher(Params const& params, cudaStream_t stream) +{ + if (!isSupported(params)) + { + int const smVersion = tensorrt_llm::common::getSMVersion(); + TLLM_LOG_WARNING( + "tensorrt_llm::kernels::cutlass_gemm_w4a16_nvfp4::cutlassGemmDispatcher [NOT DISPATCHED], " + "inputType=%d, outputType=%d, m=%d, n=%d, k=%d, sm=%d", + params.inputType, params.outputType, params.m, params.n, params.k, smVersion); + return false; + } + + return sm120::dispatch(params, stream); +} + +} // namespace cutlass_gemm_w4a16_nvfp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cutlassGemmW4A16NVFP4.h b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cutlassGemmW4A16NVFP4.h new file mode 100644 index 000000000000..ae3a2832a77e --- /dev/null +++ b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cutlassGemmW4A16NVFP4.h @@ -0,0 +1,35 @@ +/* + * Copyright (c) 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. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace cutlass_gemm_w4a16_nvfp4 +{ + +using Params = tensorrt_llm::kernels::cuda_core_gemm_w4a16_nvfp4::Params; + +bool cutlassGemmDispatcher(Params const& params, cudaStream_t stream); + +} // namespace cutlass_gemm_w4a16_nvfp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/nvfp4ScaleLayout.h b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/nvfp4ScaleLayout.h new file mode 100644 index 000000000000..03c54d80033c --- /dev/null +++ b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/nvfp4ScaleLayout.h @@ -0,0 +1,50 @@ +/* + * Copyright (c) 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. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/runtime/common.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace w4a16_nvfp4 +{ + +using SizeType32 = tensorrt_llm::runtime::SizeType32; + +static constexpr SizeType32 kScaleGranularity = 16; +static constexpr SizeType32 kScaleRowsPerTile = 128; +static constexpr SizeType32 kPackedScaleColsPerTile = 4; +static constexpr SizeType32 kScaleTileElements = 512; + +__host__ __device__ inline SizeType32 getScaleIndex(SizeType32 rowIdx, SizeType32 scaleColIdx, SizeType32 k) +{ + SizeType32 const numScaleCols = k / kScaleGranularity; + SizeType32 const numScaleColTiles = (numScaleCols + kPackedScaleColsPerTile - 1) / kPackedScaleColsPerTile; + SizeType32 const tileOffset + = ((rowIdx / kScaleRowsPerTile) * numScaleColTiles + scaleColIdx / kPackedScaleColsPerTile) + * kScaleTileElements; + return tileOffset + (rowIdx % 32) * 16 + ((rowIdx % kScaleRowsPerTile) / 32) * 4 + + scaleColIdx % kPackedScaleColsPerTile; +} + +} // namespace w4a16_nvfp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/w4a16Nvfp4Gemm.cpp b/cpp/tensorrt_llm/thop/w4a16Nvfp4Gemm.cpp index 797593a4b5db..54d8bd229b6e 100644 --- a/cpp/tensorrt_llm/thop/w4a16Nvfp4Gemm.cpp +++ b/cpp/tensorrt_llm/thop/w4a16Nvfp4Gemm.cpp @@ -15,11 +15,15 @@ */ #include "tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.h" +#include "tensorrt_llm/kernels/weightOnlyBatchedGemv/cutlassGemmW4A16NVFP4.h" +#include "tensorrt_llm/kernels/weightOnlyBatchedGemv/nvfp4ScaleLayout.h" #include "tensorrt_llm/thop/thUtils.h" #include #include +#include + using torch::Tensor; TRTLLM_NAMESPACE_BEGIN @@ -35,9 +39,36 @@ void checkActDtype(Tensor const& act) "w4a16_nvfp4_gemm only supports FP16/BF16 activations, got ", act.scalar_type()); } +int64_t padUp(int64_t value, int64_t alignment) +{ + return ((value + alignment - 1) / alignment) * alignment; +} + +size_t getCudaDataTypeSize(cudaDataType_t dataType) +{ + switch (dataType) + { + case CUDA_R_16F: + case CUDA_R_16BF: return 2; + case CUDA_R_32F: return 4; + default: return 0; + } +} + +void checkWeightScaleSize(Tensor const& weightScale, int64_t n, int64_t k) +{ + using namespace tensorrt_llm::kernels::w4a16_nvfp4; + TORCH_CHECK( + k % kScaleGranularity == 0, "K must be divisible by ", kScaleGranularity, " for W4A16 NVFP4 GEMM, got K=", k); + int64_t const expectedNumel = padUp(n, kScaleRowsPerTile) * padUp(k / kScaleGranularity, kPackedScaleColsPerTile); + TORCH_CHECK(weightScale.numel() >= expectedNumel, "weight_scale has too few elements for W4A16 NVFP4 GEMM: got ", + weightScale.numel(), ", expected at least ", expectedNumel, " for N=", n, " K=", k); +} + void w4a16Nvfp4GemmCaller( Tensor& out, Tensor const& act, Tensor const& weight, Tensor const& weightScale, Tensor const& weightScale2) { + constexpr int32_t kCudaCoreMaxM = 16; auto const m = static_cast(act.sizes()[0]); auto const k = static_cast(act.sizes()[1]); auto const n = static_cast(weight.sizes()[0]); @@ -56,8 +87,53 @@ void w4a16Nvfp4GemmCaller( tensorrt_llm::kernels::cuda_core_gemm_w4a16_nvfp4::Params params( actPtr, weightPtr, weightScalePtr, weightScale2Ptr, outPtr, m, n, k, inputType, outType); - bool const dispatched = tensorrt_llm::kernels::cuda_core_gemm_w4a16_nvfp4::cudaCoreGemmDispatcher(params, stream); - TORCH_CHECK(dispatched, "Failed to dispatch w4a16_nvfp4_gemm CUDA-core kernel"); + bool dispatched = false; + if (m <= kCudaCoreMaxM) + { + dispatched = tensorrt_llm::kernels::cuda_core_gemm_w4a16_nvfp4::cudaCoreGemmDispatcher(params, stream); + } + else + { + size_t const inputElementSize = getCudaDataTypeSize(inputType); + size_t const outputElementSize = getCudaDataTypeSize(outType); + dispatched = inputElementSize != 0 && outputElementSize != 0; + for (int32_t start = 0; dispatched && start < m; start += kCudaCoreMaxM) + { + int32_t const chunkM = std::min(kCudaCoreMaxM, m - start); + auto const* chunkActPtr + = static_cast(actPtr) + static_cast(start) * k * inputElementSize; + auto* chunkOutPtr = static_cast(outPtr) + static_cast(start) * n * outputElementSize; + tensorrt_llm::kernels::cuda_core_gemm_w4a16_nvfp4::Params chunkParams( + chunkActPtr, weightPtr, weightScalePtr, weightScale2Ptr, chunkOutPtr, chunkM, n, k, inputType, outType); + dispatched = tensorrt_llm::kernels::cuda_core_gemm_w4a16_nvfp4::cudaCoreGemmDispatcher(chunkParams, stream); + } + } + TORCH_CHECK(dispatched, "Failed to dispatch w4a16_nvfp4_gemm kernel"); +} + +void w4a16Nvfp4CutlassGemmCaller( + Tensor& out, Tensor const& act, Tensor const& weight, Tensor const& weightScale, Tensor const& weightScale2) +{ + auto const m = static_cast(act.sizes()[0]); + auto const k = static_cast(act.sizes()[1]); + auto const n = static_cast(weight.sizes()[0]); + TORCH_CHECK(weight.sizes()[1] * 2 == k, "weight shape [N, K/2] must match activation shape [M, K]"); + + auto stream = at::cuda::getCurrentCUDAStream(act.get_device()); + + auto* actPtr = static_cast(act.data_ptr()); + auto* weightPtr = static_cast(weight.data_ptr()); + auto* weightScalePtr = static_cast(weightScale.data_ptr()); + auto* weightScale2Ptr = static_cast(weightScale2.data_ptr()); + auto* outPtr = static_cast(out.data_ptr()); + + auto const inputType = convert_torch_dtype(act.scalar_type()); + auto const outType = convert_torch_dtype(out.scalar_type()); + + tensorrt_llm::kernels::cutlass_gemm_w4a16_nvfp4::Params params( + actPtr, weightPtr, weightScalePtr, weightScale2Ptr, outPtr, m, n, k, inputType, outType); + bool const dispatched = tensorrt_llm::kernels::cutlass_gemm_w4a16_nvfp4::cutlassGemmDispatcher(params, stream); + TORCH_CHECK(dispatched, "Failed to dispatch w4a16_nvfp4_cutlass_gemm kernel"); } } // namespace @@ -78,6 +154,7 @@ Tensor& w4a16_nvfp4_gemm_out(Tensor const& act, Tensor const& weight, Tensor con TORCH_CHECK(act.sizes()[0] == out.sizes()[0]); TORCH_CHECK(weight.sizes()[0] == out.sizes()[1]); TORCH_CHECK(weight.sizes()[1] * 2 == act.sizes()[1]); + checkWeightScaleSize(weightScale, weight.sizes()[0], act.sizes()[1]); TORCH_CHECK(weightScale2.numel() == 1, "weight_scale_2 must be a scalar tensor"); TORCH_CHECK(!bias.has_value(), "w4a16_nvfp4_gemm does not support bias"); TORCH_CHECK(!outDtype.has_value() || out.scalar_type() == outDtype.value()); @@ -96,6 +173,29 @@ Tensor w4a16_nvfp4_gemm(Tensor const& act, Tensor const& weight, Tensor const& w return w4a16_nvfp4_gemm_out(act, weight, weightScale, weightScale2, outDtype, bias, out); } +Tensor w4a16_nvfp4_cutlass_gemm(Tensor const& act, Tensor const& weight, Tensor const& weightScale, + Tensor const& weightScale2, std::optional outDtype, std::optional const& bias) +{ + CHECK_TH_CUDA(act); + CHECK_CONTIGUOUS(act); + checkActDtype(act); + CHECK_INPUT(weight, FLOAT4_E2M1X2); + CHECK_INPUT(weightScale, SF_DTYPE); + CHECK_INPUT(weightScale2, torch::kFloat32); + + TORCH_CHECK(act.dim() == 2 && weight.dim() == 2); + TORCH_CHECK(weight.sizes()[1] * 2 == act.sizes()[1]); + checkWeightScaleSize(weightScale, weight.sizes()[0], act.sizes()[1]); + TORCH_CHECK(weightScale2.numel() == 1, "weight_scale_2 must be a scalar tensor"); + TORCH_CHECK(!bias.has_value(), "w4a16_nvfp4_cutlass_gemm does not support bias"); + + auto const outDtypeValue = outDtype.value_or(act.scalar_type()); + std::vector outputSize = {act.sizes()[0], weight.sizes()[0]}; + Tensor out = at::empty(outputSize, act.options().dtype(outDtypeValue)); + w4a16Nvfp4CutlassGemmCaller(out, act, weight, weightScale, weightScale2); + return out; +} + } // namespace torch_ext TRTLLM_NAMESPACE_END @@ -105,9 +205,13 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) m.def( "w4a16_nvfp4_gemm(Tensor act, Tensor weight, Tensor weight_scale, Tensor weight_scale_2, ScalarType? " "out_dtype, Tensor? bias=None) -> Tensor"); + m.def( + "w4a16_nvfp4_cutlass_gemm(Tensor act, Tensor weight, Tensor weight_scale, Tensor weight_scale_2, ScalarType? " + "out_dtype, Tensor? bias=None) -> Tensor"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) { m.impl("w4a16_nvfp4_gemm", &tensorrt_llm::torch_ext::w4a16_nvfp4_gemm); + m.impl("w4a16_nvfp4_cutlass_gemm", &tensorrt_llm::torch_ext::w4a16_nvfp4_cutlass_gemm); } diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 618a6eaa4fdf..b84df0999abd 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -1864,6 +1864,25 @@ def post_load_weights(self, module: Linear): class W4A16NVFP4LinearMethod(NVFP4LinearMethod): CUDA_CORE_MAX_M: ClassVar[int] = 16 + CUTLASS3_ENV: ClassVar[str] = "TRTLLM_W4A16_NVFP4_CUTLASS3" + + def _can_use_cutlass3_w4a16_prefill(self, module: Linear, + input: torch.Tensor, m: int) -> bool: + if os.environ.get(self.CUTLASS3_ENV, "0") != "1": + return False + if m <= self.CUDA_CORE_MAX_M: + return False + if get_sm_version() not in (120, 121): + return False + if input.dtype != torch.bfloat16: + return False + if module.dtype != torch.bfloat16: + return False + if input.shape[-1] % 32 != 0: + return False + if module.weight.shape[0] % 32 != 0: + return False + return True def create_weights(self, module: Linear, in_features: int, out_features: int, bias: bool, dtype: torch.dtype): @@ -1957,8 +1976,8 @@ def apply(self, module: Linear, input: torch.Tensor, m = math.prod(input.shape[:-1]) else: m = input.shape[0] - if m > self.CUDA_CORE_MAX_M: - return super().apply(module, input, bias) + use_cutlass3_prefill = self._can_use_cutlass3_w4a16_prefill( + module, input, m) original_shape = None if input.dim() > 2: @@ -1969,7 +1988,9 @@ def apply(self, module: Linear, input: torch.Tensor, assert input.dtype == module.pre_quant_scale.dtype, "Input dtype and pre_quant_scale dtype must match" input = input * module.pre_quant_scale - output = torch.ops.trtllm.w4a16_nvfp4_gemm( + gemm_op = (torch.ops.trtllm.w4a16_nvfp4_cutlass_gemm if + use_cutlass3_prefill else torch.ops.trtllm.w4a16_nvfp4_gemm) + output = gemm_op( input, module.weight, module.weight_scale, diff --git a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py index fa78c964f9f3..0e20fb6034c4 100644 --- a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py +++ b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py @@ -13,16 +13,123 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os from types import SimpleNamespace from unittest.mock import patch +import pytest import torch +import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils from tensorrt_llm._torch.modules.embedding import LMHead -from tensorrt_llm._torch.modules.linear import W4A16NVFP4LinearMethod, get_quant_method +from tensorrt_llm._torch.modules.linear import ( + W4A16NVFP4LinearMethod, + get_quant_method, + get_sm_version, +) from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig +def _run_w4a16_cutlass3_reference_case(m: int, n: int, k: int, dtype: torch.dtype) -> None: + assert k % 32 == 0 + assert n % 32 == 0 + act, weight, weight_scale, weight_scale_2 = _make_w4a16_nvfp4_case(m, n, k, dtype) + + expected = torch.empty((m, n), device="cuda", dtype=dtype) + for start in range(0, m, 16): + stop = min(start + 16, m) + expected[start:stop, :] = torch.ops.trtllm.w4a16_nvfp4_gemm( + act[start:stop, :], + weight, + weight_scale, + weight_scale_2, + dtype, + bias=None, + ) + + actual = torch.ops.trtllm.w4a16_nvfp4_cutlass_gemm( + act, + weight, + weight_scale, + weight_scale_2, + dtype, + bias=None, + ) + torch.testing.assert_close(actual, expected, atol=0.08, rtol=0.08) + + +def _make_w4a16_nvfp4_case(m: int, n: int, k: int, dtype: torch.dtype): + act = torch.randn((m, k), device="cuda", dtype=dtype) + weight = torch.empty((n, k // 2), device="cuda", dtype=fp4_utils.float4_e2m1x2) + weight_u8 = torch.randint( + 0, + 256, + (n, k // 2), + device="cuda", + dtype=torch.uint8, + ) + weight.copy_(weight_u8.view(fp4_utils.float4_e2m1x2)) + + scale_cols = fp4_utils.pad_up(k // 16, 4) + scale_rows = fp4_utils.pad_up(n, 128) + weight_scale_linear = torch.randint( + 1, + 120, + (scale_rows, scale_cols), + device="cuda", + dtype=torch.uint8, + ) + weight_scale = torch.ops.trtllm.block_scale_interleave(weight_scale_linear).view( + fp4_utils.float4_sf_dtype + ) + weight_scale_2 = torch.ones((1,), device="cuda", dtype=torch.float32) + return act, weight, weight_scale, weight_scale_2 + + +@pytest.mark.skipif( + not torch.cuda.is_available() or get_sm_version() not in (120, 121), + reason="requires CUDA SM120/121", +) +@pytest.mark.parametrize( + "shape", + [(32, 256, 256), (128, 512, 1024), (256, 1024, 2048)], +) +def test_w4a16_nvfp4_cutlass3_bf16_matches_cuda_core(shape): + m, n, k = shape + _run_w4a16_cutlass3_reference_case(m, n, k, torch.bfloat16) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or get_sm_version() not in (120, 121), + reason="requires CUDA SM120/121", +) +def test_w4a16_nvfp4_gemm_large_m_chunks_cuda_core(): + m, n, k = 32, 64, 64 + act, weight, weight_scale, weight_scale_2 = _make_w4a16_nvfp4_case(m, n, k, torch.bfloat16) + expected = torch.empty((m, n), device="cuda", dtype=torch.bfloat16) + for start in range(0, m, 16): + stop = min(start + 16, m) + expected[start:stop, :] = torch.ops.trtllm.w4a16_nvfp4_gemm( + act[start:stop, :], + weight, + weight_scale, + weight_scale_2, + torch.bfloat16, + bias=None, + ) + + actual = torch.ops.trtllm.w4a16_nvfp4_gemm( + act, + weight, + weight_scale, + weight_scale_2, + torch.bfloat16, + bias=None, + ) + + torch.testing.assert_close(actual, expected, atol=0.08, rtol=0.08) + + def test_get_quant_method_returns_w4a16_nvfp4_linear_method(): quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) @@ -97,7 +204,7 @@ def fake_w4a16_nvfp4_gemm( assert output.shape == (2, 3, 5) -def test_w4a16_nvfp4_linear_falls_back_to_w4a4_nvfp4_for_large_m(): +def test_w4a16_nvfp4_linear_uses_chunked_w4a16_op_for_large_m(): method = W4A16NVFP4LinearMethod() input_tensor = torch.ones((17, 16), dtype=torch.bfloat16) module = SimpleNamespace( @@ -107,62 +214,194 @@ def test_w4a16_nvfp4_linear_falls_back_to_w4a4_nvfp4_for_large_m(): dtype=torch.bfloat16, out_features=3, pre_quant_scale=None, - input_scale=None, - alpha=None, - force_dynamic_quantization=False, - scaling_vector_size=16, - nvfp4_allowed_backends=["cuda_core"], - all_reduce=None, - mapping=None, ) captured = {} + def fake_w4a16_nvfp4_gemm( + input_arg, weight, weight_scale, weight_scale_2, out_dtype, bias=None + ): + captured["input"] = input_arg + captured["weight"] = weight + captured["weight_scale"] = weight_scale + captured["weight_scale_2"] = weight_scale_2 + captured["out_dtype"] = out_dtype + captured["bias"] = bias + return torch.ones((input_arg.shape[0], weight.shape[0]), dtype=out_dtype) + + def fail_fp4_quantize(*args, **kwargs): + raise AssertionError("large-M W4A16 path must not quantize activations") + + with patch("torch.ops.trtllm.w4a16_nvfp4_gemm", side_effect=fake_w4a16_nvfp4_gemm, create=True): + with patch("torch.ops.trtllm.fp4_quantize", side_effect=fail_fp4_quantize, create=True): + output = method.apply(module, input_tensor, bias=None) + + assert captured["input"] is input_tensor + assert captured["weight"] is module.weight + assert captured["weight_scale"] is module.weight_scale + assert captured["weight_scale_2"] is module.weight_scale_2 + assert captured["out_dtype"] is torch.bfloat16 + assert output.shape == (17, 3) + + +def test_w4a16_nvfp4_linear_uses_cutlass3_op_for_large_bf16_m_when_enabled(): + method = W4A16NVFP4LinearMethod() + input_tensor = torch.ones((17, 32), dtype=torch.bfloat16) + module = SimpleNamespace( + weight=torch.empty((32, 16), dtype=torch.uint8), + weight_scale=torch.empty((128 * 4,), dtype=torch.uint8), + weight_scale_2=torch.tensor([0.5], dtype=torch.float32), + dtype=torch.bfloat16, + out_features=3, + pre_quant_scale=None, + ) + captured = {} + + def fake_w4a16_nvfp4_cutlass_gemm( + input_arg, weight, weight_scale, weight_scale_2, out_dtype, bias=None + ): + captured["input"] = input_arg + captured["weight"] = weight + captured["weight_scale"] = weight_scale + captured["weight_scale_2"] = weight_scale_2 + captured["out_dtype"] = out_dtype + captured["bias"] = bias + return torch.ones((input_arg.shape[0], weight.shape[0]), dtype=out_dtype) + def fail_w4a16_gemm(*args, **kwargs): - raise AssertionError("large-M dense fallback must not call w4a16_nvfp4_gemm") - - def fake_fp4_quantize(input_arg, input_scale, scaling_vector_size, is_sf_swizzled): - captured["quant_input"] = input_arg - captured["input_scale"] = input_scale - captured["scaling_vector_size"] = scaling_vector_size - captured["is_sf_swizzled"] = is_sf_swizzled - return torch.empty( - (input_arg.shape[0], input_arg.shape[1] // 2), dtype=torch.uint8 - ), torch.empty((input_arg.shape[0], 4), dtype=torch.uint8) - - def fake_nvfp4_gemm( - act_fp4, - weight, - act_sf, - weight_scale, - alpha, - out_dtype, - output_buffer_kind=0, - allowed_backends="", - group=None, + raise AssertionError("CUTLASS3 W4A16 prefill must not call the default W4A16 op") + + def fail_fp4_quantize(*args, **kwargs): + raise AssertionError("CUTLASS3 W4A16 prefill must not quantize activations") + + with patch.dict(os.environ, {"TRTLLM_W4A16_NVFP4_CUTLASS3": "1"}): + with patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=120): + with patch( + "torch.ops.trtllm.w4a16_nvfp4_gemm", side_effect=fail_w4a16_gemm, create=True + ): + with patch( + "torch.ops.trtllm.w4a16_nvfp4_cutlass_gemm", + side_effect=fake_w4a16_nvfp4_cutlass_gemm, + create=True, + ): + with patch( + "torch.ops.trtllm.fp4_quantize", side_effect=fail_fp4_quantize, create=True + ): + output = method.apply(module, input_tensor, bias=None) + + assert captured["input"] is input_tensor + assert captured["weight"] is module.weight + assert captured["weight_scale"] is module.weight_scale + assert captured["weight_scale_2"] is module.weight_scale_2 + assert captured["out_dtype"] is torch.bfloat16 + assert captured["bias"] is None + assert output.shape == (17, 3) + + +def test_w4a16_nvfp4_linear_cutlass3_restores_high_rank_input_shape(): + method = W4A16NVFP4LinearMethod() + input_tensor = torch.ones((2, 9, 32), dtype=torch.bfloat16) + bias = torch.tensor([1.0, 2.0, 3.0], dtype=torch.bfloat16) + module = SimpleNamespace( + weight=torch.empty((32, 16), dtype=torch.uint8), + weight_scale=torch.empty((128 * 4,), dtype=torch.uint8), + weight_scale_2=torch.tensor([0.5], dtype=torch.float32), + dtype=torch.bfloat16, + out_features=3, + pre_quant_scale=None, + ) + captured = {} + + def fake_w4a16_nvfp4_cutlass_gemm( + input_arg, weight, weight_scale, weight_scale_2, out_dtype, bias=None + ): + captured["input_shape"] = input_arg.shape + return torch.ones((input_arg.shape[0], weight.shape[0]), dtype=out_dtype) + + def fail_w4a16_gemm(*args, **kwargs): + raise AssertionError("large-M CUTLASS3 path must not call the default W4A16 op") + + with patch.dict(os.environ, {"TRTLLM_W4A16_NVFP4_CUTLASS3": "1"}): + with patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=120): + with patch( + "torch.ops.trtllm.w4a16_nvfp4_gemm", side_effect=fail_w4a16_gemm, create=True + ): + with patch( + "torch.ops.trtllm.w4a16_nvfp4_cutlass_gemm", + side_effect=fake_w4a16_nvfp4_cutlass_gemm, + create=True, + ): + output = method.apply(module, input_tensor, bias=bias) + + assert captured["input_shape"] == (18, 32) + assert output.shape == (2, 9, 3) + expected = torch.tensor([2.0, 3.0, 4.0], dtype=torch.bfloat16).expand(2, 9, 3) + torch.testing.assert_close(output, expected) + + +def test_w4a16_nvfp4_cutlass3_prefill_requires_supported_shape(): + method = W4A16NVFP4LinearMethod() + input_tensor = torch.ones((17, 32), dtype=torch.bfloat16) + module = SimpleNamespace( + weight=torch.empty((6, 16), dtype=torch.uint8), + dtype=torch.bfloat16, + ) + + with patch.dict(os.environ, {"TRTLLM_W4A16_NVFP4_CUTLASS3": "1"}): + with patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=120): + assert not method._can_use_cutlass3_w4a16_prefill(module, input_tensor, m=17) + + +def test_w4a16_nvfp4_linear_cutlass3_unsupported_shape_uses_default_w4a16_op(): + method = W4A16NVFP4LinearMethod() + input_tensor = torch.ones((17, 32), dtype=torch.bfloat16) + module = SimpleNamespace( + weight=torch.empty((6, 16), dtype=torch.uint8), + weight_scale=torch.empty((128 * 4,), dtype=torch.uint8), + weight_scale_2=torch.tensor([0.5], dtype=torch.float32), + dtype=torch.bfloat16, + out_features=3, + pre_quant_scale=None, + ) + captured = {} + + def fake_w4a16_nvfp4_gemm( + input_arg, weight, weight_scale, weight_scale_2, out_dtype, bias=None ): - captured["act_fp4"] = act_fp4 + captured["input"] = input_arg captured["weight"] = weight - captured["act_sf"] = act_sf captured["weight_scale"] = weight_scale - captured["alpha"] = alpha + captured["weight_scale_2"] = weight_scale_2 captured["out_dtype"] = out_dtype - captured["allowed_backends"] = allowed_backends - captured["group"] = group - return torch.ones((act_fp4.shape[0], weight.shape[0]), dtype=out_dtype) - - with patch("torch.ops.trtllm.w4a16_nvfp4_gemm", side_effect=fail_w4a16_gemm, create=True): - with patch("torch.ops.trtllm.fp4_quantize", side_effect=fake_fp4_quantize, create=True): - with patch("torch.ops.trtllm.nvfp4_gemm", side_effect=fake_nvfp4_gemm, create=True): - output = method.apply(module, input_tensor, bias=None) - - assert captured["quant_input"] is input_tensor - assert captured["scaling_vector_size"] == 16 - assert captured["is_sf_swizzled"] is False + captured["bias"] = bias + return torch.ones((input_arg.shape[0], weight.shape[0]), dtype=out_dtype) + + def fail_cutlass_gemm(*args, **kwargs): + raise AssertionError("unsupported CUTLASS3 shape must use the default W4A16 op") + + def fail_fp4_quantize(*args, **kwargs): + raise AssertionError("unsupported CUTLASS3 W4A16 path must not quantize activations") + + with patch.dict(os.environ, {"TRTLLM_W4A16_NVFP4_CUTLASS3": "1"}): + with patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=120): + with patch( + "torch.ops.trtllm.w4a16_nvfp4_gemm", side_effect=fake_w4a16_nvfp4_gemm, create=True + ): + with patch( + "torch.ops.trtllm.w4a16_nvfp4_cutlass_gemm", + side_effect=fail_cutlass_gemm, + create=True, + ): + with patch( + "torch.ops.trtllm.fp4_quantize", side_effect=fail_fp4_quantize, create=True + ): + output = method.apply(module, input_tensor, bias=None) + + assert captured["input"] is input_tensor + assert captured["bias"] is None assert captured["weight"] is module.weight assert captured["weight_scale"] is module.weight_scale + assert captured["weight_scale_2"] is module.weight_scale_2 assert captured["out_dtype"] is torch.bfloat16 - assert captured["allowed_backends"] == "cuda_core" - assert captured["group"] is None assert output.shape == (17, 3) @@ -201,7 +440,7 @@ def test_lm_head_uses_w4a16_nvfp4_quant_method_for_packed_lm_head(): assert lm_head.weight_scale_2.shape == (1,) -def test_lm_head_w4a16_nvfp4_forward_dispatches_to_dense_op(): +def test_lm_head_w4a16_nvfp4_forward_dispatches_to_w4a16_op(): quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) lm_head = LMHead( num_embeddings=3, embedding_dim=16, dtype=torch.float16, quant_config=quant_config From 9724dbbfe71b03c08e3fbb0af3a5d510f1d7f3d2 Mon Sep 17 00:00:00 2001 From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> Date: Tue, 26 May 2026 04:57:12 +0000 Subject: [PATCH 04/25] xqa + cute_dsl rebase Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> --- cpp/tensorrt_llm/common/attentionOp.cpp | 3 +- tensorrt_llm/_torch/model_config.py | 34 ++--- .../_torch/models/modeling_nemotron_h.py | 6 +- .../_torch/modules/fused_moe/create_moe.py | 24 +++- .../fused_moe/fused_moe_cute_dsl_b12x.py | 48 ++++--- .../_torch/modules/fused_moe/quantization.py | 45 ++++-- .../test_modeling_nemotron_h_moe_quant.py | 8 +- .../moe/test_cute_dsl_b12x_moe_backend.py | 133 ++++++++++++++++-- tests/unittest/_torch/test_model_config.py | 18 ++- 9 files changed, 242 insertions(+), 77 deletions(-) diff --git a/cpp/tensorrt_llm/common/attentionOp.cpp b/cpp/tensorrt_llm/common/attentionOp.cpp index d38672283b25..c4fff526540b 100644 --- a/cpp/tensorrt_llm/common/attentionOp.cpp +++ b/cpp/tensorrt_llm/common/attentionOp.cpp @@ -232,8 +232,9 @@ bool AttentionOp::convertMMHAParamsToXQAParams(tensorrt_llm::kernels::XQAParams& { xqaParams.kv_cache_data_type = xqaParams.data_type; } + // SM120/121 ship FP8 paged-KV XQA cubins with multi-block support. if (xqaParams.kv_cache_data_type == DATA_TYPE_INT8 - || (xqaParams.kv_cache_data_type == DATA_TYPE_E4M3 && (mSM < kSM_90 || mSM >= kSM_120))) + || (xqaParams.kv_cache_data_type == DATA_TYPE_E4M3 && mSM < kSM_90)) { xqaParams.multi_block_mode = false; } diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 4af19424fa01..fcc4668c7d7b 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -319,18 +319,18 @@ def resolve_moe_backend_after_quant_config( if moe_backend.upper() != "AUTO": return moe_backend - is_nemotron_h = architecture in ("NemotronHForCausalLM", - "NemotronHPuzzleForCausalLM") is_w4a16_nvfp4 = quant_config.quant_algo in (QuantAlgo.W4A16_NVFP4, "W4A16_NVFP4") - if is_nemotron_h and is_w4a16_nvfp4 and get_sm_version() in (120, 121): - return "FLASHINFER_NVFP4SM12X" + if is_w4a16_nvfp4 and get_sm_version() in (120, 121): + return "CUTEDSL" return ModelConfig.resolve_moe_backend(moe_backend, architecture) @staticmethod - def load_modelopt_quant_config(quant_config_file, checkpoint_dir, - moe_backend, hf_quant_config=None): + def load_modelopt_quant_config(quant_config_file, + checkpoint_dir, + moe_backend, + hf_quant_config=None): with open(quant_config_file) as f: quant_config_dict = json.load(f) return ModelConfig._build_modelopt_quant_config( @@ -338,8 +338,10 @@ def load_modelopt_quant_config(quant_config_file, checkpoint_dir, moe_backend, hf_quant_config) @staticmethod - def _build_modelopt_quant_config(json_quant_configs, checkpoint_dir, - moe_backend, hf_quant_config=None): + def _build_modelopt_quant_config(json_quant_configs, + checkpoint_dir, + moe_backend, + hf_quant_config=None): """Build (quant_config, layer_quant_config) from a normalized modelopt 'quantization' inner dict. ``json_quant_configs`` should be a dict as produced by @@ -708,11 +710,6 @@ def _recursive_update_config(config: transformers.PretrainedConfig, quant_config = QuantConfig() layer_quant_config = None requested_moe_backend = kwargs.get('moe_backend', 'AUTO') -<<<<<<< HEAD -======= - moe_backend = requested_moe_backend - # Resolve AUTO to specific backend based on model architecture ->>>>>>> ff309a699b (fix cuda core w4a16) architecture = pretrained_config.architectures[ 0] if pretrained_config.architectures else "" # Use an architecture-only backend hint for quant config parsing. Some @@ -738,7 +735,10 @@ def _recursive_update_config(config: transformers.PretrainedConfig, source_file="hf_quant_config.json", ) quant_config, layer_quant_config = cls._build_modelopt_quant_config( - normalized, checkpoint_dir, moe_backend_hint, hf_quant_config=hf_quant_config) + normalized, + checkpoint_dir, + moe_backend_hint, + hf_quant_config=hf_quant_config) # quantized ckpt in other formats elif getattr(pretrained_config, "quantization_config", None) is not None: @@ -751,12 +751,6 @@ def _recursive_update_config(config: transformers.PretrainedConfig, quant_config, layer_quant_config = cls.load_quant_config_from_dtypes_json( quant_config_file, moe_backend_hint) - kwargs['moe_backend'] = cls.resolve_moe_backend( - requested_moe_backend, - architecture, - quant_config=quant_config, - ) - kwargs['moe_backend'] = cls.resolve_moe_backend_after_quant_config( requested_moe_backend, architecture, quant_config) diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index 92a6c88dbd7b..14fdec865ce4 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py @@ -169,7 +169,7 @@ def _get_nemotron_h_moe_model_config( quant_config = model_config.quant_config if (quant_config is not None and quant_config.quant_algo in (QuantAlgo.W4A16_NVFP4, "W4A16_NVFP4") - and model_config.moe_backend.upper() != "FLASHINFER_NVFP4SM12X"): + and model_config.moe_backend.upper() != "CUTEDSL"): moe_quant_config = quant_config.model_copy(deep=True) moe_quant_config.quant_algo = QuantAlgo.NVFP4 moe_quant_config.__dict__.pop("quant_mode", None) @@ -1113,8 +1113,8 @@ def __init__( # spec_config cleared. All other fields (use_cuda_graph, # moe_max_num_tokens, etc.) must be inherited so MoE layers are # configured correctly for CUDA graph capture and communication - # (e.g., DeepEP). BF16 MTP body layers cannot use the NVFP4-only - # FlashInfer MoE backend, so route those sublayers to CUTLASS. + # (e.g., DeepEP). BF16 MTP body layers cannot use NVFP4-only MoE + # backends, so route those sublayers to CUTLASS. sublayer_model_config = replace(model_config, quant_config=sublayer_quant_config, moe_backend=sublayer_moe_backend, diff --git a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py index 7cc3c0af3d80..9f5d22f1de3d 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py @@ -43,12 +43,15 @@ def get_moe_cls( elif moe_backend.upper() == "CUTEDSL": if quant_config is not None and ( quant_config.quant_mode.has_fp8_block_scales() - or quant_config.quant_mode.has_nvfp4()): - # On SM120 / SM121 + NVFP4 the cuteDSL family member is the + or quant_config.quant_mode.has_nvfp4() + or quant_config.quant_mode.has_w4a16_nvfp4()): + # On SM120 / SM121 + NVFP4/W4A16_NVFP4 the cuteDSL family member is the # hybrid CUTLASS-prefill / FlashInfer NVFP4 MoE decode backend # (CuteDslB12xFusedMoE). Prefer it when flashinfer is importable; # otherwise fall through to CuteDslFusedMoE for SM100 / SM103. - if quant_config.quant_mode.has_nvfp4(): + has_nvfp4 = quant_config.quant_mode.has_nvfp4() + has_w4a16_nvfp4 = quant_config.quant_mode.has_w4a16_nvfp4() + if has_nvfp4 or has_w4a16_nvfp4: from tensorrt_llm._utils import get_sm_version sm_version = get_sm_version() if sm_version in CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS: @@ -64,13 +67,24 @@ def get_moe_cls( except ImportError: logger.warning( "CuteDslB12xFusedMoE eligible (SM%d + NVFP4) " - "but flashinfer is not importable; using CuteDslFusedMoE.", + "but flashinfer is not importable; using %s.", sm_version, + "CutlassFusedMoE" + if has_w4a16_nvfp4 else "CuteDslFusedMoE", ) + if has_w4a16_nvfp4: + return CutlassFusedMoE + elif has_w4a16_nvfp4: + logger.warning( + "CuteDslB12xFusedMoE requires SM120/121 for W4A16_NVFP4 " + "(got SM%d). Using CutlassFusedMoE.", + sm_version, + ) + return CutlassFusedMoE return CuteDslFusedMoE else: logger.warning( - f"{layer_prefix}CuteDslFusedMoE only supports fp8_block_scales and nvfp4. " + f"{layer_prefix}CuteDslFusedMoE only supports fp8_block_scales, nvfp4, and w4a16_nvfp4. " f"Check out details in quant_config: {quant_config}. Using CutlassFusedMoE instead." ) return CutlassFusedMoE diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py index 392c076d954f..2689683768da 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py @@ -44,24 +44,25 @@ class CuteDslB12xFusedMoE(CuteDslFusedMoE): - """Hybrid CUTLASS-prefill / b12x-decode NVFP4 fused-MoE backend for SM120 / SM121. + """B12x NVFP4 fused-MoE backend for SM120 / SM121. Member of the cuteDSL backend family: the decode kernel (``flashinfer.B12xMoEWrapper.run``) is JIT-compiled CuTe DSL, so the backend slots in next to :class:`CuteDslFusedMoE` (which targets SM100 / - SM103). The hybrid prefill path still routes through the C++ CUTLASS - NVFP4 GroupGEMM via explicit :class:`CutlassFusedMoE` method calls; the - parent class on the MRO does not change which kernels execute, only - where the b12x backend sits in the family. + SM103). Plain NVFP4 prefill can route through the C++ CUTLASS NVFP4 + GroupGEMM via explicit :class:`CutlassFusedMoE` method calls; the parent + class on the MRO does not change which kernels execute, only where the + b12x backend sits in the family. Composition (see ``MOE_DEVELOPER_GUIDE.md`` for the full explainer): - - **Prefill (``m >= _PREFILL_VIA_CUTLASS_THRESHOLD``)** explicitly + - **NVFP4 prefill (``m >= _PREFILL_VIA_CUTLASS_THRESHOLD``)** explicitly invokes :class:`CutlassFusedMoE` NVFP4 GroupGEMM. The b12x kernel's 12-CTA-per-token MMA pattern is suboptimal at large ``m``. - **Decode (``m < _PREFILL_VIA_CUTLASS_THRESHOLD``)** dispatches to FlashInfer's ``B12xMoEWrapper.run`` — a kernel purpose-built for ``m=1`` / small routed-row counts. + - **W4A16_NVFP4** stays on the b12x path for both prefill and decode. NVFP4 weights are loaded via :class:`NVFP4CuteDslB12xFusedMoEMethod` (an :class:`NVFP4CutlassFusedMoEMethod` subclass returned by @@ -79,8 +80,9 @@ class CuteDslB12xFusedMoE(CuteDslFusedMoE): The backend hard-rejects EP (b12x has no dispatch / combine kernel), MoE alltoall, ``Fp4QuantizedTensor`` input, ``swiglu_gptoss_style`` biased SwiGLU, and activations outside ``{Relu2, Swiglu}``. It is - selected on the ``CUTEDSL`` MoE path when SM120 / SM121 + NVFP4 + - flashinfer-importable gates pass (see ``create_moe.get_moe_cls``). + selected on the ``CUTEDSL`` MoE path when SM120 / SM121 + NVFP4 or + W4A16_NVFP4 + flashinfer-importable gates pass (see + ``create_moe.get_moe_cls``). """ # SM versions on which the FlashInfer b12x NVFP4 MoE kernel is available. @@ -105,9 +107,9 @@ def can_implement( if sm_version not in cls._SUPPORTED_SM_VERSIONS: sm_list = "/".join(f"SM{v}" for v in sorted(cls._SUPPORTED_SM_VERSIONS)) return _warn_and_return(f"CuteDslB12xFusedMoE requires {sm_list}, got SM{sm_version}") - if quant_algo != QuantAlgo.NVFP4: + if quant_algo not in {QuantAlgo.NVFP4, QuantAlgo.W4A16_NVFP4}: return _warn_and_return( - f"CuteDslB12xFusedMoE only supports NVFP4 quantization " + f"CuteDslB12xFusedMoE only supports NVFP4 or W4A16_NVFP4 quantization " f"(got quant_algo={quant_algo})" ) if dtype_activation not in {torch.float16, torch.bfloat16}: @@ -158,7 +160,10 @@ def _get_quant_method(self): if ( self.quant_config is not None and self.quant_config.layer_quant_mode.has_any_quant(exclude_kv_cache=True) - and self.quant_config.layer_quant_mode.has_nvfp4() + and ( + self.quant_config.layer_quant_mode.has_nvfp4() + or self.quant_config.layer_quant_mode.has_w4a16_nvfp4() + ) ): from .quantization import NVFP4CuteDslB12xFusedMoEMethod @@ -167,9 +172,12 @@ def _get_quant_method(self): def _route_to_cutlass(self, x) -> bool: """Return ``True`` iff this call should fall back to the inherited - CUTLASS path (prefill chunk). ``Fp4QuantizedTensor`` inputs always - stay on the b12x path (which rejects them) so the existing error - message is preserved.""" + CUTLASS path (NVFP4 prefill chunk). ``Fp4QuantizedTensor`` inputs + always stay on the b12x path (which rejects them) so the existing + error message is preserved.""" + quant_config = getattr(self, "quant_config", None) + if quant_config is not None and quant_config.layer_quant_mode.has_w4a16_nvfp4(): + return False return isinstance(x, torch.Tensor) and x.shape[0] >= self._PREFILL_VIA_CUTLASS_THRESHOLD # ``post_load_weights`` is inherited from ``CutlassFusedMoE`` and @@ -190,12 +198,12 @@ def quantize_input( ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: """Hybrid dispatch entrypoint for activation handling. - Prefill chunks (``x.shape[0] >= _PREFILL_VIA_CUTLASS_THRESHOLD``) take - the inherited :meth:`CutlassFusedMoE.quantize_input` path so the - downstream ``run_moe`` can call CUTLASS NVFP4 GroupGEMM. Decode - chunks pass through unchanged because b12x quantizes activations - internally (consumes a bf16 / fp16 ``x`` and produces its own scale - factors). + NVFP4 prefill chunks take the inherited + :meth:`CutlassFusedMoE.quantize_input` path so the downstream + ``run_moe`` can call CUTLASS NVFP4 GroupGEMM. Decode chunks and + W4A16_NVFP4 chunks pass through unchanged because b12x quantizes + activations internally (consumes a bf16 / fp16 ``x`` and produces its + own scale factors). """ if self._route_to_cutlass(x): return CutlassFusedMoE.quantize_input( diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index e6769b241646..91e11dc8e82c 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py @@ -3020,6 +3020,10 @@ def post_load_weights(self, module: torch.nn.Module): # layers across the model share one wrapper-owned output buffer. from .fused_moe_cute_dsl_b12x import _SHARED_MOE_OUTPUT_BUF + quant_config = getattr(module, "quant_config", None) + is_w4a16_nvfp4 = (quant_config is not None + and quant_config.layer_quant_mode.has_w4a16_nvfp4()) + num_local_experts = module.w3_w1_weight.shape[0] # Tensor shapes use the *padded* per-rank dims because TP partitions # may pad ``intermediate_size`` up to a kernel-friendly boundary. @@ -3069,11 +3073,28 @@ def post_load_weights(self, module: torch.nn.Module): k=w2_in_dim, num_groups=num_local_experts) - w1_alpha_b12x = ((1.0 / module.fc31_input_scale).expand( - module.num_experts).to(torch.float32).contiguous()) - w2_alpha_b12x = ((1.0 / module.fc2_input_scale).expand( - module.num_experts).to(torch.float32).contiguous()) - fc2_input_scale_b12x = (1.0 / module.fc2_input_scale).to(torch.float32) + if is_w4a16_nvfp4: + # W4A16 path: BF16/FP16 activations multiplied by FP4 weights. + # There is no online activation quantization, so the wrapper uses + # ``w*_alpha`` only as FC1/FC2 epilogue dequant multipliers. + # Since we still un-normalize FP8 block scales by weight_scale_2 + # above, alpha is unity. ``fc2_input_scale`` is ignored in this + # FlashInfer mode. + alpha_device = module.w3_w1_weight.device + w1_alpha_b12x = torch.ones(module.num_experts, + dtype=torch.float32, + device=alpha_device) + w2_alpha_b12x = torch.ones(module.num_experts, + dtype=torch.float32, + device=alpha_device) + fc2_input_scale_b12x = None + else: + w1_alpha_b12x = ((1.0 / module.fc31_input_scale).expand( + module.num_experts).to(torch.float32).contiguous()) + w2_alpha_b12x = ((1.0 / module.fc2_input_scale).expand( + module.num_experts).to(torch.float32).contiguous()) + fc2_input_scale_b12x = (1.0 / module.fc2_input_scale).to( + torch.float32) # TRT-LLM packs 16 FP4 values per int64. flashinfer's internal # ``view(torch.float4_e2m1fn_x2)`` requires byte-contiguous storage @@ -3096,14 +3117,21 @@ def post_load_weights(self, module: torch.nn.Module): f"{ActivationType(module.activation_type).name}; " f"supported: {supported}.") + # The model config may carry the logical intermediate size while the + # NVFP4 weight tensors are padded for kernel alignment, e.g. Nano3.5 + # uses 1856 logical channels and 1920 stored channels. FlashInfer's + # CUDA-graph workspace must match the stored tensors. + b12x_intermediate_size = w2_in_dim + module.b12x_wrapper = B12xMoEWrapper( num_experts=module.num_experts, top_k=module.routing_method.experts_per_token, hidden_size=module.hidden_size, - intermediate_size=module.intermediate_size_per_partition, + intermediate_size=b12x_intermediate_size, use_cuda_graph=getattr(module, "_b12x_use_cuda_graph", False), max_num_tokens=module.moe_max_num_tokens, activation=self._ACTIVATION_MAP[module.activation_type], + quant_mode="w4a16" if is_w4a16_nvfp4 else "nvfp4", ) # Replace the wrapper's per-instance output buffer with a shared one. @@ -3124,10 +3152,11 @@ def post_load_weights(self, module: torch.nn.Module): logger.info_once( f"NVFP4CuteDslB12xFusedMoEMethod active: hidden={module.hidden_size}, " - f"intermediate={module.intermediate_size_per_partition}, " + f"intermediate={b12x_intermediate_size}, " f"experts={module.num_experts}, top_k=" f"{module.routing_method.experts_per_token}, " - f"activation={self._ACTIVATION_MAP[module.activation_type]}.", + f"activation={self._ACTIVATION_MAP[module.activation_type]}, " + f"quant_mode={'w4a16' if is_w4a16_nvfp4 else 'nvfp4'}.", key="cute_dsl_b12x_moe_active", ) diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py index 02882c834813..02c89052775c 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py @@ -76,11 +76,11 @@ def fake_create_moe(**kwargs): assert model_config.quant_config.quant_algo == QuantAlgo.W4A16_NVFP4 -def test_nemotron_h_moe_preserves_w4a16_config_for_flashinfer_sm12x(): +def test_nemotron_h_moe_preserves_w4a16_config_for_cutedsl_sm12x(): quant_config = QuantConfig( quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16, exclude_modules=["lm_head"] ) - model_config = _make_nemotron_h_moe_config(quant_config, moe_backend="FLASHINFER_NVFP4SM12X") + model_config = _make_nemotron_h_moe_config(quant_config, moe_backend="CUTEDSL") captured = {} def fake_create_moe(**kwargs): @@ -108,7 +108,7 @@ def test_nemotron_h_mtp_bf16_body_uses_cutlass_moe_backend(): mtp_hybrid_override_pattern="*E", torch_dtype=torch.bfloat16, ), - moe_backend="FLASHINFER_NVFP4SM12X", + moe_backend="CUTEDSL", quant_config=quant_config, ) captured = [] @@ -141,4 +141,4 @@ def fake_decoder_layer(**kwargs): assert sublayer_model_config.quant_config.quant_algo is None assert sublayer_model_config.moe_backend == "CUTLASS" assert model_config.quant_config.quant_algo == QuantAlgo.W4A16_NVFP4 - assert model_config.moe_backend == "FLASHINFER_NVFP4SM12X" + assert model_config.moe_backend == "CUTEDSL" diff --git a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py index 2a075f8b7c39..8fe8279209c7 100644 --- a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py @@ -23,6 +23,8 @@ SM120/SM121 hardware. """ +import sys +import types from unittest.mock import patch import pytest @@ -33,6 +35,11 @@ from tensorrt_llm._torch.modules.fused_moe.fused_moe_cute_dsl import CuteDslFusedMoE from tensorrt_llm._torch.modules.fused_moe.fused_moe_cute_dsl_b12x import CuteDslB12xFusedMoE from tensorrt_llm._torch.modules.fused_moe.fused_moe_cutlass import CutlassFusedMoE +from tensorrt_llm._torch.modules.fused_moe.quantization import ( + NVFP4CuteDslB12xFusedMoEMethod, + NVFP4CutlassFusedMoEMethod, +) +from tensorrt_llm._torch.utils import ActivationType from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig _FUSED_MOE_MODULE = "tensorrt_llm._torch.modules.fused_moe.fused_moe_cute_dsl_b12x" @@ -55,10 +62,10 @@ def test_can_implement_accepts_supported_sm_with_nvfp4(sm_version): assert reason is None -@pytest.mark.parametrize("sm_version", sorted(FlashInferNvfp4Sm12xFusedMoE._SUPPORTED_SM_VERSIONS)) +@pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) def test_can_implement_accepts_supported_sm_with_w4a16_nvfp4(sm_version): with patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=sm_version): - ok, reason = FlashInferNvfp4Sm12xFusedMoE.can_implement(QuantAlgo.W4A16_NVFP4) + ok, reason = CuteDslB12xFusedMoE.can_implement(QuantAlgo.W4A16_NVFP4) assert ok assert reason is None @@ -138,6 +145,16 @@ def test_get_moe_cls_cutedsl_returns_plain_cutedsl_on_unsupported_sm(): assert cls is CuteDslFusedMoE +def test_get_moe_cls_cutedsl_returns_cutlass_for_w4a16_nvfp4_on_unsupported_sm(): + """CUTEDSL + W4A16_NVFP4 + non-SM120/121 → CutlassFusedMoE.""" + cfg = ModelConfig() + cfg.moe_backend = "CUTEDSL" + cfg.quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) + with patch("tensorrt_llm._utils.get_sm_version", return_value=100): + cls = get_moe_cls(cfg) + assert cls is CutlassFusedMoE + + @pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) def test_get_moe_cls_cutedsl_selects_b12x_on_supported_sm(sm_version): """CUTEDSL + NVFP4 + SM120/121 + flashinfer importable → CuteDslB12xFusedMoE.""" @@ -149,6 +166,17 @@ def test_get_moe_cls_cutedsl_selects_b12x_on_supported_sm(sm_version): assert cls is CuteDslB12xFusedMoE +@pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) +def test_get_moe_cls_cutedsl_selects_b12x_for_w4a16_nvfp4_on_supported_sm(sm_version): + """CUTEDSL + W4A16_NVFP4 + SM120/121 + flashinfer importable → CuteDslB12xFusedMoE.""" + cfg = ModelConfig() + cfg.moe_backend = "CUTEDSL" + cfg.quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) + with patch("tensorrt_llm._utils.get_sm_version", return_value=sm_version): + cls = get_moe_cls(cfg) + assert cls is CuteDslB12xFusedMoE + + def test_get_moe_cls_cutedsl_falls_back_to_plain_cutedsl_when_flashinfer_missing(monkeypatch): """CUTEDSL + NVFP4 + SM120/121 + flashinfer NOT importable → CuteDslFusedMoE.""" import builtins @@ -170,15 +198,6 @@ def _raise_on_flashinfer(name, *args, **kwargs): assert cls is CuteDslFusedMoE -def test_get_moe_cls_returns_flashinfer_for_w4a16_nvfp4_on_supported_sm(): - cfg = ModelConfig() - cfg.moe_backend = "FLASHINFER_NVFP4SM12X" - cfg.quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) - with patch("tensorrt_llm._utils.get_sm_version", return_value=120): - cls = get_moe_cls(cfg) - assert cls is FlashInferNvfp4Sm12xFusedMoE - - # -------------------------------------------------------------------------- # Hybrid CUTLASS-prefill / b12x-decode dispatch predicate tests # @@ -217,6 +236,98 @@ def test_dispatch_decode_shape_takes_b12x(): assert stub._route_to_cutlass(x) is False +def test_w4a16_nvfp4_prefill_quantize_input_stays_on_b12x(): + moe = object.__new__(CuteDslB12xFusedMoE) + moe.quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) + x = torch.empty(CuteDslB12xFusedMoE._PREFILL_VIA_CUTLASS_THRESHOLD, 1024) + + with patch.object( + CutlassFusedMoE, + "quantize_input", + side_effect=AssertionError("W4A16_NVFP4 prefill must not route through CUTLASS"), + ): + out, out_sf = CuteDslB12xFusedMoE.quantize_input(moe, x) + + assert out is x + assert out_sf is None + + +def test_w4a16_nvfp4_post_load_constructs_b12x_w4a16_wrapper(monkeypatch): + class _RoutingMethod: + experts_per_token = 4 + + class _FakeB12xWrapper: + calls = [] + + def __init__(self, **kwargs): + self._moe_output = None + self.calls.append(kwargs) + + def _convert_sf_to_mma_layout(scales, *, m, k, num_groups): + return scales + + flashinfer = types.ModuleType("flashinfer") + flashinfer.B12xMoEWrapper = _FakeB12xWrapper + cute_dsl = types.ModuleType("flashinfer.cute_dsl") + utils = types.ModuleType("flashinfer.cute_dsl.utils") + utils.convert_sf_to_mma_layout = _convert_sf_to_mma_layout + monkeypatch.setitem(sys.modules, "flashinfer", flashinfer) + monkeypatch.setitem(sys.modules, "flashinfer.cute_dsl", cute_dsl) + monkeypatch.setitem(sys.modules, "flashinfer.cute_dsl.utils", utils) + + num_experts = 2 + hidden_size = 128 + logical_intermediate_size = 1856 + padded_intermediate_size = 1920 + module = torch.nn.Module() + module.num_experts = num_experts + module.hidden_size = hidden_size + module.intermediate_size_per_partition = logical_intermediate_size + module.moe_max_num_tokens = 8 + module.routing_method = _RoutingMethod() + module.activation_type = ActivationType.Swiglu + module.quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) + module.w3_w1_weight = torch.empty( + num_experts, + 2 * padded_intermediate_size, + hidden_size // 16, + dtype=torch.int64, + ) + module.w2_weight = torch.empty( + num_experts, + hidden_size, + padded_intermediate_size // 16, + dtype=torch.int64, + ) + module.w3_w1_weight_scale = torch.ones( + num_experts, + 2 * padded_intermediate_size, + hidden_size // 16, + dtype=torch.float8_e4m3fn, + ) + module.w2_weight_scale = torch.ones( + num_experts, + hidden_size, + padded_intermediate_size // 16, + dtype=torch.float8_e4m3fn, + ) + module.fc31_alpha = torch.ones(num_experts) + module.fc2_alpha = torch.ones(num_experts) + module.fc31_input_scale = torch.ones(()) + module.fc2_input_scale = torch.ones(()) + + with patch.object(NVFP4CutlassFusedMoEMethod, "post_load_weights", return_value=None): + NVFP4CuteDslB12xFusedMoEMethod().post_load_weights(module) + + assert _FakeB12xWrapper.calls + wrapper_kwargs = _FakeB12xWrapper.calls[0] + assert wrapper_kwargs.get("quant_mode") == "w4a16", wrapper_kwargs + assert wrapper_kwargs["intermediate_size"] == padded_intermediate_size + assert module._b12x_weights["fc2_input_scale"] is None + assert torch.equal(module._b12x_weights["w1_alpha"], torch.ones(num_experts)) + assert torch.equal(module._b12x_weights["w2_alpha"], torch.ones(num_experts)) + + def test_dispatch_rejects_non_tensor(): """Non-tensor inputs (e.g. Fp4QuantizedTensor) stay on the b12x path so the existing ValueError surfaces in quantize_input.""" diff --git a/tests/unittest/_torch/test_model_config.py b/tests/unittest/_torch/test_model_config.py index 912fc2a42f90..1aebdd806942 100644 --- a/tests/unittest/_torch/test_model_config.py +++ b/tests/unittest/_torch/test_model_config.py @@ -202,25 +202,33 @@ def test_load_modelopt_quant_config_respects_hf_w4a16_metadata(tmp_path): assert layer_quant_config is None -def test_auto_moe_backend_selects_flashinfer_for_nemotron_h_w4a16_sm12x(): +@pytest.mark.parametrize( + "architecture", + ["NemotronHForCausalLM", "Qwen3MoeForCausalLM", "SomeOtherForCausalLM"], +) +def test_auto_moe_backend_selects_cutedsl_for_w4a16_sm12x(architecture): quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16) with pytest.MonkeyPatch.context() as mp: mp.setattr("tensorrt_llm._torch.model_config.get_sm_version", lambda: 121) moe_backend = ModelConfig.resolve_moe_backend_after_quant_config( - "AUTO", "NemotronHForCausalLM", quant_config + "AUTO", architecture, quant_config ) - assert moe_backend == "FLASHINFER_NVFP4SM12X" + assert moe_backend == "CUTEDSL" -def test_auto_moe_backend_keeps_cutlass_for_nemotron_h_w4a16_on_other_sm(): +@pytest.mark.parametrize( + "architecture", + ["NemotronHForCausalLM", "Qwen3MoeForCausalLM", "SomeOtherForCausalLM"], +) +def test_auto_moe_backend_keeps_cutlass_for_w4a16_on_other_sm(architecture): quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16) with pytest.MonkeyPatch.context() as mp: mp.setattr("tensorrt_llm._torch.model_config.get_sm_version", lambda: 100) moe_backend = ModelConfig.resolve_moe_backend_after_quant_config( - "AUTO", "NemotronHForCausalLM", quant_config + "AUTO", architecture, quant_config ) assert moe_backend == "CUTLASS" From 58306f4f3858fd57cb945acf176f9c4f9f3b0ebc Mon Sep 17 00:00:00 2001 From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> Date: Wed, 27 May 2026 03:25:00 +0000 Subject: [PATCH 05/25] fix Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> --- cpp/kernels/xqa/mha.cu | 6 ++- .../_torch/custom_ops/cpp_custom_ops.py | 10 ++++ .../_torch/modules/fused_moe/quantization.py | 46 ++++++++++--------- tensorrt_llm/_torch/modules/linear.py | 12 +++++ .../moe/test_cute_dsl_b12x_moe_backend.py | 34 ++++++++++---- .../_torch/modules/test_w4a16_nvfp4_linear.py | 2 +- 6 files changed, 77 insertions(+), 33 deletions(-) diff --git a/cpp/kernels/xqa/mha.cu b/cpp/kernels/xqa/mha.cu index 636cd19a4a1d..a95092747271 100644 --- a/cpp/kernels/xqa/mha.cu +++ b/cpp/kernels/xqa/mha.cu @@ -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"); @@ -2538,6 +2538,7 @@ CUBIN_EXPORT __global__ // merge results from different warp groups SharedMem::XSmemBuffer* smemOutTile = mergeAndSaveOutTile(outTile, inputElemSize == 2 && cacheElemSize == 1); + bool writeOutput = !isMultiBlock; if (isMultiBlock) { static_assert(ctaShapeInWarps.y == 1, "not implemented"); @@ -2612,6 +2613,7 @@ CUBIN_EXPORT __global__ bool const isLastCta = mbsmem.isLastCta; if (isLastCta) { + writeOutput = true; MultiBlockSMem::MBBuf& mbbuf = mbsmem.storage[warpIdx.y]; SMemWarpRowMax& smemRowMax = reinterpret_cast(smem); // get row max. @@ -2700,7 +2702,7 @@ CUBIN_EXPORT __global__ smemOutTile = mergeAndSaveOutTile(mergedOutTile, false); } } - if (warpGrpIdx == 0) + if (warpGrpIdx == 0 && writeOutput) { #if SPEC_DEC copyOutputToGlobalMem(warp, &output[reqSeqOffset * nbQHeads], nbQHeads, headGrpSize, diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index ac159a009539..82d2826175ea 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -1054,6 +1054,16 @@ def _(act: torch.Tensor, del weight_scale, weight_scale_2, bias return act.new_empty((act.shape[0], weight.shape[0]), dtype=out_dtype) + @torch.library.register_fake("trtllm::w4a16_nvfp4_cutlass_gemm") + def _(act: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + weight_scale_2: torch.Tensor, + out_dtype: Optional[torch.dtype], + bias: Optional[torch.Tensor] = None): + del weight_scale, weight_scale_2, bias + return act.new_empty((act.shape[0], weight.shape[0]), dtype=out_dtype) + @torch.library.register_fake("trtllm::mla_rope_generation") def _( fused_q: torch.Tensor, diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index 91e11dc8e82c..b6f8bb305095 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py @@ -3053,16 +3053,20 @@ def post_load_weights(self, module: torch.nn.Module): w2_w_scale_2 = (module.fc2_alpha * module.fc2_input_scale).to( torch.float32) - w1_sf_fp8_norm = module.w3_w1_weight_scale.view( - torch.float8_e4m3fn).float() - w2_sf_fp8_norm = module.w2_weight_scale.view( - torch.float8_e4m3fn).float() + w1_sf_fp8_src = module.w3_w1_weight_scale.view(torch.float8_e4m3fn) + w2_sf_fp8_src = module.w2_weight_scale.view(torch.float8_e4m3fn) + w1_sf_fp8_norm = w1_sf_fp8_src.float() + w2_sf_fp8_norm = w2_sf_fp8_src.float() - # Broadcast per-expert scalar over the trailing dims (E, *). - bcast1 = w1_w_scale_2.view(-1, *([1] * (w1_sf_fp8_norm.dim() - 1))) - bcast2 = w2_w_scale_2.view(-1, *([1] * (w2_sf_fp8_norm.dim() - 1))) - w1_sf_fp8 = (w1_sf_fp8_norm * bcast1).to(torch.float8_e4m3fn) - w2_sf_fp8 = (w2_sf_fp8_norm * bcast2).to(torch.float8_e4m3fn) + if is_w4a16_nvfp4: + w1_sf_fp8 = w1_sf_fp8_src + w2_sf_fp8 = w2_sf_fp8_src + else: + # Broadcast per-expert scalar over the trailing dims (E, *). + bcast1 = w1_w_scale_2.view(-1, *([1] * (w1_sf_fp8_norm.dim() - 1))) + bcast2 = w2_w_scale_2.view(-1, *([1] * (w2_sf_fp8_norm.dim() - 1))) + w1_sf_fp8 = (w1_sf_fp8_norm * bcast1).to(torch.float8_e4m3fn) + w2_sf_fp8 = (w2_sf_fp8_norm * bcast2).to(torch.float8_e4m3fn) w1_sf_b12x = convert_sf_to_mma_layout(w1_sf_fp8, m=w3w1_out_dim, @@ -3075,18 +3079,18 @@ def post_load_weights(self, module: torch.nn.Module): if is_w4a16_nvfp4: # W4A16 path: BF16/FP16 activations multiplied by FP4 weights. - # There is no online activation quantization, so the wrapper uses - # ``w*_alpha`` only as FC1/FC2 epilogue dequant multipliers. - # Since we still un-normalize FP8 block scales by weight_scale_2 - # above, alpha is unity. ``fc2_input_scale`` is ignored in this - # FlashInfer mode. - alpha_device = module.w3_w1_weight.device - w1_alpha_b12x = torch.ones(module.num_experts, - dtype=torch.float32, - device=alpha_device) - w2_alpha_b12x = torch.ones(module.num_experts, - dtype=torch.float32, - device=alpha_device) + # FlashInfer's W4A16 packer expects the ModelOpt scale contract: + # normalized FP8 block scales plus per-expert ``weight_global_scale``. + # TRT-LLM stores the reciprocal as ``weight_scale_2``; recover the + # ModelOpt value from the already computed dequant scale. + def _to_modelopt_global_scale(scale: torch.Tensor) -> torch.Tensor: + global_scale = torch.zeros_like(scale, dtype=torch.float32) + valid = scale > 0 + global_scale[valid] = scale[valid].reciprocal() + return global_scale.contiguous() + + w1_alpha_b12x = _to_modelopt_global_scale(w1_w_scale_2) + w2_alpha_b12x = _to_modelopt_global_scale(w2_w_scale_2) fc2_input_scale_b12x = None else: w1_alpha_b12x = ((1.0 / module.fc31_input_scale).expand( diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index b84df0999abd..ab94f69346ce 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -1927,6 +1927,8 @@ def _process_weights_without_static_activation_scale( original_alpha = module.alpha had_scalar_alpha = hasattr(module, "scalar_alpha") original_scalar_alpha = getattr(module, "scalar_alpha", None) + has_weight_global_scale = bool( + getattr(module, "tmp_nvfp4_weight_scale_2_list", [])) device = module.weight_scale_2.device module.input_scale = Parameter(torch.empty([1], @@ -1943,6 +1945,8 @@ def _process_weights_without_static_activation_scale( requires_grad=False) try: process_fn(module) + if has_weight_global_scale: + self._convert_weight_global_scale_to_dequant_scale(module) finally: module.input_scale = original_input_scale module.inv_input_scale = original_inv_input_scale @@ -1952,6 +1956,14 @@ def _process_weights_without_static_activation_scale( elif hasattr(module, "scalar_alpha"): delattr(module, "scalar_alpha") + @staticmethod + def _convert_weight_global_scale_to_dequant_scale(module: Linear): + weight_scale_2 = module.weight_scale_2.data.float() + dequant_scale = torch.zeros_like(weight_scale_2) + nonzero = weight_scale_2 != 0 + dequant_scale[nonzero] = weight_scale_2[nonzero].reciprocal() + copy_weight(module.weight_scale_2, dequant_scale) + def process_weights_after_loading_vanilla(self, module: Linear): self._process_weights_without_static_activation_scale( module, diff --git a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py index 8fe8279209c7..b278e963137a 100644 --- a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py @@ -252,7 +252,7 @@ def test_w4a16_nvfp4_prefill_quantize_input_stays_on_b12x(): assert out_sf is None -def test_w4a16_nvfp4_post_load_constructs_b12x_w4a16_wrapper(monkeypatch): +def test_w4a16_nvfp4_post_load_uses_modelopt_scale_contract(monkeypatch): class _RoutingMethod: experts_per_token = 4 @@ -299,22 +299,24 @@ def _convert_sf_to_mma_layout(scales, *, m, k, num_groups): padded_intermediate_size // 16, dtype=torch.int64, ) - module.w3_w1_weight_scale = torch.ones( + w3_w1_weight_scale = torch.ones( num_experts, 2 * padded_intermediate_size, hidden_size // 16, dtype=torch.float8_e4m3fn, ) - module.w2_weight_scale = torch.ones( + w2_weight_scale = torch.ones( num_experts, hidden_size, padded_intermediate_size // 16, dtype=torch.float8_e4m3fn, ) - module.fc31_alpha = torch.ones(num_experts) - module.fc2_alpha = torch.ones(num_experts) - module.fc31_input_scale = torch.ones(()) - module.fc2_input_scale = torch.ones(()) + module.w3_w1_weight_scale = w3_w1_weight_scale.clone() + module.w2_weight_scale = w2_weight_scale.clone() + module.fc31_alpha = torch.tensor([0.25, 0.5], dtype=torch.float32) + module.fc2_alpha = torch.tensor([0.125, 0.25], dtype=torch.float32) + module.fc31_input_scale = torch.tensor(2.0, dtype=torch.float32) + module.fc2_input_scale = torch.tensor(4.0, dtype=torch.float32) with patch.object(NVFP4CutlassFusedMoEMethod, "post_load_weights", return_value=None): NVFP4CuteDslB12xFusedMoEMethod().post_load_weights(module) @@ -324,8 +326,22 @@ def _convert_sf_to_mma_layout(scales, *, m, k, num_groups): assert wrapper_kwargs.get("quant_mode") == "w4a16", wrapper_kwargs assert wrapper_kwargs["intermediate_size"] == padded_intermediate_size assert module._b12x_weights["fc2_input_scale"] is None - assert torch.equal(module._b12x_weights["w1_alpha"], torch.ones(num_experts)) - assert torch.equal(module._b12x_weights["w2_alpha"], torch.ones(num_experts)) + assert torch.equal( + module._b12x_weights["w1_weight_sf"].float(), + w3_w1_weight_scale.float(), + ) + assert torch.equal( + module._b12x_weights["w2_weight_sf"].float(), + w2_weight_scale.float(), + ) + assert torch.allclose( + module._b12x_weights["w1_alpha"], + torch.tensor([2.0, 1.0], dtype=torch.float32), + ) + assert torch.allclose( + module._b12x_weights["w2_alpha"], + torch.tensor([2.0, 1.0], dtype=torch.float32), + ) def test_dispatch_rejects_non_tensor(): diff --git a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py index 0e20fb6034c4..45305211f9c1 100644 --- a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py +++ b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py @@ -421,7 +421,7 @@ def test_w4a16_nvfp4_post_load_ignores_checkpoint_activation_scale(): assert module.input_scale is None assert module.inv_input_scale is None assert module.alpha is None - torch.testing.assert_close(module.weight_scale_2, torch.tensor([0.25], dtype=torch.float32)) + torch.testing.assert_close(module.weight_scale_2, torch.tensor([4.0], dtype=torch.float32)) assert not hasattr(module, "tmp_nvfp4_input_scales_list") assert not hasattr(module, "tmp_nvfp4_weight_scale_2_list") From c71d57635390b5068491bd126fe97c7267b5cbd0 Mon Sep 17 00:00:00 2001 From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> Date: Thu, 28 May 2026 19:22:12 +0000 Subject: [PATCH 06/25] fix pre-commit findings Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> --- tensorrt_llm/_torch/model_config.py | 8 ++++---- tensorrt_llm/models/quant_config_utils.py | 12 ++++++++---- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index fcc4668c7d7b..d01c0b82ca77 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -23,8 +23,8 @@ from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig -from tensorrt_llm.models.quant_config_utils import \ - update_quant_config_from_compressed_tensors, is_w4a16_nvfp4_hf_quant_config +from tensorrt_llm.models.quant_config_utils import ( + is_w4a16_nvfp4_hf_quant_config, update_quant_config_from_compressed_tensors) from tensorrt_llm.quantization.mode import QuantAlgo from tensorrt_llm.quantization.modelopt_config import ( is_modelopt_quant_config, read_modelopt_quant_config, @@ -366,8 +366,8 @@ def _build_modelopt_quant_config(json_quant_configs, if 'pre_quant_scale' in json_quant_configs: quant_config.pre_quant_scale = json_quant_configs['pre_quant_scale'] - if (quant_config.quant_algo in (QuantAlgo.NVFP4, "NVFP4") and - is_w4a16_nvfp4_hf_quant_config(hf_quant_config)): + if (quant_config.quant_algo in (QuantAlgo.NVFP4, "NVFP4") + and is_w4a16_nvfp4_hf_quant_config(hf_quant_config)): quant_config.quant_algo = QuantAlgo.W4A16_NVFP4 quant_config.group_size = 16 quant_config.exclude_modules = hf_quant_config.get( diff --git a/tensorrt_llm/models/quant_config_utils.py b/tensorrt_llm/models/quant_config_utils.py index e946e5e70d7a..3c1a235c595c 100644 --- a/tensorrt_llm/models/quant_config_utils.py +++ b/tensorrt_llm/models/quant_config_utils.py @@ -18,6 +18,7 @@ from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization.mode import QuantAlgo + def is_w4a16_nvfp4_hf_quant_config(hf_quant_config): if hf_quant_config is None: return False @@ -30,10 +31,13 @@ def is_w4a16_nvfp4_hf_quant_config(hf_quant_config): return False weights_quant_config = group_config.get("weights", {}) inputs_quant_config = group_config.get("input_activations") - return (hf_quant_config.get("format") == "nvfp4-pack-quantized" - and weights_quant_config.get("num_bits") == 4 - and weights_quant_config.get("group_size") == 16 - and inputs_quant_config is None) + return ( + hf_quant_config.get("format") == "nvfp4-pack-quantized" + and weights_quant_config.get("num_bits") == 4 + and weights_quant_config.get("group_size") == 16 + and inputs_quant_config is None + ) + def update_quant_config_from_compressed_tensors( quant_config: QuantConfig, hf_quant_config: Mapping[str, Any] From f0cb5b26125a21b403562ff33d7490f5dc79d776 Mon Sep 17 00:00:00 2001 From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> Date: Thu, 11 Jun 2026 16:09:43 +0000 Subject: [PATCH 07/25] [None][fix] support Qwen3.6 NVFP4 mixed precision Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> --- cpp/kernels/fmha_v2/setup.py | 4 +- tensorrt_llm/_torch/model_config.py | 42 ++- .../checkpoints/hf/qwen3_5_weight_mapper.py | 64 +++- .../_torch/models/modeling_qwen3_5.py | 87 +++-- tensorrt_llm/_torch/models/modeling_utils.py | 14 +- tensorrt_llm/_torch/modules/embedding.py | 4 +- .../modules/fused_moe/configurable_moe.py | 1 + .../_torch/modules/fused_moe/create_moe.py | 25 +- .../_torch/modules/fused_moe/quantization.py | 12 +- tensorrt_llm/_torch/modules/linear.py | 12 - .../_torch/pyexecutor/config_utils.py | 4 + tensorrt_llm/models/quant_config_utils.py | 85 +++-- .../hf/test_qwen3_5_weight_mapper.py | 64 ++++ .../moe/test_cute_dsl_b12x_moe_backend.py | 20 +- .../_torch/modules/test_w4a16_nvfp4_linear.py | 4 +- tests/unittest/_torch/test_model_config.py | 357 ++++++++++++++++++ .../models/test_quant_config_utils.py | 27 ++ 17 files changed, 713 insertions(+), 113 deletions(-) create mode 100644 tests/unittest/_torch/models/checkpoints/hf/test_qwen3_5_weight_mapper.py diff --git a/cpp/kernels/fmha_v2/setup.py b/cpp/kernels/fmha_v2/setup.py index 0d505c233bc6..7137ee4c4e05 100644 --- a/cpp/kernels/fmha_v2/setup.py +++ b/cpp/kernels/fmha_v2/setup.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2020-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2020-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -6810,7 +6810,7 @@ def enumerate_kernels(): enumerate_qmma_flash_kernels(specs, sm=120, dtype='e4m3_fp32', - head_sizes=[128, 192, 576], + head_sizes=[128, 192, 256, 576], output_dtype="bf16") if 'ENABLE_HMMA_FP32' in os.environ: diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index d01c0b82ca77..ec6d90a9a86f 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -12,8 +12,8 @@ from transformers.utils import HF_MODULES_CACHE from tensorrt_llm._torch.pyexecutor.config_utils import ( - get_qwen3_hybrid_num_attention_layers, is_nemotron_hybrid, is_qwen3_hybrid, - load_pretrained_config) + get_qwen3_hybrid_layer_types, get_qwen3_hybrid_num_attention_layers, + is_nemotron_hybrid, is_qwen3_5, is_qwen3_hybrid, load_pretrained_config) from tensorrt_llm._utils import (get_sm_version, is_sm_100f, torch_dtype_to_binding) from tensorrt_llm.bindings import LayerType as LayerTypeCpp @@ -313,14 +313,20 @@ def resolve_moe_backend(moe_backend: str, return "CUTLASS" def resolve_moe_backend_after_quant_config( - moe_backend: str, architecture: str, - quant_config: QuantConfig) -> str: + moe_backend: str, + architecture: str, + quant_config: QuantConfig, + layer_quant_config: Optional[Dict[str, QuantConfig]] = None) -> str: """Resolve AUTO moe_backend after quantization metadata is known.""" if moe_backend.upper() != "AUTO": return moe_backend is_w4a16_nvfp4 = quant_config.quant_algo in (QuantAlgo.W4A16_NVFP4, "W4A16_NVFP4") + if not is_w4a16_nvfp4 and layer_quant_config is not None: + is_w4a16_nvfp4 = any(config.quant_algo in (QuantAlgo.W4A16_NVFP4, + "W4A16_NVFP4") + for config in layer_quant_config.values()) if is_w4a16_nvfp4 and get_sm_version() in (120, 121): return "CUTEDSL" @@ -407,12 +413,15 @@ def _build_modelopt_quant_config(json_quant_configs, 'group_size', quant_config.group_size) quant_config.exclude_modules = json_quant_configs.get( 'exclude_modules', quant_config.exclude_modules) + w4a16_nvfp4_group_size = None for layer in mixed_quant_configs: layer_cfg = mixed_quant_configs[layer] config = QuantConfig() config.kv_cache_quant_algo = kv_cache_quant_algo config.quant_algo = QuantAlgo(layer_cfg['quant_algo']) config.group_size = layer_cfg.get('group_size', None) + if config.quant_algo == QuantAlgo.W4A16_NVFP4: + w4a16_nvfp4_group_size = config.group_size or 16 # AWQ-specific extras emitted by modelopt per-layer. if 'has_zero_point' in layer_cfg: config.has_zero_point = layer_cfg['has_zero_point'] @@ -420,6 +429,9 @@ def _build_modelopt_quant_config(json_quant_configs, config.pre_quant_scale = layer_cfg['pre_quant_scale'] mixed_quant_configs[layer] = config layer_quant_config = mixed_quant_configs + if w4a16_nvfp4_group_size is not None: + quant_config.quant_algo = None + quant_config.group_size = w4a16_nvfp4_group_size elif quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES: if quant_config.group_size is None: quant_config.group_size = 128 @@ -432,6 +444,24 @@ def _build_modelopt_quant_config(json_quant_configs, ] return quant_config, layer_quant_config + @staticmethod + def _add_qwen35_qkvz_bf16_excludes(json_quant_configs, + pretrained_config) -> None: + """Keep packed Qwen3.5 linear-attention qkvz on the BF16 path.""" + if pretrained_config is None or not is_qwen3_5(pretrained_config): + return + try: + layer_types = get_qwen3_hybrid_layer_types(pretrained_config) + except (ValueError, AttributeError): + return + + exclude_modules = list(json_quant_configs.get('exclude_modules') or []) + for layer_idx, layer_type in enumerate(layer_types): + if layer_type == "linear_attention": + exclude_modules.append( + f"model.layers.{layer_idx}.linear_attn.in_proj_qkvz") + json_quant_configs['exclude_modules'] = sorted(set(exclude_modules)) + @staticmethod def get_mxfp4_quant_algo(moe_backend, is_dynamic_quant=False): quant_algo = ModelConfig.override_quant_algo() @@ -726,6 +756,7 @@ def _recursive_update_config(config: transformers.PretrainedConfig, 'hf_quant_config.json'): with open(quant_config_file) as f: normalized = read_modelopt_quant_config(json.load(f)) + cls._add_qwen35_qkvz_bf16_excludes(normalized, pretrained_config) # The file is authoritative; warn if the inline copy disagrees. # Done before _build_modelopt_quant_config since the builder may # mutate ``normalized`` via ``.update`` from quant_cfg.json. @@ -752,7 +783,8 @@ def _recursive_update_config(config: transformers.PretrainedConfig, quant_config_file, moe_backend_hint) kwargs['moe_backend'] = cls.resolve_moe_backend_after_quant_config( - requested_moe_backend, architecture, quant_config) + requested_moe_backend, architecture, quant_config, + layer_quant_config) model_config = cls(pretrained_config=pretrained_config, quant_config=quant_config, diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py index fa2f161bdc4f..3cd598ca21a0 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py @@ -73,9 +73,15 @@ def _preprocess_modelopt_ckpt(self, weights: dict) -> tuple[bool, dict]: is_modelopt_ckpt = True prefix = key[: -len(".weight_scale")] if prefix not in nvfp4_prefixes: - new_key = f"{key}_inv" - # modelopt fp8_pb_wo has 2 extra singleton dimensions - tensor = tensor.squeeze(1).squeeze(-1) + # modelopt fp8_pb_wo stores per-block FP8 scales as (out, 1, 1). + # MIXED_PRECISION checkpoints (e.g. Qwen3.6 NVFP4) also use + # per-tensor scalar FP8 scales (0-D). Keep those under the + # ModelOpt weight_scale key because FP8QDQLinearMethod + # loads that spelling. + if tensor.ndim != 0: + new_key = f"{key}_inv" + if tensor.ndim == 3 and tensor.shape[1] == 1 and tensor.shape[-1] == 1: + tensor = tensor.squeeze(1).squeeze(-1) if new_key in remapped_weights: raise ValueError(f"Duplicate remapped key found: {new_key}") remapped_weights[new_key] = tensor @@ -191,6 +197,52 @@ def _dequantize_linear_attn_fp8_qkvz(self, weights: dict) -> dict: updated_weights.pop(scale_name, None) return updated_weights + def _dequantize_linear_attn_per_tensor_fp8(self, weights: dict) -> dict: + """Dequantize per-tensor FP8 linear-attn split projections to BF16. + + MIXED_PRECISION checkpoints (e.g. Qwen3.6 NVFP4) store the linear-attn + ``in_proj_*`` split projections (in_proj_qkv, in_proj_z, in_proj_q/k/v, + etc.) as per-tensor FP8: scalar ``weight_scale`` or + ``weight_scale_inv`` (and matching scalar ``input_scale``). The packer + below assumes BF16 weights or per-block FP8 with block-shaped scales, + so dequantize here and drop the scale tensors. The packed BF16 tensor + is then concatenated with the already-BF16 ``in_proj_a`` / + ``in_proj_b`` / ``in_proj_z`` weights and TP-sharded by the standard + path. + """ + target_dtype = getattr(self.config.pretrained_config, "torch_dtype", torch.bfloat16) + if target_dtype is None: + target_dtype = torch.bfloat16 + updated = dict(weights) + drop = [] + for name in list(weights): + if not name.endswith(".weight"): + continue + if self._SPLIT_PROJ_PATTERN.match(name) is None: + continue + prefix = name[: -len(".weight")] + scale_name = prefix + ".weight_scale_inv" + if scale_name not in weights: + scale_name = prefix + ".weight_scale" + if scale_name not in weights: + continue + scale = weights[scale_name] + # Per-tensor scalar scale → dequantize. Per-block (n-D) handled elsewhere. + if scale.ndim != 0: + continue + updated[name] = ( + (weights[name].to(torch.float32) * scale.to(torch.float32)) + .to(target_dtype) + .contiguous() + ) + drop.append(scale_name) + input_scale_name = name[: -len(".weight")] + ".input_scale" + if input_scale_name in weights: + drop.append(input_scale_name) + for k in drop: + updated.pop(k, None) + return updated + def _pack_split_projections(self, weights: dict) -> dict: config = self.config.pretrained_config num_k_groups = config.linear_num_key_heads @@ -306,6 +358,12 @@ def preprocess_weights(self, weights: dict) -> dict: normalized_weights = self._normalize_weight_names(weights) is_modelopt_ckpt, normalized_weights = self._preprocess_modelopt_ckpt(normalized_weights) + # MIXED_PRECISION modelopt checkpoints store the linear-attn split + # projections as per-tensor FP8 (scalar scales). Dequantize to BF16 + # before packing so the packer can run its existing BF16 path. + if is_modelopt_ckpt: + normalized_weights = self._dequantize_linear_attn_per_tensor_fp8(normalized_weights) + packed_weights = self._pack_split_projections(normalized_weights) if not is_modelopt_ckpt: packed_weights = self._dequantize_linear_attn_fp8_qkvz(packed_weights) diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_5.py b/tensorrt_llm/_torch/models/modeling_qwen3_5.py index bf83e916db29..8a84c380eb05 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_5.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_5.py @@ -51,46 +51,73 @@ def _translate_mtp_pattern(name, n_hidden_layers): return None +def _normalize_qwen35_module_path(name, n_hidden_layers, translate_mtp): + if name.startswith(_LANG_PREFIX): + name = "model." + name[len(_LANG_PREFIX) :] + if name.startswith("model.visual."): + return None + if name.startswith("mtp."): + if not translate_mtp or n_hidden_layers is None: + return None + return _translate_mtp_pattern(name, n_hidden_layers) + name = re.sub( + r"(\.mlp\.shared_expert)\.(gate_proj|up_proj)(?=$|\.|\*)", + r"\1.gate_up_proj", + name, + ) + name = re.sub(r"\.in_proj_[ab](?=$|\.|\*)", ".in_proj_ba", name) + name = re.sub(r"\.in_proj_(q|k|v|z|qkv)(?=$|\.|\*)", ".in_proj_qkvz", name) + return name + + def _normalize_qwen35_exclude_modules(model_config): """Normalize NVFP4/FP8 exclude_modules from HF naming to TRT-LLM naming. hf_quant_config.json stores exclude patterns in HF checkpoint namespace (e.g. ``model.language_model.layers.0.linear_attn*`` and ``mtp.layers.0*``), - but TRT-LLM modules use ``model.layers.0.linear_attn.in_proj_qkvz`` and - map the MTP layer to ``model.layers..*``. This - function translates the patterns so that - ``apply_quant_config_exclude_modules`` can match them. + and mixed-precision ``quantized_layers`` use the same namespace. TRT-LLM + modules use ``model.layers.0.linear_attn.in_proj_qkvz`` and map the MTP + layer to ``model.layers..*``. This function translates + the patterns so that ``apply_quant_config_exclude_modules`` and + ``apply_layerwise_quant_config`` can match them. """ qc = model_config.quant_config - if qc is None or qc.exclude_modules is None: - return - n_hidden_layers = getattr(model_config.pretrained_config, "num_hidden_layers", None) - normalized = set() - for name in qc.exclude_modules: - # Strip VLM prefix: model.language_model.X -> model.X - if name.startswith(_LANG_PREFIX): - name = "model." + name[len(_LANG_PREFIX) :] - # Drop vision tensors (not part of the language model graph) - if name.startswith("model.visual"): - continue - # Translate MTP-namespace patterns to TRT-LLM paths so the MTP - # layer (which the checkpoint stores unquantized) gets correctly - # excluded from the global NVFP4/FP8 quant_config. - if name.startswith("mtp."): - if n_hidden_layers is None: + if qc is not None and qc.exclude_modules is not None: + normalized = set() + for name in qc.exclude_modules: + normalized_name = _normalize_qwen35_module_path( + name, n_hidden_layers, translate_mtp=True + ) + if normalized_name is not None: + normalized.add(normalized_name) + qc.exclude_modules = sorted(normalized) + + if model_config.quant_config_dict is not None: + normalized_quant_config_dict = {} + for name, quant_config in model_config.quant_config_dict.items(): + normalized_name = _normalize_qwen35_module_path( + name, n_hidden_layers, translate_mtp=False + ) + if normalized_name is None: + continue + if normalized_name in normalized_quant_config_dict: + if normalized_quant_config_dict[normalized_name] != quant_config: + raise ValueError( + "Conflicting Qwen3.5 quant configs after normalizing " + f"{name} to {normalized_name}" + ) continue - translated = _translate_mtp_pattern(name, n_hidden_layers) - if translated is not None: - normalized.add(translated) - continue - # Map split projection names to packed TRT-LLM names - name = re.sub(r"\.in_proj_[ab](\b|\*)", ".in_proj_ba*", name) - name = re.sub(r"\.in_proj_(q|k|v|z|qkv)(\b|\*)", ".in_proj_qkvz*", name) - normalized.add(name) - - qc.exclude_modules = sorted(normalized) + normalized_quant_config_dict[normalized_name] = quant_config + was_frozen = getattr(model_config, "_frozen", False) + if was_frozen: + model_config._frozen = False + try: + model_config.quant_config_dict = normalized_quant_config_dict + finally: + if was_frozen: + model_config._frozen = True @register_auto_model("Qwen3_5MoeForCausalLM") diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 0fae7075ed13..8136e2db92bc 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -51,6 +51,10 @@ class MetaInitException(RuntimeError): pass +def _is_same_or_child_module_path(lhs: str, rhs: str) -> bool: + return lhs == rhs or lhs.startswith(f"{rhs}.") or rhs.startswith(f"{lhs}.") + + class MetaInitMode(TorchDispatchMode): """ Context for skip random parameter initialization @@ -374,6 +378,11 @@ def __init__(self, model: TModel, *, config: ModelConfig[TConfig], in (QuantAlgo.W4A16_NVFP4, "W4A16_NVFP4") and not config. quant_config.is_module_excluded_from_quantization("lm_head")): lm_head_quant_config = config.quant_config + elif config.quant_config_dict is not None: + for name, quant_config in config.quant_config_dict.items(): + if _is_same_or_child_module_path("lm_head", name): + lm_head_quant_config = quant_config + break if config.mapping.enable_attention_dp and not config.mapping.enable_lm_head_tp_in_adp: self.lm_head = LMHead( @@ -381,6 +390,7 @@ def __init__(self, model: TModel, *, config: ModelConfig[TConfig], hidden_size, dtype=config.pretrained_config.torch_dtype, quant_config=lm_head_quant_config, + skip_create_weights_in_init=config.skip_create_weights_in_init, ) else: if (hasattr(config, 'lora_config') @@ -405,6 +415,7 @@ def __init__(self, model: TModel, *, config: ModelConfig[TConfig], use_custom_cublas_mm=getattr(model, 'use_custom_cublas_mm', False), quant_config=lm_head_quant_config, + skip_create_weights_in_init=config.skip_create_weights_in_init, ) if self.has_custom_lm_head: @@ -425,6 +436,7 @@ def __init__(self, model: TModel, *, config: ModelConfig[TConfig], assert self.lm_head.tp_mode == self.model.embed_tokens.tp_mode, ( "lm_head and vocab embedding should use the same TP mode") self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head._weights_created = True if config.mapping.is_last_pp_rank(): self.model.keep_embed_tokens = True @@ -454,7 +466,7 @@ def apply_layerwise_quant_config(self): if isinstance(module, (MoE, VanillaMoE)): for n, q in quant_config_dict.items(): # all linear layers inside FusedMoE share the same quant config - if name in n: + if _is_same_or_child_module_path(name, n): module.quant_config = q break elif isinstance(module, Linear): diff --git a/tensorrt_llm/_torch/modules/embedding.py b/tensorrt_llm/_torch/modules/embedding.py index 0d878dee4cfd..2256e5fa30d7 100644 --- a/tensorrt_llm/_torch/modules/embedding.py +++ b/tensorrt_llm/_torch/modules/embedding.py @@ -38,6 +38,7 @@ def __init__( reduce_output: bool = True, use_custom_cublas_mm: bool = False, quant_config: Optional[QuantConfig] = None, + skip_create_weights_in_init: bool = False, ): local_in_features = embedding_dim local_out_features = num_embeddings @@ -71,6 +72,7 @@ def __init__( reduce_output=reduce_output, use_custom_cublas_mm=use_custom_cublas_mm, quant_config=quant_config, + skip_create_weights_in_init=skip_create_weights_in_init, ) if tensor_parallel_mode == TensorParallelMode.ROW: @@ -81,7 +83,7 @@ def __init__( self.num_embeddings = num_embeddings self.embedding_dim = embedding_dim - if not self.has_any_quant: + if not skip_create_weights_in_init and not self.has_any_quant: weight_shape = (self.out_features, self.in_features) self.weight = Parameter(torch.empty(weight_shape, dtype=dtype)) self.register_parameter("bias", None) diff --git a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py index 690894af84a6..79b86ded67df 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py @@ -292,6 +292,7 @@ def _create_and_sync_backend( routing_method, self.dtype, override_quant_config=override_quant_config, + layer_idx=self.layer_idx, ) backend_model_config = model_config diff --git a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py index 9f5d22f1de3d..cae5237f7d60 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py @@ -26,15 +26,31 @@ from .routing import BaseMoeRoutingMethod +def _is_same_or_child_module_path(lhs: str, rhs: str) -> bool: + return lhs == rhs or lhs.startswith(f"{rhs}.") or rhs.startswith(f"{lhs}.") + + +def _get_layer_quant_config(model_config: ModelConfig, + layer_idx: Optional[int]) -> Optional[QuantConfig]: + if layer_idx is None or model_config.quant_config_dict is None: + return None + + moe_module_name = f"model.layers.{layer_idx}.mlp.experts" + for name, quant_config in model_config.quant_config_dict.items(): + if _is_same_or_child_module_path(name, moe_module_name): + return quant_config + return None + + def get_moe_cls( model_config: ModelConfig, override_quant_config: Optional[QuantConfig] = None, layer_idx: Optional[int] = None, ) -> Type[MoE]: moe_backend = model_config.moe_backend - quant_config = model_config.quant_config - if override_quant_config is not None: - quant_config = override_quant_config + quant_config = (override_quant_config + or _get_layer_quant_config(model_config, layer_idx) + or model_config.quant_config) layer_prefix = f"[layer_idx={layer_idx}] " if layer_idx is not None else "" if moe_backend.upper() == "CUTLASS": return CutlassFusedMoE @@ -192,7 +208,8 @@ def resolve_moe_cls( ) -> Type[MoE]: moe_cls = get_moe_cls(model_config, override_quant_config, layer_idx) - effective_quant_config = override_quant_config or model_config.quant_config + effective_quant_config = (override_quant_config or _get_layer_quant_config( + model_config, layer_idx) or model_config.quant_config) has_quant = (effective_quant_config is not None and effective_quant_config.layer_quant_mode.has_any_quant( exclude_kv_cache=True)) diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index b6f8bb305095..c0913d3b2ce3 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py @@ -3081,16 +3081,8 @@ def post_load_weights(self, module: torch.nn.Module): # W4A16 path: BF16/FP16 activations multiplied by FP4 weights. # FlashInfer's W4A16 packer expects the ModelOpt scale contract: # normalized FP8 block scales plus per-expert ``weight_global_scale``. - # TRT-LLM stores the reciprocal as ``weight_scale_2``; recover the - # ModelOpt value from the already computed dequant scale. - def _to_modelopt_global_scale(scale: torch.Tensor) -> torch.Tensor: - global_scale = torch.zeros_like(scale, dtype=torch.float32) - valid = scale > 0 - global_scale[valid] = scale[valid].reciprocal() - return global_scale.contiguous() - - w1_alpha_b12x = _to_modelopt_global_scale(w1_w_scale_2) - w2_alpha_b12x = _to_modelopt_global_scale(w2_w_scale_2) + w1_alpha_b12x = w1_w_scale_2.to(torch.float32).contiguous() + w2_alpha_b12x = w2_w_scale_2.to(torch.float32).contiguous() fc2_input_scale_b12x = None else: w1_alpha_b12x = ((1.0 / module.fc31_input_scale).expand( diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index ab94f69346ce..b84df0999abd 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -1927,8 +1927,6 @@ def _process_weights_without_static_activation_scale( original_alpha = module.alpha had_scalar_alpha = hasattr(module, "scalar_alpha") original_scalar_alpha = getattr(module, "scalar_alpha", None) - has_weight_global_scale = bool( - getattr(module, "tmp_nvfp4_weight_scale_2_list", [])) device = module.weight_scale_2.device module.input_scale = Parameter(torch.empty([1], @@ -1945,8 +1943,6 @@ def _process_weights_without_static_activation_scale( requires_grad=False) try: process_fn(module) - if has_weight_global_scale: - self._convert_weight_global_scale_to_dequant_scale(module) finally: module.input_scale = original_input_scale module.inv_input_scale = original_inv_input_scale @@ -1956,14 +1952,6 @@ def _process_weights_without_static_activation_scale( elif hasattr(module, "scalar_alpha"): delattr(module, "scalar_alpha") - @staticmethod - def _convert_weight_global_scale_to_dequant_scale(module: Linear): - weight_scale_2 = module.weight_scale_2.data.float() - dequant_scale = torch.zeros_like(weight_scale_2) - nonzero = weight_scale_2 != 0 - dequant_scale[nonzero] = weight_scale_2[nonzero].reciprocal() - copy_weight(module.weight_scale_2, dequant_scale) - def process_weights_after_loading_vanilla(self, module: Linear): self._process_weights_without_static_activation_scale( module, diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index ef270040fc26..5a3181f95aa7 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -335,6 +335,10 @@ def _inherit_quantization_config(config_dict: dict, modules = _Qwen35ConfigCompat._add_qkvz_bf16_workaround( text_config, modules) quantization_config["modules_to_not_convert"] = sorted(set(modules)) + if "ignore" in quantization_config: + modules = _Qwen35ConfigCompat._add_qkvz_bf16_workaround( + text_config, list(quantization_config["ignore"])) + quantization_config["ignore"] = sorted(set(modules)) text_config["quantization_config"] = quantization_config return text_config diff --git a/tensorrt_llm/models/quant_config_utils.py b/tensorrt_llm/models/quant_config_utils.py index 3c1a235c595c..41fe12dafe3e 100644 --- a/tensorrt_llm/models/quant_config_utils.py +++ b/tensorrt_llm/models/quant_config_utils.py @@ -47,56 +47,59 @@ def update_quant_config_from_compressed_tensors( if config_groups is None: raise ValueError(f"config_groups is not set in {hf_quant_config}.") - weights_quant_config = config_groups["group_0"]["weights"] - inputs_quant_config = config_groups["group_0"]["input_activations"] - weights_quant_strategy = weights_quant_config["strategy"] - inputs_quant_strategy = inputs_quant_config["strategy"] + group_config = config_groups["group_0"] + weights_quant_config = group_config["weights"] if is_w4a16_nvfp4_hf_quant_config(hf_quant_config): quant_config.quant_algo = QuantAlgo.W4A16_NVFP4 quant_config.group_size = 16 - elif weights_quant_config["num_bits"] == 8: - if weights_quant_strategy == "channel": - if inputs_quant_strategy != "token": - raise ValueError(f"Unsupported inputs_quant_strategy: {inputs_quant_strategy}.") - quant_config.quant_algo = QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN - elif weights_quant_strategy == "block": - if inputs_quant_strategy != "group": - raise ValueError(f"Unsupported inputs_quant_strategy: {inputs_quant_strategy}.") - quant_config.quant_algo = QuantAlgo.FP8_BLOCK_SCALES - group_size = inputs_quant_config["group_size"] + else: + inputs_quant_config = group_config["input_activations"] + weights_quant_strategy = weights_quant_config["strategy"] + inputs_quant_strategy = inputs_quant_config["strategy"] - # TRT-LLM only supports group_size=128 for FP8_BLOCK_SCALES. - if group_size != 128: - raise ValueError(f"Unsupported group_size: {group_size}. Supported: 128.") - quant_config.group_size = group_size + if weights_quant_config["num_bits"] == 8: + if weights_quant_strategy == "channel": + if inputs_quant_strategy != "token": + raise ValueError(f"Unsupported inputs_quant_strategy: {inputs_quant_strategy}.") + quant_config.quant_algo = QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN + elif weights_quant_strategy == "block": + if inputs_quant_strategy != "group": + raise ValueError(f"Unsupported inputs_quant_strategy: {inputs_quant_strategy}.") + quant_config.quant_algo = QuantAlgo.FP8_BLOCK_SCALES + group_size = inputs_quant_config["group_size"] + # TRT-LLM only supports group_size=128 for FP8_BLOCK_SCALES. + if group_size != 128: + raise ValueError(f"Unsupported group_size: {group_size}. Supported: 128.") + quant_config.group_size = group_size + + else: + raise ValueError( + f"Unsupported weights_quant_strategy: {weights_quant_strategy}. " + "Supported strategies: 'channel', 'block'." + ) + elif ( + weights_quant_config["num_bits"] == 4 + and weights_quant_config.get("type") == "float" + and weights_quant_strategy == "tensor_group" + ): + # llm-compressor NVFP4: weights FP4 with FP8 per-group scales + # (group_size=16), scaled by an FP32 global scale. + if inputs_quant_strategy != "tensor_group": + raise ValueError( + f"Unsupported inputs_quant_strategy for NVFP4: {inputs_quant_strategy}." + ) + group_size = weights_quant_config["group_size"] + if group_size != 16: + raise ValueError(f"Unsupported group_size: {group_size}. Supported: 16 for NVFP4.") + quant_config.quant_algo = QuantAlgo.NVFP4 + quant_config.group_size = group_size else: raise ValueError( - f"Unsupported weights_quant_strategy: {weights_quant_strategy}. " - "Supported strategies: 'channel', 'block'." - ) - elif ( - weights_quant_config["num_bits"] == 4 - and weights_quant_config.get("type") == "float" - and weights_quant_strategy == "tensor_group" - ): - # llm-compressor NVFP4: weights FP4 with FP8 per-group scales - # (group_size=16), scaled by an FP32 global scale. - if inputs_quant_strategy != "tensor_group": - raise ValueError( - f"Unsupported inputs_quant_strategy for NVFP4: {inputs_quant_strategy}." + f"Unsupported quant_bits: {weights_quant_config['num_bits']}. " + "Supported: 8 (FP8) or 4 (NVFP4)." ) - group_size = weights_quant_config["group_size"] - if group_size != 16: - raise ValueError(f"Unsupported group_size: {group_size}. Supported: 16 for NVFP4.") - quant_config.quant_algo = QuantAlgo.NVFP4 - quant_config.group_size = group_size - else: - raise ValueError( - f"Unsupported quant_bits: {weights_quant_config['num_bits']}. " - "Supported: 8 (FP8) or 4 (NVFP4)." - ) # kv_cache_scheme (llm-compressor): FP8 per-tensor KV cache. kv_cache_scheme = hf_quant_config.get("kv_cache_scheme") diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_qwen3_5_weight_mapper.py b/tests/unittest/_torch/models/checkpoints/hf/test_qwen3_5_weight_mapper.py new file mode 100644 index 000000000000..2a5e6867b9f6 --- /dev/null +++ b/tests/unittest/_torch/models/checkpoints/hf/test_qwen3_5_weight_mapper.py @@ -0,0 +1,64 @@ +# 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 types import SimpleNamespace + +import torch + +from tensorrt_llm._torch.models.checkpoints.hf.qwen3_5_weight_mapper import Qwen3_5MoeHfWeightMapper + + +def _make_mapper(dtype=torch.bfloat16): + mapper = Qwen3_5MoeHfWeightMapper() + mapper._config = SimpleNamespace(pretrained_config=SimpleNamespace(torch_dtype=dtype)) + return mapper + + +def test_qwen35_modelopt_preprocess_preserves_scalar_fp8_scale_name(): + mapper = _make_mapper() + scale = torch.tensor(0.5, dtype=torch.float32) + + is_modelopt_ckpt, weights = mapper._preprocess_modelopt_ckpt( + { + "model.layers.0.linear_attn.out_proj.weight_scale": scale, + } + ) + + assert is_modelopt_ckpt + assert "model.layers.0.linear_attn.out_proj.weight_scale" in weights + assert weights["model.layers.0.linear_attn.out_proj.weight_scale"].shape == torch.Size([]) + + +def test_qwen35_dequantizes_per_tensor_fp8_linear_attention_projection(): + mapper = _make_mapper(dtype=torch.bfloat16) + weight_name = "model.layers.0.linear_attn.in_proj_qkv.weight" + scale_name = "model.layers.0.linear_attn.in_proj_qkv.weight_scale" + input_scale_name = "model.layers.0.linear_attn.in_proj_qkv.input_scale" + + weights = { + weight_name: torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32), + scale_name: torch.tensor(0.5, dtype=torch.float32), + input_scale_name: torch.tensor(1.0, dtype=torch.float32), + } + + updated = mapper._dequantize_linear_attn_per_tensor_fp8(weights) + + assert scale_name not in updated + assert input_scale_name not in updated + assert updated[weight_name].dtype == torch.bfloat16 + torch.testing.assert_close( + updated[weight_name], + torch.tensor([[0.5, 1.0], [1.5, 2.0]], dtype=torch.bfloat16), + ) diff --git a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py index b278e963137a..8ed41823048d 100644 --- a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py @@ -177,6 +177,22 @@ def test_get_moe_cls_cutedsl_selects_b12x_for_w4a16_nvfp4_on_supported_sm(sm_ver assert cls is CuteDslB12xFusedMoE +@pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) +def test_get_moe_cls_cutedsl_selects_b12x_for_layer_w4a16_nvfp4_on_supported_sm(sm_version): + """MIXED_PRECISION per-layer W4A16 must select the same backend that the + layer will use after apply_layerwise_quant_config(). + """ + cfg = ModelConfig() + cfg.moe_backend = "CUTEDSL" + cfg.quant_config = QuantConfig() + cfg.quant_config_dict = { + "model.layers.0.mlp.experts": QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16), + } + with patch("tensorrt_llm._utils.get_sm_version", return_value=sm_version): + cls = get_moe_cls(cfg, layer_idx=0) + assert cls is CuteDslB12xFusedMoE + + def test_get_moe_cls_cutedsl_falls_back_to_plain_cutedsl_when_flashinfer_missing(monkeypatch): """CUTEDSL + NVFP4 + SM120/121 + flashinfer NOT importable → CuteDslFusedMoE.""" import builtins @@ -336,11 +352,11 @@ def _convert_sf_to_mma_layout(scales, *, m, k, num_groups): ) assert torch.allclose( module._b12x_weights["w1_alpha"], - torch.tensor([2.0, 1.0], dtype=torch.float32), + torch.tensor([0.5, 1.0], dtype=torch.float32), ) assert torch.allclose( module._b12x_weights["w2_alpha"], - torch.tensor([2.0, 1.0], dtype=torch.float32), + torch.tensor([0.5, 1.0], dtype=torch.float32), ) diff --git a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py index 45305211f9c1..f8732b4875c9 100644 --- a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py +++ b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py @@ -405,7 +405,7 @@ def fail_fp4_quantize(*args, **kwargs): assert output.shape == (17, 3) -def test_w4a16_nvfp4_post_load_ignores_checkpoint_activation_scale(): +def test_w4a16_nvfp4_post_load_preserves_checkpoint_weight_global_scale(): method = W4A16NVFP4LinearMethod() module = SimpleNamespace( input_scale=None, @@ -421,7 +421,7 @@ def test_w4a16_nvfp4_post_load_ignores_checkpoint_activation_scale(): assert module.input_scale is None assert module.inv_input_scale is None assert module.alpha is None - torch.testing.assert_close(module.weight_scale_2, torch.tensor([4.0], dtype=torch.float32)) + torch.testing.assert_close(module.weight_scale_2, torch.tensor([0.25], dtype=torch.float32)) assert not hasattr(module, "tmp_nvfp4_input_scales_list") assert not hasattr(module, "tmp_nvfp4_weight_scale_2_list") diff --git a/tests/unittest/_torch/test_model_config.py b/tests/unittest/_torch/test_model_config.py index 1aebdd806942..6bb8533cfbf7 100644 --- a/tests/unittest/_torch/test_model_config.py +++ b/tests/unittest/_torch/test_model_config.py @@ -3,8 +3,14 @@ import pytest import torch +from torch import nn from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_qwen3_5 import _normalize_qwen35_exclude_modules +from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM +from tensorrt_llm._torch.modules.fused_moe.interface import MoE +from tensorrt_llm._torch.modules.linear import TensorParallelMode +from tensorrt_llm._torch.pyexecutor.config_utils import _Qwen35ConfigCompat from tensorrt_llm._torch.pyexecutor.model_loader import validate_and_set_kv_cache_quant from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig @@ -202,6 +208,338 @@ def test_load_modelopt_quant_config_respects_hf_w4a16_metadata(tmp_path): assert layer_quant_config is None +def test_load_modelopt_mixed_precision_w4a16_uses_layer_quant_config(tmp_path): + modelopt_quant_config = { + "quantization": { + "quant_algo": "MIXED_PRECISION", + "kv_cache_quant_algo": "FP8", + "exclude_modules": ["mtp*"], + "quantized_layers": { + "model.layers.0.linear_attn.in_proj_qkv": { + "quant_algo": "FP8", + }, + "model.layers.0.mlp.experts": { + "quant_algo": "W4A16_NVFP4", + "group_size": 16, + }, + }, + } + } + quant_config_file = tmp_path / "hf_quant_config.json" + quant_config_file.write_text(json.dumps(modelopt_quant_config), encoding="utf-8") + + quant_config, layer_quant_config = ModelConfig.load_modelopt_quant_config( + str(quant_config_file), + str(tmp_path), + moe_backend="CUTLASS", + ) + + assert quant_config.quant_algo is None + assert quant_config.group_size == 16 + assert quant_config.kv_cache_quant_algo == QuantAlgo.FP8 + assert quant_config.exclude_modules == ["mtp*"] + assert layer_quant_config is not None + assert layer_quant_config["model.layers.0.linear_attn.in_proj_qkv"].quant_algo == QuantAlgo.FP8 + assert layer_quant_config["model.layers.0.mlp.experts"].quant_algo == QuantAlgo.W4A16_NVFP4 + + +def test_qwen35_compat_adds_qkvz_excludes_to_modelopt_ignore(): + config_dict = { + "architectures": ["Qwen3_5MoeForConditionalGeneration"], + "model_type": "qwen3_5_moe", + "text_config": { + "model_type": "qwen3_5_moe_text", + "num_hidden_layers": 4, + "num_experts": 256, + "layer_types": [ + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + ], + }, + "quantization_config": { + "producer": {"name": "modelopt"}, + "quant_method": "modelopt", + "quant_algo": "MIXED_PRECISION", + "ignore": [ + "model.language_model.layers.0.linear_attn.in_proj_qkv", + "mtp.layers.0*", + ], + "quantized_layers": {}, + }, + } + + normalized = _Qwen35ConfigCompat.normalize(config_dict) + + ignore = normalized["quantization_config"]["ignore"] + assert "model.language_model.layers.0.linear_attn.in_proj_qkv" in ignore + assert "mtp.layers.0*" in ignore + assert "model.layers.0.linear_attn.in_proj_qkvz" in ignore + assert "model.layers.2.linear_attn.in_proj_qkvz" in ignore + assert "model.layers.3.linear_attn.in_proj_qkvz" in ignore + assert "model.layers.1.linear_attn.in_proj_qkvz" not in ignore + + +def test_qwen35_file_modelopt_config_adds_qkvz_excludes(tmp_path): + config_dict = { + "architectures": ["Qwen3_5MoeForConditionalGeneration"], + "model_type": "qwen3_5_moe", + "text_config": { + "model_type": "qwen3_5_moe_text", + "vocab_size": 1024, + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 4, + "num_attention_heads": 4, + "num_key_value_heads": 1, + "head_dim": 4, + "max_position_embeddings": 4096, + "torch_dtype": "bfloat16", + "num_experts": 2, + "num_experts_per_tok": 1, + "moe_intermediate_size": 8, + "shared_expert_intermediate_size": 8, + "linear_key_head_dim": 4, + "linear_value_head_dim": 4, + "linear_num_key_heads": 1, + "linear_num_value_heads": 4, + "linear_conv_kernel_dim": 4, + "layer_types": [ + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + ], + }, + } + hf_quant_config = { + "producer": {"name": "modelopt"}, + "quantization": { + "quant_algo": "MIXED_PRECISION", + "kv_cache_quant_algo": "FP8", + "exclude_modules": ["mtp.layers.0*", "mtp*"], + "quantized_layers": { + "model.language_model.layers.0.linear_attn.in_proj_qkv": { + "quant_algo": "FP8", + }, + "model.language_model.layers.0.mlp.experts": { + "quant_algo": "W4A16_NVFP4", + "group_size": 16, + }, + }, + }, + } + (tmp_path / "config.json").write_text(json.dumps(config_dict), encoding="utf-8") + (tmp_path / "hf_quant_config.json").write_text(json.dumps(hf_quant_config), encoding="utf-8") + + model_config = ModelConfig.from_pretrained( + str(tmp_path), + moe_backend="CUTLASS", + ) + + exclude_modules = model_config.quant_config.exclude_modules + assert "mtp.layers.0*" in exclude_modules + assert "mtp*" in exclude_modules + assert "model.layers.0.linear_attn.in_proj_qkvz" in exclude_modules + assert "model.layers.2.linear_attn.in_proj_qkvz" in exclude_modules + assert "model.layers.3.linear_attn.in_proj_qkvz" in exclude_modules + assert "model.layers.1.linear_attn.in_proj_qkvz" not in exclude_modules + + +def test_qwen35_normalizes_layer_quant_config_keys(): + model_config = types.SimpleNamespace( + pretrained_config=types.SimpleNamespace(num_hidden_layers=40), + quant_config=QuantConfig(exclude_modules=["mtp.layers.0*"]), + quant_config_dict={ + "model.language_model.layers.0.linear_attn.in_proj_qkv": QuantConfig( + quant_algo=QuantAlgo.FP8 + ), + "model.language_model.layers.0.linear_attn.in_proj_z": QuantConfig( + quant_algo=QuantAlgo.FP8 + ), + "model.language_model.layers.0.linear_attn.in_proj_b": QuantConfig( + quant_algo=QuantAlgo.FP8 + ), + "model.language_model.layers.0.mlp.experts": QuantConfig( + quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16 + ), + "model.language_model.layers.0.mlp.shared_expert.gate_proj": QuantConfig( + quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16 + ), + "model.language_model.layers.0.mlp.shared_expert.up_proj": QuantConfig( + quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16 + ), + "model.language_model.layers.0.mlp.shared_expert.down_proj": QuantConfig( + quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16 + ), + "model.visual.patch_embed": QuantConfig(quant_algo=QuantAlgo.FP8), + "mtp.layers.0.mlp.experts": QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4), + }, + ) + + _normalize_qwen35_exclude_modules(model_config) + + assert model_config.quant_config.exclude_modules == ["model.layers.40*"] + assert set(model_config.quant_config_dict) == { + "model.layers.0.linear_attn.in_proj_ba", + "model.layers.0.linear_attn.in_proj_qkvz", + "model.layers.0.mlp.experts", + "model.layers.0.mlp.shared_expert.down_proj", + "model.layers.0.mlp.shared_expert.gate_up_proj", + } + assert ( + model_config.quant_config_dict["model.layers.0.linear_attn.in_proj_qkvz"].quant_algo + == QuantAlgo.FP8 + ) + assert ( + model_config.quant_config_dict["model.layers.0.mlp.experts"].quant_algo + == QuantAlgo.W4A16_NVFP4 + ) + assert ( + model_config.quant_config_dict["model.layers.0.mlp.shared_expert.down_proj"].quant_algo + == QuantAlgo.W4A16_NVFP4 + ) + assert ( + model_config.quant_config_dict["model.layers.0.mlp.shared_expert.gate_up_proj"].quant_algo + == QuantAlgo.W4A16_NVFP4 + ) + + +def test_qwen35_normalizes_frozen_layer_quant_config(): + model_config = ModelConfig( + pretrained_config=types.SimpleNamespace(num_hidden_layers=40), + quant_config=QuantConfig(exclude_modules=[]), + quant_config_dict={ + "model.language_model.layers.0.mlp.shared_expert.down_proj": QuantConfig( + quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16 + ), + }, + ) + model_config._frozen = True + + _normalize_qwen35_exclude_modules(model_config) + + assert model_config._frozen + assert set(model_config.quant_config_dict) == {"model.layers.0.mlp.shared_expert.down_proj"} + assert ( + model_config.quant_config_dict["model.layers.0.mlp.shared_expert.down_proj"].quant_algo + == QuantAlgo.W4A16_NVFP4 + ) + + +class _LayerwiseQuantRoutingMoE(MoE): + @classmethod + def can_implement(cls, quant_algo, dtype_activation=torch.bfloat16, swiglu_gptoss_style=False): + return True, None + + def __init__(self): + nn.Module.__init__(self) + self.quant_config = QuantConfig() + + def create_weights(self): + pass + + def load_weights(self, weights, allow_partial_loading=False): + pass + + def quantize_input(self, x, **kwargs): + return x, None + + def run_moe(self, x, token_selected_experts, token_final_scales, x_sf=None, **kwargs): + return x + + def forward_impl(self, x, router_logits, **kwargs): + return x + + +class _LayerwiseQuantRoutingModel(nn.Module): + apply_layerwise_quant_config = DecoderModelForCausalLM.apply_layerwise_quant_config + + def __init__(self, quant_config_dict): + super().__init__() + self.model_config = types.SimpleNamespace(quant_config_dict=quant_config_dict) + self.model = nn.Module() + layer = nn.Module() + layer.mlp = nn.Module() + layer.mlp.experts = _LayerwiseQuantRoutingMoE() + layer.mlp.experts.backend = _LayerwiseQuantRoutingMoE() + self.model.layers = nn.ModuleList([layer]) + + +class _LmHeadRoutingDecoder(nn.Module): + def __pp_init__(self): + pass + + +class _TiedLmHeadRoutingDecoder(_LmHeadRoutingDecoder): + def __init__(self, hidden_size, vocab_size): + super().__init__() + self.embed_tokens = nn.Module() + self.embed_tokens.tp_size = 1 + self.embed_tokens.tp_mode = TensorParallelMode.COLUMN + self.embed_tokens.weight = nn.Parameter(torch.empty(vocab_size, hidden_size)) + + +def test_apply_layerwise_quant_config_updates_moe_inner_backend(): + quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16) + model = _LayerwiseQuantRoutingModel({"model.layers.0.mlp.experts": quant_config}) + outer = model.model.layers[0].mlp.experts + inner = outer.backend + + model.apply_layerwise_quant_config() + + assert outer.quant_config.quant_algo == QuantAlgo.W4A16_NVFP4 + assert inner.quant_config.quant_algo == QuantAlgo.W4A16_NVFP4 + + +def test_decoder_model_routes_layer_quant_config_to_lm_head_before_weight_creation(): + hidden_size = 2048 + vocab_size = 2048 + quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16) + model_config = ModelConfig( + pretrained_config=types.SimpleNamespace( + torch_dtype=torch.bfloat16, + tie_word_embeddings=False, + ), + quant_config=QuantConfig(), + quant_config_dict={"lm_head": quant_config}, + skip_create_weights_in_init=True, + ) + + model = DecoderModelForCausalLM( + model=_LmHeadRoutingDecoder(), + config=model_config, + hidden_size=hidden_size, + vocab_size=vocab_size, + ) + + assert model.lm_head.quant_config.quant_algo == QuantAlgo.W4A16_NVFP4 + assert tuple(model.lm_head.weight.shape) == (vocab_size, hidden_size // 2) + + +def test_decoder_model_preserves_tied_lm_head_with_delayed_weight_creation(): + hidden_size = 128 + vocab_size = 256 + decoder = _TiedLmHeadRoutingDecoder(hidden_size, vocab_size) + model_config = ModelConfig( + pretrained_config=types.SimpleNamespace( + torch_dtype=torch.bfloat16, + tie_word_embeddings=True, + ), + skip_create_weights_in_init=True, + ) + + model = DecoderModelForCausalLM( + model=decoder, + config=model_config, + hidden_size=hidden_size, + vocab_size=vocab_size, + ) + + assert model.lm_head.weight is model.model.embed_tokens.weight + + @pytest.mark.parametrize( "architecture", ["NemotronHForCausalLM", "Qwen3MoeForCausalLM", "SomeOtherForCausalLM"], @@ -218,6 +556,25 @@ def test_auto_moe_backend_selects_cutedsl_for_w4a16_sm12x(architecture): assert moe_backend == "CUTEDSL" +@pytest.mark.parametrize( + "architecture", + ["NemotronHForCausalLM", "Qwen3MoeForCausalLM", "SomeOtherForCausalLM"], +) +def test_auto_moe_backend_selects_cutedsl_for_layer_w4a16_sm12x(architecture): + quant_config = QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8) + layer_quant_config = { + "model.layers.0.mlp.experts": QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16), + } + + with pytest.MonkeyPatch.context() as mp: + mp.setattr("tensorrt_llm._torch.model_config.get_sm_version", lambda: 121) + moe_backend = ModelConfig.resolve_moe_backend_after_quant_config( + "AUTO", architecture, quant_config, layer_quant_config + ) + + assert moe_backend == "CUTEDSL" + + @pytest.mark.parametrize( "architecture", ["NemotronHForCausalLM", "Qwen3MoeForCausalLM", "SomeOtherForCausalLM"], diff --git a/tests/unittest/models/test_quant_config_utils.py b/tests/unittest/models/test_quant_config_utils.py index 9b2860f58114..d1921d58e978 100644 --- a/tests/unittest/models/test_quant_config_utils.py +++ b/tests/unittest/models/test_quant_config_utils.py @@ -64,6 +64,33 @@ def test_update_quant_config_from_compressed_tensors_parses_nvfp4(): assert set(quant_config.exclude_modules) == {gate_exclude, "lm_head"} +def test_update_quant_config_from_compressed_tensors_parses_w4a16_nvfp4(): + quant_config = QuantConfig() + update_quant_config_from_compressed_tensors( + quant_config, + { + "quant_method": "compressed-tensors", + "format": "nvfp4-pack-quantized", + "config_groups": { + "group_0": { + "weights": { + "num_bits": 4, + "type": "float", + "strategy": "tensor_group", + "group_size": 16, + }, + "input_activations": None, + }, + }, + "ignore": ["lm_head"], + }, + ) + + assert quant_config.quant_algo == QuantAlgo.W4A16_NVFP4 + assert quant_config.group_size == 16 + assert quant_config.exclude_modules == ["lm_head"] + + def test_update_quant_config_from_compressed_tensors_parses_fp8_block_scales(): quant_config = QuantConfig() update_quant_config_from_compressed_tensors( From 24de1aa3d84898f24dc9654f8f4b1e4c2c897e39 Mon Sep 17 00:00:00 2001 From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> Date: Thu, 2 Jul 2026 01:46:47 +0000 Subject: [PATCH 08/25] perf update Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> --- .../decoderXQAImplCommon.h | 20 +- .../decoderXQAImplJIT/decoderXQAImplJIT.cpp | 17 +- .../decoderXQAImplPrecompiled.cpp | 6 +- cpp/tensorrt_llm/thop/cublasScaledMM.cpp | 24 +- cpp/tensorrt_llm/thop/cublasScaledMMLut.h | 19 ++ cpp/tests/unit_tests/kernels/CMakeLists.txt | 3 +- .../kernels/decoderXQAImplCommonTest.cpp | 92 ++++++ .../_torch/attention_backend/trtllm.py | 15 +- tensorrt_llm/_torch/model_config.py | 23 +- .../models/checkpoints/base_weight_loader.py | 153 +++++++++- .../models/checkpoints/base_weight_mapper.py | 8 +- .../checkpoints/hf/exaone4_5_weight_mapper.py | 12 +- .../checkpoints/hf/qwen3_5_weight_mapper.py | 118 +++++--- .../models/checkpoints/hf/weight_loader.py | 33 ++- .../checkpoints/mistral/weight_mapper.py | 25 +- tensorrt_llm/_torch/models/modeling_utils.py | 43 ++- .../_torch/pyexecutor/config_utils.py | 32 +-- .../_torch/pyexecutor/mamba_cache_manager.py | 9 +- .../_torch/pyexecutor/model_engine.py | 2 + .../_torch/pyexecutor/model_loader.py | 264 ++++++++++-------- .../executor/test_mamba_cache_manager.py | 24 +- .../executor/test_pytorch_model_engine.py | 38 +++ .../hf/test_mmapped_safetensors_weights.py | 38 +++ .../hf/test_qwen3_5_weight_mapper.py | 45 ++- tests/unittest/_torch/test_model_config.py | 21 +- 25 files changed, 804 insertions(+), 280 deletions(-) create mode 100644 cpp/tests/unit_tests/kernels/decoderXQAImplCommonTest.cpp create mode 100644 tests/unittest/_torch/models/checkpoints/hf/test_mmapped_safetensors_weights.py diff --git a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplCommon.h b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplCommon.h index 60a5524de09e..cc11080825e7 100644 --- a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplCommon.h +++ b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplCommon.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2020-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. @@ -84,6 +84,20 @@ struct XQAKernelRuntimeHashKey uint32_t getKernelMTileSize( uint32_t headGrpSize, bool isSpecDec, uint32_t qSeqLen, bool isXqaJit, bool supportQGMMA, bool supportMLA); +inline uint32_t getSpecDecHmmaMTileSize(uint32_t headGrpSize, uint32_t qSeqLen) +{ + return getKernelMTileSize( + headGrpSize, /*isSpecDec=*/true, qSeqLen, /*isXqaJit=*/true, /*supportQGMMA=*/false, /*supportMLA=*/false); +} + +inline uint32_t getSpecDecHmmaTokenBlocksPerGroup(uint32_t headGrpSize, uint32_t qSeqLen) +{ + uint32_t const mTileSize = getSpecDecHmmaMTileSize(headGrpSize, qSeqLen); + TLLM_CHECK_WITH_INFO(mTileSize > 0U, "Spec-dec HMMA M tile size must be positive."); + uint32_t const headTokens = qSeqLen * headGrpSize; + return (headTokens + mTileSize - 1U) / mTileSize; +} + XQAKernelRuntimeHashKey getRuntimeHashKeyFromXQAParams(XQAParams const& xqaParams, bool isXqaJit, int SM); struct XQAKernelRuntimeHasher @@ -397,7 +411,7 @@ inline int computeMultiBlockCountForMLA(XQAParams const& xqaParams, int multipro return 1; // disable multi-block for MLA kernel for now. } -inline int computeMultiBlockCountSpecDecGMMA( +inline int computeMultiBlockCountSpecDec( XQAParams const& xqaParams, int batch_size, int multiprocessor_count, int specDecBlocks) { auto const userSpecified = tensorrt_llm::common::getEnvXqaBlocksPerSequence(); @@ -416,7 +430,7 @@ inline int computeMultiBlockCountSpecDecGMMA( return multi_block_count; } - // gridDim = dim3{specDecBlocks, multi_block, nbKVHeads * xqaParams.batch_size} + // GMMA uses specDecBlocks along grid.x. HMMA passes its token blocks per group, which is folded into grid.y. int single_block_count = specDecBlocks * num_kv_heads * batch_size; double wave_count = (double) single_block_count / (double) multiprocessor_count; diff --git a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplJIT/decoderXQAImplJIT.cpp b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplJIT/decoderXQAImplJIT.cpp index 877a780072c4..e7c11326a806 100644 --- a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplJIT/decoderXQAImplJIT.cpp +++ b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplJIT/decoderXQAImplJIT.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2020-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. @@ -420,10 +420,8 @@ void DecoderXQAImplJIT::runImpl(XQAParams const& xqaParams, KVCacheBuffer const& // MultiQueryTokens (generation_input_length > 1) need extra parameters (like qSeqLen, headGrpSize, and // mask). Input parameters for MultiQueryTokens kernels. unsigned int headGrpSize = num_q_heads_over_kv; - // Use mTileSize = 16 kernels when qSeqLen <= 16. unsigned int qSeqLen = static_cast(xqaParams.generation_input_length); - unsigned int mTileSize = qSeqLen <= 16 ? 16 : 32; - unsigned int nbTokenBlocksPerGrp = divUp(qSeqLen * headGrpSize, mTileSize); + unsigned int nbTokenBlocksPerGrp = getSpecDecHmmaTokenBlocksPerGroup(headGrpSize, qSeqLen); unsigned int maxQSeqLen = xqaParams.spec_decoding_is_generation_length_variable ? // true for ReDrafter xqaParams.spec_decoding_max_generation_length : qSeqLen; @@ -459,10 +457,11 @@ void DecoderXQAImplJIT::runImpl(XQAParams const& xqaParams, KVCacheBuffer const& appendParam(&launchParams.scratch); uint32_t multi_block = 1; - // if (xqaParams.multi_block_mode) - // { - // multi_block = computeMultiBlockCount(xqaParams, xqaParams.batch_size, multiprocessor_count); - // } + if (xqaParams.multi_block_mode) + { + multi_block = computeMultiBlockCountSpecDec( + xqaParams, xqaParams.batch_size, multiprocessor_count, nbTokenBlocksPerGrp); + } auto const gridDim = (dim3{multi_block, xqaParams.num_kv_heads * nbTokenBlocksPerGrp, xqaParams.batch_size}); dim3 const blockDim(128, 1, 2); @@ -533,7 +532,7 @@ void DecoderXQAImplJIT::runImpl(XQAParams const& xqaParams, KVCacheBuffer const& { if (isSpecDec && isGMMAKernel) { - multi_block = computeMultiBlockCountSpecDecGMMA( + multi_block = computeMultiBlockCountSpecDec( xqaParams, xqaParams.batch_size, multiprocessor_count, specDecBlocks); } else if (!isSpecDec) diff --git a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplPrecompiled.cpp b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplPrecompiled.cpp index d37c7a7a4037..9f77a8a40ded 100644 --- a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplPrecompiled.cpp +++ b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplPrecompiled.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2020-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. @@ -276,10 +276,8 @@ class XQAKernelList // MultiQueryTokens (generation_input_length > 1) need extra parameters (like qSeqLen, headGrpSize, and // mask). Input parameters for MultiQueryTokens kernels. unsigned int headGrpSize = num_q_heads_over_kv; - // Use mTileSize = 16 kernels when qSeqLen <= 16. unsigned int qSeqLen = static_cast(xqaParams.generation_input_length); - unsigned int mTileSize = qSeqLen <= 16 ? 16 : 32; - unsigned int nbTokenBlocksPerGrp = divUp(qSeqLen * headGrpSize, mTileSize); + unsigned int nbTokenBlocksPerGrp = getSpecDecHmmaTokenBlocksPerGroup(headGrpSize, qSeqLen); int const* maskPtr = xqaParams.spec_decoding_packed_mask; int const* cuQSeqLens = launchParams.cu_seq_lens; unsigned int maxQSeqLen = xqaParams.spec_decoding_is_generation_length_variable ? // true for ReDrafter diff --git a/cpp/tensorrt_llm/thop/cublasScaledMM.cpp b/cpp/tensorrt_llm/thop/cublasScaledMM.cpp index 4951db12cfd2..8383850509a2 100644 --- a/cpp/tensorrt_llm/thop/cublasScaledMM.cpp +++ b/cpp/tensorrt_llm/thop/cublasScaledMM.cpp @@ -26,6 +26,7 @@ #include "userbuffersTensor.h" #include #include +#include using torch::Tensor; @@ -189,11 +190,28 @@ void cublas_gemm_caller(torch::Tensor& out, torch::Tensor const& a, torch::Tenso cudaDataType_t scaleType = CUDA_R_32F; cublasWrapper->setGemmConfig(aType, bType, outType, /*computeType=*/scaleType); - auto const workspace_options = torch::TensorOptions().dtype(torch::kUInt8).device(a.device()); - auto workspace = torch::empty(CUBLAS_WORKSPACE_SIZE, workspace_options); - auto stream = at::cuda::getCurrentCUDAStream(a.get_device()); + // Persistent per-stream cublasLt workspace. The cublasLt kernel records a + // fixed pointer to this workspace; if we destruct the workspace storage at + // function return (the previous behavior), the CUDA caching allocator may + // hand the same block out to a later allocation while a captured CUDA + // graph still references the workspace pointer -> use-after-free that + // surfaces as free-block-tree corruption on the next allocator operation + // (e.g. FusedMoeRunner::getWorkspaceInfo's first torch::empty destructor). + // + // Keyed by (device, stream) so that concurrent GEMMs on different streams + // don't race on the same scratch bytes. Mirrors PyTorch's per-stream + // cublasLt workspace cache in at::cuda. + thread_local std::unordered_map workspace_cache; + auto stream_ptr = stream.stream(); + auto& workspace = workspace_cache[stream_ptr]; + if (!workspace.defined() || workspace.device() != a.device()) + { + auto const workspace_options = torch::TensorOptions().dtype(torch::kUInt8).device(a.device()); + workspace = torch::empty(CUBLAS_WORKSPACE_SIZE, workspace_options); + } + auto* a_ptr = static_cast(a.data_ptr()); auto* b_ptr = static_cast(b.data_ptr()); auto* out_ptr = static_cast(out.data_ptr()); diff --git a/cpp/tensorrt_llm/thop/cublasScaledMMLut.h b/cpp/tensorrt_llm/thop/cublasScaledMMLut.h index 6c1471f70ebb..534edda011d6 100644 --- a/cpp/tensorrt_llm/thop/cublasScaledMMLut.h +++ b/cpp/tensorrt_llm/thop/cublasScaledMMLut.h @@ -118,6 +118,25 @@ inline const AlgoListType fp8_algo_list = { // Llama-3.3-70B TP4 (this is the default algo on B200. Here we aim to use the same algo on GB200.) // [-algo66 -m_tile393 -m_stages36 -m_numsK1 -m_reduction0 -m_swizzle0 -m_custom1 -m_mma0 -m_cga4 -m_scheduling1] {{8, 8192, 14336}, {66, 393, 36, 1, 0, 1, 1, 4}}, + + // Qwen3.6-35B-A3B-NVFP4 on GB10 (sm_121), FP8 per-tensor projections. Output is BF16; the FP8 + // branch of find_special_algo is output-dtype-agnostic so these match. Shapes are the ACTUAL + // runtime cuBLASLt inventory from the DEBUG log (1330 calls, all previously missing) -- NOT + // derived from weight headers (TRT-LLM FUSES qkv+z and gate+up, so derived dims were wrong). + // NVJET algos chosen via cublasTest matmulFind + -T200 -Tme20p retime, validated at M=1/4/8. + // Swept against the RUNTIME cuBLASLt (CUDA 13.1, libcublasLt.so.13.2.1.1) -- NVJET tile/custom + // configs are version-specific, valid only for that cuBLASLt generation. Qwen3.6-specific keys. + // [-algo67 -m_tile6 -m_stages38 -m_numsK1 -m_reduction0 -m_swizzle0 -m_custom16 -m_mma0 -m_cga0 -m_scheduling1] + {{8, 2048, 12288}, + {67, 6, 38, 1, 0, 0, 16, 0}}, // GDN in_proj qkv+z fused (x210 in log): heur 61.5->51.3us (robust across M) + // [-algo67 -m_tile3 -m_stages38 -m_numsK1 -m_reduction0 -m_swizzle0 -m_custom114 -m_mma0 -m_cga0 -m_scheduling1] + {{8, 2048, 256}, {67, 3, 38, 1, 0, 0, 114, 0}}, // k/v proj n=256 (x280 in log): heur 8.35->6.08us (M1) + // [-algo67 -m_tile10 -m_stages37 -m_numsK1 -m_reduction0 -m_swizzle0 -m_custom28 -m_mma0 -m_cga0 -m_scheduling1] + {{8, 2048, 9216}, + {67, 10, 37, 1, 0, 0, 28, 0}}, // shared_expert gate+up fused (x70 in log): heur 49/110->24.8us (M1/-80% M4) + // [-algo67 -m_tile10 -m_stages36 -m_numsK1 -m_reduction0 -m_swizzle0 -m_custom29 -m_mma0 -m_cga0 -m_scheduling1] + {{8, 4096, 2048}, + {67, 10, 36, 1, 0, 0, 29, 0}}, // out_proj/down_proj K=4096 (x280 in log): heur 22.5->12.3us (-45%) }; } // namespace cublas_lut diff --git a/cpp/tests/unit_tests/kernels/CMakeLists.txt b/cpp/tests/unit_tests/kernels/CMakeLists.txt index 95d33e421050..4f700e95f9f5 100644 --- a/cpp/tests/unit_tests/kernels/CMakeLists.txt +++ b/cpp/tests/unit_tests/kernels/CMakeLists.txt @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & +# 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"); you may not @@ -15,6 +15,7 @@ add_gtest(banRepeatNGramsKernelsTest banRepeatNGramsKernelsTest.cpp) add_gtest(decodingKernelsTest decodingKernelTest.cpp) +add_gtest(decoderXQAImplCommonTest decoderXQAImplCommonTest.cpp) add_gtest(logitsBitmaskTest logitsBitmaskTest.cpp) macro(remove_compile_definition TARGET_NAME DEFINITION) diff --git a/cpp/tests/unit_tests/kernels/decoderXQAImplCommonTest.cpp b/cpp/tests/unit_tests/kernels/decoderXQAImplCommonTest.cpp new file mode 100644 index 000000000000..2032d2ab7f1a --- /dev/null +++ b/cpp/tests/unit_tests/kernels/decoderXQAImplCommonTest.cpp @@ -0,0 +1,92 @@ +/* + * 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. + */ + +#include "tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplCommon.h" +#include "tensorrt_llm/common/envUtils.h" + +#include + +namespace tk = tensorrt_llm::kernels; +namespace tc = tensorrt_llm::common; + +namespace +{ + +tk::XQAParams makeSpecDecXqaParams(int numKvHeads, int maxPastKvLength) +{ + tk::XQAParams params{}; + params.num_kv_heads = numKvHeads; + params.max_past_kv_length = maxPastKvLength; + return params; +} + +void skipIfXqaBlocksPerSequenceIsForced() +{ + if (tc::getEnvXqaBlocksPerSequence().has_value()) + { + GTEST_SKIP() << "TRTLLM_XQA_BLOCKS_PER_SEQUENCE overrides multi-block tuning."; + } +} + +TEST(DecoderXQAImplCommonTest, computeMultiBlockCountSpecDecPopulatesLowCtaLaunch) +{ + skipIfXqaBlocksPerSequenceIsForced(); + + int constexpr kBatchSize = 1; + int constexpr kMultiprocessorCount = 128; + int constexpr kSpecDecBlocks = 1; + auto const params = makeSpecDecXqaParams(/*numKvHeads=*/1, /*maxPastKvLength=*/4096); + + EXPECT_EQ(tk::computeMultiBlockCountSpecDec(params, kBatchSize, kMultiprocessorCount, kSpecDecBlocks), 16); +} + +TEST(DecoderXQAImplCommonTest, computeMultiBlockCountSpecDecAccountsForTokenBlocks) +{ + skipIfXqaBlocksPerSequenceIsForced(); + + int constexpr kBatchSize = 1; + int constexpr kMultiprocessorCount = 128; + int constexpr kSpecDecBlocks = 8; + auto const params = makeSpecDecXqaParams(/*numKvHeads=*/1, /*maxPastKvLength=*/4096); + + EXPECT_EQ(tk::computeMultiBlockCountSpecDec(params, kBatchSize, kMultiprocessorCount, kSpecDecBlocks), 8); +} + +TEST(DecoderXQAImplCommonTest, computeMultiBlockCountSpecDecKeepsShortHistorySingleBlock) +{ + skipIfXqaBlocksPerSequenceIsForced(); + + int constexpr kBatchSize = 1; + int constexpr kMultiprocessorCount = 128; + int constexpr kSpecDecBlocks = 1; + auto const params = makeSpecDecXqaParams(/*numKvHeads=*/1, /*maxPastKvLength=*/1024); + + EXPECT_EQ(tk::computeMultiBlockCountSpecDec(params, kBatchSize, kMultiprocessorCount, kSpecDecBlocks), 1); +} + +TEST(DecoderXQAImplCommonTest, getSpecDecHmmaMTileSizeUsesHeadTokens) +{ + EXPECT_EQ(tk::getSpecDecHmmaMTileSize(/*headGrpSize=*/8, /*qSeqLen=*/4), 32U); +} + +TEST(DecoderXQAImplCommonTest, getSpecDecHmmaTokenBlocksPerGroupMatchesCompiledTile) +{ + EXPECT_EQ(tk::getSpecDecHmmaTokenBlocksPerGroup(/*headGrpSize=*/8, /*qSeqLen=*/4), 1U); + EXPECT_EQ(tk::getSpecDecHmmaTokenBlocksPerGroup(/*headGrpSize=*/16, /*qSeqLen=*/4), 2U); +} + +} // namespace diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 66899f3d5b26..62b4252ee47f 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -486,13 +486,14 @@ def prepare(self) -> None: self.kv_cache_block_offsets, self.request_ids, self.beam_width, self.num_contexts, self.num_seqs) - error_message = ( - f"The max KV cache length of input sequences ({self.kv_lens[:self.num_seqs].max()}) " - f"exceeds the KV cache manager's maximum supported length " - f"({self.kv_cache_manager.max_seq_len}).") - - assert self.kv_lens[:self.num_seqs].max( - ) <= self.kv_cache_manager.max_seq_len, error_message + if self.num_seqs > 0: + max_kv_len = self.kv_lens[:self.num_seqs].max() + error_message = ( + f"The max KV cache length of input sequences ({max_kv_len}) " + f"exceeds the KV cache manager's maximum supported length " + f"({self.kv_cache_manager.max_seq_len}).") + + assert max_kv_len <= self.kv_cache_manager.max_seq_len, error_message # Also prepare draft KV cache block offsets if draft_kv_cache_manager exists if self.draft_kv_cache_manager is not None: diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index ec6d90a9a86f..d529bf743fca 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -12,8 +12,8 @@ from transformers.utils import HF_MODULES_CACHE from tensorrt_llm._torch.pyexecutor.config_utils import ( - get_qwen3_hybrid_layer_types, get_qwen3_hybrid_num_attention_layers, - is_nemotron_hybrid, is_qwen3_5, is_qwen3_hybrid, load_pretrained_config) + get_qwen3_hybrid_num_attention_layers, is_nemotron_hybrid, is_qwen3_hybrid, + load_pretrained_config) from tensorrt_llm._utils import (get_sm_version, is_sm_100f, torch_dtype_to_binding) from tensorrt_llm.bindings import LayerType as LayerTypeCpp @@ -444,24 +444,6 @@ def _build_modelopt_quant_config(json_quant_configs, ] return quant_config, layer_quant_config - @staticmethod - def _add_qwen35_qkvz_bf16_excludes(json_quant_configs, - pretrained_config) -> None: - """Keep packed Qwen3.5 linear-attention qkvz on the BF16 path.""" - if pretrained_config is None or not is_qwen3_5(pretrained_config): - return - try: - layer_types = get_qwen3_hybrid_layer_types(pretrained_config) - except (ValueError, AttributeError): - return - - exclude_modules = list(json_quant_configs.get('exclude_modules') or []) - for layer_idx, layer_type in enumerate(layer_types): - if layer_type == "linear_attention": - exclude_modules.append( - f"model.layers.{layer_idx}.linear_attn.in_proj_qkvz") - json_quant_configs['exclude_modules'] = sorted(set(exclude_modules)) - @staticmethod def get_mxfp4_quant_algo(moe_backend, is_dynamic_quant=False): quant_algo = ModelConfig.override_quant_algo() @@ -756,7 +738,6 @@ def _recursive_update_config(config: transformers.PretrainedConfig, 'hf_quant_config.json'): with open(quant_config_file) as f: normalized = read_modelopt_quant_config(json.load(f)) - cls._add_qwen35_qkvz_bf16_excludes(normalized, pretrained_config) # The file is authoritative; warn if the inline copy disagrees. # Done before _build_modelopt_quant_config since the builder may # mutate ``normalized`` via ``.update`` from quant_cfg.json. diff --git a/tensorrt_llm/_torch/models/checkpoints/base_weight_loader.py b/tensorrt_llm/_torch/models/checkpoints/base_weight_loader.py index 10d72a3a86c7..0a0cd4e8eca6 100644 --- a/tensorrt_llm/_torch/models/checkpoints/base_weight_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/base_weight_loader.py @@ -1,9 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import re import threading from abc import ABC, abstractmethod -from typing import Any, Dict, Iterator, Tuple, Union +from typing import Any, Callable, Dict, Iterator, List, Tuple, Union + +import safetensors from tensorrt_llm.mapping import Mapping @@ -91,6 +94,154 @@ def mark_consumed(self, prefix: str) -> int: return len(keys_to_delete) +class MmappedSafetensorsWeights: + """ + Memory-mapped safetensors checkpoint access for integrated GPU systems. + + Opens checkpoint shards via ``safetensors.safe_open`` without loading full + tensors into RAM. ``mark_consumed`` drops index entries after weights are + copied into the model to reduce peak memory on unified-memory machines. + """ + + def __init__(self, weight_files: List[str]): + self._files = weight_files + self._handles: list = [] + self._key_locations: dict[str, tuple[int, str]] = {} + self._lock = threading.Lock() + for file_idx, path in enumerate(weight_files): + handle = safetensors.safe_open(path, framework="pt", device="cpu") + self._handles.append(handle) + for key in handle.keys(): + if key in self._key_locations: + raise RuntimeError( + f"Duplicate tensor name {key!r} in checkpoint files.") + self._key_locations[key] = (file_idx, key) + + def _get_slice(self, key: str): + file_idx, tensor_name = self._key_locations[key] + return self._handles[file_idx].get_slice(tensor_name) + + def __getitem__(self, key: str) -> Any: + return self._get_slice(key) + + def __contains__(self, key: str) -> bool: + return key in self._key_locations + + def __len__(self) -> int: + return len(self._key_locations) + + def __iter__(self) -> Iterator[str]: + with self._lock: + return iter(list(self._key_locations.keys())) + + def keys(self): + with self._lock: + return list(self._key_locations.keys()) + + def values(self): + with self._lock: + return [self._get_slice(k) for k in self._key_locations.keys()] + + def items(self) -> Iterator[Tuple[str, Any]]: + with self._lock: + keys = list(self._key_locations.keys()) + return [(k, self._get_slice(k)) for k in keys] + + def get(self, key: str, default: Any = None) -> Any: + if key in self._key_locations: + return self._get_slice(key) + return default + + def mark_consumed(self, prefix: str) -> int: + with self._lock: + keys_to_delete = [ + k for k in self._key_locations if k.startswith(prefix + ".") + ] + for key in keys_to_delete: + del self._key_locations[key] + return len(keys_to_delete) + + def remap_keys(self, key_mapping: dict[str, + str]) -> "MmappedSafetensorsWeights": + remapped = object.__new__(MmappedSafetensorsWeights) + remapped._files = self._files + remapped._handles = self._handles + remapped._lock = threading.Lock() + remapped._key_locations = {} + for old_key, loc in self._key_locations.items(): + new_key = key_mapping.get(old_key, old_key) + if new_key in remapped._key_locations: + raise RuntimeError( + f"Duplicate tensor name {new_key!r} after key remap.") + remapped._key_locations[new_key] = loc + return remapped + + def transform_keys( + self, transform_fn: Callable[[str], + str]) -> "MmappedSafetensorsWeights": + key_mapping = { + old_key: transform_fn(old_key) + for old_key in self._key_locations + } + return self.remap_keys(key_mapping) + + def rename_by_regex( + self, pattern_mapping: dict[str, + str]) -> "MmappedSafetensorsWeights": + key_mapping = {} + matched_keys = set() + for key in self._key_locations: + new_key = key + for pattern, replacement in pattern_mapping.items(): + if re.match(pattern, key): + new_key = re.sub(pattern, replacement, key) + matched_keys.add(key) + break + key_mapping[key] = new_key + return self.remap_keys(key_mapping) + + +def rename_weight_keys_with_regex( + weights: Union[Dict[str, Any], ConsumableWeightsDict, + MmappedSafetensorsWeights], + pattern_mapping: dict[str, str], +) -> Union[Dict[str, Any], ConsumableWeightsDict, MmappedSafetensorsWeights]: + if isinstance(weights, MmappedSafetensorsWeights): + return weights.rename_by_regex(pattern_mapping) + + is_consumable = isinstance(weights, ConsumableWeightsDict) + renamed_weights = {} + matched_keys = set() + for key in list(weights.keys()): + new_key = key + for pattern, replacement in pattern_mapping.items(): + if re.match(pattern, key): + new_key = re.sub(pattern, replacement, key) + matched_keys.add(key) + break + renamed_weights[new_key] = weights[key] + if is_consumable: + return ConsumableWeightsDict(renamed_weights) + return renamed_weights + + +def remap_weight_keys( + weights: Union[Dict[str, Any], ConsumableWeightsDict, + MmappedSafetensorsWeights], + key_mapping: dict[str, str], +) -> Union[Dict[str, Any], ConsumableWeightsDict, MmappedSafetensorsWeights]: + if isinstance(weights, MmappedSafetensorsWeights): + return weights.remap_keys(key_mapping) + + is_consumable = isinstance(weights, ConsumableWeightsDict) + renamed_weights = {} + for key in weights.keys(): + renamed_weights[key_mapping.get(key, key)] = weights[key] + if is_consumable: + return ConsumableWeightsDict(renamed_weights) + return renamed_weights + + class BaseWeightLoader(ABC): @abstractmethod diff --git a/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py index 3ac76d6809be..06be9134bd7a 100644 --- a/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py @@ -81,8 +81,12 @@ def rename_by_params_map(self, params_map: dict[str, str], """ import re - from tensorrt_llm._torch.models.checkpoints.base_weight_loader import \ - ConsumableWeightsDict + from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ( + ConsumableWeightsDict, MmappedSafetensorsWeights, + rename_weight_keys_with_regex) + + if isinstance(weights, MmappedSafetensorsWeights): + return rename_weight_keys_with_regex(weights, params_map) # Check if input is a ConsumableWeightsDict to preserve the type is_consumable = isinstance(weights, ConsumableWeightsDict) diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/exaone4_5_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/exaone4_5_weight_mapper.py index 7e315acced94..0a4637488629 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/exaone4_5_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/exaone4_5_weight_mapper.py @@ -13,7 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict +from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ( + ConsumableWeightsDict, + MmappedSafetensorsWeights, +) from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper from tensorrt_llm._torch.models.modeling_utils import register_mapper @@ -22,6 +25,13 @@ class Exaone4_5HfWeightMapper(HfWeightMapper): def preprocess_weights(self, weights: dict): """Rename HF checkpoint prefixes; supports plain dict and ConsumableWeightsDict.""" + if isinstance(weights, MmappedSafetensorsWeights): + return weights.transform_keys( + lambda key: key.replace("model.visual.", "visual.") + if key.startswith("model.visual.") + else key + ) + is_consumable = isinstance(weights, ConsumableWeightsDict) renamed = {} for key, value in weights.items(): diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py index 3cd598ca21a0..2018feedec4f 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py @@ -32,8 +32,11 @@ class Qwen3_5MoeHfWeightMapper(Qwen3NextHfWeightMapper): tensors. Qwen3.5 checkpoints store them as separate in_proj_qkv + z (or fully split q/k/v/z) and b + a tensors. This mapper packs them into the grouped-interleaved layout that TRT-LLM expects. - For FP8 checkpoints, the packed qkvz tensor is then dequantized to - bf16 as a temporary workaround for TP loading + For FP8 checkpoints with per-tensor scales, qkv/z weights are rescaled + to one packed-module scale before packing + (handled in _rescale_linear_attn_qkvz_per_tensor_fp8). + Existing per-block FP8 checkpoints still dequantize the packed qkvz + tensor to bf16 as a temporary workaround for TP loading (handled in _dequantize_linear_attn_fp8_qkvz). 3. MoE expert tensors (handled in handle_special_instance_module): @@ -197,50 +200,79 @@ def _dequantize_linear_attn_fp8_qkvz(self, weights: dict) -> dict: updated_weights.pop(scale_name, None) return updated_weights - def _dequantize_linear_attn_per_tensor_fp8(self, weights: dict) -> dict: - """Dequantize per-tensor FP8 linear-attn split projections to BF16. - - MIXED_PRECISION checkpoints (e.g. Qwen3.6 NVFP4) store the linear-attn - ``in_proj_*`` split projections (in_proj_qkv, in_proj_z, in_proj_q/k/v, - etc.) as per-tensor FP8: scalar ``weight_scale`` or - ``weight_scale_inv`` (and matching scalar ``input_scale``). The packer - below assumes BF16 weights or per-block FP8 with block-shaped scales, - so dequantize here and drop the scale tensors. The packed BF16 tensor - is then concatenated with the already-BF16 ``in_proj_a`` / - ``in_proj_b`` / ``in_proj_z`` weights and TP-sharded by the standard - path. + def _rescale_linear_attn_qkvz_per_tensor_fp8(self, weights: dict) -> dict: + """Rescale per-tensor FP8 qkv/z shards to one packed qkvz scale. + + Qwen3.6 MIXED_PRECISION checkpoints store ``in_proj_qkv`` and + ``in_proj_z`` as separate per-tensor FP8 projections with scalar + ``weight_scale`` / ``input_scale`` metadata. TRT-LLM packs them into a + single ``in_proj_qkvz`` Linear, which has one scalar ``weight_scale`` + and one scalar ``input_scale``. Match the fused-FP8 Linear loader: pick + the max scale and requantize each split weight into that shared scale + before packing. """ - target_dtype = getattr(self.config.pretrained_config, "torch_dtype", torch.bfloat16) - if target_dtype is None: - target_dtype = torch.bfloat16 updated = dict(weights) - drop = [] - for name in list(weights): - if not name.endswith(".weight"): + + qkvz_groups = defaultdict(dict) + for name, tensor in weights.items(): + match = self._SPLIT_PROJ_PATTERN.match(name) + if match is None: + continue + prefix, projection_name, suffix = match.groups() + if projection_name not in {"qkv", "q", "k", "v", "z"}: continue - if self._SPLIT_PROJ_PATTERN.match(name) is None: + if suffix not in {"weight", "weight_scale", "input_scale"}: continue - prefix = name[: -len(".weight")] - scale_name = prefix + ".weight_scale_inv" - if scale_name not in weights: - scale_name = prefix + ".weight_scale" - if scale_name not in weights: + qkvz_groups[(prefix, projection_name)][suffix] = (name, tensor) + + prefixes = {prefix for prefix, _ in qkvz_groups} + for prefix in prefixes: + projection_names = { + name for group_prefix, name in qkvz_groups if group_prefix == prefix + } + if "qkv" in projection_names: + shard_names = ["qkv", "z"] + else: + shard_names = ["q", "k", "v", "z"] + if not all((prefix, name) in qkvz_groups for name in shard_names): + continue + + shard_groups = [qkvz_groups[(prefix, name)] for name in shard_names] + if not all("weight" in group and "weight_scale" in group for group in shard_groups): continue - scale = weights[scale_name] - # Per-tensor scalar scale → dequantize. Per-block (n-D) handled elsewhere. - if scale.ndim != 0: + + weight_scales = [group["weight_scale"][1].reshape([]) for group in shard_groups] + if not all(scale.ndim == 0 for scale in weight_scales): continue - updated[name] = ( - (weights[name].to(torch.float32) * scale.to(torch.float32)) - .to(target_dtype) - .contiguous() - ) - drop.append(scale_name) - input_scale_name = name[: -len(".weight")] + ".input_scale" - if input_scale_name in weights: - drop.append(input_scale_name) - for k in drop: - updated.pop(k, None) + max_weight_scale = torch.stack( + [scale.to(torch.float32) for scale in weight_scales] + ).max() + + input_scale_groups = [group for group in shard_groups if "input_scale" in group] + if input_scale_groups: + if len(input_scale_groups) != len(shard_groups): + raise ValueError( + f"Expected input_scale for every split qkvz shard under {prefix}" + ) + input_scales = [group["input_scale"][1].reshape([]) for group in shard_groups] + if not all(scale.ndim == 0 for scale in input_scales): + raise ValueError( + f"Expected scalar input_scale for every split qkvz shard under {prefix}" + ) + max_input_scale = torch.stack( + [scale.to(torch.float32) for scale in input_scales] + ).max() + for group in shard_groups: + updated.pop(group["input_scale"][0], None) + updated[f"{prefix}.in_proj_qkvz.input_scale"] = max_input_scale + + for group, weight_scale in zip(shard_groups, weight_scales): + weight_name, weight = group["weight"] + rescaled_weight = weight.to(torch.float32).mul(weight_scale.to(torch.float32)) + rescaled_weight = rescaled_weight.div(max_weight_scale.to(rescaled_weight.device)) + updated[weight_name] = rescaled_weight.to(torch.float8_e4m3fn).contiguous() + updated.pop(group["weight_scale"][0], None) + updated[f"{prefix}.in_proj_qkvz.weight_scale"] = max_weight_scale return updated def _pack_split_projections(self, weights: dict) -> dict: @@ -359,10 +391,10 @@ def preprocess_weights(self, weights: dict) -> dict: is_modelopt_ckpt, normalized_weights = self._preprocess_modelopt_ckpt(normalized_weights) # MIXED_PRECISION modelopt checkpoints store the linear-attn split - # projections as per-tensor FP8 (scalar scales). Dequantize to BF16 - # before packing so the packer can run its existing BF16 path. + # qkv/z projections as per-tensor FP8 (scalar scales). Rescale them to + # one packed-module scale before packing. if is_modelopt_ckpt: - normalized_weights = self._dequantize_linear_attn_per_tensor_fp8(normalized_weights) + normalized_weights = self._rescale_linear_attn_qkvz_per_tensor_fp8(normalized_weights) packed_weights = self._pack_split_projections(normalized_weights) if not is_modelopt_ckpt: diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py index 264104c6f2e1..ed57ae748e82 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py @@ -25,11 +25,12 @@ from mpi4py import MPI as _MPI from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ( - BaseWeightLoader, ConsumableWeightsDict) + BaseWeightLoader, ConsumableWeightsDict, MmappedSafetensorsWeights) from tensorrt_llm._torch.models.modeling_utils import ( register_checkpoint_weight_loader, run_concurrently) -from tensorrt_llm._utils import (ENABLE_MULTI_DEVICE, local_mpi_barrier, - local_mpi_comm, local_mpi_rank, local_mpi_size) +from tensorrt_llm._utils import (ENABLE_MULTI_DEVICE, is_device_integrated, + local_mpi_barrier, local_mpi_comm, + local_mpi_rank, local_mpi_size) from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping @@ -77,12 +78,19 @@ def load_weights(self, if len(filtered_weight_files) > 0: weight_files = filtered_weight_files if weight_files: + num_layers = int(os.environ.get("TLLM_OVERRIDE_LAYER_NUM", "0")) + if is_device_integrated(): + logger.info( + "Integrated GPU: using mmap-backed safetensors loading " + "(no prefetch, no parallel shard load).") + local_mpi_barrier() + return MmappedSafetensorsWeights(weight_files) + # Prefetch the weight files to CPU memory if the size is less than 90% of the available memory. # This is a heuristic to avoid prefetching files that are too large and causing file cache thrashing. prefetch_size = sum(os.path.getsize(file) for file in weight_files) # If the layer number is overridden, it indicates that only a subset of layers are loaded. # Prefetching all layers is unnecessary. - num_layers = int(os.environ.get("TLLM_OVERRIDE_LAYER_NUM", "0")) enable_prefetch = (prefetch_size < self._get_local_available_host_memory() * 0.9 and num_layers == 0) @@ -107,14 +115,20 @@ def load_weights(self, weight_files = glob.glob(f"{checkpoint_dir}/*.pth") if weight_files: + num_workers = 1 if is_device_integrated() else None return self._load_weights_in_parallel( - weight_files, self._load_bin_or_path_file, - "Loading bin weights in parallel") + weight_files, + self._load_bin_or_path_file, + "Loading bin weights in parallel", + num_workers=num_workers) raise RuntimeError(f"No weight files found in {checkpoint_dir}.") - def _load_weights_in_parallel(self, weight_files: List[str], load_func, - description: str) -> ConsumableWeightsDict: + def _load_weights_in_parallel(self, + weight_files: List[str], + load_func, + description: str, + num_workers=None) -> ConsumableWeightsDict: """ Load weight files in parallel using the specified loading function. @@ -133,7 +147,8 @@ def _load_weights_in_parallel(self, weight_files: List[str], load_func, # specifically the comma right after the w is important to make it a tuple. run_concurrently(load_func, [(w, ) for w in weight_files], reduce_func=weights.update, - pbar=pbar) + pbar=pbar, + num_workers=num_workers) return ConsumableWeightsDict(weights) diff --git a/tensorrt_llm/_torch/models/checkpoints/mistral/weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/mistral/weight_mapper.py index dd6e0332b849..c4e5b9936b3c 100644 --- a/tensorrt_llm/_torch/models/checkpoints/mistral/weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/mistral/weight_mapper.py @@ -45,7 +45,30 @@ def __init__(self): # Adapted from: # https://github.com/vllm-project/vllm/blob/883b42896a9ed9791750d721fad26005b7569eba/vllm/model_executor/models/llama.py#L657 def rename_by_params_map(self, params_map: dict[str, str], weights: dict) -> dict: - from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict + from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ( + ConsumableWeightsDict, + MmappedSafetensorsWeights, + remap_weight_keys, + ) + + if isinstance(weights, MmappedSafetensorsWeights): + key_mapping = {} + for key in weights.keys(): + new_key = key + modules = key.split(".") + num_modules = len(modules) + for i in range(num_modules): + item = modules[i] + next_item = modules[i + 1] if i < num_modules - 1 else None + + combined_item = f"{item}.{next_item}" if next_item is not None else None + + if combined_item in params_map: + new_key = new_key.replace(combined_item, params_map[combined_item]) + elif item in params_map: + new_key = new_key.replace(item, params_map[item]) + key_mapping[key] = new_key + return remap_weight_keys(weights, key_mapping) # Check if input is a ConsumableWeightsDict to preserve the type is_consumable = isinstance(weights, ConsumableWeightsDict) diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 8136e2db92bc..82ab5a569b9d 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -12,7 +12,7 @@ from torch.utils._pytree import tree_any_only from tqdm import tqdm -from tensorrt_llm._utils import local_mpi_rank +from tensorrt_llm._utils import is_device_integrated, local_mpi_rank from tensorrt_llm.lora_manager import HfLoraLoader from tensorrt_llm.models.convert_utils import split_matrix_tp @@ -831,8 +831,12 @@ def rename_weights_with_regex(pattern_mapping: Dict[str, str], weights: Dict): """ import re - from tensorrt_llm._torch.models.checkpoints.base_weight_loader import \ - ConsumableWeightsDict + from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ( + ConsumableWeightsDict, rename_weight_keys_with_regex) + + if isinstance(weights, ConsumableWeightsDict) or hasattr( + weights, "rename_by_regex"): + return rename_weight_keys_with_regex(weights, pattern_mapping) # Check if input is a ConsumableWeightsDict to preserve the type is_consumable = isinstance(weights, ConsumableWeightsDict) @@ -874,6 +878,29 @@ def filter_weights(prefix, weights: Dict): return result +def materialize_meta_parameters(module: nn.Module) -> None: + """Materialize meta parameters on CUDA for integrated GPU lazy loading.""" + for key, param in list(module._parameters.items()): + if param is None: + continue + if param.is_meta: + module._parameters[key] = nn.Parameter( + torch.empty_like(param, device='cuda'), + requires_grad=False, + ) + elif param.device.type != 'cuda': + module._parameters[key] = nn.Parameter(param.cuda(), + requires_grad=False) + + +def _should_load_weights_serially() -> bool: + env_flag = os.environ.get("TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL", + "False") + return is_device_integrated() or env_flag in [ + "True", "true", "1", "yes", "y" + ] + + def run_concurrently(func, args_list, reduce_func=None, @@ -945,6 +972,8 @@ def _load_weights_impl(model: Union[nn.Module, DecoderModelForCausalLM], def load_single_module(name, module): torch.cuda.set_device(device_id) if len(module._parameters) > 0: + if is_device_integrated(): + materialize_meta_parameters(module) # skip load weights if module is in skip_modules if any(skip_module in name for skip_module in skip_modules): return @@ -1029,8 +1058,7 @@ def load_single_module(name, module): if hasattr(weights, 'mark_consumed'): weights.mark_consumed(name) - if os.environ.get("TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL", - "False") in ["True", "true", "1", "yes", "y"]: + if _should_load_weights_serially(): for name, module in tqdm(list( model.named_modules(remove_duplicate=False)), desc="Loading weights"): @@ -1081,6 +1109,8 @@ def _load_weights_impl_v2(model: Union[nn.Module, DecoderModelForCausalLM], def load_single_module(name, module): torch.cuda.set_device(device_id) if len(module._parameters) > 0: + if is_device_integrated(): + materialize_meta_parameters(module) if weight_mapper.should_skip_module(name): return @@ -1144,8 +1174,7 @@ def load_single_module(name, module): if hasattr(weights, 'mark_consumed'): weights.mark_consumed(name) - if os.environ.get("TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL", - "False") in ["True", "true", "1", "yes", "y"]: + if _should_load_weights_serially(): for name, module in tqdm(list( model.named_modules(remove_duplicate=False)), desc="Loading weights"): diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index 5a3181f95aa7..eb38482cabb7 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -1,6 +1,5 @@ import dataclasses import re -from types import SimpleNamespace from typing import List, Optional import torch @@ -317,12 +316,7 @@ def _extract_text_config(config_dict: dict) -> dict: @staticmethod def _inherit_quantization_config(config_dict: dict, text_config: dict) -> dict: - """Copy top-level quantization_config into text_config with name normalization. - - Also adds a temporary workaround that keeps packed linear-attention - in_proj_qkvz on the bf16 path until FP8 block-scale TP loading is - fixed for that layout. - """ + """Copy top-level quantization_config into text_config with name normalization.""" if "quantization_config" in text_config: return text_config if "quantization_config" not in config_dict: @@ -332,12 +326,10 @@ def _inherit_quantization_config(config_dict: dict, if "modules_to_not_convert" in quantization_config: modules = _Qwen35ConfigCompat._normalize_exclude_modules( quantization_config["modules_to_not_convert"]) - modules = _Qwen35ConfigCompat._add_qkvz_bf16_workaround( - text_config, modules) quantization_config["modules_to_not_convert"] = sorted(set(modules)) if "ignore" in quantization_config: - modules = _Qwen35ConfigCompat._add_qkvz_bf16_workaround( - text_config, list(quantization_config["ignore"])) + modules = _Qwen35ConfigCompat._normalize_exclude_modules( + list(quantization_config["ignore"])) quantization_config["ignore"] = sorted(set(modules)) text_config["quantization_config"] = quantization_config return text_config @@ -361,24 +353,6 @@ def _normalize_exclude_modules(modules: list[str]) -> list[str]: normalized.add(name) return sorted(normalized) - @staticmethod - def _add_qkvz_bf16_workaround(text_config: dict, - modules: list[str]) -> list[str]: - """Keep packed linear-attention qkvz on bf16 path for all linear-attention layers. - - Temporary until FP8 block-scale TP loading is fixed for this layout. - """ - try: - layer_types = get_qwen3_hybrid_layer_types( - SimpleNamespace(**text_config)) - except (ValueError, AttributeError): - return modules - for layer_idx, layer_type in enumerate(layer_types): - if layer_type == "linear_attention": - modules.append( - f"model.layers.{layer_idx}.linear_attn.in_proj_qkvz") - return modules - @staticmethod def _flatten_rope(text_config: dict) -> dict: """Flatten rope_parameters into top-level rope_theta / partial_rotary_factor / rope_scaling. diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 5aeeeadfe861..50bb73cf78be 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -1484,11 +1484,10 @@ def get_num_available_tokens(self, stats = self.impl.get_kv_cache_stats() rs_free = stats.num_free_blocks_per_window_size.get( LinearCacheType.RECURRENT_STATES.value, 0) - # Reserve 1 block for the always-allocated last block (corner case - # / final live state) so we don't promise more tokens than the - # pool can actually back at allocation time. - usable_rs_blocks = max(0, rs_free - 1) - rs_token_cap = usable_rs_blocks * interval + # Each recurrent-state block backs one context stop position. + # For prompts shorter than the snapshot interval, the only stop + # position is the final/live state, so one free block is enough. + rs_token_cap = max(0, rs_free) * interval result = min(result, rs_token_cap) return max(result, 0) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 45aafe381b27..41b91d8a0d33 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1207,6 +1207,8 @@ def _create_warmup_request( token_num_upper_bound=num_tokens, max_num_draft_tokens=self.max_total_draft_tokens) available_blocks = kv_cache_manager.get_num_free_blocks() + if num_tokens <= 0 and num_gen_requests == 0: + return None if num_tokens > self.max_num_tokens or num_tokens > available_tokens: return None diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 14d813a99dfd..bd66a3e9cff9 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -3,13 +3,14 @@ import os import traceback import warnings +from contextlib import nullcontext from typing import Callable, Optional, Tuple import torch from tensorrt_llm._torch.models.checkpoints.base_checkpoint_loader import ( AutoCheckpointMapper, BaseCheckpointLoader) -from tensorrt_llm._utils import str_dtype_to_torch +from tensorrt_llm._utils import is_device_integrated, str_dtype_to_torch from tensorrt_llm.llmapi.llm_args import (ExecutorMemoryType, ModelExpressConfig, TorchLlmArgs) from tensorrt_llm.llmapi.llm_utils import apply_model_defaults_to_llm_args @@ -348,137 +349,170 @@ def load( is_meta_init = False memo = dict() + defer_weight_allocation = (is_meta_init and is_device_integrated() + and load_format == LoadFormat.AUTO) + if defer_weight_allocation: + logger.info( + "Integrated GPU (unified memory): deferring CUDA weight " + "allocation until checkpoint load to reduce peak memory.") + + def materialize_cuda_tensor(t: torch.Tensor) -> torch.Tensor: + if t not in memo: + if t.device == torch.device('meta'): + cuda_t = torch.empty_like(t, device='cuda') + elif not t.is_cuda: + cuda_t = t.cuda() + else: + cuda_t = t + memo[t] = cuda_t + memo[cuda_t] = cuda_t + return memo[t] if self.model_weights_memory_tag is not None: # Allocate buffers to the outer virtual_memory_scope, # but parameters (weights) to the dedicated inner virtual_memory_scope. - def allocate_buffer_on_cuda(t: torch.Tensor): - if t not in memo: - if t.device == torch.device('meta'): - cuda_t = torch.empty_like(t, device='cuda') - else: - cuda_t = t.cuda() - memo[t] = cuda_t - memo[cuda_t] = cuda_t - return memo[t] - - _apply_to_buffers_only(model, allocate_buffer_on_cuda) + _apply_to_buffers_only(model, materialize_cuda_tensor) - need_initialized_weights = load_format not in (LoadFormat.AUTO, - LoadFormat.DUMMY) + if not defer_weight_allocation: + need_initialized_weights = load_format not in ( + LoadFormat.AUTO, LoadFormat.DUMMY) - def allocate_weights_on_cuda(t: torch.Tensor): - if t not in memo: - cuda_t = torch.empty_like(t, device='cuda') - if t.device != torch.device('meta') and ( - need_initialized_weights or is_meta_init): - if t.is_cuda: - memory_type_map = { - ExecutorMemoryType.MODEL_WEIGHTS_MAIN: - ExecutorMemoryType.MODEL_ENGINE_MAIN, - ExecutorMemoryType.MODEL_WEIGHTS_DRAFT: - ExecutorMemoryType.MODEL_ENGINE_DRAFT, - } - - warnings.warn( - f"A weight tensor of shape {t.shape} is already allocated on CUDA device before " - f"the weight allocation stage. This will cause extra CUDA memory usage in the " - f"'{memory_type_map[self.model_weights_memory_tag]}' scope." - ) - cuda_t.copy_(t) - memo[t] = cuda_t - memo[cuda_t] = cuda_t - return memo[t] - - with virtual_memory_scope( - self.model_weights_memory_tag, - self.model_weights_restore_mode) as pool: - model._apply(allocate_weights_on_cuda) - self._weight_pool_proxy = pool + def allocate_weights_on_cuda(t: torch.Tensor): + if t not in memo: + cuda_t = torch.empty_like(t, device='cuda') + if t.device != torch.device('meta') and ( + need_initialized_weights or is_meta_init): + if t.is_cuda: + memory_type_map = { + ExecutorMemoryType.MODEL_WEIGHTS_MAIN: + ExecutorMemoryType.MODEL_ENGINE_MAIN, + ExecutorMemoryType.MODEL_WEIGHTS_DRAFT: + ExecutorMemoryType.MODEL_ENGINE_DRAFT, + } + + warnings.warn( + f"A weight tensor of shape {t.shape} is already allocated on CUDA device before " + f"the weight allocation stage. This will cause extra CUDA memory usage in the " + f"'{memory_type_map[self.model_weights_memory_tag]}' scope." + ) + cuda_t.copy_(t) + memo[t] = cuda_t + memo[cuda_t] = cuda_t + return memo[t] + + with virtual_memory_scope( + self.model_weights_memory_tag, + self.model_weights_restore_mode) as pool: + model._apply(allocate_weights_on_cuda) + self._weight_pool_proxy = pool elif is_meta_init: + if defer_weight_allocation: + _apply_to_buffers_only(model, materialize_cuda_tensor) + else: - def init_meta_tensor(t: torch.Tensor): - if t.device != torch.device('meta'): - return t + def init_meta_tensor(t: torch.Tensor): + if t.device != torch.device('meta'): + return t - if t not in memo: - memo[t] = torch.empty_like(t, device='cuda') - return memo[t] + if t not in memo: + memo[t] = torch.empty_like(t, device='cuda') + return memo[t] - model._apply(init_meta_tensor) + model._apply(init_meta_tensor) - # Ensure everything is at least on CUDA - # No-op if worked as expected - model.to("cuda") + if not defer_weight_allocation: + # Ensure everything is at least on CUDA + # No-op if worked as expected + model.to("cuda") del memo - rank_model_storage = get_rank_model_storage(model) - logger.info( - f"Use {rank_model_storage / (1024**3):.2f} GB for model weights." - ) - weights_preloaded = False - if load_format == LoadFormat.AUTO: - # Pass model= so format-specific loaders (e.g. MX) can - # write weights directly into parameter buffers via P2P. - # Generic loaders ignore model=; loaders that can consume a - # live module reference (MX) use it for direct writes. - load_weights_kwargs: dict = { - "mapping": self.mapping, - "model": model, - } - - if hasattr(model, 'llm_checkpoint_dir'): - weights = checkpoint_loader.load_weights( - model.llm_checkpoint_dir, **load_weights_kwargs) - else: - weights = checkpoint_loader.load_weights( - checkpoint_dir, **load_weights_kwargs) - - # When MX P2P succeeds, weights are already in model params. - # A non-empty dict contains size-mismatched tensors that - # should be merged via the standard disk pipeline. - weights_preloaded = checkpoint_loader.is_weights_preloaded() - self.weight_mapper = checkpoint_loader.get_initialized_weight_mapper( - model, config) - - if weights: - self._call_load_weights(model.load_weights, weights, - self.weight_mapper) - - if self.spec_config is not None and self.spec_config.spec_dec_mode.need_load_draft_weights( - ): - weights = checkpoint_loader.load_weights( - self.spec_config.speculative_model, - mapping=self.mapping) - - draft_model_arch = model.draft_config.pretrained_config.architectures[ - 0] - draft_weight_mapper = AutoCheckpointMapper.get( - checkpoint_loader.checkpoint_format, draft_model_arch) - draft_weight_mapper.init_model_and_config( - model.draft_model, model.draft_config) - - self._call_load_weights(model.load_draft_weights, weights, - draft_weight_mapper) - - elif load_format == LoadFormat.DUMMY: - self.weight_mapper = checkpoint_loader.get_initialized_weight_mapper( - model, config) - initialize_dummy_weights(model) - if self.spec_config is not None and self.spec_config.spec_dec_mode.need_load_draft_weights( - ): - model.draft_model.load_weights_from_target_model(model) - - elif load_format == LoadFormat.VISION_ONLY: - # Vision weights are already loaded within the model. + if not defer_weight_allocation: + rank_model_storage = get_rank_model_storage(model) logger.info( - "LoadFormat.VISION_ONLY: skipping weight loading; using preloaded vision weights." + f"Use {rank_model_storage / (1024**3):.2f} GB for model weights." ) - else: - raise NotImplementedError( - f"No load support for load format: {load_format}") + weight_vm_ctx = nullcontext() + if defer_weight_allocation and self.model_weights_memory_tag is not None: + weight_vm_ctx = virtual_memory_scope( + self.model_weights_memory_tag, + self.model_weights_restore_mode) + + weights_preloaded = False + with weight_vm_ctx as weight_pool: + if defer_weight_allocation and self.model_weights_memory_tag is not None: + self._weight_pool_proxy = weight_pool + + if load_format == LoadFormat.AUTO: + # Pass model= so format-specific loaders (e.g. MX) can + # write weights directly into parameter buffers via P2P. + # Generic loaders ignore model=; loaders that can consume a + # live module reference (MX) use it for direct writes. + load_weights_kwargs: dict = { + "mapping": self.mapping, + "model": model, + } + + if hasattr(model, 'llm_checkpoint_dir'): + weights = checkpoint_loader.load_weights( + model.llm_checkpoint_dir, **load_weights_kwargs) + else: + weights = checkpoint_loader.load_weights( + checkpoint_dir, **load_weights_kwargs) + + # When MX P2P succeeds, weights are already in model params. + # A non-empty dict contains size-mismatched tensors that + # should be merged via the standard disk pipeline. + weights_preloaded = checkpoint_loader.is_weights_preloaded() + self.weight_mapper = checkpoint_loader.get_initialized_weight_mapper( + model, config) + + if weights: + self._call_load_weights(model.load_weights, weights, + self.weight_mapper) + + if self.spec_config is not None and self.spec_config.spec_dec_mode.need_load_draft_weights( + ): + weights = checkpoint_loader.load_weights( + self.spec_config.speculative_model, + mapping=self.mapping) + + draft_model_arch = model.draft_config.pretrained_config.architectures[ + 0] + draft_weight_mapper = AutoCheckpointMapper.get( + checkpoint_loader.checkpoint_format, + draft_model_arch) + draft_weight_mapper.init_model_and_config( + model.draft_model, model.draft_config) + + self._call_load_weights(model.load_draft_weights, + weights, draft_weight_mapper) + + elif load_format == LoadFormat.DUMMY: + self.weight_mapper = checkpoint_loader.get_initialized_weight_mapper( + model, config) + initialize_dummy_weights(model) + if self.spec_config is not None and self.spec_config.spec_dec_mode.need_load_draft_weights( + ): + model.draft_model.load_weights_from_target_model(model) + + elif load_format == LoadFormat.VISION_ONLY: + # Vision weights are already loaded within the model. + logger.info( + "LoadFormat.VISION_ONLY: skipping weight loading; using preloaded vision weights." + ) + + else: + raise NotImplementedError( + f"No load support for load format: {load_format}") + + if defer_weight_allocation: + model.to("cuda") + rank_model_storage = get_rank_model_storage(model) + logger.info( + f"Use {rank_model_storage / (1024**3):.2f} GB for model weights " + f"after integrated GPU lazy allocation.") checkpoint_loader.post_load_apply( model, weights_preloaded=weights_preloaded) diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 033dfa9f7949..af4bde27d28f 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -16,7 +16,7 @@ CppMambaHybridCacheManager, PythonMambaCacheManager, ) -from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp +from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp, KVCacheManager from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.bindings.internal.batch_manager import LinearCacheType from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MTPDecodingConfig @@ -442,6 +442,28 @@ def test_cpp_hybrid_recurrent_pool_floor_with_block_reuse(): ) +def test_cpp_hybrid_available_tokens_allows_live_state_only(monkeypatch): + """One recurrent-state block can hold a prompt shorter than the snapshot interval.""" + + interval = 8192 + token_num_upper_bound = 4095 + mgr = object.__new__(CppMambaHybridCacheManager) + mgr.linear_attention_metadata = SimpleNamespace(states_snapshot_interval=interval) + mgr.impl = SimpleNamespace( + get_kv_cache_stats=lambda: SimpleNamespace( + num_free_blocks_per_window_size={LinearCacheType.RECURRENT_STATES.value: 1} + ) + ) + + monkeypatch.setattr( + KVCacheManager, + "get_num_available_tokens", + lambda self, token_num_upper_bound, max_num_draft_tokens=0, **kwargs: token_num_upper_bound, + ) + + assert mgr.get_num_available_tokens(token_num_upper_bound) == token_num_upper_bound + + # --------------------------------------------------------------------------- # CppMambaHybridCacheManager: rank with zero local mamba layers # diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index db6873d04cc3..560668b4f04d 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -1,5 +1,6 @@ import unittest from dataclasses import dataclass +from types import SimpleNamespace from unittest.mock import Mock import torch @@ -144,6 +145,43 @@ def create_model_engine_and_kvcache(llm_args: TorchLlmArgs = None, class PyTorchModelEngineTestCase(unittest.TestCase): + def test_create_warmup_request_rejects_zero_token_batch(self) -> None: + """A zero-token warmup must not create an empty scheduled batch.""" + + class ZeroTokenKvCacheManager: + tokens_per_block = 32 + + def get_num_available_tokens(self, **kwargs): + return 0 + + def get_num_free_blocks(self): + return 1 + + def add_dummy_requests(self, *args, **kwargs): + raise AssertionError( + "zero-token warmup should not allocate dummy requests") + + model_engine = SimpleNamespace( + kv_cache_manager_key=ResourceManagerType.KV_CACHE_MANAGER, + spec_config=None, + max_total_draft_tokens=0, + max_draft_loop_tokens=0, + max_beam_width=1, + max_num_tokens=4096, + batch_size=1, + max_seq_len=4096, + use_mrope=False, + _get_draft_kv_cache_manager=lambda resource_manager: None, + _get_num_extra_decoding_steps=lambda: 0, + ) + resource_manager = ResourceManager( + {ResourceManagerType.KV_CACHE_MANAGER: ZeroTokenKvCacheManager()}) + + warmup_request = PyTorchModelEngine._create_warmup_request( + model_engine, resource_manager, num_tokens=0, num_gen_requests=0) + + self.assertIsNone(warmup_request) + def test_pad_generation_requests(self) -> None: model_engine, kv_cache_manager = create_model_engine_and_kvcache() resource_manager = ResourceManager( diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_mmapped_safetensors_weights.py b/tests/unittest/_torch/models/checkpoints/hf/test_mmapped_safetensors_weights.py new file mode 100644 index 000000000000..6205ba49f97f --- /dev/null +++ b/tests/unittest/_torch/models/checkpoints/hf/test_mmapped_safetensors_weights.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch +from safetensors.torch import save_file + +from tensorrt_llm._torch.models.checkpoints.base_weight_loader import MmappedSafetensorsWeights +from tensorrt_llm._torch.models.modeling_utils import materialize_meta_parameters + + +def test_mmapped_safetensors_weights_lazy_load_and_mark_consumed(tmp_path): + weight_path = tmp_path / "model.safetensors" + save_file({"layer.weight": torch.ones(2, 3)}, weight_path) + + weights = MmappedSafetensorsWeights([str(weight_path)]) + assert "layer.weight" in weights + assert len(weights) == 1 + + tensor = weights["layer.weight"] + assert hasattr(tensor, "get_shape") + + deleted = weights.mark_consumed("layer") + assert deleted == 1 + assert len(weights) == 0 + + +def test_materialize_meta_parameters(): + module = torch.nn.Linear(4, 2, bias=False) + module.weight = torch.nn.Parameter( + torch.empty(2, 4, device="meta"), + requires_grad=False, + ) + + materialize_meta_parameters(module) + + assert module.weight.is_cuda + assert not module.weight.is_meta + assert module.weight.shape == (2, 4) diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_qwen3_5_weight_mapper.py b/tests/unittest/_torch/models/checkpoints/hf/test_qwen3_5_weight_mapper.py index 2a5e6867b9f6..4c3535747c79 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_qwen3_5_weight_mapper.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_qwen3_5_weight_mapper.py @@ -41,24 +41,55 @@ def test_qwen35_modelopt_preprocess_preserves_scalar_fp8_scale_name(): assert weights["model.layers.0.linear_attn.out_proj.weight_scale"].shape == torch.Size([]) -def test_qwen35_dequantizes_per_tensor_fp8_linear_attention_projection(): +def test_qwen35_rescales_per_tensor_fp8_linear_attention_qkvz_projection(): mapper = _make_mapper(dtype=torch.bfloat16) + mapper._config.pretrained_config.linear_num_key_heads = 1 + mapper._config.pretrained_config.linear_num_value_heads = 1 + mapper._config.pretrained_config.linear_key_head_dim = 2 + mapper._config.pretrained_config.linear_value_head_dim = 2 weight_name = "model.layers.0.linear_attn.in_proj_qkv.weight" scale_name = "model.layers.0.linear_attn.in_proj_qkv.weight_scale" input_scale_name = "model.layers.0.linear_attn.in_proj_qkv.input_scale" + z_weight_name = "model.layers.0.linear_attn.in_proj_z.weight" + z_scale_name = "model.layers.0.linear_attn.in_proj_z.weight_scale" + z_input_scale_name = "model.layers.0.linear_attn.in_proj_z.input_scale" weights = { - weight_name: torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32), + weight_name: torch.tensor( + [ + [1.0, 2.0], + [3.0, 4.0], + [1.0, 2.0], + [3.0, 4.0], + [1.0, 2.0], + [3.0, 4.0], + ], + dtype=torch.float8_e4m3fn, + ), scale_name: torch.tensor(0.5, dtype=torch.float32), - input_scale_name: torch.tensor(1.0, dtype=torch.float32), + input_scale_name: torch.tensor(2.0, dtype=torch.float32), + z_weight_name: torch.tensor([[5.0, 6.0], [7.0, 8.0]], dtype=torch.float8_e4m3fn), + z_scale_name: torch.tensor(1.0, dtype=torch.float32), + z_input_scale_name: torch.tensor(1.0, dtype=torch.float32), } - updated = mapper._dequantize_linear_attn_per_tensor_fp8(weights) + updated = mapper._rescale_linear_attn_qkvz_per_tensor_fp8(weights) + packed = mapper._pack_split_projections(updated) assert scale_name not in updated assert input_scale_name not in updated - assert updated[weight_name].dtype == torch.bfloat16 + assert z_scale_name not in updated + assert z_input_scale_name not in updated + assert packed["model.layers.0.linear_attn.in_proj_qkvz.weight"].dtype == torch.float8_e4m3fn torch.testing.assert_close( - updated[weight_name], - torch.tensor([[0.5, 1.0], [1.5, 2.0]], dtype=torch.bfloat16), + packed["model.layers.0.linear_attn.in_proj_qkvz.weight_scale"], + torch.tensor(1.0, dtype=torch.float32), + ) + torch.testing.assert_close( + packed["model.layers.0.linear_attn.in_proj_qkvz.input_scale"], + torch.tensor(2.0, dtype=torch.float32), + ) + torch.testing.assert_close( + packed["model.layers.0.linear_attn.in_proj_qkvz.weight"][:6].to(torch.float32), + weights[weight_name].to(torch.float32) * 0.5, ) diff --git a/tests/unittest/_torch/test_model_config.py b/tests/unittest/_torch/test_model_config.py index 6bb8533cfbf7..71d570a7a3b4 100644 --- a/tests/unittest/_torch/test_model_config.py +++ b/tests/unittest/_torch/test_model_config.py @@ -243,7 +243,7 @@ def test_load_modelopt_mixed_precision_w4a16_uses_layer_quant_config(tmp_path): assert layer_quant_config["model.layers.0.mlp.experts"].quant_algo == QuantAlgo.W4A16_NVFP4 -def test_qwen35_compat_adds_qkvz_excludes_to_modelopt_ignore(): +def test_qwen35_compat_normalizes_modelopt_ignore_without_qkvz_override(): config_dict = { "architectures": ["Qwen3_5MoeForConditionalGeneration"], "model_type": "qwen3_5_moe", @@ -263,7 +263,7 @@ def test_qwen35_compat_adds_qkvz_excludes_to_modelopt_ignore(): "quant_method": "modelopt", "quant_algo": "MIXED_PRECISION", "ignore": [ - "model.language_model.layers.0.linear_attn.in_proj_qkv", + "model.language_model.layers.0.linear_attn.out_proj", "mtp.layers.0*", ], "quantized_layers": {}, @@ -273,15 +273,14 @@ def test_qwen35_compat_adds_qkvz_excludes_to_modelopt_ignore(): normalized = _Qwen35ConfigCompat.normalize(config_dict) ignore = normalized["quantization_config"]["ignore"] - assert "model.language_model.layers.0.linear_attn.in_proj_qkv" in ignore - assert "mtp.layers.0*" in ignore - assert "model.layers.0.linear_attn.in_proj_qkvz" in ignore - assert "model.layers.2.linear_attn.in_proj_qkvz" in ignore - assert "model.layers.3.linear_attn.in_proj_qkvz" in ignore + assert ignore == ["model.layers.0.linear_attn.out_proj"] + assert "model.layers.0.linear_attn.in_proj_qkvz" not in ignore + assert "model.layers.2.linear_attn.in_proj_qkvz" not in ignore + assert "model.layers.3.linear_attn.in_proj_qkvz" not in ignore assert "model.layers.1.linear_attn.in_proj_qkvz" not in ignore -def test_qwen35_file_modelopt_config_adds_qkvz_excludes(tmp_path): +def test_qwen35_file_modelopt_config_preserves_excludes_without_qkvz_override(tmp_path): config_dict = { "architectures": ["Qwen3_5MoeForConditionalGeneration"], "model_type": "qwen3_5_moe", @@ -341,9 +340,9 @@ def test_qwen35_file_modelopt_config_adds_qkvz_excludes(tmp_path): exclude_modules = model_config.quant_config.exclude_modules assert "mtp.layers.0*" in exclude_modules assert "mtp*" in exclude_modules - assert "model.layers.0.linear_attn.in_proj_qkvz" in exclude_modules - assert "model.layers.2.linear_attn.in_proj_qkvz" in exclude_modules - assert "model.layers.3.linear_attn.in_proj_qkvz" in exclude_modules + assert "model.layers.0.linear_attn.in_proj_qkvz" not in exclude_modules + assert "model.layers.2.linear_attn.in_proj_qkvz" not in exclude_modules + assert "model.layers.3.linear_attn.in_proj_qkvz" not in exclude_modules assert "model.layers.1.linear_attn.in_proj_qkvz" not in exclude_modules From d3fe03536f33d769ee9edfdb7326b16766c84234 Mon Sep 17 00:00:00 2001 From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> Date: Thu, 16 Jul 2026 19:54:44 +0000 Subject: [PATCH 09/25] Fix semantic conflicts after main merge Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> --- cpp/tensorrt_llm/thop/cublasScaledMM.cpp | 20 -------- .../_torch/modules/fused_moe/create_moe.py | 50 +++++++++++++++---- tensorrt_llm/models/quant_config_utils.py | 20 +++++++- .../hf/test_qwen3_5_weight_mapper.py | 18 +++++-- .../moe/test_cute_dsl_b12x_moe_backend.py | 4 +- 5 files changed, 74 insertions(+), 38 deletions(-) diff --git a/cpp/tensorrt_llm/thop/cublasScaledMM.cpp b/cpp/tensorrt_llm/thop/cublasScaledMM.cpp index 40d2f5dbaf94..dea6f51363e5 100644 --- a/cpp/tensorrt_llm/thop/cublasScaledMM.cpp +++ b/cpp/tensorrt_llm/thop/cublasScaledMM.cpp @@ -220,26 +220,6 @@ void cublas_gemm_caller(torch::Tensor& out, torch::Tensor const& a, torch::Tenso auto stream = at::cuda::getCurrentCUDAStream(a.get_device()); auto const& workspace = getWorkspaceTensor(a.device(), stream.stream()); - // Persistent per-stream cublasLt workspace. The cublasLt kernel records a - // fixed pointer to this workspace; if we destruct the workspace storage at - // function return (the previous behavior), the CUDA caching allocator may - // hand the same block out to a later allocation while a captured CUDA - // graph still references the workspace pointer -> use-after-free that - // surfaces as free-block-tree corruption on the next allocator operation - // (e.g. FusedMoeRunner::getWorkspaceInfo's first torch::empty destructor). - // - // Keyed by (device, stream) so that concurrent GEMMs on different streams - // don't race on the same scratch bytes. Mirrors PyTorch's per-stream - // cublasLt workspace cache in at::cuda. - thread_local std::unordered_map workspace_cache; - auto stream_ptr = stream.stream(); - auto& workspace = workspace_cache[stream_ptr]; - if (!workspace.defined() || workspace.device() != a.device()) - { - auto const workspace_options = torch::TensorOptions().dtype(torch::kUInt8).device(a.device()); - workspace = torch::empty(CUBLAS_WORKSPACE_SIZE, workspace_options); - } - auto* a_ptr = static_cast(a.data_ptr()); auto* b_ptr = static_cast(b.data_ptr()); auto* out_ptr = static_cast(out.data_ptr()); diff --git a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py index 1f3e9b112dd0..d4267287b0ae 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py @@ -28,6 +28,22 @@ from .routing import BaseMoeRoutingMethod +def _is_same_or_child_module_path(lhs: str, rhs: str) -> bool: + return lhs == rhs or lhs.startswith(f"{rhs}.") or rhs.startswith(f"{lhs}.") + + +def _get_layer_quant_config(model_config: ModelConfig, + layer_idx: Optional[int]) -> Optional[QuantConfig]: + if layer_idx is None or model_config.quant_config_dict is None: + return None + + moe_module_name = f"model.layers.{layer_idx}.mlp.experts" + for name, quant_config in model_config.quant_config_dict.items(): + if _is_same_or_child_module_path(name, moe_module_name): + return quant_config + return None + + def _get_pretrained_megamoe_capability_args( model_config: ModelConfig) -> Dict[str, Optional[object]]: """Extract dtype / hidden / intermediate kwargs for MegaMoE @@ -59,9 +75,9 @@ def get_moe_cls( layer_idx: Optional[int] = None, ) -> Type[MoE]: moe_backend = model_config.moe_backend - quant_config = model_config.quant_config - if override_quant_config is not None: - quant_config = override_quant_config + quant_config = (override_quant_config + or _get_layer_quant_config(model_config, layer_idx) + or model_config.quant_config) layer_prefix = f"[layer_idx={layer_idx}] " if layer_idx is not None else "" if moe_backend.upper() == "MARLIN": # Marlin MoE is a Hopper-specific NVFP4 W4A16 backend. Require nvfp4 @@ -74,14 +90,18 @@ def get_moe_cls( elif moe_backend.upper() == "VANILLA": return VanillaMoE elif moe_backend.upper() == "CUTEDSL": + has_w4a16_nvfp4 = (quant_config is not None + and quant_config.quant_algo == QuantAlgo.W4A16_NVFP4) if quant_config is not None and ( quant_config.quant_mode.has_fp8_block_scales() - or quant_config.quant_mode.has_nvfp4()): - # On SM120 / SM121 + NVFP4 the cuteDSL family member is the + or quant_config.quant_mode.has_nvfp4() or has_w4a16_nvfp4): + # On SM120 / SM121 + NVFP4/W4A16_NVFP4 the cuteDSL family member is the # hybrid CUTLASS-prefill / FlashInfer NVFP4 MoE decode backend # (CuteDslB12xFusedMoE). Prefer it when flashinfer is importable; # otherwise fall through to CuteDslFusedMoE for SM100 / SM103. - if quant_config.quant_mode.has_nvfp4(): + has_nvfp4 = (quant_config.quant_mode.has_nvfp4() + and not has_w4a16_nvfp4) + if has_nvfp4 or has_w4a16_nvfp4: from tensorrt_llm._utils import get_sm_version sm_version = get_sm_version() if sm_version in CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS: @@ -97,13 +117,24 @@ def get_moe_cls( except ImportError: logger.warning( "CuteDslB12xFusedMoE eligible (SM%d + NVFP4) " - "but flashinfer is not importable; using CuteDslFusedMoE.", + "but flashinfer is not importable; using %s.", sm_version, + "CutlassFusedMoE" + if has_w4a16_nvfp4 else "CuteDslFusedMoE", ) + if has_w4a16_nvfp4: + return CutlassFusedMoE + elif has_w4a16_nvfp4: + logger.warning( + "CuteDslB12xFusedMoE requires SM120/121 for W4A16_NVFP4 " + "(got SM%d). Using CutlassFusedMoE.", + sm_version, + ) + return CutlassFusedMoE return CuteDslFusedMoE else: logger.warning( - f"{layer_prefix}CuteDslFusedMoE only supports fp8_block_scales and nvfp4. " + f"{layer_prefix}CuteDslFusedMoE only supports fp8_block_scales, nvfp4, and w4a16_nvfp4. " f"Check out details in quant_config: {quant_config}. Using CutlassFusedMoE instead." ) return CutlassFusedMoE @@ -212,7 +243,8 @@ def resolve_moe_cls( ) -> Type[MoE]: moe_cls = get_moe_cls(model_config, override_quant_config, layer_idx) - effective_quant_config = override_quant_config or model_config.quant_config + effective_quant_config = (override_quant_config or _get_layer_quant_config( + model_config, layer_idx) or model_config.quant_config) has_quant = (effective_quant_config is not None and effective_quant_config.layer_quant_mode.has_any_quant( exclude_kv_cache=True)) diff --git a/tensorrt_llm/models/quant_config_utils.py b/tensorrt_llm/models/quant_config_utils.py index 09d50a1727e3..e008833e4d97 100644 --- a/tensorrt_llm/models/quant_config_utils.py +++ b/tensorrt_llm/models/quant_config_utils.py @@ -40,9 +40,24 @@ def update_quant_config_from_compressed_tensors( weights_quant_config = group_config["weights"] inputs_quant_config = group_config["input_activations"] weights_quant_strategy = weights_quant_config["strategy"] - inputs_quant_strategy = inputs_quant_config["strategy"] - if weights_quant_config["num_bits"] == 8: + if inputs_quant_config is None: + is_w4a16_nvfp4 = ( + hf_quant_config.get("format") == "nvfp4-pack-quantized" + and weights_quant_config["num_bits"] == 4 + and weights_quant_config.get("type") == "float" + and weights_quant_strategy == "tensor_group" + and weights_quant_config.get("group_size") == 16 + ) + if not is_w4a16_nvfp4: + raise ValueError( + "input_activations=None is only supported for W4A16 NVFP4 " + "with 4-bit float weights and group_size=16." + ) + quant_config.quant_algo = QuantAlgo.W4A16_NVFP4 + quant_config.group_size = 16 + elif weights_quant_config["num_bits"] == 8: + inputs_quant_strategy = inputs_quant_config["strategy"] if weights_quant_strategy == "channel": if inputs_quant_strategy != "token": raise ValueError(f"Unsupported inputs_quant_strategy: {inputs_quant_strategy}.") @@ -68,6 +83,7 @@ def update_quant_config_from_compressed_tensors( and weights_quant_config.get("type") == "float" and weights_quant_strategy == "tensor_group" ): + inputs_quant_strategy = inputs_quant_config["strategy"] # llm-compressor NVFP4: weights FP4 with FP8 per-group scales # (group_size=16), scaled by an FP32 global scale. if inputs_quant_strategy != "tensor_group": diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_qwen3_5_weight_mapper.py b/tests/unittest/_torch/models/checkpoints/hf/test_qwen3_5_weight_mapper.py index 4c3535747c79..9796217a2382 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_qwen3_5_weight_mapper.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_qwen3_5_weight_mapper.py @@ -18,11 +18,15 @@ import torch from tensorrt_llm._torch.models.checkpoints.hf.qwen3_5_weight_mapper import Qwen3_5MoeHfWeightMapper +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig def _make_mapper(dtype=torch.bfloat16): mapper = Qwen3_5MoeHfWeightMapper() - mapper._config = SimpleNamespace(pretrained_config=SimpleNamespace(torch_dtype=dtype)) + mapper._config = SimpleNamespace( + pretrained_config=SimpleNamespace(torch_dtype=dtype), + quant_config_dict={}, + ) return mapper @@ -30,19 +34,23 @@ def test_qwen35_modelopt_preprocess_preserves_scalar_fp8_scale_name(): mapper = _make_mapper() scale = torch.tensor(0.5, dtype=torch.float32) - is_modelopt_ckpt, weights = mapper._preprocess_modelopt_ckpt( + weights, is_modelopt_pb_wo = mapper._normalize_scale_names( { "model.layers.0.linear_attn.out_proj.weight_scale": scale, - } + }, + QuantAlgo.MIXED_PRECISION, ) - assert is_modelopt_ckpt + assert not is_modelopt_pb_wo assert "model.layers.0.linear_attn.out_proj.weight_scale" in weights assert weights["model.layers.0.linear_attn.out_proj.weight_scale"].shape == torch.Size([]) def test_qwen35_rescales_per_tensor_fp8_linear_attention_qkvz_projection(): mapper = _make_mapper(dtype=torch.bfloat16) + mapper._config.quant_config_dict = { + "model.layers.0.linear_attn.in_proj_qkvz": QuantConfig(quant_algo=QuantAlgo.FP8), + } mapper._config.pretrained_config.linear_num_key_heads = 1 mapper._config.pretrained_config.linear_num_value_heads = 1 mapper._config.pretrained_config.linear_key_head_dim = 2 @@ -73,7 +81,7 @@ def test_qwen35_rescales_per_tensor_fp8_linear_attention_qkvz_projection(): z_input_scale_name: torch.tensor(1.0, dtype=torch.float32), } - updated = mapper._rescale_linear_attn_qkvz_per_tensor_fp8(weights) + updated = mapper._requantize_linear_attn_fp8_qkvz(weights) packed = mapper._pack_split_projections(updated) assert scale_name not in updated diff --git a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py index 8ed41823048d..5d2ce75f435c 100644 --- a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py @@ -334,8 +334,8 @@ def _convert_sf_to_mma_layout(scales, *, m, k, num_groups): module.fc31_input_scale = torch.tensor(2.0, dtype=torch.float32) module.fc2_input_scale = torch.tensor(4.0, dtype=torch.float32) - with patch.object(NVFP4CutlassFusedMoEMethod, "post_load_weights", return_value=None): - NVFP4CuteDslB12xFusedMoEMethod().post_load_weights(module) + with patch.object(NVFP4CutlassFusedMoEMethod, "transform_weights", return_value=None): + NVFP4CuteDslB12xFusedMoEMethod().transform_weights(module) assert _FakeB12xWrapper.calls wrapper_kwargs = _FakeB12xWrapper.calls[0] From 44fd1f0e06d846a28032117f947370e66b8cd2f8 Mon Sep 17 00:00:00 2001 From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> Date: Thu, 16 Jul 2026 20:59:58 +0000 Subject: [PATCH 10/25] Remove unrelated changes from feature branch Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> --- .../models/checkpoints/base_weight_loader.py | 153 +----------------- .../models/checkpoints/base_weight_mapper.py | 8 +- .../checkpoints/hf/exaone4_5_weight_mapper.py | 12 +- .../checkpoints/mistral/weight_mapper.py | 25 +-- .../_torch/pyexecutor/model_engine.py | 2 - .../hf/test_mmapped_safetensors_weights.py | 38 ----- 6 files changed, 5 insertions(+), 233 deletions(-) delete mode 100644 tests/unittest/_torch/models/checkpoints/hf/test_mmapped_safetensors_weights.py diff --git a/tensorrt_llm/_torch/models/checkpoints/base_weight_loader.py b/tensorrt_llm/_torch/models/checkpoints/base_weight_loader.py index 0a0cd4e8eca6..10d72a3a86c7 100644 --- a/tensorrt_llm/_torch/models/checkpoints/base_weight_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/base_weight_loader.py @@ -1,12 +1,9 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import re import threading from abc import ABC, abstractmethod -from typing import Any, Callable, Dict, Iterator, List, Tuple, Union - -import safetensors +from typing import Any, Dict, Iterator, Tuple, Union from tensorrt_llm.mapping import Mapping @@ -94,154 +91,6 @@ def mark_consumed(self, prefix: str) -> int: return len(keys_to_delete) -class MmappedSafetensorsWeights: - """ - Memory-mapped safetensors checkpoint access for integrated GPU systems. - - Opens checkpoint shards via ``safetensors.safe_open`` without loading full - tensors into RAM. ``mark_consumed`` drops index entries after weights are - copied into the model to reduce peak memory on unified-memory machines. - """ - - def __init__(self, weight_files: List[str]): - self._files = weight_files - self._handles: list = [] - self._key_locations: dict[str, tuple[int, str]] = {} - self._lock = threading.Lock() - for file_idx, path in enumerate(weight_files): - handle = safetensors.safe_open(path, framework="pt", device="cpu") - self._handles.append(handle) - for key in handle.keys(): - if key in self._key_locations: - raise RuntimeError( - f"Duplicate tensor name {key!r} in checkpoint files.") - self._key_locations[key] = (file_idx, key) - - def _get_slice(self, key: str): - file_idx, tensor_name = self._key_locations[key] - return self._handles[file_idx].get_slice(tensor_name) - - def __getitem__(self, key: str) -> Any: - return self._get_slice(key) - - def __contains__(self, key: str) -> bool: - return key in self._key_locations - - def __len__(self) -> int: - return len(self._key_locations) - - def __iter__(self) -> Iterator[str]: - with self._lock: - return iter(list(self._key_locations.keys())) - - def keys(self): - with self._lock: - return list(self._key_locations.keys()) - - def values(self): - with self._lock: - return [self._get_slice(k) for k in self._key_locations.keys()] - - def items(self) -> Iterator[Tuple[str, Any]]: - with self._lock: - keys = list(self._key_locations.keys()) - return [(k, self._get_slice(k)) for k in keys] - - def get(self, key: str, default: Any = None) -> Any: - if key in self._key_locations: - return self._get_slice(key) - return default - - def mark_consumed(self, prefix: str) -> int: - with self._lock: - keys_to_delete = [ - k for k in self._key_locations if k.startswith(prefix + ".") - ] - for key in keys_to_delete: - del self._key_locations[key] - return len(keys_to_delete) - - def remap_keys(self, key_mapping: dict[str, - str]) -> "MmappedSafetensorsWeights": - remapped = object.__new__(MmappedSafetensorsWeights) - remapped._files = self._files - remapped._handles = self._handles - remapped._lock = threading.Lock() - remapped._key_locations = {} - for old_key, loc in self._key_locations.items(): - new_key = key_mapping.get(old_key, old_key) - if new_key in remapped._key_locations: - raise RuntimeError( - f"Duplicate tensor name {new_key!r} after key remap.") - remapped._key_locations[new_key] = loc - return remapped - - def transform_keys( - self, transform_fn: Callable[[str], - str]) -> "MmappedSafetensorsWeights": - key_mapping = { - old_key: transform_fn(old_key) - for old_key in self._key_locations - } - return self.remap_keys(key_mapping) - - def rename_by_regex( - self, pattern_mapping: dict[str, - str]) -> "MmappedSafetensorsWeights": - key_mapping = {} - matched_keys = set() - for key in self._key_locations: - new_key = key - for pattern, replacement in pattern_mapping.items(): - if re.match(pattern, key): - new_key = re.sub(pattern, replacement, key) - matched_keys.add(key) - break - key_mapping[key] = new_key - return self.remap_keys(key_mapping) - - -def rename_weight_keys_with_regex( - weights: Union[Dict[str, Any], ConsumableWeightsDict, - MmappedSafetensorsWeights], - pattern_mapping: dict[str, str], -) -> Union[Dict[str, Any], ConsumableWeightsDict, MmappedSafetensorsWeights]: - if isinstance(weights, MmappedSafetensorsWeights): - return weights.rename_by_regex(pattern_mapping) - - is_consumable = isinstance(weights, ConsumableWeightsDict) - renamed_weights = {} - matched_keys = set() - for key in list(weights.keys()): - new_key = key - for pattern, replacement in pattern_mapping.items(): - if re.match(pattern, key): - new_key = re.sub(pattern, replacement, key) - matched_keys.add(key) - break - renamed_weights[new_key] = weights[key] - if is_consumable: - return ConsumableWeightsDict(renamed_weights) - return renamed_weights - - -def remap_weight_keys( - weights: Union[Dict[str, Any], ConsumableWeightsDict, - MmappedSafetensorsWeights], - key_mapping: dict[str, str], -) -> Union[Dict[str, Any], ConsumableWeightsDict, MmappedSafetensorsWeights]: - if isinstance(weights, MmappedSafetensorsWeights): - return weights.remap_keys(key_mapping) - - is_consumable = isinstance(weights, ConsumableWeightsDict) - renamed_weights = {} - for key in weights.keys(): - renamed_weights[key_mapping.get(key, key)] = weights[key] - if is_consumable: - return ConsumableWeightsDict(renamed_weights) - return renamed_weights - - class BaseWeightLoader(ABC): @abstractmethod diff --git a/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py index 06be9134bd7a..3ac76d6809be 100644 --- a/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py @@ -81,12 +81,8 @@ def rename_by_params_map(self, params_map: dict[str, str], """ import re - from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ( - ConsumableWeightsDict, MmappedSafetensorsWeights, - rename_weight_keys_with_regex) - - if isinstance(weights, MmappedSafetensorsWeights): - return rename_weight_keys_with_regex(weights, params_map) + from tensorrt_llm._torch.models.checkpoints.base_weight_loader import \ + ConsumableWeightsDict # Check if input is a ConsumableWeightsDict to preserve the type is_consumable = isinstance(weights, ConsumableWeightsDict) diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/exaone4_5_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/exaone4_5_weight_mapper.py index 0a4637488629..7e315acced94 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/exaone4_5_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/exaone4_5_weight_mapper.py @@ -13,10 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ( - ConsumableWeightsDict, - MmappedSafetensorsWeights, -) +from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper from tensorrt_llm._torch.models.modeling_utils import register_mapper @@ -25,13 +22,6 @@ class Exaone4_5HfWeightMapper(HfWeightMapper): def preprocess_weights(self, weights: dict): """Rename HF checkpoint prefixes; supports plain dict and ConsumableWeightsDict.""" - if isinstance(weights, MmappedSafetensorsWeights): - return weights.transform_keys( - lambda key: key.replace("model.visual.", "visual.") - if key.startswith("model.visual.") - else key - ) - is_consumable = isinstance(weights, ConsumableWeightsDict) renamed = {} for key, value in weights.items(): diff --git a/tensorrt_llm/_torch/models/checkpoints/mistral/weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/mistral/weight_mapper.py index c4e5b9936b3c..dd6e0332b849 100644 --- a/tensorrt_llm/_torch/models/checkpoints/mistral/weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/mistral/weight_mapper.py @@ -45,30 +45,7 @@ def __init__(self): # Adapted from: # https://github.com/vllm-project/vllm/blob/883b42896a9ed9791750d721fad26005b7569eba/vllm/model_executor/models/llama.py#L657 def rename_by_params_map(self, params_map: dict[str, str], weights: dict) -> dict: - from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ( - ConsumableWeightsDict, - MmappedSafetensorsWeights, - remap_weight_keys, - ) - - if isinstance(weights, MmappedSafetensorsWeights): - key_mapping = {} - for key in weights.keys(): - new_key = key - modules = key.split(".") - num_modules = len(modules) - for i in range(num_modules): - item = modules[i] - next_item = modules[i + 1] if i < num_modules - 1 else None - - combined_item = f"{item}.{next_item}" if next_item is not None else None - - if combined_item in params_map: - new_key = new_key.replace(combined_item, params_map[combined_item]) - elif item in params_map: - new_key = new_key.replace(item, params_map[item]) - key_mapping[key] = new_key - return remap_weight_keys(weights, key_mapping) + from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict # Check if input is a ConsumableWeightsDict to preserve the type is_consumable = isinstance(weights, ConsumableWeightsDict) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 6dcc8adc9f1d..62a442972830 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1772,8 +1772,6 @@ def _create_warmup_request( token_num_upper_bound=num_tokens, max_num_draft_tokens=self.max_total_draft_tokens) available_blocks = kv_cache_manager.get_num_free_blocks() - if num_tokens <= 0 and num_gen_requests == 0: - return None if num_tokens > self.max_num_tokens or num_tokens > available_tokens: return None diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_mmapped_safetensors_weights.py b/tests/unittest/_torch/models/checkpoints/hf/test_mmapped_safetensors_weights.py deleted file mode 100644 index 6205ba49f97f..000000000000 --- a/tests/unittest/_torch/models/checkpoints/hf/test_mmapped_safetensors_weights.py +++ /dev/null @@ -1,38 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -import torch -from safetensors.torch import save_file - -from tensorrt_llm._torch.models.checkpoints.base_weight_loader import MmappedSafetensorsWeights -from tensorrt_llm._torch.models.modeling_utils import materialize_meta_parameters - - -def test_mmapped_safetensors_weights_lazy_load_and_mark_consumed(tmp_path): - weight_path = tmp_path / "model.safetensors" - save_file({"layer.weight": torch.ones(2, 3)}, weight_path) - - weights = MmappedSafetensorsWeights([str(weight_path)]) - assert "layer.weight" in weights - assert len(weights) == 1 - - tensor = weights["layer.weight"] - assert hasattr(tensor, "get_shape") - - deleted = weights.mark_consumed("layer") - assert deleted == 1 - assert len(weights) == 0 - - -def test_materialize_meta_parameters(): - module = torch.nn.Linear(4, 2, bias=False) - module.weight = torch.nn.Parameter( - torch.empty(2, 4, device="meta"), - requires_grad=False, - ) - - materialize_meta_parameters(module) - - assert module.weight.is_cuda - assert not module.weight.is_meta - assert module.weight.shape == (2, 4) From bc2fb1a50dceabd0ae5c5272e732908ce5d7b41c Mon Sep 17 00:00:00 2001 From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> Date: Tue, 21 Jul 2026 02:38:16 +0000 Subject: [PATCH 11/25] use marlin nvfp4 kernel Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> --- cpp/tensorrt_llm/kernels/CMakeLists.txt | 10 +- .../w4a16_nvfp4_gemm_sm120.cuh | 183 -------- .../kernels/marlin/marlin_nvfp4_gemm.cu | 6 +- .../kernels/marlin/marlin_nvfp4_template.h | 5 +- .../kernels/marlin/marlin_repack.cu | 6 +- .../cutlassGemmW4A16NVFP4.cu | 59 --- .../cutlassGemmW4A16NVFP4.h | 35 -- cpp/tensorrt_llm/thop/w4a16Nvfp4Gemm.cpp | 53 --- .../_torch/custom_ops/cpp_custom_ops.py | 10 - tensorrt_llm/_torch/model_config.py | 25 +- .../_torch/models/modeling_nemotron_h.py | 59 ++- tensorrt_llm/_torch/modules/linear.py | 367 ++++++++------- .../defs/accuracy/test_llm_api_pytorch.py | 15 +- .../test_lists/test-db/l0_gb10.yml | 1 + .../test_lists/test-db/l0_rtx_pro_6000.yml | 1 + .../modeling/test_modeling_qwen3_5_quant.py | 124 +++++ .../_torch/modules/test_w4a16_nvfp4_linear.py | 422 ++++++++++-------- tests/unittest/_torch/test_model_config.py | 55 +++ .../thop/serial/test_nvfp4_marlin_gemm.py | 7 +- 19 files changed, 728 insertions(+), 715 deletions(-) delete mode 100644 cpp/tensorrt_llm/kernels/cutlass_kernels/w4a16_nvfp4_gemm/w4a16_nvfp4_gemm_sm120.cuh delete mode 100644 cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cutlassGemmW4A16NVFP4.cu delete mode 100644 cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cutlassGemmW4A16NVFP4.h create mode 100644 tests/unittest/_torch/modeling/test_modeling_qwen3_5_quant.py diff --git a/cpp/tensorrt_llm/kernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/CMakeLists.txt index fcea8829442b..c17a3fea1a9f 100644 --- a/cpp/tensorrt_llm/kernels/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/CMakeLists.txt @@ -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"); you may not @@ -61,7 +61,7 @@ list(FILTER SRC_CPP EXCLUDE REGEX "mhcKernels/.*") list(FILTER SRC_CU EXCLUDE REGEX "mhcKernels/.*") list(FILTER SRC_CPP EXCLUDE REGEX "compressorKernels/.*") list(FILTER SRC_CU EXCLUDE REGEX "compressorKernels/.*") -# Marlin is built as its own Hopper-only OBJECT library below. +# Marlin is built as its own architecture-scoped OBJECT library below. list(FILTER SRC_CPP EXCLUDE REGEX "marlin/.*") list(FILTER SRC_CU EXCLUDE REGEX "marlin/.*") @@ -82,8 +82,8 @@ if(FAST_BUILD) STATUS "FAST_BUILD enabled for kernels: using -O1 for CUDA compilation") endif() -# Marlin NVFP4: Hopper-only OBJECT library. Pinned to sm_90 so the global -# CMAKE_CUDA_ARCHITECTURES doesn't propagate. +# Marlin NVFP4: build the existing Hopper kernels and the dense W4A16/repack +# kernels used on SM12x. Unsupported device passes emit empty kernel stubs. file(GLOB_RECURSE MARLIN_SRC "marlin/*.cu" "marlin/*.cpp") if(MARLIN_SRC) add_library(marlin_src OBJECT ${MARLIN_SRC}) @@ -95,7 +95,7 @@ if(MARLIN_SRC) $ ) target_link_libraries(marlin_src PRIVATE trtllm_gen_fmha_interface) - set_cuda_architectures(marlin_src 90) + set_cuda_architectures(marlin_src 90 120f) endif() add_library( diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/w4a16_nvfp4_gemm/w4a16_nvfp4_gemm_sm120.cuh b/cpp/tensorrt_llm/kernels/cutlass_kernels/w4a16_nvfp4_gemm/w4a16_nvfp4_gemm_sm120.cuh deleted file mode 100644 index 0d2a3caeaac6..000000000000 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/w4a16_nvfp4_gemm/w4a16_nvfp4_gemm_sm120.cuh +++ /dev/null @@ -1,183 +0,0 @@ -/* - * Copyright (c) 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. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "cutlass/cutlass.h" -#include "cutlass/epilogue/thread/linear_combination.h" -#include "cutlass/epilogue/thread/scale_type.h" -#include "cutlass/gemm/device/gemm.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/kernels/weightOnlyBatchedGemv/cutlassGemmW4A16NVFP4.h" -#include "tensorrt_llm/kernels/weightOnlyBatchedGemv/nvfp4ScaleLayout.h" - -#include -#include -#include - -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace cutlass_gemm_w4a16_nvfp4 -{ -namespace sm120 -{ -namespace -{ - -using SizeType32 = tensorrt_llm::runtime::SizeType32; - -__device__ float loadNvfp4Weight(uint8_t const* weight, __nv_fp8_e4m3 const* weightScale, float const weightGlobalScale, - SizeType32 nIdx, SizeType32 kIdx, SizeType32 k) -{ - size_t const packedOffset = (static_cast(nIdx) * static_cast(k) + static_cast(kIdx)) / 2; - uint8_t const packed = weight[packedOffset]; - __nv_fp4_storage_t const nibble = (kIdx % 2 == 0) ? (packed & 0x0FU) : (packed >> 4U); - half const weightValue = static_cast(__nv_cvt_fp4_to_halfraw(nibble, __NV_E2M1)); - - SizeType32 const scaleIdx = w4a16_nvfp4::getScaleIndex(nIdx, kIdx / w4a16_nvfp4::kScaleGranularity, k); - float const scale = static_cast(weightScale[scaleIdx]); - return __half2float(weightValue) * scale * weightGlobalScale; -} - -__global__ void dequantizeWeightToBf16ColumnMajor(__nv_bfloat16* dequantizedWeight, uint8_t const* weight, - __nv_fp8_e4m3 const* weightScale, float const* weightGlobalScale, SizeType32 n, SizeType32 k) -{ - size_t const total = static_cast(n) * static_cast(k); - float const globalScale = weightGlobalScale[0]; - for (size_t idx - = static_cast(blockIdx.x) * static_cast(blockDim.x) + static_cast(threadIdx.x); - idx < total; idx += static_cast(gridDim.x) * static_cast(blockDim.x)) - { - SizeType32 const nIdx = static_cast(idx / static_cast(k)); - SizeType32 const kIdx = static_cast(idx % static_cast(k)); - size_t const dequantizedOffset = static_cast(kIdx) + static_cast(nIdx) * static_cast(k); - dequantizedWeight[dequantizedOffset] - = __float2bfloat16(loadNvfp4Weight(weight, weightScale, globalScale, nIdx, kIdx, k)); - } -} - -void dequantizeWeight(Params const& params, __nv_bfloat16* dequantizedWeight, cudaStream_t stream) -{ - constexpr SizeType32 kBlockSize = 256; - constexpr SizeType32 kMaxGridSize = 65535; - size_t const total = static_cast(params.n) * static_cast(params.k); - auto const gridSize = static_cast( - std::min((total + static_cast(kBlockSize) - 1) / static_cast(kBlockSize), - static_cast(kMaxGridSize))); - dequantizeWeightToBf16ColumnMajor<<>>(dequantizedWeight, - reinterpret_cast(params.weight), reinterpret_cast<__nv_fp8_e4m3 const*>(params.weightScale), - params.weightGlobalScale, params.n, params.k); - TLLM_CUDA_CHECK(cudaGetLastError()); -} - -bool runCutlassBf16Gemm(Params const& params, __nv_bfloat16 const* dequantizedWeight, cudaStream_t stream) -{ - using ElementA = cutlass::bfloat16_t; - using ElementB = cutlass::bfloat16_t; - using ElementOutput = cutlass::bfloat16_t; - using ElementAccumulator = float; - using ElementCompute = float; - using LayoutA = cutlass::layout::RowMajor; - using LayoutB = cutlass::layout::ColumnMajor; - using LayoutOutput = cutlass::layout::RowMajor; - using ThreadblockShape = cutlass::gemm::GemmShape<128, 128, 32>; - using WarpShape = cutlass::gemm::GemmShape<64, 64, 32>; - using InstructionShape = cutlass::gemm::GemmShape<16, 8, 16>; - using EpilogueOp - = cutlass::epilogue::thread::LinearCombination::value, - ElementAccumulator, ElementCompute, cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling>; - using Gemm = cutlass::gemm::device::Gemm, 4, 8, 8>; - - typename Gemm::Arguments arguments({params.m, params.n, params.k}, - {reinterpret_cast(params.act), params.k}, - {reinterpret_cast(dequantizedWeight), params.k}, - {reinterpret_cast(params.output), params.n}, - {reinterpret_cast(params.output), params.n}, {ElementCompute(1), ElementCompute(0)}, 1); - - Gemm gemm; - cutlass::Status status = gemm.can_implement(arguments); - if (status != cutlass::Status::kSuccess) - { - TLLM_LOG_WARNING("w4a16_nvfp4 transient BF16 CUTLASS GEMM cannot implement shape m=%d n=%d k=%d: %s", params.m, - params.n, params.k, cutlass::cutlassGetStatusString(status)); - return false; - } - - size_t const workspaceBytes = Gemm::get_workspace_size(arguments); - void* workspace = nullptr; - if (workspaceBytes > 0) - { - TLLM_CUDA_CHECK(cudaMallocAsync(&workspace, workspaceBytes, stream)); - } - - status = gemm.initialize(arguments, workspace, stream); - if (status == cutlass::Status::kSuccess) - { - status = gemm.run(stream); - } - - if (workspace != nullptr) - { - TLLM_CUDA_CHECK(cudaFreeAsync(workspace, stream)); - } - - if (status != cutlass::Status::kSuccess) - { - TLLM_LOG_WARNING("w4a16_nvfp4 transient BF16 CUTLASS GEMM failed for shape m=%d n=%d k=%d: %s", params.m, - params.n, params.k, cutlass::cutlassGetStatusString(status)); - return false; - } - TLLM_CUDA_CHECK(cudaGetLastError()); - return true; -} - -bool runTransientDequantCutlassGemm(Params const& params, cudaStream_t stream) -{ - void* dequantizedWeight = nullptr; - size_t const dequantizedBytes - = static_cast(params.n) * static_cast(params.k) * sizeof(__nv_bfloat16); - TLLM_CUDA_CHECK(cudaMallocAsync(&dequantizedWeight, dequantizedBytes, stream)); - dequantizeWeight(params, reinterpret_cast<__nv_bfloat16*>(dequantizedWeight), stream); - bool const dispatched - = runCutlassBf16Gemm(params, reinterpret_cast<__nv_bfloat16 const*>(dequantizedWeight), stream); - TLLM_CUDA_CHECK(cudaFreeAsync(dequantizedWeight, stream)); - return dispatched; -} - -} // namespace - -inline bool dispatch(Params const& params, cudaStream_t stream) -{ - if (params.inputType == CUDA_R_16BF && params.outputType == CUDA_R_16BF) - { - return runTransientDequantCutlassGemm(params, stream); - } - return false; -} - -} // namespace sm120 -} // namespace cutlass_gemm_w4a16_nvfp4 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_gemm.cu b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_gemm.cu index f9cbbe3d64ab..9bd26c9f210f 100644 --- a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_gemm.cu +++ b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_gemm.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -327,8 +327,8 @@ void marlinNvfp4Gemm(void const* act_bf16, void const* weight, void* output, voi bool use_fp32_reduce, cudaStream_t stream) { int const sm = tensorrt_llm::common::getSMVersion(); - TLLM_CHECK_WITH_INFO( - sm >= 90 && sm < 100, "Marlin NVFP4 GEMM is only supported on Hopper (SM 9.x); current SM = %d", sm); + TLLM_CHECK_WITH_INFO((sm >= 90 && sm < 100) || (sm >= 120 && sm < 130), + "Marlin NVFP4 GEMM is only supported on SM 9.x and SM 12.x; current SM = %d", sm); int dev; cudaGetDevice(&dev); diff --git a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_template.h b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_template.h index 9fdc9d2a6d38..e059b60d8a02 100644 --- a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_template.h +++ b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_template.h @@ -49,8 +49,9 @@ namespace MARLIN_NAMESPACE_NAME { -// Empty kernel stub for non-Hopper device passes; see marlin.cuh. -#if defined(__CUDA_ARCH__) && !(__CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000) +// Empty kernel stub outside the architectures supported by dense Marlin. +#if defined(__CUDA_ARCH__) \ + && !((__CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000) || (__CUDA_ARCH__ >= 1200 && __CUDA_ARCH__ < 1300)) template = 90 && sm < 100, "Marlin NVFP4 repack is only supported on Hopper (SM 9.x); current SM = %d", sm); + TLLM_CHECK_WITH_INFO((sm >= 90 && sm < 100) || (sm >= 120 && sm < 130), + "Marlin NVFP4 repack is only supported on SM 9.x and SM 12.x; current SM = %d", sm); int blocks; int dev; diff --git a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cutlassGemmW4A16NVFP4.cu b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cutlassGemmW4A16NVFP4.cu deleted file mode 100644 index 2c42d8cf8c7d..000000000000 --- a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cutlassGemmW4A16NVFP4.cu +++ /dev/null @@ -1,59 +0,0 @@ -/* - * Copyright (c) 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. - * 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. - */ - -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/kernels/cutlass_kernels/w4a16_nvfp4_gemm/w4a16_nvfp4_gemm_sm120.cuh" -#include "tensorrt_llm/kernels/weightOnlyBatchedGemv/cutlassGemmW4A16NVFP4.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace cutlass_gemm_w4a16_nvfp4 -{ -namespace -{ - -bool isSupported(Params const& params) -{ - int const smVersion = tensorrt_llm::common::getSMVersion(); - return (smVersion == 120 || smVersion == 121) && params.inputType == CUDA_R_16BF && params.outputType == CUDA_R_16BF - && params.m > 16 && params.n > 0 && params.k > 0 && params.n % 32 == 0 && params.k % 32 == 0 - && params.weightScale != nullptr && params.weightGlobalScale != nullptr; -} - -} // namespace - -bool cutlassGemmDispatcher(Params const& params, cudaStream_t stream) -{ - if (!isSupported(params)) - { - int const smVersion = tensorrt_llm::common::getSMVersion(); - TLLM_LOG_WARNING( - "tensorrt_llm::kernels::cutlass_gemm_w4a16_nvfp4::cutlassGemmDispatcher [NOT DISPATCHED], " - "inputType=%d, outputType=%d, m=%d, n=%d, k=%d, sm=%d", - params.inputType, params.outputType, params.m, params.n, params.k, smVersion); - return false; - } - - return sm120::dispatch(params, stream); -} - -} // namespace cutlass_gemm_w4a16_nvfp4 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cutlassGemmW4A16NVFP4.h b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cutlassGemmW4A16NVFP4.h deleted file mode 100644 index ae3a2832a77e..000000000000 --- a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cutlassGemmW4A16NVFP4.h +++ /dev/null @@ -1,35 +0,0 @@ -/* - * Copyright (c) 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. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace cutlass_gemm_w4a16_nvfp4 -{ - -using Params = tensorrt_llm::kernels::cuda_core_gemm_w4a16_nvfp4::Params; - -bool cutlassGemmDispatcher(Params const& params, cudaStream_t stream); - -} // namespace cutlass_gemm_w4a16_nvfp4 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/w4a16Nvfp4Gemm.cpp b/cpp/tensorrt_llm/thop/w4a16Nvfp4Gemm.cpp index 54d8bd229b6e..d5faf50da520 100644 --- a/cpp/tensorrt_llm/thop/w4a16Nvfp4Gemm.cpp +++ b/cpp/tensorrt_llm/thop/w4a16Nvfp4Gemm.cpp @@ -15,7 +15,6 @@ */ #include "tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.h" -#include "tensorrt_llm/kernels/weightOnlyBatchedGemv/cutlassGemmW4A16NVFP4.h" #include "tensorrt_llm/kernels/weightOnlyBatchedGemv/nvfp4ScaleLayout.h" #include "tensorrt_llm/thop/thUtils.h" @@ -111,31 +110,6 @@ void w4a16Nvfp4GemmCaller( TORCH_CHECK(dispatched, "Failed to dispatch w4a16_nvfp4_gemm kernel"); } -void w4a16Nvfp4CutlassGemmCaller( - Tensor& out, Tensor const& act, Tensor const& weight, Tensor const& weightScale, Tensor const& weightScale2) -{ - auto const m = static_cast(act.sizes()[0]); - auto const k = static_cast(act.sizes()[1]); - auto const n = static_cast(weight.sizes()[0]); - TORCH_CHECK(weight.sizes()[1] * 2 == k, "weight shape [N, K/2] must match activation shape [M, K]"); - - auto stream = at::cuda::getCurrentCUDAStream(act.get_device()); - - auto* actPtr = static_cast(act.data_ptr()); - auto* weightPtr = static_cast(weight.data_ptr()); - auto* weightScalePtr = static_cast(weightScale.data_ptr()); - auto* weightScale2Ptr = static_cast(weightScale2.data_ptr()); - auto* outPtr = static_cast(out.data_ptr()); - - auto const inputType = convert_torch_dtype(act.scalar_type()); - auto const outType = convert_torch_dtype(out.scalar_type()); - - tensorrt_llm::kernels::cutlass_gemm_w4a16_nvfp4::Params params( - actPtr, weightPtr, weightScalePtr, weightScale2Ptr, outPtr, m, n, k, inputType, outType); - bool const dispatched = tensorrt_llm::kernels::cutlass_gemm_w4a16_nvfp4::cutlassGemmDispatcher(params, stream); - TORCH_CHECK(dispatched, "Failed to dispatch w4a16_nvfp4_cutlass_gemm kernel"); -} - } // namespace Tensor& w4a16_nvfp4_gemm_out(Tensor const& act, Tensor const& weight, Tensor const& weightScale, @@ -173,29 +147,6 @@ Tensor w4a16_nvfp4_gemm(Tensor const& act, Tensor const& weight, Tensor const& w return w4a16_nvfp4_gemm_out(act, weight, weightScale, weightScale2, outDtype, bias, out); } -Tensor w4a16_nvfp4_cutlass_gemm(Tensor const& act, Tensor const& weight, Tensor const& weightScale, - Tensor const& weightScale2, std::optional outDtype, std::optional const& bias) -{ - CHECK_TH_CUDA(act); - CHECK_CONTIGUOUS(act); - checkActDtype(act); - CHECK_INPUT(weight, FLOAT4_E2M1X2); - CHECK_INPUT(weightScale, SF_DTYPE); - CHECK_INPUT(weightScale2, torch::kFloat32); - - TORCH_CHECK(act.dim() == 2 && weight.dim() == 2); - TORCH_CHECK(weight.sizes()[1] * 2 == act.sizes()[1]); - checkWeightScaleSize(weightScale, weight.sizes()[0], act.sizes()[1]); - TORCH_CHECK(weightScale2.numel() == 1, "weight_scale_2 must be a scalar tensor"); - TORCH_CHECK(!bias.has_value(), "w4a16_nvfp4_cutlass_gemm does not support bias"); - - auto const outDtypeValue = outDtype.value_or(act.scalar_type()); - std::vector outputSize = {act.sizes()[0], weight.sizes()[0]}; - Tensor out = at::empty(outputSize, act.options().dtype(outDtypeValue)); - w4a16Nvfp4CutlassGemmCaller(out, act, weight, weightScale, weightScale2); - return out; -} - } // namespace torch_ext TRTLLM_NAMESPACE_END @@ -205,13 +156,9 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) m.def( "w4a16_nvfp4_gemm(Tensor act, Tensor weight, Tensor weight_scale, Tensor weight_scale_2, ScalarType? " "out_dtype, Tensor? bias=None) -> Tensor"); - m.def( - "w4a16_nvfp4_cutlass_gemm(Tensor act, Tensor weight, Tensor weight_scale, Tensor weight_scale_2, ScalarType? " - "out_dtype, Tensor? bias=None) -> Tensor"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) { m.impl("w4a16_nvfp4_gemm", &tensorrt_llm::torch_ext::w4a16_nvfp4_gemm); - m.impl("w4a16_nvfp4_cutlass_gemm", &tensorrt_llm::torch_ext::w4a16_nvfp4_cutlass_gemm); } diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index a187617c56aa..e8363dab2aa0 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -1166,16 +1166,6 @@ def _(act: torch.Tensor, del weight_scale, weight_scale_2, bias return act.new_empty((act.shape[0], weight.shape[0]), dtype=out_dtype) - @torch.library.register_fake("trtllm::w4a16_nvfp4_cutlass_gemm") - def _(act: torch.Tensor, - weight: torch.Tensor, - weight_scale: torch.Tensor, - weight_scale_2: torch.Tensor, - out_dtype: Optional[torch.dtype], - bias: Optional[torch.Tensor] = None): - del weight_scale, weight_scale_2, bias - return act.new_empty((act.shape[0], weight.shape[0]), dtype=out_dtype) - @torch.library.register_fake("trtllm::marlin_nvfp4_gemm") def _(mat_a: torch.Tensor, mat_b: torch.Tensor, diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 29b8b8fd390e..2f9c346fc0da 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -328,9 +328,11 @@ def is_generation_model(model_architectures: Optional[List[str]], # once ModelType is used in pytorch flow. @staticmethod - def resolve_moe_backend(moe_backend: str, - architecture: str, - quant_config: Optional[QuantConfig] = None) -> str: + def resolve_moe_backend( + moe_backend: str, + architecture: str, + quant_config: Optional[QuantConfig] = None, + layer_quant_config: Optional[Dict[str, QuantConfig]] = None) -> str: """Resolve AUTO moe_backend to a specific backend based on model architecture. Args: @@ -345,6 +347,15 @@ def resolve_moe_backend(moe_backend: str, if moe_backend.upper() != "AUTO": return moe_backend + is_w4a16_nvfp4 = (quant_config is not None and quant_config.quant_algo + in (QuantAlgo.W4A16_NVFP4, "W4A16_NVFP4")) + if not is_w4a16_nvfp4 and layer_quant_config is not None: + is_w4a16_nvfp4 = any(config.quant_algo in (QuantAlgo.W4A16_NVFP4, + "W4A16_NVFP4") + for config in layer_quant_config.values()) + if is_w4a16_nvfp4 and get_sm_version() in (120, 121): + return "CUTEDSL" + if architecture in _DEEPSEEK_V4_ARCHITECTURES: sm_version = get_sm_version() if 100 <= sm_version < 120: @@ -439,12 +450,15 @@ def _build_modelopt_quant_config(json_quant_configs, checkpoint_dir, quant_config.exclude_modules = json_quant_configs.get( 'exclude_modules', quant_config.exclude_modules) + w4a16_nvfp4_group_size = None for layer in mixed_quant_configs: layer_cfg = mixed_quant_configs[layer] config = QuantConfig() config.kv_cache_quant_algo = kv_cache_quant_algo config.quant_algo = QuantAlgo(layer_cfg['quant_algo']) config.group_size = layer_cfg.get('group_size', None) + if config.quant_algo == QuantAlgo.W4A16_NVFP4: + w4a16_nvfp4_group_size = config.group_size or 16 # AWQ-specific extras emitted by modelopt per-layer. if 'has_zero_point' in layer_cfg: config.has_zero_point = layer_cfg['has_zero_point'] @@ -452,6 +466,9 @@ def _build_modelopt_quant_config(json_quant_configs, checkpoint_dir, config.pre_quant_scale = layer_cfg['pre_quant_scale'] mixed_quant_configs[layer] = config layer_quant_config = mixed_quant_configs + if w4a16_nvfp4_group_size is not None: + quant_config.quant_algo = None + quant_config.group_size = w4a16_nvfp4_group_size elif quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES: if quant_config.group_size is None: quant_config.group_size = 128 @@ -1156,7 +1173,7 @@ def _recursive_update_config(config: transformers.PretrainedConfig, requested_moe_backend, architecture, quant_config=quant_config, - ) + layer_quant_config=layer_quant_config) if architecture in _DEEPSEEK_V4_ARCHITECTURES: layer_quant_config = cls._set_deepseek_v4_routed_moe_quant_config( diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index 5c2e60c4706a..1013a898aa5b 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py @@ -162,6 +162,29 @@ def forward( **kwargs) +def _get_nemotron_h_moe_model_config( + model_config: ModelConfig[PretrainedConfig], + layer_idx: int) -> ModelConfig[PretrainedConfig]: + # Per-expert mixed precision config is more specific than the global config. + if model_config.quant_config_dict is not None: + experts_prefix = f"model.layers.{layer_idx}.mixer.experts." + for key, cfg in model_config.quant_config_dict.items(): + if key.startswith(experts_prefix): + return replace(model_config, quant_config=cfg) + + quant_config = model_config.quant_config + if (quant_config is not None and quant_config.quant_algo + in (QuantAlgo.W4A16_NVFP4, "W4A16_NVFP4") + and model_config.moe_backend.upper() != "CUTEDSL"): + moe_quant_config = quant_config.model_copy(deep=True) + moe_quant_config.quant_algo = QuantAlgo.NVFP4 + moe_quant_config.__dict__.pop("quant_mode", None) + moe_quant_config.__dict__.pop("layer_quant_mode", None) + return replace(model_config, quant_config=moe_quant_config) + + return model_config + + # Ref code: https://huggingface.co/nvidia/Nemotron-Nano-3-30B-A3.5B-dev-1024/blob/main/modeling_nemotron_h.py#L818 class NemotronHMOE(nn.Module): @@ -247,18 +270,8 @@ def _moe(name): moe_backend=model_config.moe_backend, ) - # For MIXED_PRECISION models, the global quant_config has quant_algo=MIXED_PRECISION - # which maps to QuantMode(0) (no quant). This would cause the MoE backend to select - # UnquantizedFusedMoEMethod and allocate BF16 weight buffers, causing a shape mismatch - # when loading NVFP4/W4A8_NVFP4_FP8 quantized expert weights. - # Look up the per-expert quant config from quant_config_dict and use it for create_moe. - moe_model_config = model_config - if model_config.quant_config_dict is not None: - experts_prefix = f"model.layers.{layer_idx}.mixer.experts." - for key, cfg in model_config.quant_config_dict.items(): - if key.startswith(experts_prefix): - moe_model_config = replace(model_config, quant_config=cfg) - break + moe_model_config = _get_nemotron_h_moe_model_config( + model_config, layer_idx) # Setup MoE experts. self.experts = create_moe( @@ -1210,14 +1223,20 @@ def __init__( sublayer_quant_config = self._get_mtp_sublayer_quant_config( model_config, self.layer_idx) + sublayer_moe_backend = model_config.moe_backend + if (sublayer_quant_config is None + or sublayer_quant_config.quant_algo is None): + sublayer_moe_backend = "CUTLASS" # Create a model_config copy with quant_config overridden and # spec_config cleared. All other fields (use_cuda_graph, - # moe_backend, moe_max_num_tokens, etc.) must be inherited - # so MoE layers are configured correctly for CUDA graph - # capture and communication (e.g., DeepEP). + # moe_max_num_tokens, etc.) must be inherited so MoE layers are + # configured correctly for CUDA graph capture and communication + # (e.g., DeepEP). BF16 MTP body layers cannot use NVFP4-only MoE + # backends, so route those sublayers to CUTLASS. sublayer_model_config = replace(model_config, quant_config=sublayer_quant_config, + moe_backend=sublayer_moe_backend, spec_config=None) self.layers[str(step_rel_idx)] = NemotronHMTPDecoderLayer( @@ -1237,14 +1256,16 @@ def _get_mtp_sublayer_quant_config(self, model_config: NemotronHModelConfig, layer_idx: int): """ Get quantization config for MTP sublayer. - The MTP layer in the nvfp4 checkpoint is unquantized. Because the TRTLLM - moe_backend only supports fp8/fp4 quantization, we need to override - the quant_config for the MTP layer. + The Nano3.5 NVFP4 W4A16 checkpoint stores the MTP body tensors in BF16. + The shared MTP head still receives the checkpoint-backed lm_head, so its + logits path keeps the lm_head precision instead of inheriting this + sublayer override. """ from tensorrt_llm.models.modeling_utils import QuantConfig quant_config = model_config.quant_config - # MTP layers are always unquantized, force quant_algo=None + # This checkpoint's MTP body is unquantized, so force quant_algo=None + # only for the MTP sublayers constructed here. if quant_config is None: return None return QuantConfig( diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 89acfc65af55..befb5285d9db 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -384,6 +384,9 @@ def load_weights(self, def transform_weights(self, module: Linear) -> None: return None + def cache_derived_state(self, module: Linear) -> None: + return None + def post_load_weights(self, module: Linear) -> None: self.transform_weights(module) @@ -1902,16 +1905,9 @@ def transform_weights(self, module: Linear) -> None: class W4A16NVFP4LinearMethod(NVFP4LinearMethod): - CUDA_CORE_MAX_M: ClassVar[int] = 16 - CUTLASS3_ENV: ClassVar[str] = "TRTLLM_W4A16_NVFP4_CUTLASS3" + """W4A16 NVFP4 linear with a small-M CUDA-core fast path.""" - def _can_use_cutlass3_w4a16_prefill(self, module: Linear, - input: torch.Tensor, m: int) -> bool: - return (os.environ.get(self.CUTLASS3_ENV, "0") == "1" - and m > self.CUDA_CORE_MAX_M and get_sm_version() in (120, 121) - and input.dtype == torch.bfloat16 - and module.dtype == torch.bfloat16 and input.shape[-1] % 32 == 0 - and module.weight.shape[0] % 32 == 0) + CUDA_CORE_MAX_M: ClassVar[int] = 16 def create_weights(self, module: Linear, in_features: int, out_features: int, bias: bool, dtype: torch.dtype): @@ -1995,16 +1991,34 @@ def process_weights_after_loading_fused_gate_up_linear( module, super().process_weights_after_loading_fused_gate_up_linear) - def apply(self, module: Linear, input: torch.Tensor, - bias: Optional[torch.Tensor]): - if input.dim() == 1: - m = 1 - elif input.dim() > 2: - m = math.prod(input.shape[:-1]) - else: - m = input.shape[0] - use_cutlass3_prefill = self._can_use_cutlass3_w4a16_prefill( - module, input, m) + def transform_weights(self, module: Linear) -> None: + # Keep the checkpoint layout for the CUDA-core path and materialize the + # much smaller linear scale view once for Triton dequantization. + LinearMethodBase.transform_weights(self, module) + self.cache_derived_state(module) + + def cache_derived_state(self, module: Linear) -> None: + pad_rows = fp4_utils.pad_up(module.out_features, 128) + pad_cols = fp4_utils.pad_up( + module.in_features // module.scaling_vector_size, 4) + scale_swizzled = module.weight_scale.data.view( + fp4_utils.float4_sf_dtype).reshape(pad_rows, pad_cols) + module._w4a16_weight_scale_linear = ( + torch.ops.trtllm.block_scale_interleave_reverse( + scale_swizzled).reshape(-1)) + + @staticmethod + def _prepare_input(module: Linear, input: torch.Tensor): + if isinstance(input, (Fp4QuantizedTensor, tuple)): + raise RuntimeError( + "W4A16NVFP4LinearMethod requires a high-precision input; " + "disable upstream FP4 fusion") + + if input.dtype == torch.float8_e4m3fn: + assert module.inv_input_scale is not None, \ + "W4A16NVFP4LinearMethod: FP8 input requires static inv_input_scale" + input = (input.to(module.dtype) / module.inv_input_scale).to( + module.dtype) original_shape = None if input.dim() > 2: @@ -2015,28 +2029,190 @@ def apply(self, module: Linear, input: torch.Tensor, assert input.dtype == module.pre_quant_scale.dtype, \ "Input dtype and pre_quant_scale dtype must match" input = input * module.pre_quant_scale + return input, original_shape - gemm_op = (torch.ops.trtllm.w4a16_nvfp4_cutlass_gemm if - use_cutlass3_prefill else torch.ops.trtllm.w4a16_nvfp4_gemm) - output = gemm_op( - input, - module.weight, - module.weight_scale, - module.weight_scale_2, - module.dtype, - bias=None, - ) - - if output.shape[-1] > module.out_features: - output = output[..., :module.out_features].contiguous() + @classmethod + def _can_use_cuda_core(cls, module: Linear, input: torch.Tensor) -> bool: + return (input.dim() == 2 and 0 < input.shape[0] <= cls.CUDA_CORE_MAX_M + and get_sm_version() in (120, 121) + and input.dtype in (torch.float16, torch.bfloat16) + and module.dtype in (torch.float16, torch.bfloat16) + and input.is_contiguous() and input.shape[1] % 32 == 0 + and module.weight.shape[0] % 2 == 0 + and module.weight.shape[1] * 2 == input.shape[1] + and hasattr(torch.ops.trtllm, "w4a16_nvfp4_gemm")) + @staticmethod + def _restore_output(output: torch.Tensor, original_shape, + bias: Optional[torch.Tensor]): if original_shape is not None: output = output.reshape(*original_shape[:-1], output.shape[-1]) - if bias is not None: output = output + bias return output + def apply(self, module: Linear, input: torch.Tensor, + bias: Optional[torch.Tensor]): + input, original_shape = self._prepare_input(module, input) + if self._can_use_cuda_core(module, input): + output = torch.ops.trtllm.w4a16_nvfp4_gemm( + input, + module.weight, + module.weight_scale, + module.weight_scale_2, + module.dtype, + bias=None, + ) + return self._restore_output(output, original_shape, bias) + + from tensorrt_llm._torch.modules.fused_moe.triton_dequant_nvfp4 import \ + dequant_nvfp4_2d_triton + weight_deq = dequant_nvfp4_2d_triton( + module.weight.view(torch.uint8), + module._w4a16_weight_scale_linear, + module.weight_scale_2, + target_dtype=module.dtype, + sf_vec_size=module.scaling_vector_size, + ) + + if getattr(module, "use_custom_cublas_mm", False): + output_buffer_kind = ( + int(BufferKind.NCCL_WINDOW) + if self.supports_nccl_symmetric_memory_window_output + and getattr(module, "all_reduce", None) is not None + and module.all_reduce.uses_nccl_symmetric_memory_window() else + int(BufferKind.DEFAULT)) + mapping = getattr(module, "mapping", None) + group = (mapping.tp_group + if output_buffer_kind == int(BufferKind.NCCL_WINDOW) + and mapping is not None else None) + output = torch.ops.trtllm.cublas_mm( + input, + weight_deq.t(), + bias, + out_dtype=None, + output_buffer_kind=output_buffer_kind, + group=group, + ) + else: + output = F.linear(input, weight_deq, bias) + return self._restore_output(output, original_shape, bias=None) + + +class MarlinNVFP4LinearMethod(W4A16NVFP4LinearMethod): + """W4A16 NVFP4 linear backed by Marlin.""" + + @staticmethod + def is_supported(module: Linear) -> bool: + return (get_sm_version() in (120, 121) + and getattr(module, "dtype", None) == torch.bfloat16 + and not getattr(module, "use_fused_gemm_allreduce", False) + and hasattr(torch.ops.trtllm, "marlin_nvfp4_gemm") + and hasattr(torch.ops.trtllm, "gptq_marlin_repack")) + + def transform_weights(self, module: Linear) -> None: + from tensorrt_llm.quantization.utils import marlin_utils + + weight = module.weight.data + weight_scale = module.weight_scale.data + size_n = module.out_features + size_k = module.in_features + group_size = module.scaling_vector_size + + assert size_k % group_size == 0, ( + f"size_k {size_k} must be divisible by group_size {group_size}") + + size_k_pad = fp4_utils.pad_up(size_k, 64) + size_n_pad = fp4_utils.pad_up(size_n, 128) + num_groups = size_k // group_size + scale_unswizzled = torch.ops.trtllm.block_scale_interleave_reverse( + weight_scale.view(size_n_pad, -1)) + scale_2d = scale_unswizzled[:size_n, :num_groups] + + if size_k_pad != size_k or size_n_pad != size_n: + num_groups_pad = size_k_pad // group_size + weight = F.pad(weight, + (0, + (size_k_pad - size_k) // 2, 0, size_n_pad - size_n)) + # Marlin's S0E5M3 fast dequantization requires a non-zero scale for + # zero-weight K padding. 0x08 is the smallest normal E4M3 value. + scale_2d = F.pad(scale_2d, (0, num_groups_pad - num_groups), + value=0x08) + scale_2d = F.pad(scale_2d, (0, 0, 0, size_n_pad - size_n), value=0) + + qweight_int32 = weight.view(torch.int32).T.contiguous() + perm = torch.empty(0, dtype=torch.int32, device=weight.device) + marlin_weight = torch.ops.trtllm.gptq_marlin_repack( + b_q_weight=qweight_int32, + perm=perm, + size_k=size_k_pad, + size_n=size_n_pad, + num_bits=4, + is_a_8bit=False, + ) + + scale_2d = scale_2d.view(torch.float8_e4m3fn).T.contiguous() + marlin_scale = marlin_utils.marlin_permute_scales(scale_2d.to( + torch.half), + size_k_pad, + size_n_pad, + group_size=group_size) + marlin_scale = marlin_utils.nvfp4_marlin_process_scales(marlin_scale) + + weight_scale_2 = module.weight_scale_2.data + if (weight_scale_2.numel() == 0 + or not torch.isfinite(weight_scale_2).all() + or weight_scale_2.item() == 0): + weight_scale_2 = torch.tensor([1.0], + dtype=torch.float32, + device=weight.device) + weight_global_scale = marlin_utils.nvfp4_marlin_process_global_scale( + weight_scale_2.to(torch.bfloat16)) + + module.weight = Parameter(marlin_weight, requires_grad=False) + module.weight_scale = Parameter(marlin_scale, requires_grad=False) + module.weight_global_scale = Parameter(weight_global_scale, + requires_grad=False) + self.cache_derived_state(module) + + def cache_derived_state(self, module: Linear) -> None: + module._marlin_size_k = fp4_utils.pad_up(module.in_features, 64) + module._marlin_size_n = fp4_utils.pad_up(module.out_features, 128) + + def apply(self, module: Linear, input: torch.Tensor, + bias: Optional[torch.Tensor]): + input, original_shape = self._prepare_input(module, input) + size_k = module.in_features + size_n = module.out_features + size_k_pad = getattr(module, "_marlin_size_k", size_k) + size_n_pad = getattr(module, "_marlin_size_n", size_n) + + input_bf16 = input.bfloat16() + if size_k_pad != size_k: + input_bf16 = F.pad(input_bf16, (0, size_k_pad - size_k)) + output = torch.ops.trtllm.marlin_nvfp4_gemm( + input_bf16, + module.weight, + scale_a=None, + scale_b=module.weight_scale, + alpha=None, + weight_global_scale=module.weight_global_scale, + bias=None, + out_dtype=module.dtype, + size_n=size_n_pad, + size_k=size_k_pad, + output_buffer_kind=int(BufferKind.DEFAULT), + ) + if size_n_pad != size_n: + output = output[..., :size_n].contiguous() + return self._restore_output(output, original_shape, bias) + + def apply_linear_allreduce(self, module: Linear, input: torch.Tensor, + bias: Optional[torch.Tensor], tp_rank: int, + tp_group: List[int]): + raise RuntimeError( + "MarlinNVFP4LinearMethod does not support apply_linear_allreduce") + class W4A8NVFP4FP8LinearMethod(LinearMethodBase): @@ -2882,124 +3058,6 @@ def apply(self, module: Linear, input: torch.Tensor, return output -class MarlinNVFP4LinearMethod(NVFP4LinearMethod): - """NVFP4 Linear method backed by the Marlin W4A16 kernel (Hopper only).""" - - def transform_weights(self, module: Linear) -> None: - from tensorrt_llm.quantization.utils import marlin_utils - - weight = module.weight.data - weight_scale = module.weight_scale.data - size_n = module.out_features - size_k = module.in_features - group_size = module.scaling_vector_size # 16 - - assert size_k % group_size == 0, ( - f"size_k {size_k} must be divisible by group_size {group_size}") - - size_k_pad = fp4_utils.pad_up(size_k, 64) - size_n_pad = fp4_utils.pad_up(size_n, 128) - - num_groups = size_k // group_size - n_padded = size_n_pad - scale_unswizzled = torch.ops.trtllm.block_scale_interleave_reverse( - weight_scale.view(n_padded, -1)) - # [size_n, num_groups] block scales; uint8 storage (reverse interleave), - # reinterpreted as E4M3 after any padding. Pad in uint8 since F.pad does - # not support float8. - scale_2d = scale_unswizzled[:size_n, :num_groups] - - if size_k_pad != size_k or size_n_pad != size_n: - num_groups_pad = size_k_pad // group_size - # weight: [N, K/2] uint8 -> [N_pad, K_pad/2] (FP4 zero == 0.0) - weight = F.pad(weight, - (0, - (size_k_pad - size_k) // 2, 0, size_n_pad - size_n)) - # scales: [N, num_groups] -> [N_pad, num_groups_pad]. - # The Marlin S0E5M3 fast-dequant is NOT zero-safe: a zero scale on - # a (zero-weight) padded K-group still corrupts that tile's output. - # Since K is the contraction dim, one bad group-scale poisons every - # output row, so padded K-groups must carry a valid non-zero fp8 - # scale -- use the smallest-normal e4m3 value (0x08), matching the - # quantizer's own zero-block scale. N-row padding is sliced off in - # ``apply`` and can stay zero. - scale_2d = F.pad(scale_2d, (0, num_groups_pad - num_groups), - value=0x08) - scale_2d = F.pad(scale_2d, (0, 0, 0, size_n_pad - size_n), value=0) - - qweight_int32 = weight.view( - torch.int32).T.contiguous() # [K_pad/4, N_pad] - perm = torch.empty(0, dtype=torch.int32, device=weight.device) - marlin_weight = torch.ops.trtllm.gptq_marlin_repack( - b_q_weight=qweight_int32, - perm=perm, - size_k=size_k_pad, - size_n=size_n_pad, - num_bits=4, - is_a_8bit=False, - ) - - scale_2d = scale_2d.view( - torch.float8_e4m3fn).T.contiguous() # [num_groups_pad, N_pad] - marlin_scale = marlin_utils.marlin_permute_scales(scale_2d.to( - torch.half), - size_k_pad, - size_n_pad, - group_size=group_size) - marlin_scale = marlin_utils.nvfp4_marlin_process_scales(marlin_scale) - - ws2 = module.weight_scale_2.data - if ws2.numel() == 0 or not torch.isfinite(ws2).all() or ws2.item() == 0: - ws2 = torch.tensor([1.0], dtype=torch.float32, device=weight.device) - weight_global_scale = marlin_utils.nvfp4_marlin_process_global_scale( - ws2.to(torch.bfloat16)) - - module.weight = Parameter(marlin_weight, requires_grad=False) - module.weight_scale = Parameter(marlin_scale, requires_grad=False) - module.weight_global_scale = Parameter(weight_global_scale, - requires_grad=False) - # Padded GEMM dims consumed by ``apply``; default to the real sizes. - module._marlin_size_k = size_k_pad - module._marlin_size_n = size_n_pad - - def apply(self, module: Linear, input: torch.Tensor, - bias: Optional[torch.Tensor]): - assert is_nvfp4_marlin_enabled() - size_k = module.in_features - size_n = module.out_features - # Set by transform_weights; equal to size_k/size_n when 64-aligned. - size_k_pad = getattr(module, "_marlin_size_k", size_k) - size_n_pad = getattr(module, "_marlin_size_n", size_n) - - x = input.bfloat16() - if size_k_pad != size_k: - x = F.pad(x, (0, size_k_pad - size_k)) - output = torch.ops.trtllm.marlin_nvfp4_gemm( - x, - module.weight, - scale_a=None, - scale_b=module.weight_scale, - alpha=None, - weight_global_scale=module.weight_global_scale, - bias=None, - out_dtype=module.dtype, - size_n=size_n_pad, - size_k=size_k_pad, - output_buffer_kind=int(BufferKind.DEFAULT), - ) - if size_n_pad != size_n: - output = output[..., :size_n].contiguous() - if bias is not None: - output = output + bias - return output - - def apply_linear_allreduce(self, module: Linear, input: torch.Tensor, - bias: Optional[torch.Tensor], tp_rank: int, - tp_group: List[int]): - raise RuntimeError( - "MarlinNVFP4LinearMethod does not support apply_linear_allreduce") - - def _mxfp8_cutlass_op_available() -> bool: """Cached check for whether the CUTLASS MXFP8xMXFP8 GEMM op is compiled in. @@ -3350,7 +3408,11 @@ def __init__( self.create_weights() def get_quant_method(self, quant_config: Optional[QuantConfig] = None): - return get_quant_method(quant_config) + quant_method = get_quant_method(quant_config) + if (type(quant_method) is W4A16NVFP4LinearMethod + and MarlinNVFP4LinearMethod.is_supported(self)): + return MarlinNVFP4LinearMethod() + return quant_method def create_weights(self): if self._weights_created: @@ -3549,6 +3611,7 @@ def transform_weights(self) -> None: self._weights_transformed = True def cache_derived_state(self) -> None: + self.quant_method.cache_derived_state(self) self._weights_transformed = True def post_load_weights(self) -> None: diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 2305c0852664..ebeecd40542c 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -6270,19 +6270,20 @@ class TestQwen3_6_35B_A3B(LlmapiAccuracyTestHarness): @pytest.mark.parametrize("moe_backend", ["TRTLLM", "CUTEDSL"]) def test_nvfp4(self, moe_backend, mocker): - # Qwen3.6-35B-A3B NVFP4 MoE checkpoint. The TRTLLM-Gen / CuteDSL NVFP4 - # MoE backends only support the SM100 family (B200/B300); RTX 6000 - # (SM120) uses a different MoE path, so restrict this test to SM100/103. - if get_sm_version() not in (100, 103): - pytest.skip("Qwen3.6-35B-A3B NVFP4 MoE test runs on SM100/103 only") + sm_version = get_sm_version() + if moe_backend == "TRTLLM" and sm_version not in (100, 103): + pytest.skip("TRTLLM NVFP4 MoE supports SM100/103 only") + if moe_backend == "CUTEDSL" and sm_version not in (100, 103, 120, 121): + pytest.skip("CuteDSL NVFP4 MoE supports SM100/103/120/121 only") if not os.path.exists(self.MODEL_PATH): pytest.skip(f"Model directory {self.MODEL_PATH} does not exist") + max_batch_size = 32 if sm_version in (120, 121) else 128 kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.8, enable_block_reuse=False) cuda_graph_config = CudaGraphConfig(enable_padding=True, - max_batch_size=128) + max_batch_size=max_batch_size) moe_config = MoeConfig(backend=moe_backend) with LLM(self.MODEL_PATH, @@ -6290,7 +6291,7 @@ def test_nvfp4(self, moe_backend, mocker): tensor_parallel_size=1, moe_expert_parallel_size=1, max_seq_len=4096, - max_batch_size=128, + max_batch_size=max_batch_size, kv_cache_config=kv_cache_config, cuda_graph_config=cuda_graph_config, moe_config=moe_config) as llm: diff --git a/tests/integration/test_lists/test-db/l0_gb10.yml b/tests/integration/test_lists/test-db/l0_gb10.yml index efc6a64530f9..3efa15b293ca 100644 --- a/tests/integration/test_lists/test-db/l0_gb10.yml +++ b/tests/integration/test_lists/test-db/l0_gb10.yml @@ -40,3 +40,4 @@ l0_gb10: - unittest/_torch/modeling -k "modeling_out_of_tree" - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e8_k1_h512_i512-seq=8-dtype=torch.float16-backend=CUTLASS-quant=NVFP4-routing=Renormalize] - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=NVFP4-routing=Renormalize] + - accuracy/test_llm_api_pytorch.py::TestQwen3_6_35B_A3B::test_nvfp4[CUTEDSL] diff --git a/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml b/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml index f7372a8c4224..6b4793e191a1 100644 --- a/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml +++ b/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml @@ -37,6 +37,7 @@ l0_rtx_pro_6000: - accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_bf16 - accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_fp8 - accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_nvfp4 + - accuracy/test_llm_api_pytorch.py::TestQwen3_6_35B_A3B::test_nvfp4[CUTEDSL] - condition: ranges: diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen3_5_quant.py b/tests/unittest/_torch/modeling/test_modeling_qwen3_5_quant.py new file mode 100644 index 000000000000..c3fbc4f4f37b --- /dev/null +++ b/tests/unittest/_torch/modeling/test_modeling_qwen3_5_quant.py @@ -0,0 +1,124 @@ +# 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 types import SimpleNamespace +from unittest.mock import patch + +import torch +from torch import nn + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_qwen3_5 import ( + _normalize_qwen35_exclude_modules, + _normalize_qwen35_quant_config_dict, +) +from tensorrt_llm._torch.models.modeling_qwen3_next import Qwen3NextSparseMoeBlock +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + + +def test_qwen36_normalizes_modelopt_quantized_layer_paths_on_sm121(): + w4a16_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16) + fp8_config = QuantConfig(quant_algo=QuantAlgo.FP8) + model_config = ModelConfig( + pretrained_config=SimpleNamespace(num_hidden_layers=40), + quant_config=QuantConfig( + exclude_modules=[ + "model.language_model.layers.0.linear_attn.in_proj_qkv", + "mtp.layers.0*", + ] + ), + quant_config_dict={ + "model.language_model.layers.0.linear_attn.in_proj_qkv": fp8_config, + "model.language_model.layers.0.linear_attn.in_proj_z": fp8_config, + "model.language_model.layers.0.mlp.experts": w4a16_config, + "model.language_model.layers.0.mlp.shared_expert.gate_proj": w4a16_config, + "model.language_model.layers.0.mlp.shared_expert.up_proj": w4a16_config, + "model.language_model.layers.0.mlp.shared_expert.down_proj": w4a16_config, + "model.visual.patch_embed": fp8_config, + "mtp.layers.0.mlp.experts": w4a16_config, + "lm_head": w4a16_config, + }, + ) + + with patch("tensorrt_llm._torch.models.modeling_qwen3_5.get_sm_version", return_value=121): + _normalize_qwen35_exclude_modules(model_config) + _normalize_qwen35_quant_config_dict(model_config) + + assert model_config.quant_config.exclude_modules == [ + "*linear_attn.conv1d", + "lm_head", + "model.layers.0.linear_attn.in_proj_qkvz*", + "model.layers.40*", + ] + assert set(model_config.quant_config_dict) == { + "model.layers.0.linear_attn.in_proj_qkvz", + "model.layers.0.mlp.experts", + "model.layers.0.mlp.shared_expert.down_proj", + "model.layers.0.mlp.shared_expert.gate_proj", + "model.layers.0.mlp.shared_expert.up_proj", + "model.layers.40.mlp.experts", + } + for name, quant_config in model_config.quant_config_dict.items(): + if name.endswith("in_proj_qkvz"): + assert quant_config.quant_algo == QuantAlgo.FP8 + else: + assert quant_config.quant_algo == QuantAlgo.W4A16_NVFP4 + + +def test_qwen36_sparse_moe_uses_layer_w4a16_quant_config(): + experts_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16) + model_config = ModelConfig( + pretrained_config=SimpleNamespace( + hidden_size=16, + intermediate_size=32, + moe_intermediate_size=8, + num_experts=4, + num_experts_per_tok=1, + shared_expert_intermediate_size=8, + torch_dtype=torch.bfloat16, + model_type="qwen3_5_moe_text", + mlp_bias=False, + ), + moe_backend="CUTEDSL", + quant_config=QuantConfig(), + quant_config_dict={ + "model.layers.0.mlp.experts": experts_config, + }, + ) + captured = {} + + def fake_create_moe(**kwargs): + captured.update(kwargs) + return nn.Identity() + + with ( + patch( + "tensorrt_llm._torch.models.modeling_qwen3_next.create_moe", side_effect=fake_create_moe + ), + patch( + "tensorrt_llm._torch.models.modeling_qwen3_next.AllReduce", + side_effect=lambda **kwargs: nn.Identity(), + ), + patch( + "tensorrt_llm._torch.models.modeling_qwen3_next.GatedMLP", + side_effect=lambda **kwargs: nn.Identity(), + ), + patch("torch.cuda.Event", side_effect=lambda: object()), + ): + Qwen3NextSparseMoeBlock(model_config=model_config, aux_stream=None, layer_idx=0) + + assert captured["model_config"] is model_config + assert captured["override_quant_config"] is experts_config + assert captured["layer_idx"] == 0 diff --git a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py index f8732b4875c9..0a700e415caa 100644 --- a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py +++ b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os from types import SimpleNamespace from unittest.mock import patch @@ -23,6 +22,8 @@ import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils from tensorrt_llm._torch.modules.embedding import LMHead from tensorrt_llm._torch.modules.linear import ( + Linear, + MarlinNVFP4LinearMethod, W4A16NVFP4LinearMethod, get_quant_method, get_sm_version, @@ -30,35 +31,64 @@ from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig -def _run_w4a16_cutlass3_reference_case(m: int, n: int, k: int, dtype: torch.dtype) -> None: - assert k % 32 == 0 - assert n % 32 == 0 - act, weight, weight_scale, weight_scale_2 = _make_w4a16_nvfp4_case(m, n, k, dtype) - - expected = torch.empty((m, n), device="cuda", dtype=dtype) - for start in range(0, m, 16): - stop = min(start + 16, m) - expected[start:stop, :] = torch.ops.trtllm.w4a16_nvfp4_gemm( - act[start:stop, :], - weight, - weight_scale, - weight_scale_2, - dtype, - bias=None, - ) +def _run_w4a16_marlin_reference_case(m: int, n: int, k: int) -> None: + act, weight, weight_scale, weight_scale_2 = _make_w4a16_nvfp4_case(m, n, k, torch.bfloat16) + expected = torch.ops.trtllm.w4a16_nvfp4_gemm( + act, + weight, + weight_scale, + weight_scale_2, + torch.bfloat16, + bias=None, + ) - actual = torch.ops.trtllm.w4a16_nvfp4_cutlass_gemm( + linear = Linear( + k, + n, + bias=False, + dtype=torch.bfloat16, + quant_config=QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4), + reduce_output=False, + ).cuda() + assert isinstance(linear.quant_method, MarlinNVFP4LinearMethod) + linear.weight.data.copy_(weight) + linear.weight_scale.data.copy_(weight_scale) + linear.weight_scale_2.data.copy_(weight_scale_2) + linear.transform_weights() + actual = linear(act) + torch.testing.assert_close(actual, expected, atol=0.75, rtol=0.02) + + +def _run_w4a16_triton_reference_case(m: int, n: int, k: int) -> None: + act, weight, weight_scale, weight_scale_2 = _make_w4a16_nvfp4_case(m, n, k, torch.float16) + expected = torch.ops.trtllm.w4a16_nvfp4_gemm( act, weight, weight_scale, weight_scale_2, - dtype, + torch.float16, bias=None, ) + + linear = Linear( + k, + n, + bias=False, + dtype=torch.float16, + quant_config=QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4), + reduce_output=False, + ).cuda() + assert type(linear.quant_method) is W4A16NVFP4LinearMethod + linear.weight.data.copy_(weight) + linear.weight_scale.data.copy_(weight_scale) + linear.weight_scale_2.data.copy_(weight_scale_2) + linear.transform_weights() + actual = linear(act) torch.testing.assert_close(actual, expected, atol=0.08, rtol=0.08) def _make_w4a16_nvfp4_case(m: int, n: int, k: int, dtype: torch.dtype): + torch.manual_seed(m + n + k) act = torch.randn((m, k), device="cuda", dtype=dtype) weight = torch.empty((n, k // 2), device="cuda", dtype=fp4_utils.float4_e2m1x2) weight_u8 = torch.randint( @@ -72,9 +102,10 @@ def _make_w4a16_nvfp4_case(m: int, n: int, k: int, dtype: torch.dtype): scale_cols = fp4_utils.pad_up(k // 16, 4) scale_rows = fp4_utils.pad_up(n, 128) + # E4M3 bit patterns in [0x30, 0x40] represent scales from 0.5 to 2.0. weight_scale_linear = torch.randint( - 1, - 120, + 0x30, + 0x41, (scale_rows, scale_cols), device="cuda", dtype=torch.uint8, @@ -92,11 +123,25 @@ def _make_w4a16_nvfp4_case(m: int, n: int, k: int, dtype: torch.dtype): ) @pytest.mark.parametrize( "shape", - [(32, 256, 256), (128, 512, 1024), (256, 1024, 2048)], + [ + (1, 256, 256), + (4, 160, 288), + (4, 256, 256), + (32, 256, 256), + (128, 512, 1024), + ], ) -def test_w4a16_nvfp4_cutlass3_bf16_matches_cuda_core(shape): +def test_w4a16_nvfp4_marlin_bf16_matches_cuda_core(shape): m, n, k = shape - _run_w4a16_cutlass3_reference_case(m, n, k, torch.bfloat16) + _run_w4a16_marlin_reference_case(m, n, k) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or get_sm_version() not in (120, 121), + reason="requires CUDA SM120/121", +) +def test_w4a16_nvfp4_triton_fallback_matches_cuda_core(): + _run_w4a16_triton_reference_case(32, 64, 64) @pytest.mark.skipif( @@ -135,19 +180,19 @@ def test_get_quant_method_returns_w4a16_nvfp4_linear_method(): method = get_quant_method(quant_config) - assert isinstance(method, W4A16NVFP4LinearMethod) + assert type(method) is W4A16NVFP4LinearMethod def test_w4a16_nvfp4_linear_uses_high_precision_activation_without_fp4_quantize(): method = W4A16NVFP4LinearMethod() - input_tensor = torch.ones((2, 4), dtype=torch.bfloat16) - bias = torch.tensor([1.0, 2.0, 3.0], dtype=torch.bfloat16) + input_tensor = torch.ones((2, 32), dtype=torch.bfloat16) + bias = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.bfloat16) module = SimpleNamespace( - weight=torch.empty((5, 2), dtype=torch.uint8), + weight=torch.empty((4, 16), dtype=torch.uint8), weight_scale=torch.empty((128 * 4,), dtype=torch.uint8), weight_scale_2=torch.tensor([0.25], dtype=torch.float32), dtype=torch.bfloat16, - out_features=3, + out_features=4, pre_quant_scale=None, ) captured = {} @@ -166,9 +211,12 @@ def fake_w4a16_nvfp4_gemm( def fail_fp4_quantize(*args, **kwargs): raise AssertionError("W4A16 NVFP4 must not quantize activations") - with patch("torch.ops.trtllm.w4a16_nvfp4_gemm", side_effect=fake_w4a16_nvfp4_gemm, create=True): - with patch("torch.ops.trtllm.fp4_quantize", side_effect=fail_fp4_quantize, create=True): - output = method.apply(module, input_tensor, bias) + with ( + patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=120), + patch("torch.ops.trtllm.w4a16_nvfp4_gemm", side_effect=fake_w4a16_nvfp4_gemm, create=True), + patch("torch.ops.trtllm.fp4_quantize", side_effect=fail_fp4_quantize, create=True), + ): + output = method.apply(module, input_tensor, bias) assert captured["input"] is input_tensor assert captured["weight"] is module.weight @@ -176,161 +224,151 @@ def fail_fp4_quantize(*args, **kwargs): assert captured["weight_scale_2"] is module.weight_scale_2 assert captured["out_dtype"] is torch.bfloat16 assert captured["bias"] is None - expected = torch.tensor([[2.0, 3.0, 4.0], [2.0, 3.0, 4.0]], dtype=torch.bfloat16) + expected = torch.tensor([[2.0, 3.0, 4.0, 5.0], [2.0, 3.0, 4.0, 5.0]], dtype=torch.bfloat16) torch.testing.assert_close(output, expected) def test_w4a16_nvfp4_linear_restores_high_rank_input_shape(): method = W4A16NVFP4LinearMethod() - input_tensor = torch.ones((2, 3, 4), dtype=torch.float16) + input_tensor = torch.ones((2, 3, 32), dtype=torch.float16) module = SimpleNamespace( - weight=torch.empty((7, 2), dtype=torch.uint8), + weight=torch.empty((8, 16), dtype=torch.uint8), weight_scale=torch.empty((128 * 4,), dtype=torch.uint8), weight_scale_2=torch.tensor([0.5], dtype=torch.float32), dtype=torch.float16, - out_features=5, + out_features=8, pre_quant_scale=None, ) def fake_w4a16_nvfp4_gemm( input_arg, weight, weight_scale, weight_scale_2, out_dtype, bias=None ): - assert input_arg.shape == (6, 4) + assert input_arg.shape == (6, 32) return torch.ones((input_arg.shape[0], weight.shape[0]), dtype=out_dtype) - with patch("torch.ops.trtllm.w4a16_nvfp4_gemm", side_effect=fake_w4a16_nvfp4_gemm, create=True): + with ( + patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=120), + patch("torch.ops.trtllm.w4a16_nvfp4_gemm", side_effect=fake_w4a16_nvfp4_gemm, create=True), + ): output = method.apply(module, input_tensor, bias=None) - assert output.shape == (2, 3, 5) + assert output.shape == (2, 3, 8) -def test_w4a16_nvfp4_linear_uses_chunked_w4a16_op_for_large_m(): +def test_w4a16_nvfp4_linear_uses_triton_dequant_for_large_m(): method = W4A16NVFP4LinearMethod() - input_tensor = torch.ones((17, 16), dtype=torch.bfloat16) + input_tensor = torch.ones((17, 32), dtype=torch.bfloat16) + bias = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.bfloat16) module = SimpleNamespace( - weight=torch.empty((5, 8), dtype=torch.uint8), + weight=torch.empty((4, 16), dtype=torch.uint8), weight_scale=torch.empty((128 * 4,), dtype=torch.uint8), + _w4a16_weight_scale_linear=torch.empty((128 * 4,), dtype=torch.uint8), weight_scale_2=torch.tensor([0.5], dtype=torch.float32), dtype=torch.bfloat16, - out_features=3, + out_features=4, + scaling_vector_size=16, pre_quant_scale=None, + use_custom_cublas_mm=False, ) captured = {} - def fake_w4a16_nvfp4_gemm( - input_arg, weight, weight_scale, weight_scale_2, out_dtype, bias=None - ): - captured["input"] = input_arg + def fake_dequant(weight, weight_scale, weight_scale_2, **kwargs): captured["weight"] = weight captured["weight_scale"] = weight_scale captured["weight_scale_2"] = weight_scale_2 - captured["out_dtype"] = out_dtype - captured["bias"] = bias - return torch.ones((input_arg.shape[0], weight.shape[0]), dtype=out_dtype) - - def fail_fp4_quantize(*args, **kwargs): - raise AssertionError("large-M W4A16 path must not quantize activations") + captured.update(kwargs) + return torch.ones((4, 32), dtype=torch.bfloat16) - with patch("torch.ops.trtllm.w4a16_nvfp4_gemm", side_effect=fake_w4a16_nvfp4_gemm, create=True): - with patch("torch.ops.trtllm.fp4_quantize", side_effect=fail_fp4_quantize, create=True): - output = method.apply(module, input_tensor, bias=None) + def fail_w4a16_gemm(*args, **kwargs): + raise AssertionError("large-M W4A16 must use Triton dequantization") + + with ( + patch("torch.ops.trtllm.w4a16_nvfp4_gemm", side_effect=fail_w4a16_gemm, create=True), + patch( + "tensorrt_llm._torch.modules.fused_moe.triton_dequant_nvfp4.dequant_nvfp4_2d_triton", + side_effect=fake_dequant, + ), + ): + output = method.apply(module, input_tensor, bias=bias) - assert captured["input"] is input_tensor - assert captured["weight"] is module.weight - assert captured["weight_scale"] is module.weight_scale + assert captured["weight"].data_ptr() == module.weight.data_ptr() + assert captured["weight_scale"] is module._w4a16_weight_scale_linear assert captured["weight_scale_2"] is module.weight_scale_2 - assert captured["out_dtype"] is torch.bfloat16 - assert output.shape == (17, 3) + assert captured["target_dtype"] is torch.bfloat16 + assert captured["sf_vec_size"] == 16 + expected = torch.tensor([33.0, 34.0, 35.0, 36.0], dtype=torch.bfloat16).expand(17, 4) + torch.testing.assert_close(output, expected) -def test_w4a16_nvfp4_linear_uses_cutlass3_op_for_large_bf16_m_when_enabled(): - method = W4A16NVFP4LinearMethod() - input_tensor = torch.ones((17, 32), dtype=torch.bfloat16) +def test_w4a16_nvfp4_linear_uses_marlin_op_after_weight_transform(): + method = MarlinNVFP4LinearMethod() + input_tensor = torch.ones((1, 32), dtype=torch.bfloat16) module = SimpleNamespace( - weight=torch.empty((32, 16), dtype=torch.uint8), - weight_scale=torch.empty((128 * 4,), dtype=torch.uint8), - weight_scale_2=torch.tensor([0.5], dtype=torch.float32), + weight=torch.empty((8, 16), dtype=torch.int32), + weight_scale=torch.empty((2, 128), dtype=torch.float8_e4m3fn), + weight_global_scale=torch.tensor([0.5], dtype=torch.bfloat16), dtype=torch.bfloat16, + in_features=32, out_features=3, pre_quant_scale=None, + _marlin_size_k=32, + _marlin_size_n=32, ) captured = {} - def fake_w4a16_nvfp4_cutlass_gemm( - input_arg, weight, weight_scale, weight_scale_2, out_dtype, bias=None - ): + def fake_marlin_nvfp4_gemm(input_arg, weight, **kwargs): captured["input"] = input_arg captured["weight"] = weight - captured["weight_scale"] = weight_scale - captured["weight_scale_2"] = weight_scale_2 - captured["out_dtype"] = out_dtype - captured["bias"] = bias - return torch.ones((input_arg.shape[0], weight.shape[0]), dtype=out_dtype) + captured.update(kwargs) + return torch.ones((input_arg.shape[0], kwargs["size_n"]), dtype=kwargs["out_dtype"]) def fail_w4a16_gemm(*args, **kwargs): - raise AssertionError("CUTLASS3 W4A16 prefill must not call the default W4A16 op") - - def fail_fp4_quantize(*args, **kwargs): - raise AssertionError("CUTLASS3 W4A16 prefill must not quantize activations") - - with patch.dict(os.environ, {"TRTLLM_W4A16_NVFP4_CUTLASS3": "1"}): - with patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=120): - with patch( - "torch.ops.trtllm.w4a16_nvfp4_gemm", side_effect=fail_w4a16_gemm, create=True - ): - with patch( - "torch.ops.trtllm.w4a16_nvfp4_cutlass_gemm", - side_effect=fake_w4a16_nvfp4_cutlass_gemm, - create=True, - ): - with patch( - "torch.ops.trtllm.fp4_quantize", side_effect=fail_fp4_quantize, create=True - ): - output = method.apply(module, input_tensor, bias=None) + raise AssertionError("Marlin W4A16 must not call the default W4A16 op") + + with patch("torch.ops.trtllm.w4a16_nvfp4_gemm", side_effect=fail_w4a16_gemm, create=True): + with patch( + "torch.ops.trtllm.marlin_nvfp4_gemm", + side_effect=fake_marlin_nvfp4_gemm, + create=True, + ): + output = method.apply(module, input_tensor, bias=None) assert captured["input"] is input_tensor assert captured["weight"] is module.weight - assert captured["weight_scale"] is module.weight_scale - assert captured["weight_scale_2"] is module.weight_scale_2 assert captured["out_dtype"] is torch.bfloat16 assert captured["bias"] is None - assert output.shape == (17, 3) + assert captured["scale_b"] is module.weight_scale + assert captured["weight_global_scale"] is module.weight_global_scale + assert output.shape == (1, 3) -def test_w4a16_nvfp4_linear_cutlass3_restores_high_rank_input_shape(): - method = W4A16NVFP4LinearMethod() +def test_w4a16_nvfp4_linear_marlin_restores_high_rank_input_shape(): + method = MarlinNVFP4LinearMethod() input_tensor = torch.ones((2, 9, 32), dtype=torch.bfloat16) bias = torch.tensor([1.0, 2.0, 3.0], dtype=torch.bfloat16) module = SimpleNamespace( - weight=torch.empty((32, 16), dtype=torch.uint8), - weight_scale=torch.empty((128 * 4,), dtype=torch.uint8), - weight_scale_2=torch.tensor([0.5], dtype=torch.float32), + weight=torch.empty((8, 16), dtype=torch.int32), + weight_scale=torch.empty((2, 128), dtype=torch.float8_e4m3fn), + weight_global_scale=torch.tensor([0.5], dtype=torch.bfloat16), dtype=torch.bfloat16, + in_features=32, out_features=3, pre_quant_scale=None, + _marlin_size_k=32, + _marlin_size_n=32, ) captured = {} - def fake_w4a16_nvfp4_cutlass_gemm( - input_arg, weight, weight_scale, weight_scale_2, out_dtype, bias=None - ): + def fake_marlin_nvfp4_gemm(input_arg, weight, **kwargs): captured["input_shape"] = input_arg.shape - return torch.ones((input_arg.shape[0], weight.shape[0]), dtype=out_dtype) + return torch.ones((input_arg.shape[0], kwargs["size_n"]), dtype=kwargs["out_dtype"]) - def fail_w4a16_gemm(*args, **kwargs): - raise AssertionError("large-M CUTLASS3 path must not call the default W4A16 op") - - with patch.dict(os.environ, {"TRTLLM_W4A16_NVFP4_CUTLASS3": "1"}): - with patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=120): - with patch( - "torch.ops.trtllm.w4a16_nvfp4_gemm", side_effect=fail_w4a16_gemm, create=True - ): - with patch( - "torch.ops.trtllm.w4a16_nvfp4_cutlass_gemm", - side_effect=fake_w4a16_nvfp4_cutlass_gemm, - create=True, - ): - output = method.apply(module, input_tensor, bias=bias) + with patch( + "torch.ops.trtllm.marlin_nvfp4_gemm", + side_effect=fake_marlin_nvfp4_gemm, + create=True, + ): + output = method.apply(module, input_tensor, bias=bias) assert captured["input_shape"] == (18, 32) assert output.shape == (2, 9, 3) @@ -338,71 +376,93 @@ def fail_w4a16_gemm(*args, **kwargs): torch.testing.assert_close(output, expected) -def test_w4a16_nvfp4_cutlass3_prefill_requires_supported_shape(): - method = W4A16NVFP4LinearMethod() - input_tensor = torch.ones((17, 32), dtype=torch.bfloat16) +def test_w4a16_nvfp4_linear_marlin_applies_pre_quant_scale_once(): + method = MarlinNVFP4LinearMethod() + input_tensor = torch.ones((1, 32), dtype=torch.bfloat16) module = SimpleNamespace( - weight=torch.empty((6, 16), dtype=torch.uint8), + weight=torch.empty((8, 16), dtype=torch.int32), + weight_scale=torch.empty((2, 128), dtype=torch.float8_e4m3fn), + weight_global_scale=torch.tensor([0.5], dtype=torch.bfloat16), dtype=torch.bfloat16, + in_features=32, + out_features=3, + pre_quant_scale=torch.full((32,), 2.0, dtype=torch.bfloat16), + _marlin_size_k=32, + _marlin_size_n=32, ) + captured = {} - with patch.dict(os.environ, {"TRTLLM_W4A16_NVFP4_CUTLASS3": "1"}): - with patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=120): - assert not method._can_use_cutlass3_w4a16_prefill(module, input_tensor, m=17) + def fake_marlin_nvfp4_gemm(input_arg, weight, **kwargs): + captured["input"] = input_arg + return torch.ones((input_arg.shape[0], kwargs["size_n"]), dtype=kwargs["out_dtype"]) + with patch( + "torch.ops.trtllm.marlin_nvfp4_gemm", + side_effect=fake_marlin_nvfp4_gemm, + create=True, + ): + method.apply(module, input_tensor, bias=None) -def test_w4a16_nvfp4_linear_cutlass3_unsupported_shape_uses_default_w4a16_op(): - method = W4A16NVFP4LinearMethod() - input_tensor = torch.ones((17, 32), dtype=torch.bfloat16) - module = SimpleNamespace( - weight=torch.empty((6, 16), dtype=torch.uint8), - weight_scale=torch.empty((128 * 4,), dtype=torch.uint8), - weight_scale_2=torch.tensor([0.5], dtype=torch.float32), - dtype=torch.bfloat16, - out_features=3, - pre_quant_scale=None, - ) - captured = {} + torch.testing.assert_close(captured["input"], input_tensor * module.pre_quant_scale) - def fake_w4a16_nvfp4_gemm( - input_arg, weight, weight_scale, weight_scale_2, out_dtype, bias=None + +@pytest.mark.parametrize( + ("dtype", "use_fused_gemm_allreduce"), + [ + (torch.float16, False), + (torch.bfloat16, True), + ], +) +def test_w4a16_nvfp4_marlin_selection_requires_supported_module(dtype, use_fused_gemm_allreduce): + module = SimpleNamespace(dtype=dtype, use_fused_gemm_allreduce=use_fused_gemm_allreduce) + + with ( + patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=120), + patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), + patch("torch.ops.trtllm.gptq_marlin_repack", create=True), ): - captured["input"] = input_arg - captured["weight"] = weight - captured["weight_scale"] = weight_scale - captured["weight_scale_2"] = weight_scale_2 - captured["out_dtype"] = out_dtype - captured["bias"] = bias - return torch.ones((input_arg.shape[0], weight.shape[0]), dtype=out_dtype) + assert not MarlinNVFP4LinearMethod.is_supported(module) - def fail_cutlass_gemm(*args, **kwargs): - raise AssertionError("unsupported CUTLASS3 shape must use the default W4A16 op") - def fail_fp4_quantize(*args, **kwargs): - raise AssertionError("unsupported CUTLASS3 W4A16 path must not quantize activations") - - with patch.dict(os.environ, {"TRTLLM_W4A16_NVFP4_CUTLASS3": "1"}): - with patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=120): - with patch( - "torch.ops.trtllm.w4a16_nvfp4_gemm", side_effect=fake_w4a16_nvfp4_gemm, create=True - ): - with patch( - "torch.ops.trtllm.w4a16_nvfp4_cutlass_gemm", - side_effect=fail_cutlass_gemm, - create=True, - ): - with patch( - "torch.ops.trtllm.fp4_quantize", side_effect=fail_fp4_quantize, create=True - ): - output = method.apply(module, input_tensor, bias=None) +def test_w4a16_nvfp4_linear_selects_marlin_for_supported_module(): + quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) - assert captured["input"] is input_tensor - assert captured["bias"] is None - assert captured["weight"] is module.weight - assert captured["weight_scale"] is module.weight_scale - assert captured["weight_scale_2"] is module.weight_scale_2 - assert captured["out_dtype"] is torch.bfloat16 - assert output.shape == (17, 3) + with ( + patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=120), + patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), + patch("torch.ops.trtllm.gptq_marlin_repack", create=True), + ): + linear = Linear( + 32, + 32, + bias=False, + dtype=torch.bfloat16, + quant_config=quant_config, + reduce_output=False, + ) + + assert isinstance(linear.quant_method, MarlinNVFP4LinearMethod) + assert isinstance(linear.quant_method, W4A16NVFP4LinearMethod) + + +def test_w4a16_nvfp4_linear_keeps_default_method_for_fp16(): + quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) + + with ( + patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=120), + patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), + patch("torch.ops.trtllm.gptq_marlin_repack", create=True), + ): + linear = Linear( + 32, + 32, + bias=False, + dtype=torch.float16, + quant_config=quant_config, + reduce_output=False, + ) + + assert type(linear.quant_method) is W4A16NVFP4LinearMethod def test_w4a16_nvfp4_post_load_preserves_checkpoint_weight_global_scale(): @@ -443,9 +503,12 @@ def test_lm_head_uses_w4a16_nvfp4_quant_method_for_packed_lm_head(): def test_lm_head_w4a16_nvfp4_forward_dispatches_to_w4a16_op(): quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) lm_head = LMHead( - num_embeddings=3, embedding_dim=16, dtype=torch.float16, quant_config=quant_config + num_embeddings=4, + embedding_dim=32, + dtype=torch.float16, + quant_config=quant_config, ) - input_tensor = torch.ones((2, 16), dtype=torch.float16) + input_tensor = torch.ones((2, 32), dtype=torch.float16) captured = {} def fake_w4a16_nvfp4_gemm( @@ -459,7 +522,10 @@ def fake_w4a16_nvfp4_gemm( captured["bias"] = bias return torch.ones((input_arg.shape[0], weight.shape[0]), dtype=out_dtype) - with patch("torch.ops.trtllm.w4a16_nvfp4_gemm", side_effect=fake_w4a16_nvfp4_gemm, create=True): + with ( + patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=120), + patch("torch.ops.trtllm.w4a16_nvfp4_gemm", side_effect=fake_w4a16_nvfp4_gemm, create=True), + ): output = lm_head(input_tensor) assert captured["input"] is input_tensor @@ -468,4 +534,4 @@ def fake_w4a16_nvfp4_gemm( assert captured["weight_scale_2"] is lm_head.weight_scale_2 assert captured["out_dtype"] is torch.float16 assert captured["bias"] is None - assert output.shape == (2, 3) + assert output.shape == (2, 4) diff --git a/tests/unittest/_torch/test_model_config.py b/tests/unittest/_torch/test_model_config.py index 486c150f2040..ff04936063b8 100644 --- a/tests/unittest/_torch/test_model_config.py +++ b/tests/unittest/_torch/test_model_config.py @@ -42,6 +42,61 @@ def make_pretrained_config( ) +def test_qwen36_modelopt_mixed_precision_preserves_w4a16_layers(tmp_path): + quant_config_file = tmp_path / "hf_quant_config.json" + quant_config_file.write_text( + json.dumps( + { + "quantization": { + "quant_algo": "MIXED_PRECISION", + "kv_cache_quant_algo": "FP8", + "exclude_modules": ["mtp.layers.0*"], + "quantized_layers": { + "model.language_model.layers.0.linear_attn.out_proj": { + "quant_algo": "FP8", + }, + "model.language_model.layers.0.mlp.experts": { + "quant_algo": "W4A16_NVFP4", + "group_size": 16, + }, + }, + }, + } + ), + encoding="utf-8", + ) + + quant_config, layer_quant_config = ModelConfig.load_modelopt_quant_config( + str(quant_config_file), str(tmp_path), moe_backend="CUTLASS" + ) + + assert quant_config.quant_algo is None + assert quant_config.group_size == 16 + assert quant_config.kv_cache_quant_algo == QuantAlgo.FP8 + assert layer_quant_config is not None + experts_config = layer_quant_config["model.language_model.layers.0.mlp.experts"] + assert experts_config.quant_algo == QuantAlgo.W4A16_NVFP4 + assert experts_config.group_size == 16 + + +def test_qwen36_auto_moe_backend_uses_layer_w4a16_on_sm121(monkeypatch): + layer_quant_config = { + "model.language_model.layers.0.mlp.experts": QuantConfig( + quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16 + ), + } + monkeypatch.setattr("tensorrt_llm._torch.model_config.get_sm_version", lambda: 121) + + moe_backend = ModelConfig.resolve_moe_backend( + "AUTO", + "Qwen3_5MoeForConditionalGeneration", + quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8), + layer_quant_config=layer_quant_config, + ) + + assert moe_backend == "CUTEDSL" + + @pytest.mark.parametrize( "num_key_value_heads", [ diff --git a/tests/unittest/_torch/thop/serial/test_nvfp4_marlin_gemm.py b/tests/unittest/_torch/thop/serial/test_nvfp4_marlin_gemm.py index 765fe516bd19..0deace4bc7f8 100644 --- a/tests/unittest/_torch/thop/serial/test_nvfp4_marlin_gemm.py +++ b/tests/unittest/_torch/thop/serial/test_nvfp4_marlin_gemm.py @@ -24,7 +24,7 @@ import torch from parameterized import parameterized -from utils.util import skip_non_hopper_unittest, unittest_name_func +from utils.util import getSMVersion, unittest_name_func import tensorrt_llm # noqa: F401 # registers torch.ops.tensorrt_llm / torch.ops.trtllm ops @@ -65,7 +65,10 @@ class TestNvfp4MarlinGemm(unittest.TestCase): list(product([1024, 2048], [1024, 2048], [1, 8, 128], [16], [1.0, 2.0], ["nvfp4", "bf16"])), name_func=unittest_name_func, ) - @skip_non_hopper_unittest + @unittest.skipUnless( + getSMVersion() == 90 or 120 <= getSMVersion() < 130, + "Marlin NVFP4 GEMM requires SM90 or SM12x", + ) def test_nvfp4_marlin_gemm( self, input_dim, output_dim, batch_size, sf_vec_size, alpha, act_dtype ): From 12f12065b4965257d4b54d919a4c87bf502e9650 Mon Sep 17 00:00:00 2001 From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> Date: Tue, 21 Jul 2026 04:44:17 +0000 Subject: [PATCH 12/25] fix: correct W4A16 checkpoint handling on SM12x Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> --- .../hf/nemotron_h_weight_mapper.py | 79 +++++++++++-------- .../_torch/models/modeling_qwen3_5.py | 28 ++++--- .../_torch/modules/fused_moe/create_moe.py | 7 ++ tensorrt_llm/_torch/modules/mlp.py | 9 +-- .../modeling/test_modeling_qwen3_5_quant.py | 15 +++- .../hf/test_nemotron_h_weight_mapper.py | 76 ++++++++++++------ .../hf/test_qwen3_5_weight_mapper.py | 17 ++++ .../moe/test_cute_dsl_b12x_moe_backend.py | 25 ++++++ .../_torch/modules/test_w4a16_nvfp4_linear.py | 23 ++++++ 9 files changed, 197 insertions(+), 82 deletions(-) diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py index ea009727e1ce..77fa808267c9 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py @@ -50,20 +50,27 @@ def _split_mamba2_mixer_in_proj(w: torch.Tensor) -> torch.Tensor: d_state = config.ssm_state_size nheads = config.mamba_num_heads - def _canonicalize_quant_key(key: str) -> str: - replacements = { - ".weight_packed": ".weight", - ".weight_global_scale": ".weight_scale_2", - ".input_global_scale": ".input_scale_2", - } - for suffix, replacement in replacements.items(): - if key.endswith(suffix): - return f"{key[:-len(suffix)]}{replacement}" - return key + def _invert_compressed_tensors_scale(value) -> torch.Tensor: + value = value[...] if not isinstance(value, torch.Tensor) else value + value = value.to(torch.float32) + return torch.where(value > 0, value.reciprocal(), + torch.zeros_like(value)).contiguous() + + def _canonicalize_quant_weight(key: str, value): + if key.endswith(".weight_packed"): + return f"{key[:-len('.weight_packed')]}.weight", value + if key.endswith(".weight_global_scale"): + key = f"{key[:-len('.weight_global_scale')]}.weight_scale_2" + return key, _invert_compressed_tensors_scale(value) + if key.endswith(".input_global_scale"): + key = f"{key[:-len('.input_global_scale')]}.input_scale" + return key, _invert_compressed_tensors_scale(value) + return key, value new_weights = {} for name, _ in weights.items(): key = name + value = weights[name] # change backbone root name to model if "backbone" in key: @@ -85,29 +92,28 @@ def _canonicalize_quant_key(key: str) -> str: if "A_log" in key: key = key.replace("A_log", "A") - key = _canonicalize_quant_key(key) + key, value = _canonicalize_quant_weight(key, value) if ("mixer.in_proj" in key or "mixer.out_proj" in key) and "_scale" in key: # Special handing for nvfp4 Mamba2 mixer in_proj.weight_scale. if is_nvfp4 and "in_proj.weight_scale_2" not in key and "in_proj.weight_scale" in key: - new_weights[key] = _split_mamba2_mixer_in_proj( - weights[name]) + new_weights[key] = _split_mamba2_mixer_in_proj(value) else: - new_weights[key] = weights[name] + new_weights[key] = value elif "A" in key: - w = split(weights[name], tp_size, tp_rank) + w = split(value, tp_size, tp_rank) w = w.to(torch.float32) # Avoid extra temporaries: one fp32 cast, then in-place exp/neg. w.exp_() w.neg_() new_weights[key] = w elif "D" in key: - w = split(weights[name], tp_size, tp_rank) + w = split(value, tp_size, tp_rank) w = w.to(torch.float32) new_weights[key] = w elif "dt_bias" in key: - w = split(weights[name], tp_size, tp_rank) + w = split(value, tp_size, tp_rank) w = w.to(torch.float32) new_weights[key] = w elif "mixer.in_proj" in key: @@ -116,9 +122,9 @@ def _canonicalize_quant_key(key: str) -> str: # ``weight_scale``, ``weight_scale_2``, …) under ``mixer.in_proj.*`` # — those are scalars / 1-D scales and must not go through the # Mamba2 split rearrangement. - new_weights[key] = _split_mamba2_mixer_in_proj(weights[name]) + new_weights[key] = _split_mamba2_mixer_in_proj(value) elif "conv1d" in key: - w = weights[name] + w = value # removing dim(1) because we are using Linear to store conv1d weights if "weight" in key: w = w.squeeze(1) @@ -136,12 +142,12 @@ def _canonicalize_quant_key(key: str) -> str: w = torch.concat(w).contiguous() new_weights[key] = w elif "mixer.norm.weight" in key: - w = split(weights[name], tp_size, tp_rank) + w = split(value, tp_size, tp_rank) new_weights[key] = w # Remap MoE expert weights. elif "mixer.experts." in key: if self.config.moe_backend == 'VANILLA': - new_weights[key] = weights[name] + new_weights[key] = value else: # HF transformers 5.x exposes routed MoE experts as fused # tensors stacked along dim 0 ([num_experts, ...]) under keys @@ -150,7 +156,7 @@ def _canonicalize_quant_key(key: str) -> str: # contrast, stores per-expert keys (``experts.{i}.up_proj``). # The VANILLA FusedMoE loader expects per-expert keys, so # unfuse the 3D HF format here before the standard rename. - val = weights[name] + val = value m = re.match(r"(.*\.mixer\.experts)\.(up_proj|down_proj)$", key) is_hf_fused = (m is not None @@ -173,34 +179,37 @@ def _canonicalize_quant_key(key: str) -> str: elif "up_proj" in key: w1_key = key.replace("up_proj", "w1") w3_key = key.replace("up_proj", "w3") - # Don't need to handle with input_scale and weight_scale_2 since they are scalar for fp8 and nvfp4 models. - if "input_scale" in key or "weight_scale_2" in key or "input_quantizer" in key or "weight_quantizer" in key: - new_weights[w3_key] = weights[name] - new_weights[w1_key] = weights[name] + # Per-tensor quantization parameters are shared by w1 + # and the empty w3 projection. + if ("input_scale" in key or "weight_scale_2" in key + or "input_quantizer" in key + or "weight_quantizer" in key): + new_weights[w3_key] = value + new_weights[w1_key] = value elif "weight_scale" in key: # NVFP4 case. - if weights[name].shape: + if value.shape: # w3 weight (gate_proj) scale should be empty for Nemotron-H MoE model. # Use [:0] to keep the same input dimension as the other weights. # The w3 weight_scale shape should be [0, input_dim]. - new_weights[w3_key] = weights[name][:0] - new_weights[w1_key] = weights[name] + new_weights[w3_key] = value[:0] + new_weights[w1_key] = value # FP8 case. else: - new_weights[w3_key] = weights[name] - new_weights[w1_key] = weights[name] + new_weights[w3_key] = value + new_weights[w1_key] = value else: # w3 weight (gate_proj) should be empty for Nemotron-H MoE model. # Use [:0] to keep the same input dimension as the other weights. # The w3 weight shape should be [0, input_dim]. - new_weights[w3_key] = weights[name][:0] - new_weights[w1_key] = weights[name] + new_weights[w3_key] = value[:0] + new_weights[w1_key] = value elif "down_proj" in key: key = key.replace("down_proj", "w2") - new_weights[key] = weights[name] + new_weights[key] = value else: raise ValueError(f"Unknown MoE weight: {key}") else: - new_weights[key] = weights[name] + new_weights[key] = value return new_weights diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_5.py b/tensorrt_llm/_torch/models/modeling_qwen3_5.py index 6fc43129a353..f1556e61b3db 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_5.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_5.py @@ -391,12 +391,11 @@ def _lm_head_nvfp4_enabled(model_config): """Whether the checkpoint's quantized lm_head should stay quantized. ModelOpt MIXED_PRECISION exports for Qwen3.5/3.6 quantize lm_head to - W4A16_NVFP4 (packed FP4 weight + per-group FP8 scales). On SM100/103 the - NVFP4 (W4A4) Linear path can consume it directly, cutting the lm_head - GEMM's weight traffic 4x vs the bf16 dequant fallback -- the decode - lm_head is purely weight-bandwidth-bound. Conditions mirror what the - quantized LMHead supports (see LMHead.__init__ guards) plus the paths - that bypass the Linear machinery entirely: + W4A16_NVFP4 (packed FP4 weight + per-group FP8 scales). On SM100/103 the + NVFP4 (W4A4) Linear path consumes it; on SM120/121 the W4A16 Linear path + consumes it directly. Both avoid the bf16 dequant fallback and cut lm_head + weight traffic 4x. Conditions mirror what the quantized LMHead supports + (see LMHead.__init__ guards) plus paths that bypass Linear entirely: - tie_word_embeddings shares the weight with the embedding lookup, which needs a dense bf16 weight; @@ -415,7 +414,7 @@ def _lm_head_nvfp4_enabled(model_config): return ( cfg is not None and cfg.quant_algo == QuantAlgo.W4A16_NVFP4 - and get_sm_version() in (100, 103) + and get_sm_version() in (100, 103, 120, 121) and not getattr(pretrained, "tie_word_embeddings", False) and not mapping.enable_attention_dp and getattr(pretrained, "vocab_size", 0) % mapping.tp_size == 0 @@ -507,10 +506,11 @@ def _normalize_qwen35_quant_config_dict(model_config, keep_lm_head_quant=False): shared scale (_requantize_linear_attn_fp8_qkvz). Incomplete or non-FP8 sets get no fused entry, and the mapper dequantizes them to bf16 instead. - The ``lm_head`` entry is promoted W4A16_NVFP4 -> NVFP4 when - ``keep_lm_head_quant`` (see _lm_head_nvfp4_enabled) and dropped otherwise: - a leftover entry would make DecoderModelForCausalLM build a quantized - LMHead whose weights the mapper had already dequantized to bf16. + When ``keep_lm_head_quant`` is true (see _lm_head_nvfp4_enabled), the + ``lm_head`` entry is promoted W4A16_NVFP4 -> NVFP4 on SM100/103 and remains + W4A16_NVFP4 on SM120/121. It is dropped otherwise: a leftover entry would + make DecoderModelForCausalLM build a quantized LMHead whose weights the + mapper had already dequantized to bf16. """ qcd = getattr(model_config, "quant_config_dict", None) if not qcd: @@ -528,7 +528,9 @@ def _normalize_qwen35_quant_config_dict(model_config, keep_lm_head_quant=False): continue if name == "lm_head": if keep_lm_head_quant: - normalized[name] = cfg.model_copy(update={"quant_algo": QuantAlgo.NVFP4}) + if convert_to_nvfp4: + cfg = cfg.model_copy(update={"quant_algo": QuantAlgo.NVFP4}) + normalized[name] = cfg else: # Make the fallback visible: the checkpoint quantizes lm_head # but this configuration can't keep it quantized (see @@ -536,7 +538,7 @@ def _normalize_qwen35_quant_config_dict(model_config, keep_lm_head_quant=False): logger.info( f"lm_head quant entry ({cfg.quant_algo}) dropped: " "unsupported configuration for quantized LMHead " - "(requires SM100/103, untied embeddings, no attention-DP, " + "(requires SM100/103/120/121, untied embeddings, no attention-DP, " "vocab divisible by tp_size); lm_head runs bf16" ) continue diff --git a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py index d4267287b0ae..3b0fc2ed8028 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py @@ -105,6 +105,13 @@ def get_moe_cls( from tensorrt_llm._utils import get_sm_version sm_version = get_sm_version() if sm_version in CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS: + mapping = model_config.mapping + if mapping.moe_ep_size > 1 or mapping.dp_size > 1: + logger.info( + "CuteDslB12xFusedMoE does not support expert " + "parallelism or attention-DP/all-to-all; selecting " + "CutlassFusedMoE.") + return CutlassFusedMoE try: import flashinfer # noqa: F401 logger.info( diff --git a/tensorrt_llm/_torch/modules/mlp.py b/tensorrt_llm/_torch/modules/mlp.py index dcd2cc6ebfe5..0d06d670ee82 100644 --- a/tensorrt_llm/_torch/modules/mlp.py +++ b/tensorrt_llm/_torch/modules/mlp.py @@ -10,7 +10,8 @@ from ..model_config import ModelConfig from ..peft.lora.layer import LoraLayer, LoraModuleType from ..utils import Fp4QuantizedTensor, gelu_tanh, relu2 -from .linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig +from .linear import (Linear, TensorParallelMode, WeightMode, + WeightsLoadingConfig, is_static_nvfp4_input_eligible) class MLP(nn.Module): @@ -99,10 +100,8 @@ def create_weights(self): self.up_proj.create_weights() self.down_proj.create_weights() - has_nvfp4 = hasattr(self.down_proj, - 'has_nvfp4') and self.down_proj.has_nvfp4 + has_static_nvfp4_input = is_static_nvfp4_input_eligible(self.down_proj) has_kernel = hasattr(torch.ops.trtllm, 'fused_relu2_quantize') - has_scale = hasattr(self.down_proj, 'input_scale') is_relu2 = self.activation is relu2 # The fused relu2+fp4_quantize kernel body is guarded by # ``__CUDA_ARCH__ >= 1000`` (see fusedActivationQuant.cu). On pre-SM100 @@ -110,7 +109,7 @@ def create_weights(self): # quantize in the downstream linear layer. is_sm100_or_later = get_sm_version() >= 100 - self._use_fused_relu2_quant = (has_nvfp4 and has_kernel and has_scale + self._use_fused_relu2_quant = (has_static_nvfp4_input and has_kernel and is_relu2 and is_sm100_or_later) # Static eligibility for the fused GELU(tanh) CuteDSL epilogue (mirrors diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen3_5_quant.py b/tests/unittest/_torch/modeling/test_modeling_qwen3_5_quant.py index c3fbc4f4f37b..404fa703b9d9 100644 --- a/tests/unittest/_torch/modeling/test_modeling_qwen3_5_quant.py +++ b/tests/unittest/_torch/modeling/test_modeling_qwen3_5_quant.py @@ -21,6 +21,7 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_qwen3_5 import ( + _lm_head_nvfp4_enabled, _normalize_qwen35_exclude_modules, _normalize_qwen35_quant_config_dict, ) @@ -32,7 +33,11 @@ def test_qwen36_normalizes_modelopt_quantized_layer_paths_on_sm121(): w4a16_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16) fp8_config = QuantConfig(quant_algo=QuantAlgo.FP8) model_config = ModelConfig( - pretrained_config=SimpleNamespace(num_hidden_layers=40), + pretrained_config=SimpleNamespace( + num_hidden_layers=40, + tie_word_embeddings=False, + vocab_size=248320, + ), quant_config=QuantConfig( exclude_modules=[ "model.language_model.layers.0.linear_attn.in_proj_qkv", @@ -53,12 +58,13 @@ def test_qwen36_normalizes_modelopt_quantized_layer_paths_on_sm121(): ) with patch("tensorrt_llm._torch.models.modeling_qwen3_5.get_sm_version", return_value=121): - _normalize_qwen35_exclude_modules(model_config) - _normalize_qwen35_quant_config_dict(model_config) + keep_lm_head_quant = _lm_head_nvfp4_enabled(model_config) + assert keep_lm_head_quant + _normalize_qwen35_exclude_modules(model_config, keep_lm_head_quant=keep_lm_head_quant) + _normalize_qwen35_quant_config_dict(model_config, keep_lm_head_quant=keep_lm_head_quant) assert model_config.quant_config.exclude_modules == [ "*linear_attn.conv1d", - "lm_head", "model.layers.0.linear_attn.in_proj_qkvz*", "model.layers.40*", ] @@ -69,6 +75,7 @@ def test_qwen36_normalizes_modelopt_quantized_layer_paths_on_sm121(): "model.layers.0.mlp.shared_expert.gate_proj", "model.layers.0.mlp.shared_expert.up_proj", "model.layers.40.mlp.experts", + "lm_head", } for name, quant_config in model_config.quant_config_dict.items(): if name.endswith("in_proj_qkvz"): diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_nemotron_h_weight_mapper.py b/tests/unittest/_torch/models/checkpoints/hf/test_nemotron_h_weight_mapper.py index 3e14fab36cb4..1fe27820f8f6 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_nemotron_h_weight_mapper.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_nemotron_h_weight_mapper.py @@ -14,6 +14,7 @@ # limitations under the License. from types import SimpleNamespace +from typing import Optional import torch @@ -25,7 +26,9 @@ from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig -def _make_mapper() -> NemotronHHfWeightMapper: +def _make_mapper( + quant_algo: Optional[QuantAlgo] = None, +) -> NemotronHHfWeightMapper: mapper = NemotronHHfWeightMapper() mapper._config = ModelConfig( pretrained_config=SimpleNamespace( @@ -33,54 +36,77 @@ def _make_mapper() -> NemotronHHfWeightMapper: mamba_num_heads=1, n_groups=1, num_hidden_layers=52, + quantization_config={ + "producer": {"name": "modelopt", "version": "0.37.0"}, + "quant_method": "modelopt", + }, ssm_state_size=1, ), mapping=Mapping(), moe_backend="CUTLASS", - quant_config=QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4), + quant_config=QuantConfig(quant_algo=quant_algo), ) return mapper -def test_nemotron_h_mapper_canonicalizes_w4a16_nvfp4_checkpoint_keys(): +def test_nemotron_h_mapper_preserves_smart_panda_w4a16_lm_head_weights(): mapper = _make_mapper() + weight_scale_2 = torch.tensor(0.291 / (448 * 6), dtype=torch.float32) weights = { - "lm_head.weight_packed": torch.empty((8, 4), dtype=torch.uint8), + "lm_head.weight": torch.empty((8, 4), dtype=torch.uint8), "lm_head.weight_scale": torch.empty((8, 1), dtype=torch.float8_e4m3fn), - "lm_head.weight_global_scale": torch.tensor(0.25, dtype=torch.float32), + "lm_head.weight_scale_2": weight_scale_2, } mapped = mapper.preprocess_weights(weights) - assert "lm_head.weight" in mapped - assert "lm_head.weight_scale" in mapped - assert "lm_head.weight_scale_2" in mapped - assert "lm_head.input_scale" in mapped - assert "lm_head.weight_packed" not in mapped - assert "lm_head.weight_global_scale" not in mapped - assert mapped["lm_head.weight"] is weights["lm_head.weight_packed"] - torch.testing.assert_close( - mapped["lm_head.weight_scale_2"], torch.tensor(4.0, dtype=torch.float32) - ) - torch.testing.assert_close( - mapped["lm_head.input_scale"], torch.tensor([1.0], dtype=torch.float32) - ) + assert mapped["lm_head.weight"] is weights["lm_head.weight"] + assert mapped["lm_head.weight_scale"] is weights["lm_head.weight_scale"] + assert mapped["lm_head.weight_scale_2"] is weight_scale_2 + assert "lm_head.input_scale" not in mapped -def test_nemotron_h_mapper_handles_scalar_w4a16_nvfp4_moe_global_scales(): +def test_nemotron_h_mapper_remaps_smart_panda_w4a16_moe_weights(): mapper = _make_mapper() - prefix = "backbone.layers.1.mixer.experts.0.up_proj" + up_prefix = "backbone.layers.1.mixer.experts.0.up_proj" + down_prefix = "backbone.layers.1.mixer.experts.0.down_proj" + up_weight_scale_2 = torch.tensor(0.134 / (448 * 6), dtype=torch.float32) + down_weight_scale_2 = torch.tensor(0.214 / (448 * 6), dtype=torch.float32) weights = { - f"{prefix}.weight_packed": torch.empty((8, 4), dtype=torch.uint8), - f"{prefix}.weight_scale": torch.empty((8, 1), dtype=torch.float8_e4m3fn), - f"{prefix}.weight_global_scale": torch.tensor(0.25, dtype=torch.float32), + f"{up_prefix}.weight": torch.empty((8, 4), dtype=torch.uint8), + f"{up_prefix}.weight_scale": torch.empty((8, 1), dtype=torch.float8_e4m3fn), + f"{up_prefix}.weight_scale_2": up_weight_scale_2, + f"{down_prefix}.weight": torch.empty((8, 4), dtype=torch.uint8), + f"{down_prefix}.weight_scale": torch.empty((8, 1), dtype=torch.float8_e4m3fn), + f"{down_prefix}.weight_scale_2": down_weight_scale_2, } mapped = mapper.preprocess_weights(weights) assert "model.layers.1.mixer.experts.0.w1.weight" in mapped assert "model.layers.1.mixer.experts.0.w3.weight" in mapped - assert "model.layers.1.mixer.experts.0.w1.weight_scale_2" in mapped - assert "model.layers.1.mixer.experts.0.w3.weight_scale_2" in mapped + assert "model.layers.1.mixer.experts.0.w2.weight" in mapped assert mapped["model.layers.1.mixer.experts.0.w3.weight"].shape == (0, 4) assert mapped["model.layers.1.mixer.experts.0.w3.weight_scale_2"].shape == () + assert mapped["model.layers.1.mixer.experts.0.w1.weight_scale_2"] is up_weight_scale_2 + assert mapped["model.layers.1.mixer.experts.0.w3.weight_scale_2"] is up_weight_scale_2 + assert mapped["model.layers.1.mixer.experts.0.w2.weight_scale_2"] is down_weight_scale_2 + assert not any(key.endswith(".input_scale") for key in mapped) + + +def test_nemotron_h_mapper_converts_compressed_tensors_global_scale(): + mapper = _make_mapper(QuantAlgo.W4A16_NVFP4) + global_scale = torch.tensor(9362.2861328125, dtype=torch.float32) + weights = { + "lm_head.weight_packed": torch.empty((8, 4), dtype=torch.uint8), + "lm_head.weight_scale": torch.empty((8, 1), dtype=torch.float8_e4m3fn), + "lm_head.weight_global_scale": global_scale, + } + + mapped = mapper.preprocess_weights(weights) + + assert mapped["lm_head.weight"] is weights["lm_head.weight_packed"] + assert "lm_head.weight_packed" not in mapped + assert "lm_head.weight_global_scale" not in mapped + assert "lm_head.input_scale" not in mapped + torch.testing.assert_close(mapped["lm_head.weight_scale_2"], global_scale.reciprocal()) diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_qwen3_5_weight_mapper.py b/tests/unittest/_torch/models/checkpoints/hf/test_qwen3_5_weight_mapper.py index 9796217a2382..e9a644845e05 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_qwen3_5_weight_mapper.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_qwen3_5_weight_mapper.py @@ -46,6 +46,23 @@ def test_qwen35_modelopt_preprocess_preserves_scalar_fp8_scale_name(): assert weights["model.layers.0.linear_attn.out_proj.weight_scale"].shape == torch.Size([]) +def test_qwen36_preserves_quantized_modelopt_lm_head(): + mapper = _make_mapper() + mapper._config.quant_config_dict = { + "lm_head": QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16), + } + weights = { + "lm_head.weight": torch.empty((8, 4), dtype=torch.uint8), + "lm_head.weight_scale": torch.empty((8, 1), dtype=torch.float8_e4m3fn), + "lm_head.weight_scale_2": torch.tensor(0.00011189778888365254, dtype=torch.float32), + "lm_head.input_scale": torch.tensor(0.02771577425301075, dtype=torch.float32), + } + + mapped = mapper._dequantize_lm_head_nvfp4(weights) + + assert mapped is weights + + def test_qwen35_rescales_per_tensor_fp8_linear_attention_qkvz_projection(): mapper = _make_mapper(dtype=torch.bfloat16) mapper._config.quant_config_dict = { diff --git a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py index 5d2ce75f435c..61c90629f071 100644 --- a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py @@ -40,6 +40,7 @@ NVFP4CutlassFusedMoEMethod, ) from tensorrt_llm._torch.utils import ActivationType +from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig _FUSED_MOE_MODULE = "tensorrt_llm._torch.modules.fused_moe.fused_moe_cute_dsl_b12x" @@ -193,6 +194,30 @@ def test_get_moe_cls_cutedsl_selects_b12x_for_layer_w4a16_nvfp4_on_supported_sm( assert cls is CuteDslB12xFusedMoE +@pytest.mark.parametrize( + "mapping", + [ + Mapping(world_size=2, tp_size=2, moe_tp_size=1, moe_ep_size=2), + Mapping( + world_size=2, + tp_size=2, + enable_attention_dp=True, + dwdp_size=2, + dwdp_rank=0, + ), + ], +) +def test_get_moe_cls_cutedsl_falls_back_to_cutlass_for_distributed_b12x(mapping): + cfg = ModelConfig(mapping=mapping) + cfg.moe_backend = "CUTEDSL" + cfg.quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) + + with patch("tensorrt_llm._utils.get_sm_version", return_value=120): + cls = get_moe_cls(cfg) + + assert cls is CutlassFusedMoE + + def test_get_moe_cls_cutedsl_falls_back_to_plain_cutedsl_when_flashinfer_missing(monkeypatch): """CUTEDSL + NVFP4 + SM120/121 + flashinfer NOT importable → CuteDslFusedMoE.""" import builtins diff --git a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py index 0a700e415caa..86e2a7d9e2a2 100644 --- a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py +++ b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py @@ -20,6 +20,7 @@ import torch import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils +from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.modules.embedding import LMHead from tensorrt_llm._torch.modules.linear import ( Linear, @@ -28,6 +29,8 @@ get_quant_method, get_sm_version, ) +from tensorrt_llm._torch.modules.mlp import MLP +from tensorrt_llm._torch.utils import relu2 from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig @@ -183,6 +186,26 @@ def test_get_quant_method_returns_w4a16_nvfp4_linear_method(): assert type(method) is W4A16NVFP4LinearMethod +def test_w4a16_nvfp4_mlp_disables_relu2_fp4_fusion_without_input_scale(): + model_config = ModelConfig(quant_config=QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4)) + + with patch("tensorrt_llm._torch.modules.mlp.get_sm_version", return_value=121): + mlp = MLP( + hidden_size=32, + intermediate_size=64, + bias=False, + activation=relu2, + dtype=torch.bfloat16, + config=model_config, + reduce_output=False, + ) + mlp.create_weights() + + assert mlp.down_proj.has_nvfp4 + assert mlp.down_proj.input_scale is None + assert not mlp._use_fused_relu2_quant + + def test_w4a16_nvfp4_linear_uses_high_precision_activation_without_fp4_quantize(): method = W4A16NVFP4LinearMethod() input_tensor = torch.ones((2, 32), dtype=torch.bfloat16) From 12cf7e24847e948ff8587628ea9a6e4e2294f35d Mon Sep 17 00:00:00 2001 From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> Date: Tue, 21 Jul 2026 04:57:10 +0000 Subject: [PATCH 13/25] fix: guard fused ReLU2 quantization scale Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> --- tensorrt_llm/_torch/modules/mlp.py | 5 +- .../_torch/modules/test_w4a16_nvfp4_linear.py | 63 ++++++++++++++++++- 2 files changed, 66 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mlp.py b/tensorrt_llm/_torch/modules/mlp.py index 0d06d670ee82..de14c160dea1 100644 --- a/tensorrt_llm/_torch/modules/mlp.py +++ b/tensorrt_llm/_torch/modules/mlp.py @@ -147,7 +147,10 @@ def forward( x_up = self.up_proj(x) - if self._use_fused_relu2_quant: + # Weight loading may replace the quantization method after + # create_weights(), so do not rely on the cached eligibility alone. + if (self._use_fused_relu2_quant + and is_static_nvfp4_input_eligible(self.down_proj)): x_act = self._fused_relu2_quant(x_up) else: x_act = self.activation(x_up) diff --git a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py index 86e2a7d9e2a2..3594777e4cc2 100644 --- a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py +++ b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py @@ -189,7 +189,10 @@ def test_get_quant_method_returns_w4a16_nvfp4_linear_method(): def test_w4a16_nvfp4_mlp_disables_relu2_fp4_fusion_without_input_scale(): model_config = ModelConfig(quant_config=QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4)) - with patch("tensorrt_llm._torch.modules.mlp.get_sm_version", return_value=121): + with ( + patch("tensorrt_llm._torch.modules.mlp.get_sm_version", return_value=121), + patch("torch.ops.trtllm.fused_relu2_quantize", create=True), + ): mlp = MLP( hidden_size=32, intermediate_size=64, @@ -206,6 +209,64 @@ def test_w4a16_nvfp4_mlp_disables_relu2_fp4_fusion_without_input_scale(): assert not mlp._use_fused_relu2_quant +def test_nvfp4_mlp_enables_relu2_fp4_fusion_with_static_input_scale(): + model_config = ModelConfig(quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4)) + + with ( + patch("tensorrt_llm._torch.modules.mlp.get_sm_version", return_value=121), + patch("torch.ops.trtllm.fused_relu2_quantize", create=True), + ): + mlp = MLP( + hidden_size=32, + intermediate_size=64, + bias=False, + activation=relu2, + dtype=torch.bfloat16, + config=model_config, + reduce_output=False, + ) + mlp.create_weights() + + assert mlp.down_proj.input_scale is not None + assert mlp._use_fused_relu2_quant + + +def test_w4a16_nvfp4_mlp_rechecks_relu2_fp4_fusion_before_forward(): + model_config = ModelConfig(quant_config=QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4)) + + with ( + patch("tensorrt_llm._torch.modules.mlp.get_sm_version", return_value=121), + patch("torch.ops.trtllm.fused_relu2_quantize", create=True), + ): + mlp = MLP( + hidden_size=32, + intermediate_size=64, + bias=False, + activation=relu2, + dtype=torch.bfloat16, + config=model_config, + reduce_output=False, + ) + mlp.create_weights() + + # Simulate eligibility cached before weight loading replaced the linear + # method with W4A16, whose high-precision activation has no input_scale. + mlp._use_fused_relu2_quant = True + x_up = torch.tensor([[-2.0, 3.0]], dtype=torch.bfloat16) + with ( + patch.object(mlp.up_proj, "forward", return_value=x_up), + patch.object(mlp.down_proj, "forward", side_effect=lambda x: x), + patch.object( + MLP, + "_fused_relu2_quant", + side_effect=AssertionError("missing input_scale must use unfused ReLU2"), + ), + ): + output = mlp(torch.empty((1, 32), dtype=torch.bfloat16)) + + torch.testing.assert_close(output, relu2(x_up)) + + def test_w4a16_nvfp4_linear_uses_high_precision_activation_without_fp4_quantize(): method = W4A16NVFP4LinearMethod() input_tensor = torch.ones((2, 32), dtype=torch.bfloat16) From 67d78ff14fa396df55fcc6c6d616be823b162fee Mon Sep 17 00:00:00 2001 From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> Date: Tue, 21 Jul 2026 05:05:47 +0000 Subject: [PATCH 14/25] update document Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_nemotron_h.py | 7 +++---- tensorrt_llm/_torch/modules/fused_moe/quantization.py | 3 +-- .../models/checkpoints/hf/test_nemotron_h_weight_mapper.py | 4 ++-- 3 files changed, 6 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index 1013a898aa5b..4e9a00be81e2 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py @@ -1256,10 +1256,9 @@ def _get_mtp_sublayer_quant_config(self, model_config: NemotronHModelConfig, layer_idx: int): """ Get quantization config for MTP sublayer. - The Nano3.5 NVFP4 W4A16 checkpoint stores the MTP body tensors in BF16. - The shared MTP head still receives the checkpoint-backed lm_head, so its - logits path keeps the lm_head precision instead of inheriting this - sublayer override. + The MTP body tensors are stored in BF16. The shared MTP head still + receives the checkpoint-backed lm_head, so its logits path keeps the + lm_head precision instead of inheriting this sublayer override. """ from tensorrt_llm.models.modeling_utils import QuantConfig diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index 952d9fe9c851..4df1bbcc5bde 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py @@ -3388,8 +3388,7 @@ def transform_weights(self, module: torch.nn.Module) -> None: f"supported: {supported}.") # The model config may carry the logical intermediate size while the - # NVFP4 weight tensors are padded for kernel alignment, e.g. Nano3.5 - # uses 1856 logical channels and 1920 stored channels. FlashInfer's + # NVFP4 weight tensors are padded for kernel alignment. FlashInfer's # CUDA-graph workspace must match the stored tensors. b12x_intermediate_size = w2_in_dim diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_nemotron_h_weight_mapper.py b/tests/unittest/_torch/models/checkpoints/hf/test_nemotron_h_weight_mapper.py index 1fe27820f8f6..16525abe8d74 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_nemotron_h_weight_mapper.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_nemotron_h_weight_mapper.py @@ -49,7 +49,7 @@ def _make_mapper( return mapper -def test_nemotron_h_mapper_preserves_smart_panda_w4a16_lm_head_weights(): +def test_nemotron_h_mapper_preserves_w4a16_lm_head_weights_without_input_scale(): mapper = _make_mapper() weight_scale_2 = torch.tensor(0.291 / (448 * 6), dtype=torch.float32) weights = { @@ -66,7 +66,7 @@ def test_nemotron_h_mapper_preserves_smart_panda_w4a16_lm_head_weights(): assert "lm_head.input_scale" not in mapped -def test_nemotron_h_mapper_remaps_smart_panda_w4a16_moe_weights(): +def test_nemotron_h_mapper_remaps_w4a16_moe_weights_without_input_scale(): mapper = _make_mapper() up_prefix = "backbone.layers.1.mixer.experts.0.up_proj" down_prefix = "backbone.layers.1.mixer.experts.0.down_proj" From 897d5549b66702649d1ffca9b9d618bd6d91ad00 Mon Sep 17 00:00:00 2001 From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> Date: Tue, 21 Jul 2026 15:22:17 +0000 Subject: [PATCH 15/25] fix nvinfer Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> --- .../kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.h | 2 -- 1 file changed, 2 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.h b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.h index 2f525c50c9b7..30b4a91da32b 100644 --- a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.h +++ b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.h @@ -19,8 +19,6 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/runtime/common.h" -#include - #include #include #include From 7bc2719c4f055b6679917f6e9a45360c1c4ee45d Mon Sep 17 00:00:00 2001 From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> Date: Fri, 24 Jul 2026 01:14:08 +0000 Subject: [PATCH 16/25] remove cuda_core Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> --- .../cudaCoreGemmW4A16NVFP4.cu | 311 ---------------- .../cudaCoreGemmW4A16NVFP4.h | 69 ---- .../weightOnlyBatchedGemv/nvfp4ScaleLayout.h | 50 --- cpp/tensorrt_llm/thop/CMakeLists.txt | 1 - cpp/tensorrt_llm/thop/w4a16Nvfp4Gemm.cpp | 164 --------- .../_torch/custom_ops/cpp_custom_ops.py | 10 - tensorrt_llm/_torch/modules/linear.py | 43 +-- .../_torch/modules/test_w4a16_nvfp4_linear.py | 339 +++++++++++------- 8 files changed, 214 insertions(+), 773 deletions(-) delete mode 100644 cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.cu delete mode 100644 cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.h delete mode 100644 cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/nvfp4ScaleLayout.h delete mode 100644 cpp/tensorrt_llm/thop/w4a16Nvfp4Gemm.cpp diff --git a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.cu b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.cu deleted file mode 100644 index 984179f7393b..000000000000 --- a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.cu +++ /dev/null @@ -1,311 +0,0 @@ -/* - * Copyright (c) 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. - * 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. - */ - -#include "cutlass/numeric_conversion.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h" -#include "tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.h" -#include "tensorrt_llm/kernels/weightOnlyBatchedGemv/nvfp4ScaleLayout.h" - -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace cuda_core_gemm_w4a16_nvfp4 -{ -namespace -{ - -template -__device__ float toFloat(T value) -{ - return static_cast(value); -} - -template <> -__device__ float toFloat(half value) -{ - return __half2float(value); -} - -template <> -__device__ float toFloat<__nv_bfloat16>(__nv_bfloat16 value) -{ - return __bfloat162float(value); -} - -} // namespace - -template -__device__ void cudaCoreGemmImpl(ActType const* __restrict__ act, __nv_fp4_e2m1 const* __restrict__ weight, - ScaleType const* __restrict__ weightScale, float const weightGlobalScale, OutputType* __restrict__ output, - SizeType32 m, SizeType32 n, SizeType32 k) -{ - using VecType = int4; - using ScaleVecType = __nv_fp8x2_e4m3; - using CvtWeightType = - typename tensorrt_llm::kernels::cutlass_kernels::TllmToCutlassTypeAdapter<__nv_fp4_e2m1>::type; - using Converter = cutlass::NumericArrayConverter; - using CvtSrcType = typename Converter::source_type; - using CvtResType = typename Converter::result_type; - - static constexpr SizeType32 kStepK = 32; - static constexpr SizeType32 kStepKScale = kStepK / w4a16_nvfp4::kScaleGranularity; - static constexpr SizeType32 kTileK = kStepK * kBlockSize; - static constexpr SizeType32 kCvtCount = static_cast(sizeof(VecType) / sizeof(CvtSrcType)); - - static_assert(kStepK % w4a16_nvfp4::kScaleGranularity == 0); - - auto const tileIdM = static_cast(blockIdx.x * kTileM); - auto const tileIdN = static_cast(blockIdx.y * kTileN); - auto const tid = static_cast(threadIdx.x); - (void) m; - - float tileAct[kStepK]; - float tileWeight[kTileN * kStepK]; - float tileWeightScale[kTileN * kStepKScale]; - float acc[kTileM * kTileN]; - -#pragma unroll - for (SizeType32 i = 0; i < kTileM * kTileN; ++i) - { - acc[i] = 0.0F; - } - - act += tileIdM * k; - weight += tileIdN * k / 2; - output += tileIdM * n + tileIdN; - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaGridDependencySynchronize(); -#endif - - for (SizeType32 idxK = tid * kStepK; idxK < k; idxK += kTileK) - { -#pragma unroll - for (SizeType32 j = 0; j < kTileN; ++j) - { - auto tileWeightQuantized = reinterpret_cast(weight + (j * k + idxK) / 2)[0]; -#pragma unroll - for (SizeType32 cvtIdx = 0; cvtIdx < kCvtCount; ++cvtIdx) - { - reinterpret_cast(tileWeight)[j * kCvtCount + cvtIdx] - = Converter::convert(reinterpret_cast(&tileWeightQuantized)[cvtIdx]); - } - } - -#pragma unroll - for (SizeType32 j = 0; j < kTileN; ++j) - { - SizeType32 const rowIdx = tileIdN + j; - SizeType32 const colIdx = idxK / w4a16_nvfp4::kScaleGranularity; - SizeType32 const dstIdx = w4a16_nvfp4::getScaleIndex(rowIdx, colIdx, k); - auto const tileWeightScaleFp8x2 = reinterpret_cast(weightScale + dstIdx)[0]; - char2 const tmp = reinterpret_cast(tileWeightScaleFp8x2); - tileWeightScale[j * kStepKScale + 0] = static_cast(reinterpret_cast<__nv_fp8_e4m3 const&>(tmp.x)); - tileWeightScale[j * kStepKScale + 1] = static_cast(reinterpret_cast<__nv_fp8_e4m3 const&>(tmp.y)); - } - -#pragma unroll - for (SizeType32 i = 0; i < kTileM; ++i) - { -#pragma unroll - for (SizeType32 l = 0; l < kStepK; ++l) - { - tileAct[l] = toFloat(act[i * k + idxK + l]); - } - -#pragma unroll - for (SizeType32 j = 0; j < kTileN; ++j) - { -#pragma unroll - for (SizeType32 l = 0; l < kStepK; ++l) - { - float const scaledWeight = tileWeight[j * kStepK + l] - * tileWeightScale[j * kStepKScale + l / w4a16_nvfp4::kScaleGranularity] * weightGlobalScale; - acc[i * kTileN + j] = fma(tileAct[l], scaledWeight, acc[i * kTileN + j]); - } - } - } - } - - using WarpReduce = cub::WarpReduce; - static constexpr SizeType32 kWarpSize = 32; - static constexpr SizeType32 kWarpNum = kBlockSize / kWarpSize; - SizeType32 const warpId = tid / kWarpSize; - SizeType32 const laneId = tid % kWarpSize; - __shared__ float shmem[kTileM * kTileN * kWarpNum]; - __shared__ typename WarpReduce::TempStorage tempStorage[kWarpNum]; - -#pragma unroll - for (SizeType32 mi = 0; mi < kTileM; ++mi) - { -#pragma unroll - for (SizeType32 ni = 0; ni < kTileN; ++ni) - { - float const val = WarpReduce(tempStorage[warpId]).Sum(acc[mi * kTileN + ni]); - if (laneId == 0) - { - shmem[mi * kTileN + ni + warpId * kTileM * kTileN] = val; - } - } - } - __syncthreads(); - - for (SizeType32 ii = tid; ii < kTileM * kTileN; ii += kBlockSize) - { - SizeType32 const mid = ii / kTileN; - SizeType32 const nid = ii % kTileN; - float val = 0.0F; -#pragma unroll - for (SizeType32 jj = 0; jj < kWarpNum; ++jj) - { - val += shmem[jj * kTileM * kTileN + ii]; - } - output[mid * n + nid] = static_cast(val); - } - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaTriggerProgrammaticLaunchCompletion(); -#endif -} - -template -__global__ void cudaCoreGemm(ActType const* __restrict__ act, __nv_fp4_e2m1 const* __restrict__ weight, - ScaleType const* __restrict__ weightScale, float const* weightGlobalScale, OutputType* __restrict__ output, - SizeType32 m, SizeType32 n, SizeType32 k) -{ - cudaCoreGemmImpl( - act, weight, weightScale, weightGlobalScale[0], output, m, n, k); -} - -template -void cudaCoreGemmKernel(Params const& params, cudaStream_t stream) -{ - dim3 const block(kBlockSize); - dim3 const grid(params.m / kTileM, params.n / kTileN); - cudaCoreGemm<<>>( - reinterpret_cast(params.act), reinterpret_cast<__nv_fp4_e2m1 const*>(params.weight), - reinterpret_cast(params.weightScale), params.weightGlobalScale, - reinterpret_cast(params.output), params.m, params.n, params.k); - TLLM_CUDA_CHECK(cudaGetLastError()); -} - -template -bool cudaCoreGemmTemplateCaller(Params const& params, cudaStream_t stream) -{ - constexpr int kCudaCoreGemmTemplateMaxM = 16; - if (params.m == kTileM) - { - cudaCoreGemmKernel(params, stream); - return true; - } - if constexpr (kTileM < kCudaCoreGemmTemplateMaxM) - { - return cudaCoreGemmTemplateCaller( - params, stream); - } - return false; -} - -template -bool cudaCoreGemmLauncher(Params const& params, cudaStream_t stream) -{ - constexpr int kDefaultTileN = 2; - constexpr int kWideTileN = 4; - constexpr int kMaxGridDimY = 65535; - if (params.n / kDefaultTileN <= kMaxGridDimY) - { - return cudaCoreGemmTemplateCaller(params, stream); - } - if (params.n % kWideTileN == 0 && params.n / kWideTileN <= kMaxGridDimY) - { - return cudaCoreGemmTemplateCaller(params, stream); - } - return false; -} - -template -bool dispatchOutputType(Params const& params, cudaStream_t stream) -{ - if (params.outputType == CUDA_R_16F) - { - return cudaCoreGemmLauncher(params, stream); - } - if (params.outputType == CUDA_R_16BF) - { - return cudaCoreGemmLauncher(params, stream); - } - if (params.outputType == CUDA_R_32F) - { - return cudaCoreGemmLauncher(params, stream); - } - return false; -} - -bool cudaCoreGemmDispatcher(Params const& params, cudaStream_t stream) -{ - bool dispatched = true; - int const smVersion = tensorrt_llm::common::getSMVersion(); - if (smVersion != 120 && smVersion != 121) - { - dispatched = false; - } - else if (params.n % 2 != 0 || params.k % 32 != 0) - { - dispatched = false; - } - else if (params.m < 1 || params.m > 16) - { - dispatched = false; - } - else if (params.weightScale == nullptr || params.weightGlobalScale == nullptr) - { - dispatched = false; - } - else if (params.inputType == CUDA_R_16F) - { - dispatched = dispatchOutputType(params, stream); - } - else if (params.inputType == CUDA_R_16BF) - { - dispatched = dispatchOutputType<__nv_bfloat16>(params, stream); - } - else - { - dispatched = false; - } - - if (!dispatched) - { - TLLM_LOG_WARNING( - "tensorrt_llm::kernels::cuda_core_gemm_w4a16_nvfp4::cudaCoreGemmDispatcher [NOT DISPATCHED], " - "inputType=%d, outputType=%d, m=%d, n=%d, k=%d, sm=%d", - params.inputType, params.outputType, params.m, params.n, params.k, smVersion); - } - return dispatched; -} - -} // namespace cuda_core_gemm_w4a16_nvfp4 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.h b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.h deleted file mode 100644 index 30b4a91da32b..000000000000 --- a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.h +++ /dev/null @@ -1,69 +0,0 @@ -/* - * Copyright (c) 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. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/common/config.h" -#include "tensorrt_llm/runtime/common.h" - -#include -#include -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace cuda_core_gemm_w4a16_nvfp4 -{ -using SizeType32 = tensorrt_llm::runtime::SizeType32; - -struct Params -{ - void const* act; - void const* weight; - void const* weightScale; - float const* weightGlobalScale; - void* output; - SizeType32 m, n, k; - cudaDataType_t inputType; - cudaDataType_t outputType; - - Params(void const* act_, void const* weight_, void const* weightScale_, float const* weightGlobalScale_, - void* output_, SizeType32 m_, SizeType32 n_, SizeType32 k_, cudaDataType_t inputType_, - cudaDataType_t outputType_) - : act(act_) - , weight(weight_) - , weightScale(weightScale_) - , weightGlobalScale(weightGlobalScale_) - , output(output_) - , m(m_) - , n(n_) - , k(k_) - , inputType(inputType_) - , outputType(outputType_) - { - } -}; - -bool cudaCoreGemmDispatcher(Params const& params, cudaStream_t stream); - -} // namespace cuda_core_gemm_w4a16_nvfp4 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/nvfp4ScaleLayout.h b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/nvfp4ScaleLayout.h deleted file mode 100644 index 03c54d80033c..000000000000 --- a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/nvfp4ScaleLayout.h +++ /dev/null @@ -1,50 +0,0 @@ -/* - * Copyright (c) 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. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/common/config.h" -#include "tensorrt_llm/runtime/common.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace w4a16_nvfp4 -{ - -using SizeType32 = tensorrt_llm::runtime::SizeType32; - -static constexpr SizeType32 kScaleGranularity = 16; -static constexpr SizeType32 kScaleRowsPerTile = 128; -static constexpr SizeType32 kPackedScaleColsPerTile = 4; -static constexpr SizeType32 kScaleTileElements = 512; - -__host__ __device__ inline SizeType32 getScaleIndex(SizeType32 rowIdx, SizeType32 scaleColIdx, SizeType32 k) -{ - SizeType32 const numScaleCols = k / kScaleGranularity; - SizeType32 const numScaleColTiles = (numScaleCols + kPackedScaleColsPerTile - 1) / kPackedScaleColsPerTile; - SizeType32 const tileOffset - = ((rowIdx / kScaleRowsPerTile) * numScaleColTiles + scaleColIdx / kPackedScaleColsPerTile) - * kScaleTileElements; - return tileOffset + (rowIdx % 32) * 16 + ((rowIdx % kScaleRowsPerTile) / 32) * 4 - + scaleColIdx % kPackedScaleColsPerTile; -} - -} // namespace w4a16_nvfp4 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index a684fd046e4d..b95426f580fe 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -57,7 +57,6 @@ add_library( marlinNvfp4MoeMM.cpp marlinRepack.cpp cudaScaledMM.cpp - w4a16Nvfp4Gemm.cpp dynamicDecodeOp.cpp fmhaPackMaskOp.cpp fp8Op.cpp diff --git a/cpp/tensorrt_llm/thop/w4a16Nvfp4Gemm.cpp b/cpp/tensorrt_llm/thop/w4a16Nvfp4Gemm.cpp deleted file mode 100644 index d5faf50da520..000000000000 --- a/cpp/tensorrt_llm/thop/w4a16Nvfp4Gemm.cpp +++ /dev/null @@ -1,164 +0,0 @@ -/* - * Copyright (c) 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. - * 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. - */ - -#include "tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmW4A16NVFP4.h" -#include "tensorrt_llm/kernels/weightOnlyBatchedGemv/nvfp4ScaleLayout.h" -#include "tensorrt_llm/thop/thUtils.h" - -#include -#include - -#include - -using torch::Tensor; - -TRTLLM_NAMESPACE_BEGIN - -namespace torch_ext -{ -namespace -{ - -void checkActDtype(Tensor const& act) -{ - TORCH_CHECK(act.scalar_type() == torch::kFloat16 || act.scalar_type() == torch::kBFloat16, - "w4a16_nvfp4_gemm only supports FP16/BF16 activations, got ", act.scalar_type()); -} - -int64_t padUp(int64_t value, int64_t alignment) -{ - return ((value + alignment - 1) / alignment) * alignment; -} - -size_t getCudaDataTypeSize(cudaDataType_t dataType) -{ - switch (dataType) - { - case CUDA_R_16F: - case CUDA_R_16BF: return 2; - case CUDA_R_32F: return 4; - default: return 0; - } -} - -void checkWeightScaleSize(Tensor const& weightScale, int64_t n, int64_t k) -{ - using namespace tensorrt_llm::kernels::w4a16_nvfp4; - TORCH_CHECK( - k % kScaleGranularity == 0, "K must be divisible by ", kScaleGranularity, " for W4A16 NVFP4 GEMM, got K=", k); - int64_t const expectedNumel = padUp(n, kScaleRowsPerTile) * padUp(k / kScaleGranularity, kPackedScaleColsPerTile); - TORCH_CHECK(weightScale.numel() >= expectedNumel, "weight_scale has too few elements for W4A16 NVFP4 GEMM: got ", - weightScale.numel(), ", expected at least ", expectedNumel, " for N=", n, " K=", k); -} - -void w4a16Nvfp4GemmCaller( - Tensor& out, Tensor const& act, Tensor const& weight, Tensor const& weightScale, Tensor const& weightScale2) -{ - constexpr int32_t kCudaCoreMaxM = 16; - auto const m = static_cast(act.sizes()[0]); - auto const k = static_cast(act.sizes()[1]); - auto const n = static_cast(weight.sizes()[0]); - TORCH_CHECK(weight.sizes()[1] * 2 == k, "weight shape [N, K/2] must match activation shape [M, K]"); - - auto stream = at::cuda::getCurrentCUDAStream(act.get_device()); - - auto* actPtr = static_cast(act.data_ptr()); - auto* weightPtr = static_cast(weight.data_ptr()); - auto* weightScalePtr = static_cast(weightScale.data_ptr()); - auto* weightScale2Ptr = static_cast(weightScale2.data_ptr()); - auto* outPtr = static_cast(out.data_ptr()); - - auto const inputType = convert_torch_dtype(act.scalar_type()); - auto const outType = convert_torch_dtype(out.scalar_type()); - - tensorrt_llm::kernels::cuda_core_gemm_w4a16_nvfp4::Params params( - actPtr, weightPtr, weightScalePtr, weightScale2Ptr, outPtr, m, n, k, inputType, outType); - bool dispatched = false; - if (m <= kCudaCoreMaxM) - { - dispatched = tensorrt_llm::kernels::cuda_core_gemm_w4a16_nvfp4::cudaCoreGemmDispatcher(params, stream); - } - else - { - size_t const inputElementSize = getCudaDataTypeSize(inputType); - size_t const outputElementSize = getCudaDataTypeSize(outType); - dispatched = inputElementSize != 0 && outputElementSize != 0; - for (int32_t start = 0; dispatched && start < m; start += kCudaCoreMaxM) - { - int32_t const chunkM = std::min(kCudaCoreMaxM, m - start); - auto const* chunkActPtr - = static_cast(actPtr) + static_cast(start) * k * inputElementSize; - auto* chunkOutPtr = static_cast(outPtr) + static_cast(start) * n * outputElementSize; - tensorrt_llm::kernels::cuda_core_gemm_w4a16_nvfp4::Params chunkParams( - chunkActPtr, weightPtr, weightScalePtr, weightScale2Ptr, chunkOutPtr, chunkM, n, k, inputType, outType); - dispatched = tensorrt_llm::kernels::cuda_core_gemm_w4a16_nvfp4::cudaCoreGemmDispatcher(chunkParams, stream); - } - } - TORCH_CHECK(dispatched, "Failed to dispatch w4a16_nvfp4_gemm kernel"); -} - -} // namespace - -Tensor& w4a16_nvfp4_gemm_out(Tensor const& act, Tensor const& weight, Tensor const& weightScale, - Tensor const& weightScale2, std::optional outDtype, std::optional const& bias, Tensor& out) -{ - CHECK_TH_CUDA(act); - CHECK_CONTIGUOUS(act); - checkActDtype(act); - CHECK_INPUT(weight, FLOAT4_E2M1X2); - CHECK_INPUT(weightScale, SF_DTYPE); - CHECK_INPUT(weightScale2, torch::kFloat32); - CHECK_TH_CUDA(out); - CHECK_CONTIGUOUS(out); - - TORCH_CHECK(act.dim() == 2 && weight.dim() == 2 && out.dim() == 2); - TORCH_CHECK(act.sizes()[0] == out.sizes()[0]); - TORCH_CHECK(weight.sizes()[0] == out.sizes()[1]); - TORCH_CHECK(weight.sizes()[1] * 2 == act.sizes()[1]); - checkWeightScaleSize(weightScale, weight.sizes()[0], act.sizes()[1]); - TORCH_CHECK(weightScale2.numel() == 1, "weight_scale_2 must be a scalar tensor"); - TORCH_CHECK(!bias.has_value(), "w4a16_nvfp4_gemm does not support bias"); - TORCH_CHECK(!outDtype.has_value() || out.scalar_type() == outDtype.value()); - - w4a16Nvfp4GemmCaller(out, act, weight, weightScale, weightScale2); - return out; -} - -Tensor w4a16_nvfp4_gemm(Tensor const& act, Tensor const& weight, Tensor const& weightScale, Tensor const& weightScale2, - std::optional outDtype, std::optional const& bias) -{ - TORCH_CHECK(act.dim() == 2 && weight.dim() == 2); - auto const outDtypeValue = outDtype.value_or(act.scalar_type()); - std::vector outputSize = {act.sizes()[0], weight.sizes()[0]}; - Tensor out = at::empty(outputSize, act.options().dtype(outDtypeValue)); - return w4a16_nvfp4_gemm_out(act, weight, weightScale, weightScale2, outDtype, bias, out); -} - -} // namespace torch_ext - -TRTLLM_NAMESPACE_END - -TORCH_LIBRARY_FRAGMENT(trtllm, m) -{ - m.def( - "w4a16_nvfp4_gemm(Tensor act, Tensor weight, Tensor weight_scale, Tensor weight_scale_2, ScalarType? " - "out_dtype, Tensor? bias=None) -> Tensor"); -} - -TORCH_LIBRARY_IMPL(trtllm, CUDA, m) -{ - m.impl("w4a16_nvfp4_gemm", &tensorrt_llm::torch_ext::w4a16_nvfp4_gemm); -} diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index e8363dab2aa0..e86993ef8f67 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -1156,16 +1156,6 @@ def _(mat_a: torch.Tensor, n = mat_b.shape[0] return mat_a.new_empty((m, n), dtype=out_dtype) - @torch.library.register_fake("trtllm::w4a16_nvfp4_gemm") - def _(act: torch.Tensor, - weight: torch.Tensor, - weight_scale: torch.Tensor, - weight_scale_2: torch.Tensor, - out_dtype: Optional[torch.dtype], - bias: Optional[torch.Tensor] = None): - del weight_scale, weight_scale_2, bias - return act.new_empty((act.shape[0], weight.shape[0]), dtype=out_dtype) - @torch.library.register_fake("trtllm::marlin_nvfp4_gemm") def _(mat_a: torch.Tensor, mat_b: torch.Tensor, diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index befb5285d9db..1e35f02684a4 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -1431,7 +1431,10 @@ def apply(self, module: Linear, input: torch.Tensor, # Use unified interface - supports CUTLASS, cuBLASLt, CuteDSL # Convert list to comma-separated string for torch.compile compatibility - allowed_backends_str = ','.join(module.nvfp4_allowed_backends) + sm_version = get_sm_version() + use_marlin = (90 <= sm_version < 100 and module.dtype == torch.bfloat16) + allowed_backends_str = ('marlin' if use_marlin else ','.join( + module.nvfp4_allowed_backends)) output_buffer_kind = ( int(BufferKind.NCCL_WINDOW) if self.supports_nccl_symmetric_memory_window_output @@ -1443,7 +1446,7 @@ def apply(self, module: Linear, input: torch.Tensor, and module.mapping is not None else None) # Fuse bias inside the GEMM op when N is unpadded and the output is a # plain buffer; otherwise fall back to post-op `out + bias` below. - fuse_bias_in_gemm = (bias is not None + fuse_bias_in_gemm = (bias is not None and not use_marlin and output_buffer_kind == int(BufferKind.DEFAULT) and module.weight.shape[0] == module.out_features) output = torch.ops.trtllm.nvfp4_gemm( @@ -1905,9 +1908,7 @@ def transform_weights(self, module: Linear) -> None: class W4A16NVFP4LinearMethod(NVFP4LinearMethod): - """W4A16 NVFP4 linear with a small-M CUDA-core fast path.""" - - CUDA_CORE_MAX_M: ClassVar[int] = 16 + """W4A16 NVFP4 linear using on-the-fly weight dequantization.""" def create_weights(self, module: Linear, in_features: int, out_features: int, bias: bool, dtype: torch.dtype): @@ -1992,8 +1993,7 @@ def process_weights_after_loading_fused_gate_up_linear( super().process_weights_after_loading_fused_gate_up_linear) def transform_weights(self, module: Linear) -> None: - # Keep the checkpoint layout for the CUDA-core path and materialize the - # much smaller linear scale view once for Triton dequantization. + # Materialize the smaller linear scale view once for Triton dequantization. LinearMethodBase.transform_weights(self, module) self.cache_derived_state(module) @@ -2031,17 +2031,6 @@ def _prepare_input(module: Linear, input: torch.Tensor): input = input * module.pre_quant_scale return input, original_shape - @classmethod - def _can_use_cuda_core(cls, module: Linear, input: torch.Tensor) -> bool: - return (input.dim() == 2 and 0 < input.shape[0] <= cls.CUDA_CORE_MAX_M - and get_sm_version() in (120, 121) - and input.dtype in (torch.float16, torch.bfloat16) - and module.dtype in (torch.float16, torch.bfloat16) - and input.is_contiguous() and input.shape[1] % 32 == 0 - and module.weight.shape[0] % 2 == 0 - and module.weight.shape[1] * 2 == input.shape[1] - and hasattr(torch.ops.trtllm, "w4a16_nvfp4_gemm")) - @staticmethod def _restore_output(output: torch.Tensor, original_shape, bias: Optional[torch.Tensor]): @@ -2054,17 +2043,6 @@ def _restore_output(output: torch.Tensor, original_shape, def apply(self, module: Linear, input: torch.Tensor, bias: Optional[torch.Tensor]): input, original_shape = self._prepare_input(module, input) - if self._can_use_cuda_core(module, input): - output = torch.ops.trtllm.w4a16_nvfp4_gemm( - input, - module.weight, - module.weight_scale, - module.weight_scale_2, - module.dtype, - bias=None, - ) - return self._restore_output(output, original_shape, bias) - from tensorrt_llm._torch.modules.fused_moe.triton_dequant_nvfp4 import \ dequant_nvfp4_2d_triton weight_deq = dequant_nvfp4_2d_triton( @@ -2104,7 +2082,8 @@ class MarlinNVFP4LinearMethod(W4A16NVFP4LinearMethod): @staticmethod def is_supported(module: Linear) -> bool: - return (get_sm_version() in (120, 121) + sm_version = get_sm_version() + return ((90 <= sm_version < 100 or sm_version in (120, 121)) and getattr(module, "dtype", None) == torch.bfloat16 and not getattr(module, "use_fused_gemm_allreduce", False) and hasattr(torch.ops.trtllm, "marlin_nvfp4_gemm") @@ -3311,7 +3290,9 @@ def __init__( nvfp4_allowed_backends: List of backends to consider for NVFP4 GEMM auto-selection. Default (via config): ['cutlass', 'cublaslt', 'cuda_core'] - excludes cutedsl for faster build. Add 'cutedsl' for extreme performance at the cost of longer build time. - Valid backends: 'cutlass', 'cublaslt', 'cutedsl', 'cuda_core'. + Valid backends: 'cutlass', 'cublaslt', 'cutedsl', 'cuda_core', 'marlin'. + NVFP4 uses Marlin by default on Hopper. W4A16 BF16 linear + layers use Marlin by default on Hopper and SM120/121. Configure via nvfp4_gemm_config.allowed_backends in extra_llm_api_options.yaml. """ from ..distributed import AllReduce diff --git a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py index 3594777e4cc2..61133312a5a7 100644 --- a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py +++ b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py @@ -25,6 +25,7 @@ from tensorrt_llm._torch.modules.linear import ( Linear, MarlinNVFP4LinearMethod, + NVFP4LinearMethod, W4A16NVFP4LinearMethod, get_quant_method, get_sm_version, @@ -36,14 +37,6 @@ def _run_w4a16_marlin_reference_case(m: int, n: int, k: int) -> None: act, weight, weight_scale, weight_scale_2 = _make_w4a16_nvfp4_case(m, n, k, torch.bfloat16) - expected = torch.ops.trtllm.w4a16_nvfp4_gemm( - act, - weight, - weight_scale, - weight_scale_2, - torch.bfloat16, - bias=None, - ) linear = Linear( k, @@ -58,36 +51,25 @@ def _run_w4a16_marlin_reference_case(m: int, n: int, k: int) -> None: linear.weight_scale.data.copy_(weight_scale) linear.weight_scale_2.data.copy_(weight_scale_2) linear.transform_weights() - actual = linear(act) - torch.testing.assert_close(actual, expected, atol=0.75, rtol=0.02) - - -def _run_w4a16_triton_reference_case(m: int, n: int, k: int) -> None: - act, weight, weight_scale, weight_scale_2 = _make_w4a16_nvfp4_case(m, n, k, torch.float16) - expected = torch.ops.trtllm.w4a16_nvfp4_gemm( - act, - weight, - weight_scale, - weight_scale_2, - torch.float16, - bias=None, - ) - linear = Linear( - k, - n, - bias=False, - dtype=torch.float16, - quant_config=QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4), - reduce_output=False, - ).cuda() - assert type(linear.quant_method) is W4A16NVFP4LinearMethod - linear.weight.data.copy_(weight) - linear.weight_scale.data.copy_(weight_scale) - linear.weight_scale_2.data.copy_(weight_scale_2) - linear.transform_weights() + with patch.object(MarlinNVFP4LinearMethod, "is_supported", return_value=False): + reference = Linear( + k, + n, + bias=False, + dtype=torch.bfloat16, + quant_config=QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4), + reduce_output=False, + ).cuda() + assert type(reference.quant_method) is W4A16NVFP4LinearMethod + reference.weight.data.copy_(weight) + reference.weight_scale.data.copy_(weight_scale) + reference.weight_scale_2.data.copy_(weight_scale_2) + reference.transform_weights() + + expected = reference(act) actual = linear(act) - torch.testing.assert_close(actual, expected, atol=0.08, rtol=0.08) + torch.testing.assert_close(actual, expected, atol=0.75, rtol=0.02) def _make_w4a16_nvfp4_case(m: int, n: int, k: int, dtype: torch.dtype): @@ -121,8 +103,8 @@ def _make_w4a16_nvfp4_case(m: int, n: int, k: int, dtype: torch.dtype): @pytest.mark.skipif( - not torch.cuda.is_available() or get_sm_version() not in (120, 121), - reason="requires CUDA SM120/121", + not torch.cuda.is_available() or get_sm_version() not in (90, 120, 121), + reason="requires CUDA SM90 or SM120/121", ) @pytest.mark.parametrize( "shape", @@ -134,56 +116,126 @@ def _make_w4a16_nvfp4_case(m: int, n: int, k: int, dtype: torch.dtype): (128, 512, 1024), ], ) -def test_w4a16_nvfp4_marlin_bf16_matches_cuda_core(shape): +def test_w4a16_nvfp4_marlin_bf16_matches_triton(shape): m, n, k = shape _run_w4a16_marlin_reference_case(m, n, k) -@pytest.mark.skipif( - not torch.cuda.is_available() or get_sm_version() not in (120, 121), - reason="requires CUDA SM120/121", -) -def test_w4a16_nvfp4_triton_fallback_matches_cuda_core(): - _run_w4a16_triton_reference_case(32, 64, 64) +def test_get_quant_method_returns_w4a16_nvfp4_linear_method(): + quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) + method = get_quant_method(quant_config) -@pytest.mark.skipif( - not torch.cuda.is_available() or get_sm_version() not in (120, 121), - reason="requires CUDA SM120/121", + assert type(method) is W4A16NVFP4LinearMethod + + +@pytest.mark.parametrize( + ("sm_version", "dtype", "expected_backends"), + [ + (90, torch.bfloat16, "marlin"), + (90, torch.float16, "cutlass,cublaslt,cuda_core"), + (120, torch.bfloat16, "cutlass,cublaslt,cuda_core"), + (121, torch.bfloat16, "cutlass,cublaslt,cuda_core"), + ], ) -def test_w4a16_nvfp4_gemm_large_m_chunks_cuda_core(): - m, n, k = 32, 64, 64 - act, weight, weight_scale, weight_scale_2 = _make_w4a16_nvfp4_case(m, n, k, torch.bfloat16) - expected = torch.empty((m, n), device="cuda", dtype=torch.bfloat16) - for start in range(0, m, 16): - stop = min(start + 16, m) - expected[start:stop, :] = torch.ops.trtllm.w4a16_nvfp4_gemm( - act[start:stop, :], - weight, - weight_scale, - weight_scale_2, - torch.bfloat16, - bias=None, +def test_nvfp4_linear_uses_architecture_default_backend(sm_version, dtype, expected_backends): + method = NVFP4LinearMethod() + input_tensor = torch.ones((2, 32), dtype=dtype) + module = SimpleNamespace( + weight=torch.empty((4, 16), dtype=torch.uint8), + weight_scale=torch.empty((128 * 4,), dtype=torch.uint8), + out_features=4, + dtype=dtype, + nvfp4_allowed_backends=["cutlass", "cublaslt", "cuda_core"], + all_reduce=None, + mapping=None, + ) + act_fp4 = torch.empty((2, 16), dtype=torch.uint8) + act_sf = torch.empty((128 * 4,), dtype=torch.uint8) + alpha = torch.ones((1,), dtype=torch.float32) + captured = {} + + def fake_nvfp4_gemm(*args, **kwargs): + captured["allowed_backends"] = kwargs["allowed_backends"] + return torch.ones((2, 4), dtype=dtype) + + with ( + patch.object(method, "_input_prepare", return_value=(act_fp4, act_sf, alpha)), + patch( + "tensorrt_llm._torch.modules.linear.get_sm_version", + return_value=sm_version, + ), + patch( + "torch.ops.trtllm.nvfp4_gemm", + side_effect=fake_nvfp4_gemm, + create=True, + ), + ): + output = method.apply(module, input_tensor, bias=None) + + assert captured["allowed_backends"] == expected_backends + assert output.shape == (2, 4) + + +@pytest.mark.parametrize("sm_version", [90, 120, 121]) +def test_nvfp4_linear_keeps_normal_quant_method(sm_version): + quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4) + + with patch( + "tensorrt_llm._torch.modules.linear.get_sm_version", + return_value=sm_version, + ): + linear = Linear( + 32, + 32, + bias=False, + dtype=torch.bfloat16, + quant_config=quant_config, + reduce_output=False, ) - actual = torch.ops.trtllm.w4a16_nvfp4_gemm( - act, - weight, - weight_scale, - weight_scale_2, - torch.bfloat16, - bias=None, - ) + assert type(linear.quant_method) is NVFP4LinearMethod - torch.testing.assert_close(actual, expected, atol=0.08, rtol=0.08) +def test_nvfp4_linear_hopper_marlin_applies_bias_as_post_op(): + method = NVFP4LinearMethod() + input_tensor = torch.ones((2, 32), dtype=torch.bfloat16) + bias = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.bfloat16) + module = SimpleNamespace( + weight=torch.empty((4, 16), dtype=torch.uint8), + weight_scale=torch.empty((128 * 4,), dtype=torch.uint8), + out_features=4, + dtype=torch.bfloat16, + nvfp4_allowed_backends=["cutlass", "cublaslt", "cuda_core"], + all_reduce=None, + mapping=None, + ) + act_fp4 = torch.empty((2, 16), dtype=torch.uint8) + act_sf = torch.empty((128 * 4,), dtype=torch.uint8) + alpha = torch.ones((1,), dtype=torch.float32) + captured = {} -def test_get_quant_method_returns_w4a16_nvfp4_linear_method(): - quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) + def fake_nvfp4_gemm(*args, **kwargs): + captured.update(kwargs) + return torch.ones((2, 4), dtype=torch.bfloat16) - method = get_quant_method(quant_config) + with ( + patch.object(method, "_input_prepare", return_value=(act_fp4, act_sf, alpha)), + patch( + "tensorrt_llm._torch.modules.linear.get_sm_version", + return_value=90, + ), + patch( + "torch.ops.trtllm.nvfp4_gemm", + side_effect=fake_nvfp4_gemm, + create=True, + ), + ): + output = method.apply(module, input_tensor, bias=bias) - assert type(method) is W4A16NVFP4LinearMethod + assert captured["allowed_backends"] == "marlin" + assert captured["bias"] is None + torch.testing.assert_close(output, torch.ones((2, 4), dtype=torch.bfloat16) + bias) def test_w4a16_nvfp4_mlp_disables_relu2_fp4_fusion_without_input_scale(): @@ -274,41 +326,44 @@ def test_w4a16_nvfp4_linear_uses_high_precision_activation_without_fp4_quantize( module = SimpleNamespace( weight=torch.empty((4, 16), dtype=torch.uint8), weight_scale=torch.empty((128 * 4,), dtype=torch.uint8), + _w4a16_weight_scale_linear=torch.empty((128 * 4,), dtype=torch.uint8), weight_scale_2=torch.tensor([0.25], dtype=torch.float32), dtype=torch.bfloat16, out_features=4, + scaling_vector_size=16, pre_quant_scale=None, + use_custom_cublas_mm=False, ) captured = {} - def fake_w4a16_nvfp4_gemm( - input_arg, weight, weight_scale, weight_scale_2, out_dtype, bias=None - ): - captured["input"] = input_arg + def fake_dequant(weight, weight_scale, weight_scale_2, **kwargs): captured["weight"] = weight captured["weight_scale"] = weight_scale captured["weight_scale_2"] = weight_scale_2 - captured["out_dtype"] = out_dtype - captured["bias"] = bias - return torch.ones((input_arg.shape[0], weight.shape[0]), dtype=out_dtype) + captured.update(kwargs) + return torch.ones((4, 32), dtype=torch.bfloat16) def fail_fp4_quantize(*args, **kwargs): raise AssertionError("W4A16 NVFP4 must not quantize activations") with ( - patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=120), - patch("torch.ops.trtllm.w4a16_nvfp4_gemm", side_effect=fake_w4a16_nvfp4_gemm, create=True), + patch( + "tensorrt_llm._torch.modules.fused_moe.triton_dequant_nvfp4.dequant_nvfp4_2d_triton", + side_effect=fake_dequant, + ), patch("torch.ops.trtllm.fp4_quantize", side_effect=fail_fp4_quantize, create=True), ): output = method.apply(module, input_tensor, bias) - assert captured["input"] is input_tensor - assert captured["weight"] is module.weight - assert captured["weight_scale"] is module.weight_scale + assert captured["weight"].data_ptr() == module.weight.data_ptr() + assert captured["weight_scale"] is module._w4a16_weight_scale_linear assert captured["weight_scale_2"] is module.weight_scale_2 - assert captured["out_dtype"] is torch.bfloat16 - assert captured["bias"] is None - expected = torch.tensor([[2.0, 3.0, 4.0, 5.0], [2.0, 3.0, 4.0, 5.0]], dtype=torch.bfloat16) + assert captured["target_dtype"] is torch.bfloat16 + assert captured["sf_vec_size"] == 16 + expected = torch.tensor( + [[33.0, 34.0, 35.0, 36.0], [33.0, 34.0, 35.0, 36.0]], + dtype=torch.bfloat16, + ) torch.testing.assert_close(output, expected) @@ -318,30 +373,30 @@ def test_w4a16_nvfp4_linear_restores_high_rank_input_shape(): module = SimpleNamespace( weight=torch.empty((8, 16), dtype=torch.uint8), weight_scale=torch.empty((128 * 4,), dtype=torch.uint8), + _w4a16_weight_scale_linear=torch.empty((128 * 4,), dtype=torch.uint8), weight_scale_2=torch.tensor([0.5], dtype=torch.float32), dtype=torch.float16, out_features=8, + scaling_vector_size=16, pre_quant_scale=None, + use_custom_cublas_mm=False, ) - def fake_w4a16_nvfp4_gemm( - input_arg, weight, weight_scale, weight_scale_2, out_dtype, bias=None - ): - assert input_arg.shape == (6, 32) - return torch.ones((input_arg.shape[0], weight.shape[0]), dtype=out_dtype) + def fake_dequant(*args, **kwargs): + return torch.ones((8, 32), dtype=torch.float16) - with ( - patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=120), - patch("torch.ops.trtllm.w4a16_nvfp4_gemm", side_effect=fake_w4a16_nvfp4_gemm, create=True), + with patch( + "tensorrt_llm._torch.modules.fused_moe.triton_dequant_nvfp4.dequant_nvfp4_2d_triton", + side_effect=fake_dequant, ): output = method.apply(module, input_tensor, bias=None) assert output.shape == (2, 3, 8) -def test_w4a16_nvfp4_linear_uses_triton_dequant_for_large_m(): +def test_w4a16_nvfp4_linear_uses_triton_dequant(): method = W4A16NVFP4LinearMethod() - input_tensor = torch.ones((17, 32), dtype=torch.bfloat16) + input_tensor = torch.ones((2, 32), dtype=torch.bfloat16) bias = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.bfloat16) module = SimpleNamespace( weight=torch.empty((4, 16), dtype=torch.uint8), @@ -363,15 +418,9 @@ def fake_dequant(weight, weight_scale, weight_scale_2, **kwargs): captured.update(kwargs) return torch.ones((4, 32), dtype=torch.bfloat16) - def fail_w4a16_gemm(*args, **kwargs): - raise AssertionError("large-M W4A16 must use Triton dequantization") - - with ( - patch("torch.ops.trtllm.w4a16_nvfp4_gemm", side_effect=fail_w4a16_gemm, create=True), - patch( - "tensorrt_llm._torch.modules.fused_moe.triton_dequant_nvfp4.dequant_nvfp4_2d_triton", - side_effect=fake_dequant, - ), + with patch( + "tensorrt_llm._torch.modules.fused_moe.triton_dequant_nvfp4.dequant_nvfp4_2d_triton", + side_effect=fake_dequant, ): output = method.apply(module, input_tensor, bias=bias) @@ -380,7 +429,7 @@ def fail_w4a16_gemm(*args, **kwargs): assert captured["weight_scale_2"] is module.weight_scale_2 assert captured["target_dtype"] is torch.bfloat16 assert captured["sf_vec_size"] == 16 - expected = torch.tensor([33.0, 34.0, 35.0, 36.0], dtype=torch.bfloat16).expand(17, 4) + expected = torch.tensor([33.0, 34.0, 35.0, 36.0], dtype=torch.bfloat16).expand(2, 4) torch.testing.assert_close(output, expected) @@ -406,16 +455,12 @@ def fake_marlin_nvfp4_gemm(input_arg, weight, **kwargs): captured.update(kwargs) return torch.ones((input_arg.shape[0], kwargs["size_n"]), dtype=kwargs["out_dtype"]) - def fail_w4a16_gemm(*args, **kwargs): - raise AssertionError("Marlin W4A16 must not call the default W4A16 op") - - with patch("torch.ops.trtllm.w4a16_nvfp4_gemm", side_effect=fail_w4a16_gemm, create=True): - with patch( - "torch.ops.trtllm.marlin_nvfp4_gemm", - side_effect=fake_marlin_nvfp4_gemm, - create=True, - ): - output = method.apply(module, input_tensor, bias=None) + with patch( + "torch.ops.trtllm.marlin_nvfp4_gemm", + side_effect=fake_marlin_nvfp4_gemm, + create=True, + ): + output = method.apply(module, input_tensor, bias=None) assert captured["input"] is input_tensor assert captured["weight"] is module.weight @@ -508,11 +553,15 @@ def test_w4a16_nvfp4_marlin_selection_requires_supported_module(dtype, use_fused assert not MarlinNVFP4LinearMethod.is_supported(module) -def test_w4a16_nvfp4_linear_selects_marlin_for_supported_module(): +@pytest.mark.parametrize("sm_version", [90, 120, 121]) +def test_w4a16_nvfp4_linear_selects_marlin_by_default(sm_version): quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) with ( - patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=120), + patch( + "tensorrt_llm._torch.modules.linear.get_sm_version", + return_value=sm_version, + ), patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), patch("torch.ops.trtllm.gptq_marlin_repack", create=True), ): @@ -529,6 +578,26 @@ def test_w4a16_nvfp4_linear_selects_marlin_for_supported_module(): assert isinstance(linear.quant_method, W4A16NVFP4LinearMethod) +def test_w4a16_nvfp4_linear_uses_default_method_on_sm100(): + quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) + + with ( + patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=100), + patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), + patch("torch.ops.trtllm.gptq_marlin_repack", create=True), + ): + linear = Linear( + 32, + 32, + bias=False, + dtype=torch.bfloat16, + quant_config=quant_config, + reduce_output=False, + ) + + assert type(linear.quant_method) is W4A16NVFP4LinearMethod + + def test_w4a16_nvfp4_linear_keeps_default_method_for_fp16(): quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) @@ -584,7 +653,7 @@ def test_lm_head_uses_w4a16_nvfp4_quant_method_for_packed_lm_head(): assert lm_head.weight_scale_2.shape == (1,) -def test_lm_head_w4a16_nvfp4_forward_dispatches_to_w4a16_op(): +def test_lm_head_w4a16_nvfp4_forward_uses_triton_dequant(): quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) lm_head = LMHead( num_embeddings=4, @@ -593,29 +662,25 @@ def test_lm_head_w4a16_nvfp4_forward_dispatches_to_w4a16_op(): quant_config=quant_config, ) input_tensor = torch.ones((2, 32), dtype=torch.float16) + lm_head._w4a16_weight_scale_linear = torch.empty((128 * 4,), dtype=torch.uint8) captured = {} - def fake_w4a16_nvfp4_gemm( - input_arg, weight, weight_scale, weight_scale_2, out_dtype, bias=None - ): - captured["input"] = input_arg + def fake_dequant(weight, weight_scale, weight_scale_2, **kwargs): captured["weight"] = weight captured["weight_scale"] = weight_scale captured["weight_scale_2"] = weight_scale_2 - captured["out_dtype"] = out_dtype - captured["bias"] = bias - return torch.ones((input_arg.shape[0], weight.shape[0]), dtype=out_dtype) + captured.update(kwargs) + return torch.ones((4, 32), dtype=torch.float16) - with ( - patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=120), - patch("torch.ops.trtllm.w4a16_nvfp4_gemm", side_effect=fake_w4a16_nvfp4_gemm, create=True), + with patch( + "tensorrt_llm._torch.modules.fused_moe.triton_dequant_nvfp4.dequant_nvfp4_2d_triton", + side_effect=fake_dequant, ): output = lm_head(input_tensor) - assert captured["input"] is input_tensor - assert captured["weight"] is lm_head.weight - assert captured["weight_scale"] is lm_head.weight_scale + assert captured["weight"].data_ptr() == lm_head.weight.data_ptr() + assert captured["weight_scale"] is lm_head._w4a16_weight_scale_linear assert captured["weight_scale_2"] is lm_head.weight_scale_2 - assert captured["out_dtype"] is torch.float16 - assert captured["bias"] is None + assert captured["target_dtype"] is torch.float16 + assert captured["sf_vec_size"] == 16 assert output.shape == (2, 4) From 4e8170a2d2dfe5d2194c1afffc36d2bd94147d02 Mon Sep 17 00:00:00 2001 From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> Date: Mon, 3 Aug 2026 05:39:46 +0000 Subject: [PATCH 17/25] remove cuda core nvfp4 Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> --- cpp/tensorrt_llm/nanobind/bindings.cpp | 1 + cpp/tensorrt_llm/thop/cublasScaledMMLut.h | 19 - .../_torch/custom_ops/torch_custom_ops.py | 9 +- tensorrt_llm/_torch/model_config.py | 46 ++- .../_torch/models/modeling_deepseekv3.py | 5 +- .../_torch/models/modeling_deepseekv4.py | 2 +- .../_torch/models/modeling_exaone_moe.py | 2 +- tensorrt_llm/_torch/models/modeling_glm.py | 2 +- .../_torch/models/modeling_nemotron_h.py | 31 +- .../_torch/models/modeling_qwen3_moe.py | 3 +- .../_torch/models/modeling_qwen3_next.py | 5 +- tensorrt_llm/_torch/modules/attention.py | 17 +- .../modules/fused_moe/configurable_moe.py | 18 +- .../_torch/modules/fused_moe/create_moe.py | 28 +- .../_torch/modules/fused_moe/interface.py | 5 + .../_torch/modules/fused_moe/quantization.py | 8 +- tensorrt_llm/_torch/modules/gated_mlp.py | 14 +- tensorrt_llm/_torch/modules/linear.py | 86 +++-- tensorrt_llm/_torch/modules/mlp.py | 21 +- .../_torch/pyexecutor/mamba_cache_manager.py | 9 +- tensorrt_llm/llmapi/llm_utils.py | 19 +- tensorrt_llm/models/quant_config_utils.py | 2 +- .../test_modeling_nemotron_h_moe_quant.py | 45 ++- .../_torch/modules/moe/test_moe_backend.py | 73 +++- .../_torch/modules/test_w4a16_nvfp4_linear.py | 343 +++++++++++++++++- tests/unittest/_torch/test_model_config.py | 112 ++++++ .../parallel/test_dense_gemm_act_fusion.py | 9 +- .../llmapi/test_kv_cache_dtype_override.py | 39 ++ .../models/test_quant_config_utils.py | 1 - 29 files changed, 830 insertions(+), 144 deletions(-) diff --git a/cpp/tensorrt_llm/nanobind/bindings.cpp b/cpp/tensorrt_llm/nanobind/bindings.cpp index e0e670bb483c..a2054dbd7217 100644 --- a/cpp/tensorrt_llm/nanobind/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/bindings.cpp @@ -277,6 +277,7 @@ NB_MODULE(TRTLLM_NB_MODULE, m) .def_prop_ro("has_w4a8_mxfp4_mxfp8", &tc::QuantMode::hasW4a8Mxfp4Mxfp8) .def_prop_ro("has_w4a16_mxfp4", &tc::QuantMode::hasW4a16Mxfp4) + .def_prop_ro("has_kv_cache_quant", &tc::QuantMode::hasKvCacheQuant) .def_static("from_description", &tc::QuantMode::fromDescription, nb::arg("quantize_weights"), nb::arg("quantize_activations"), nb::arg("per_token"), nb::arg("per_channel"), nb::arg("per_group"), diff --git a/cpp/tensorrt_llm/thop/cublasScaledMMLut.h b/cpp/tensorrt_llm/thop/cublasScaledMMLut.h index 534edda011d6..6c1471f70ebb 100644 --- a/cpp/tensorrt_llm/thop/cublasScaledMMLut.h +++ b/cpp/tensorrt_llm/thop/cublasScaledMMLut.h @@ -118,25 +118,6 @@ inline const AlgoListType fp8_algo_list = { // Llama-3.3-70B TP4 (this is the default algo on B200. Here we aim to use the same algo on GB200.) // [-algo66 -m_tile393 -m_stages36 -m_numsK1 -m_reduction0 -m_swizzle0 -m_custom1 -m_mma0 -m_cga4 -m_scheduling1] {{8, 8192, 14336}, {66, 393, 36, 1, 0, 1, 1, 4}}, - - // Qwen3.6-35B-A3B-NVFP4 on GB10 (sm_121), FP8 per-tensor projections. Output is BF16; the FP8 - // branch of find_special_algo is output-dtype-agnostic so these match. Shapes are the ACTUAL - // runtime cuBLASLt inventory from the DEBUG log (1330 calls, all previously missing) -- NOT - // derived from weight headers (TRT-LLM FUSES qkv+z and gate+up, so derived dims were wrong). - // NVJET algos chosen via cublasTest matmulFind + -T200 -Tme20p retime, validated at M=1/4/8. - // Swept against the RUNTIME cuBLASLt (CUDA 13.1, libcublasLt.so.13.2.1.1) -- NVJET tile/custom - // configs are version-specific, valid only for that cuBLASLt generation. Qwen3.6-specific keys. - // [-algo67 -m_tile6 -m_stages38 -m_numsK1 -m_reduction0 -m_swizzle0 -m_custom16 -m_mma0 -m_cga0 -m_scheduling1] - {{8, 2048, 12288}, - {67, 6, 38, 1, 0, 0, 16, 0}}, // GDN in_proj qkv+z fused (x210 in log): heur 61.5->51.3us (robust across M) - // [-algo67 -m_tile3 -m_stages38 -m_numsK1 -m_reduction0 -m_swizzle0 -m_custom114 -m_mma0 -m_cga0 -m_scheduling1] - {{8, 2048, 256}, {67, 3, 38, 1, 0, 0, 114, 0}}, // k/v proj n=256 (x280 in log): heur 8.35->6.08us (M1) - // [-algo67 -m_tile10 -m_stages37 -m_numsK1 -m_reduction0 -m_swizzle0 -m_custom28 -m_mma0 -m_cga0 -m_scheduling1] - {{8, 2048, 9216}, - {67, 10, 37, 1, 0, 0, 28, 0}}, // shared_expert gate+up fused (x70 in log): heur 49/110->24.8us (M1/-80% M4) - // [-algo67 -m_tile10 -m_stages36 -m_numsK1 -m_reduction0 -m_swizzle0 -m_custom29 -m_mma0 -m_cga0 -m_scheduling1] - {{8, 4096, 2048}, - {67, 10, 36, 1, 0, 0, 29, 0}}, // out_proj/down_proj K=4096 (x280 in log): heur 22.5->12.3us (-45%) }; } // namespace cublas_lut diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index f3b0dc0476d3..f404bfb800f8 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -1024,8 +1024,7 @@ def get_valid_tactics(self, inputs: List[torch.Tensor], tactics = [] act_fp4, weight, act_sf, weight_scale, alpha = inputs - # Add Marlin tactics (SM90 Hopper only) — users must opt-in explicitly - # by listing "marlin" in ``allowed_backends``. + # Add Marlin tactics (SM90 Hopper only) when selected by the caller. if self._is_backend_allowed("marlin"): marlin_runner = MarlinNVFP4Runner(self.output_buffer_kind, self.output_dtype) @@ -1153,8 +1152,8 @@ def forward( ) -> torch.Tensor: # Handle fallback tactic on cache miss if tactic == -1: - # Prefer marlin on Hopper (SM90) when explicitly allowed, cutlass - # otherwise, falling back to whatever backend is available. + # Prefer Marlin on Hopper (SM90) when allowed, Cutlass otherwise, + # falling back to whichever backend is available. assert len( self.allowed_backends) > 0, "No allowed backends available" sm_version = get_sm_version() @@ -1236,7 +1235,7 @@ def nvfp4_gemm( allowed_backends: Comma-separated list of backends to consider for auto-selection. Default: "cutlass,cublaslt,cuda_core" (excludes cutedsl for faster build) Add 'cutedsl' for extreme performance at the cost of longer build time. - Valid backends: 'cutlass', 'cublaslt', 'cutedsl', 'cuda_core'. + Valid backends: 'cutlass', 'cublaslt', 'cutedsl', 'cuda_core', 'marlin'. Returns: Output tensor [m, n] with dtype=output_dtype diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 36b5d818994a..062d26f33083 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -349,6 +349,12 @@ def resolve_moe_backend( if moe_backend.upper() != "AUTO": return moe_backend + if architecture in _DEEPSEEK_V4_ARCHITECTURES: + sm_version = get_sm_version() + if 100 <= sm_version < 120: + return "TRTLLM" + return "CUTLASS" + is_w4a16_nvfp4 = (quant_config is not None and quant_config.quant_algo in (QuantAlgo.W4A16_NVFP4, "W4A16_NVFP4")) if not is_w4a16_nvfp4 and layer_quant_config is not None: @@ -358,11 +364,6 @@ def resolve_moe_backend( if is_w4a16_nvfp4 and get_sm_version() in (120, 121): return "CUTEDSL" - if architecture in _DEEPSEEK_V4_ARCHITECTURES: - sm_version = get_sm_version() - if 100 <= sm_version < 120: - return "TRTLLM" - if architecture == "GptOssForCausalLM": sm_version = get_sm_version() # Select the best performing backend based on SM version @@ -382,17 +383,21 @@ def resolve_moe_backend( return "CUTLASS" @staticmethod - def load_modelopt_quant_config(quant_config_file, checkpoint_dir, - moe_backend): + def load_modelopt_quant_config(quant_config_file, + checkpoint_dir, + moe_backend, + hf_quant_config=None): with open(quant_config_file) as f: quant_config_dict = json.load(f) return ModelConfig._build_modelopt_quant_config( read_modelopt_quant_config(quant_config_dict), checkpoint_dir, - moe_backend) + moe_backend, hf_quant_config) @staticmethod - def _build_modelopt_quant_config(json_quant_configs, checkpoint_dir, - moe_backend): + def _build_modelopt_quant_config(json_quant_configs, + checkpoint_dir, + moe_backend, + hf_quant_config=None): """Build (quant_config, layer_quant_config) from a normalized modelopt 'quantization' inner dict. ``json_quant_configs`` should be a dict as produced by @@ -417,6 +422,15 @@ def _build_modelopt_quant_config(json_quant_configs, checkpoint_dir, if 'pre_quant_scale' in json_quant_configs: quant_config.pre_quant_scale = json_quant_configs['pre_quant_scale'] + if (quant_config.quant_algo == QuantAlgo.NVFP4 + and hf_quant_config is not None and + hf_quant_config.get("quant_method") == "compressed-tensors"): + inline_quant_config = quant_config.model_copy(deep=True) + update_quant_config_from_compressed_tensors(inline_quant_config, + hf_quant_config) + if inline_quant_config.quant_algo == QuantAlgo.W4A16_NVFP4: + quant_config = inline_quant_config + if quant_config.quant_algo == QuantAlgo.MIXED_PRECISION: json_extended_quant_configs: dict = {} # See tests/unittest/llmapi/test_llm_quant.py @@ -1152,6 +1166,8 @@ def _recursive_update_config(config: transformers.PretrainedConfig, 'hf_quant_config.json'): with open(quant_config_file) as f: normalized = read_modelopt_quant_config(json.load(f)) + modelopt_declares_quant_algo = normalized.get( + "quant_algo") is not None # The file is authoritative; warn if the inline copy disagrees. # Done before _build_modelopt_quant_config since the builder may # mutate ``normalized`` via ``.update`` from quant_cfg.json. @@ -1165,10 +1181,16 @@ def _recursive_update_config(config: transformers.PretrainedConfig, cls._has_deepseek_v4_layer_only_modelopt_quant_config( quant_config_file)) quant_config, layer_quant_config = cls._build_modelopt_quant_config( - normalized, checkpoint_dir, moe_backend_hint) + normalized, + checkpoint_dir, + moe_backend_hint, + hf_quant_config=getattr(pretrained_config, + "quantization_config", None)) hf_quant_config = getattr(pretrained_config, "quantization_config", None) - if quant_config.quant_algo is None and hf_quant_config is not None: + if (quant_config.quant_algo is None + and not modelopt_declares_quant_algo + and hf_quant_config is not None): hf_quant_config, hf_layer_quant_config = cls.load_hf_quant_config( hf_quant_config, moe_backend_hint, diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index c95aac6c8147..a08926786773 100755 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -1530,7 +1530,8 @@ def _run_MoE(hidden_states, hidden_states_fp4, do_finalize): not (self.fusion_config.POST_MOE_FUSION and hidden_states.shape[0] <= self.moe_allreduce.max_token and self.model_config.moe_backend == "TRTLLM" - and self.mlp.experts.has_nvfp4 and self.is_p2p_supported)) + and self.mlp.experts.has_nvfp4_activation_quantization + and self.is_p2p_supported)) hidden_states = _run_MoE(hidden_states, hidden_states_fp4=None, @@ -1602,7 +1603,7 @@ def forward_mlp( ) -> Tuple[torch.Tensor, torch.Tensor]: if self.fusion_config.PRE_MLP_FUSION: - if self.mlp.gate_up_proj.has_nvfp4: + if self.mlp.gate_up_proj.has_nvfp4_activation_quantization: act_fp4, act_sf, residual = self.allreduce( hidden_states, all_reduce_params=AllReduceParams( diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 99c12ba07168..a3216ecb5d53 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -2153,7 +2153,7 @@ def _run_MoE(hidden_states, hidden_states_fp4, do_finalize, input_ids): self.fusion_config.POST_MOE_FUSION and hidden_states.shape[0] <= self.moe_allreduce.max_token and self.model_config.moe_backend == "TRTLLM" - and self.mlp.experts.has_nvfp4 + and self.mlp.experts.has_nvfp4_activation_quantization and self.is_p2p_supported ) ) diff --git a/tensorrt_llm/_torch/models/modeling_exaone_moe.py b/tensorrt_llm/_torch/models/modeling_exaone_moe.py index 40ae3653d6e0..d9307e36a85c 100644 --- a/tensorrt_llm/_torch/models/modeling_exaone_moe.py +++ b/tensorrt_llm/_torch/models/modeling_exaone_moe.py @@ -383,7 +383,7 @@ def _run_moe(hidden_states, hidden_states_fp4, do_finalize): self.fusion_config.POST_MOE_FUSION and hidden_states.shape[0] <= self.moe_allreduce.max_token and self.model_config.moe_backend == "TRTLLM" - and self.mlp.experts.has_nvfp4 + and self.mlp.experts.has_nvfp4_activation_quantization and self.is_p2p_supported ) ) diff --git a/tensorrt_llm/_torch/models/modeling_glm.py b/tensorrt_llm/_torch/models/modeling_glm.py index 2572ea548e48..0835319ba54c 100644 --- a/tensorrt_llm/_torch/models/modeling_glm.py +++ b/tensorrt_llm/_torch/models/modeling_glm.py @@ -720,7 +720,7 @@ def _run_MoE(hidden_states, hidden_states_fp4, do_finalize): hidden_states.shape[0] <= self.moe_allreduce.max_token and self.fusion_config.POST_MOE_FUSION and self.model_config.moe_backend == "TRTLLM" - and self.mlp.experts.has_nvfp4 + and self.mlp.experts.has_nvfp4_activation_quantization and self.is_p2p_supported ) ) diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index 22de4f0ba67d..94bfb22bcbcf 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py @@ -34,6 +34,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.lora_helper import LoraConfig from tensorrt_llm.models.modeling_utils import QuantAlgo # noqa: E402 +from tensorrt_llm.models.modeling_utils import QuantConfig from ..attention_backend import AttentionMetadata from ..distributed import AllReduce, AllReduceFusionOp, AllReduceParams @@ -42,6 +43,7 @@ from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding from ..modules.fused_moe import MoEWeightLoadingMode, create_moe +from ..modules.fused_moe.create_moe import _get_layer_quant_config from ..modules.fused_moe.fused_moe_cutlass import CutlassFusedMoE from ..modules.fused_moe.quantization import (NVFP4CutlassFusedMoEMethod, W4A16NVFP4CutlassFusedMoEMethod) @@ -162,30 +164,26 @@ def forward( **kwargs) -def _get_nemotron_h_moe_model_config( +def _get_nemotron_h_moe_quant_config( model_config: ModelConfig[PretrainedConfig], - layer_idx: int) -> ModelConfig[PretrainedConfig]: + layer_idx: int) -> QuantConfig | None: # Per-expert mixed precision config is more specific than the global config. - if model_config.quant_config_dict is not None: - experts_prefix = f"model.layers.{layer_idx}.mixer.experts." - for key, cfg in model_config.quant_config_dict.items(): - if key.startswith(experts_prefix): - return replace(model_config, quant_config=cfg) + layer_quant_config = _get_layer_quant_config(model_config, layer_idx, + "mixer.experts") + if layer_quant_config is not None: + return layer_quant_config quant_config = model_config.quant_config if (quant_config is not None and quant_config.quant_algo in (QuantAlgo.W4A16_NVFP4, "W4A16_NVFP4") and model_config.moe_backend.upper() != "CUTEDSL"): - moe_quant_config = quant_config.model_copy(deep=True) - moe_quant_config.quant_algo = QuantAlgo.NVFP4 - moe_quant_config.__dict__.pop("quant_mode", None) - moe_quant_config.__dict__.pop("layer_quant_mode", None) - return replace(model_config, quant_config=moe_quant_config) + values = quant_config.model_dump() + values["quant_algo"] = QuantAlgo.NVFP4 + return QuantConfig.model_validate(values) - return model_config + return None -# Ref code: https://huggingface.co/nvidia/Nemotron-Nano-3-30B-A3.5B-dev-1024/blob/main/modeling_nemotron_h.py#L818 class NemotronHMOE(nn.Module): def __init__( @@ -270,7 +268,7 @@ def _moe(name): moe_backend=model_config.moe_backend, ) - moe_model_config = _get_nemotron_h_moe_model_config( + moe_quant_config = _get_nemotron_h_moe_quant_config( model_config, layer_idx) # Setup MoE experts. @@ -282,7 +280,8 @@ def _moe(name): aux_stream_dict=aux_stream_dict, dtype=config.torch_dtype, reduce_results=self.reduce_results, - model_config=moe_model_config, + model_config=model_config, + override_quant_config=moe_quant_config, layer_idx=self.layer_idx, weight_loading_mode=MoEWeightLoadingMode.VANILLA, bias=self.mlp_bias, diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_moe.py b/tensorrt_llm/_torch/models/modeling_qwen3_moe.py index 571e3fe503c0..601098acf402 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_moe.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_moe.py @@ -263,7 +263,8 @@ def forward( self.fusion_config.POST_MOE_FUSION and hidden_states.shape[0] <= self.moe_allreduce.max_token and self.model_config.moe_backend == 'TRTLLM' - and self.mlp.experts.has_nvfp4 and self.is_p2p_supported) + and self.mlp.experts.has_nvfp4_activation_quantization + and self.is_p2p_supported) hidden_states = self.mlp( hidden_states, diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_next.py b/tensorrt_llm/_torch/models/modeling_qwen3_next.py index 15047b48a4d5..0fea1ba422ca 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_next.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_next.py @@ -379,9 +379,8 @@ def __init__(self, mlp: GatedMLP, mapping): self.mlp = mlp self.mapping = mapping self.enable_attention_dp = mapping.enable_attention_dp - # Provide a dummy `experts` attribute so that - # `self.mlp.experts.has_nvfp4` checks in decoder forward don't crash. - self.experts = SimpleNamespace(has_nvfp4=False) + # Match the activation-quantization capability exposed by MoE experts. + self.experts = SimpleNamespace(has_nvfp4_activation_quantization=False) def forward( self, diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index e902a892af3e..3ec3fbe71c39 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -23,8 +23,9 @@ from ..model_config import ModelConfig from ..peft.lora.layer import LoraLayer, LoraModuleType from ..utils import (Fp4QuantizedTensor, get_model_extra_attrs, - is_nvfp4_marlin_enabled, is_torch_compiling) -from .linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig + is_torch_compiling) +from .linear import (Linear, TensorParallelMode, WeightMode, + WeightsLoadingConfig, is_static_nvfp4_input_eligible) from .multi_stream_utils import maybe_execute_in_parallel from .rotary_embedding import MRotaryEmbedding, RotaryEmbedding @@ -606,7 +607,7 @@ def __init__( attn_cls = get_attention_backend(self.attn_backend, sparse_params=sparse_params) - self.is_marlin_enabled: bool = is_nvfp4_marlin_enabled() + self.is_marlin_enabled = False # These two modules are mutually exclusive - either splitted_qkv_lora or fused_qkv_lora will be used, # but never both at the same time. splitted_qkv_lora handles Q,K,V separately while fused_qkv_lora @@ -689,7 +690,9 @@ def create_weights(self): self.attn.update_quant_config(self.quant_config) self.o_proj.create_weights() - self.has_quant_scale = (self.o_proj.has_fp8_qdq or self.o_proj.has_nvfp4 + self.is_marlin_enabled = self.o_proj.uses_marlin_nvfp4 + self.has_quant_scale = (self.o_proj.has_fp8_qdq + or self.o_proj.has_nvfp4_activation_quantization or self.o_proj.has_fp8_block_scales or self.o_proj.has_fp8_rowwise or self.o_proj.has_w4a8_nvfp4_fp8) @@ -754,6 +757,12 @@ def _use_quantize_output(self): if self.o_proj.force_dynamic_quantization: return False + # Producing FP4 output requires a calibrated activation scale on the + # consumer. Weight-only W4A16 has NVFP4 weights but consumes BF16/FP16. + if (self.o_proj.has_nvfp4 + and not is_static_nvfp4_input_eligible(self.o_proj)): + return False + # If no quant is applied, no need to quantize the output if self.quant_config is not None and not self.quant_config.layer_quant_mode.has_any_quant( exclude_kv_cache=True): diff --git a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py index 5bfb368baf7d..029de1d4302b 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py @@ -160,6 +160,11 @@ def __init__( override_quant_config: Optional["QuantConfig"] = None, **kwargs, ): + from tensorrt_llm._torch.modules.fused_moe.create_moe import _get_effective_moe_quant_config + + effective_quant_config = _get_effective_moe_quant_config( + model_config, override_quant_config, layer_idx + ) super().__init__( routing_method=routing_method, num_experts=num_experts, @@ -172,8 +177,7 @@ def __init__( layer_idx=layer_idx, # ConfigurableMoE needs correct layer_idx for EPLB initialization **kwargs, ) - if override_quant_config is not None: - self.quant_config = override_quant_config + self.quant_config = effective_quant_config # Store model_config and aux_stream_dict for later use (e.g., backend setter) self.model_config = model_config @@ -186,7 +190,7 @@ def __init__( self._create_and_sync_backend( model_config=model_config, routing_method=routing_method, - override_quant_config=override_quant_config, + effective_quant_config=effective_quant_config, **kwargs, ) @@ -269,7 +273,7 @@ def _create_and_sync_backend( *, model_config: ModelConfig, routing_method: BaseMoeRoutingMethod, - override_quant_config: Optional["QuantConfig"], + effective_quant_config: Optional["QuantConfig"], **kwargs, ) -> None: """Build the MoE backend, mirror EPLB attrs, then create weights. @@ -295,14 +299,14 @@ def _create_and_sync_backend( model_config, routing_method, self.dtype, - override_quant_config=override_quant_config, + override_quant_config=effective_quant_config, layer_idx=self.layer_idx, ) backend_model_config = model_config - if override_quant_config is not None: + if effective_quant_config is not model_config.quant_config: backend_model_config = copy.deepcopy(model_config) - backend_model_config.quant_config = override_quant_config + backend_model_config.quant_config = effective_quant_config with self._temporarily_skip_weight_creation(backend_model_config): backend = create_moe_backend( diff --git a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py index ee44e46fb771..2876a4922066 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py @@ -31,18 +31,30 @@ def _is_same_or_child_module_path(lhs: str, rhs: str) -> bool: return lhs == rhs or lhs.startswith(f"{rhs}.") or rhs.startswith(f"{lhs}.") -def _get_layer_quant_config(model_config: ModelConfig, - layer_idx: Optional[int]) -> Optional[QuantConfig]: +def _get_layer_quant_config( + model_config: ModelConfig, + layer_idx: Optional[int], + expert_module_suffix: str = "mlp.experts") -> Optional[QuantConfig]: if layer_idx is None or model_config.quant_config_dict is None: return None - moe_module_name = f"model.layers.{layer_idx}.mlp.experts" + moe_module_name = f"model.layers.{layer_idx}.{expert_module_suffix}" for name, quant_config in model_config.quant_config_dict.items(): if _is_same_or_child_module_path(name, moe_module_name): return quant_config return None +def _get_effective_moe_quant_config( + model_config: ModelConfig, + override_quant_config: Optional[QuantConfig] = None, + layer_idx: Optional[int] = None, +) -> Optional[QuantConfig]: + return (override_quant_config + or _get_layer_quant_config(model_config, layer_idx) + or model_config.quant_config) + + def _get_pretrained_megamoe_capability_args( model_config: ModelConfig) -> Dict[str, Optional[object]]: """Extract dtype / hidden / intermediate kwargs for MegaMoE @@ -74,9 +86,9 @@ def get_moe_cls( layer_idx: Optional[int] = None, ) -> Type[MoE]: moe_backend = model_config.moe_backend - quant_config = (override_quant_config - or _get_layer_quant_config(model_config, layer_idx) - or model_config.quant_config) + quant_config = _get_effective_moe_quant_config(model_config, + override_quant_config, + layer_idx) layer_prefix = f"[layer_idx={layer_idx}] " if layer_idx is not None else "" if moe_backend.upper() == "MARLIN": # Marlin MoE is a Hopper-specific NVFP4 W4A16 backend. Layers without @@ -254,8 +266,8 @@ def resolve_moe_cls( ) -> Type[MoE]: moe_cls = get_moe_cls(model_config, override_quant_config, layer_idx) - effective_quant_config = (override_quant_config or _get_layer_quant_config( - model_config, layer_idx) or model_config.quant_config) + effective_quant_config = _get_effective_moe_quant_config( + model_config, override_quant_config, layer_idx) has_quant = (effective_quant_config is not None and effective_quant_config.layer_quant_mode.has_any_quant( exclude_kv_cache=True)) diff --git a/tensorrt_llm/_torch/modules/fused_moe/interface.py b/tensorrt_llm/_torch/modules/fused_moe/interface.py index 1e7ecd06845f..dc082eef24c6 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/interface.py +++ b/tensorrt_llm/_torch/modules/fused_moe/interface.py @@ -1078,6 +1078,11 @@ def has_nvfp4(self): return self.quant_config is not None and self.quant_config.layer_quant_mode.has_nvfp4( ) + @property + def has_nvfp4_activation_quantization(self): + assert self._weights_created + return self.quant_method.quantizes_nvfp4_activations + @property def has_w4a8_nvfp4_fp8(self): assert self._weights_created diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index c5c01221e40a..52ddf2cf299e 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py @@ -236,6 +236,9 @@ class FusedMoEMethodBase(ABC): weight_alignment: int = 1 """int: Required byte alignment for MoE weight tensors.""" + quantizes_nvfp4_activations: bool = False + """Whether this method converts high-precision activations to NVFP4.""" + eplb_support_status: EplbSupportStatus = EplbSupportStatus.NOT_SUPPORTED """EplbSupportStatus: Online EPLB support status for this quantization method. @@ -2095,6 +2098,7 @@ class NVFP4FusedMoEMethod(FusedMoEMethodBase): Base class for NVFP4 fused MoE methods for all backends. """ eplb_support_status = EplbSupportStatus.SUPPORTED + quantizes_nvfp4_activations = True # Whether raw per-expert block-scale staging is an EPLB migration # target. Children that migrate derived formats and free the raw @@ -3076,6 +3080,8 @@ class W4A16NVFP4CutlassFusedMoEMethod(NVFP4CutlassFusedMoEMethod): into a static [E_total, N, K] workspace, then runs the bf16 ``fused_moe``. """ + quantizes_nvfp4_activations = False + def process_weights_after_loading(self, module: torch.nn.Module): super().process_weights_after_loading(module) @@ -3084,8 +3090,6 @@ def process_weights_after_loading(self, module: torch.nn.Module): # block_scale_interleave_reverse accepts. def _unswizzle_inplace(scale_param: torch.nn.Parameter): sf_view = scale_param.data.view(float4_sf_dtype) - E, pad_rows, pad_cols = (sf_view.shape[0], sf_view.shape[1], - sf_view.shape[2]) linear = torch.ops.trtllm.block_scale_interleave_reverse(sf_view) scale_param.data.view(float4_sf_dtype).copy_(linear) diff --git a/tensorrt_llm/_torch/modules/gated_mlp.py b/tensorrt_llm/_torch/modules/gated_mlp.py index d5787ccd0044..a9655ce45f2b 100644 --- a/tensorrt_llm/_torch/modules/gated_mlp.py +++ b/tensorrt_llm/_torch/modules/gated_mlp.py @@ -12,7 +12,8 @@ from ..model_config import ModelConfig from ..peft.lora.layer import LoraLayer, LoraModuleType from ..utils import Fp4QuantizedTensor -from .linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig +from .linear import (Linear, TensorParallelMode, WeightMode, + WeightsLoadingConfig, is_static_nvfp4_input_eligible) from .swiglu import swiglu @@ -219,7 +220,8 @@ def _can_fuse_gate_up_swiglu(self): - gate_up_proj has no bias (bias not supported in fused kernel) """ return (self.use_cute_dsl_blockscaling_mm and self.activation == F.silu - and self._is_plain_swiglu() and self.gate_up_proj.has_nvfp4 + and self._is_plain_swiglu() + and self.gate_up_proj.has_nvfp4_activation_quantization and not self.gate_up_proj.has_bias) def _can_fuse_gate_up_swiglu_fp4out(self): @@ -232,13 +234,7 @@ def _can_fuse_gate_up_swiglu_fp4out(self): """ if not self._can_fuse_gate_up_swiglu(): return False - if not self.down_proj.has_nvfp4: - return False - if self.down_proj.force_dynamic_quantization: - return False - if self.down_proj.input_scale is None: - return False - return True + return is_static_nvfp4_input_eligible(self.down_proj) def _fused_gate_up_swiglu(self, x, fp4_out=False): """Fused FC1 GEMM + SwiGLU using CuteDSL dense kernel. diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index a3d27fe53516..dab6cac766a5 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -31,7 +31,6 @@ from ..._utils import get_sm_version, is_sm_100f from ...models.modeling_utils import QuantConfig from ..utils import (Fp4QuantizedTensor, get_model_extra_attrs, - is_nvfp4_marlin_enabled, replace_parameter_and_save_metadata, unswizzle_sf) @@ -99,6 +98,29 @@ def flip(cls, mode): return cls.ROW if mode == cls.COLUMN else cls.COLUMN +def quant_config_has_nvfp4_activation_quantization( + quant_config: Optional[QuantConfig]) -> bool: + """Whether the quantization algorithm converts activations to NVFP4.""" + return (quant_config is not None + and quant_config.layer_quant_mode.has_nvfp4() + and quant_config.quant_algo != QuantAlgo.W4A16_NVFP4) + + +_DEFAULT_NVFP4_ALLOWED_BACKENDS = ('cutlass', 'cublaslt', 'cuda_core') + + +def _uses_hopper_marlin_nvfp4(module) -> bool: + """Whether a regular NVFP4 linear should use Marlin on Hopper.""" + allowed_backends = tuple(getattr(module, "nvfp4_allowed_backends", ())) + marlin_selected = (allowed_backends == _DEFAULT_NVFP4_ALLOWED_BACKENDS + or "marlin" in allowed_backends) + return (marlin_selected and 90 <= get_sm_version() < 100 + and getattr(module, "dtype", None) == torch.bfloat16 + and not getattr(module, "use_fused_gemm_allreduce", False) + and hasattr(torch.ops.trtllm, "marlin_nvfp4_gemm") + and hasattr(torch.ops.trtllm, "gptq_marlin_repack")) + + def load_weight_shard( weight, tensor_parallel_size: int = 1, @@ -354,6 +376,7 @@ class LinearMethodBase(ABC): # window buffer. apply() reads this ClassVar to derive output_buffer_kind # internally; callers do not pass output_buffer_kind as a parameter. supports_nccl_symmetric_memory_window_output: ClassVar[bool] = False + quantizes_nvfp4_activations: ClassVar[bool] = False @abstractmethod def create_weights(self, module: Linear, in_features: int, @@ -1338,6 +1361,7 @@ def transform_weights(self, module: Linear) -> None: class NVFP4LinearMethod(LinearMethodBase): supports_nccl_symmetric_memory_window_output: ClassVar[bool] = True + quantizes_nvfp4_activations: ClassVar[bool] = True # Temporary workaround which will be resolved by TRTLLM-11958 # When True, use tunable_fp4_quantize (AutoTuner selects TRTLLM vs @@ -1489,8 +1513,7 @@ def apply(self, module: Linear, input: torch.Tensor, # Use unified interface - supports CUTLASS, cuBLASLt, CuteDSL # Convert list to comma-separated string for torch.compile compatibility - sm_version = get_sm_version() - use_marlin = (90 <= sm_version < 100 and module.dtype == torch.bfloat16) + use_marlin = _uses_hopper_marlin_nvfp4(module) allowed_backends_str = ('marlin' if use_marlin else ','.join( module.nvfp4_allowed_backends)) output_buffer_kind = ( @@ -1977,6 +2000,8 @@ def transform_weights(self, module: Linear) -> None: class W4A16NVFP4LinearMethod(NVFP4LinearMethod): """W4A16 NVFP4 linear using on-the-fly weight dequantization.""" + quantizes_nvfp4_activations: ClassVar[bool] = False + def create_weights(self, module: Linear, in_features: int, out_features: int, bias: bool, dtype: torch.dtype): module.scaling_vector_size = 16 @@ -2070,9 +2095,13 @@ def cache_derived_state(self, module: Linear) -> None: module.in_features // module.scaling_vector_size, 4) scale_swizzled = module.weight_scale.data.view( fp4_utils.float4_sf_dtype).reshape(pad_rows, pad_cols) - module._w4a16_weight_scale_linear = ( - torch.ops.trtllm.block_scale_interleave_reverse( - scale_swizzled).reshape(-1)) + scale_linear = torch.ops.trtllm.block_scale_interleave_reverse( + scale_swizzled).reshape(-1) + buffer_name = "_w4a16_weight_scale_linear" + if buffer_name in module._buffers: + module._buffers[buffer_name] = scale_linear + else: + module.register_buffer(buffer_name, scale_linear, persistent=False) @staticmethod def _prepare_input(module: Linear, input: torch.Tensor): @@ -3312,10 +3341,7 @@ def get_quant_method(quant_config: Optional[QuantConfig] = None): if quant_config.layer_quant_mode.has_nvfp4(): if quant_config.quant_algo == QuantAlgo.NVFP4_ARC: return NVFP4ARCLinearMethod() - elif is_nvfp4_marlin_enabled(): - return MarlinNVFP4LinearMethod() - else: - return NVFP4LinearMethod() + return NVFP4LinearMethod() if quant_config.layer_quant_mode.has_w4a8_nvfp4_fp8(): return W4A8NVFP4FP8LinearMethod() if quant_config.layer_quant_mode.has_w4a8_mxfp4_fp8(): @@ -3486,8 +3512,8 @@ def __init__( in_features_aligned = self.in_features % 128 == 0 out_features_aligned = self.out_features % 64 == 0 tp_valid = self.tp_mode is not None and self.tp_mode == TensorParallelMode.ROW and self.tp_size > 1 - quant_valid = self.quant_config is not None and self.quant_config.layer_quant_mode.has_nvfp4( - ) + quant_valid = quant_config_has_nvfp4_activation_quantization( + self.quant_config) device_supported = get_sm_version() >= 100 enable_gemm_allreduce_fusion_env = (os.environ.get( @@ -3514,8 +3540,8 @@ def __init__( def get_quant_method(self, quant_config: Optional[QuantConfig] = None): quant_method = get_quant_method(quant_config) - if (type(quant_method) is W4A16NVFP4LinearMethod - and MarlinNVFP4LinearMethod.is_supported(self)): + use_marlin = type(quant_method) is W4A16NVFP4LinearMethod + if use_marlin and MarlinNVFP4LinearMethod.is_supported(self): return MarlinNVFP4LinearMethod() return quant_method @@ -3683,6 +3709,12 @@ def create_weights(self): if self._weights_created: return + # Mixed-precision loading may replace quant_config after __init__. + # Weight-only W4A16 must not retain the activation-quantized fused path. + if not quant_config_has_nvfp4_activation_quantization( + self.quant_config): + self.use_fused_gemm_allreduce = False + self.rebuild_tensor_metadata = {} self.quant_method = self.get_quant_method(self.quant_config) @@ -3723,6 +3755,18 @@ def has_nvfp4(self): return self.quant_config is not None and self.quant_config.layer_quant_mode.has_nvfp4( ) + @property + def has_nvfp4_activation_quantization(self): + assert self._weights_created + return self.quant_method.quantizes_nvfp4_activations + + @property + def uses_marlin_nvfp4(self): + assert self._weights_created + return (isinstance(self.quant_method, MarlinNVFP4LinearMethod) + or (type(self.quant_method) is NVFP4LinearMethod + and _uses_hopper_marlin_nvfp4(self))) + @property def has_weight_only_quant(self): assert self._weights_created @@ -3895,16 +3939,16 @@ def is_static_nvfp4_input_eligible(linear) -> bool: """Whether `linear` consumes a static (calibrated) NVFP4 input, making it eligible to have its input-quantize folded into a producing RMSNorm. - Eligible iff the Linear has NVFP4 weights, a calibrated (static) - `input_scale`, no AWQ `pre_quant_scale`, and is not forced to dynamic - quantization. This is the single canonical definition shared by every - NVFP4-fold site (the layer-boundary / dense folds in modeling_deepseekv3.py - and the q_a_layernorm -> q_b_proj fold in attention.py's MLA) so the gate - cannot drift between them. + Eligible iff the Linear quantizes activations to NVFP4, has a calibrated + (static) `input_scale`, has no AWQ `pre_quant_scale`, and is not forced to + dynamic quantization. This is the single canonical definition shared by + every NVFP4-fold site (the layer-boundary / dense folds in + modeling_deepseekv3.py and the q_a_layernorm -> q_b_proj fold in + attention.py's MLA) so the gate cannot drift between them. """ if linear is None: return False - return (getattr(linear, "has_nvfp4", False) + return (getattr(linear, "has_nvfp4_activation_quantization", False) and not getattr(linear, "force_dynamic_quantization", False) and getattr(linear, "input_scale", None) is not None and getattr(linear, "pre_quant_scale", None) is None) diff --git a/tensorrt_llm/_torch/modules/mlp.py b/tensorrt_llm/_torch/modules/mlp.py index de14c160dea1..42aae9956219 100644 --- a/tensorrt_llm/_torch/modules/mlp.py +++ b/tensorrt_llm/_torch/modules/mlp.py @@ -135,11 +135,14 @@ def forward( # may be downgraded to unquantized after create_weights (e.g. LTX-2 # quant-exclusion), so re-check the NVFP4 _input_prepare at runtime (a # torch.compile trace-time guard, not a per-step cost); else fall back to eager. - if self._use_fused_gelu and hasattr( - getattr(self.up_proj, "quant_method", None), "_input_prepare"): - if self._use_fused_gelu_fp4out and hasattr( - getattr(self.down_proj, "quant_method", None), - "_input_prepare"): + if (self._use_fused_gelu + and self.up_proj.has_nvfp4_activation_quantization + and hasattr(getattr(self.up_proj, "quant_method", None), + "_input_prepare")): + if (self._use_fused_gelu_fp4out + and is_static_nvfp4_input_eligible(self.down_proj) + and hasattr(getattr(self.down_proj, "quant_method", None), + "_input_prepare")): m = self._token_count(x) return self.down_proj( self._fused_gelu(x, fp4_out=m >= MLP._FP4OUT_MIN_M)) @@ -168,16 +171,14 @@ def _gelu_fusion_eligibility(self) -> Tuple[bool, bool]: applied in forward (quant_method can be downgraded after this). """ if (self.activation is not gelu_tanh - or get_sm_version() not in (100, 103) - or not getattr(self.up_proj, "has_nvfp4", False)): + or get_sm_version() not in (100, 103) or not getattr( + self.up_proj, "has_nvfp4_activation_quantization", False)): return False, False bf16_ok = hasattr(torch.ops.trtllm, "cute_dsl_nvfp4_dense_gemm_gelu_blackwell") fp4_ok = (bf16_ok and hasattr( torch.ops.trtllm, "cute_dsl_nvfp4_dense_gemm_gelu_fp4out_blackwell") - and getattr(self.down_proj, "has_nvfp4", False) - and not self.down_proj.force_dynamic_quantization - and self.down_proj.input_scale is not None) + and is_static_nvfp4_input_eligible(self.down_proj)) return bf16_ok, fp4_ok @staticmethod diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index e8f4b33fe51e..a670074dd23b 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -2297,10 +2297,11 @@ def get_num_available_tokens(self, stats = self.impl.get_kv_cache_stats() rs_free = stats.num_free_blocks_per_window_size.get( LinearCacheType.RECURRENT_STATES.value, 0) - # Each recurrent-state block backs one context stop position. - # For prompts shorter than the snapshot interval, the only stop - # position is the final/live state, so one free block is enough. - rs_token_cap = max(0, rs_free) * interval + # Reserve 1 block for the always-allocated last block (corner case + # / final live state) so we don't promise more tokens than the + # pool can actually back at allocation time. + usable_rs_blocks = max(0, rs_free - 1) + rs_token_cap = usable_rs_blocks * interval result = min(result, rs_token_cap) return max(result, 0) diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index 63b360f5584b..07fdd41acc22 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -160,11 +160,14 @@ def _update_from_hf_quant_config(self) -> bool: # Cross-check against inline config.json.quantization_config if any. # Done before _apply_modelopt_quant_config since the apply step # mutates ``normalized`` via ``.pop()``. + inline_quant_config = None try: with open(f"{self._model_dir}/config.json", "r") as f: + inline_quant_config = json.load(f).get( + "quantization_config") warn_if_inline_diverges( normalized, - json.load(f).get("quantization_config"), + inline_quant_config, source_file="hf_quant_config.json", ) except FileNotFoundError: @@ -180,8 +183,22 @@ def _update_from_hf_quant_config(self) -> bool: raise ValueError( "Pre-quantized checkpoint must have quant_algo.") else: + normalized_quant_algo = normalized.get("quant_algo") self._apply_modelopt_quant_config(normalized, explicit_kv_cache_quant_algo) + if (normalized_quant_algo == QuantAlgo.NVFP4 + and inline_quant_config is not None + and inline_quant_config.get("quant_method") + == "compressed-tensors"): + effective_quant_config = self.llm_args.quant_config.model_copy( + deep=True) + update_quant_config_from_compressed_tensors( + effective_quant_config, inline_quant_config) + if effective_quant_config.quant_algo == QuantAlgo.W4A16_NVFP4: + logger.info( + "Using W4A16_NVFP4 activation semantics from the " + "inline compressed-tensors config.") + self.llm_args.quant_config = effective_quant_config return True hf_config_path = f"{self._model_dir}/config.json" diff --git a/tensorrt_llm/models/quant_config_utils.py b/tensorrt_llm/models/quant_config_utils.py index e008833e4d97..fd638423629c 100644 --- a/tensorrt_llm/models/quant_config_utils.py +++ b/tensorrt_llm/models/quant_config_utils.py @@ -38,7 +38,7 @@ def update_quant_config_from_compressed_tensors( ) group_config = next(iter(config_groups.values())) weights_quant_config = group_config["weights"] - inputs_quant_config = group_config["input_activations"] + inputs_quant_config = group_config.get("input_activations") weights_quant_strategy = weights_quant_config["strategy"] if inputs_quant_config is None: diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py index 02c89052775c..1e0adb064b3b 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py @@ -52,12 +52,13 @@ def test_nemotron_h_moe_uses_w4a4_nvfp4_expert_config_for_w4a16_checkpoint(): quant_config = QuantConfig( quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16, exclude_modules=["lm_head"] ) - quant_config.mamba_ssm_cache_dtype = torch.float32 + quant_config.mamba_ssm_cache_dtype = "float32" + _ = quant_config.quant_mode model_config = _make_nemotron_h_moe_config(quant_config) captured = {} def fake_create_moe(**kwargs): - captured["model_config"] = kwargs["model_config"] + captured.update(kwargs) return nn.Identity() with patch( @@ -67,12 +68,14 @@ def fake_create_moe(**kwargs): aux_stream_dict = {AuxStreamType.MoeShared: None} NemotronHMOE(model_config=model_config, layer_idx=1, aux_stream_dict=aux_stream_dict) - moe_quant_config = captured["model_config"].quant_config + moe_quant_config = captured["override_quant_config"] assert moe_quant_config is not quant_config assert moe_quant_config.quant_algo == QuantAlgo.NVFP4 + assert moe_quant_config.quant_mode.has_nvfp4() assert moe_quant_config.group_size == 16 assert moe_quant_config.exclude_modules == ["lm_head"] - assert moe_quant_config.mamba_ssm_cache_dtype == torch.float32 + assert moe_quant_config.mamba_ssm_cache_dtype == "float32" + assert captured["model_config"] is model_config assert model_config.quant_config.quant_algo == QuantAlgo.W4A16_NVFP4 @@ -84,7 +87,7 @@ def test_nemotron_h_moe_preserves_w4a16_config_for_cutedsl_sm12x(): captured = {} def fake_create_moe(**kwargs): - captured["model_config"] = kwargs["model_config"] + captured.update(kwargs) return nn.Identity() with patch( @@ -94,9 +97,35 @@ def fake_create_moe(**kwargs): aux_stream_dict = {AuxStreamType.MoeShared: None} NemotronHMOE(model_config=model_config, layer_idx=1, aux_stream_dict=aux_stream_dict) - moe_quant_config = captured["model_config"].quant_config - assert moe_quant_config is quant_config - assert moe_quant_config.quant_algo == QuantAlgo.W4A16_NVFP4 + assert captured["override_quant_config"] is None + assert captured["model_config"] is model_config + + +def test_nemotron_h_moe_uses_mixer_expert_layer_quant_config(): + global_quant_config = QuantConfig() + layer_quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16) + model_config = _make_nemotron_h_moe_config(global_quant_config) + model_config.quant_config_dict = { + "model.layers.1.mixer.experts.0.up_proj": layer_quant_config, + } + captured = {} + + def fake_create_moe(**kwargs): + captured.update(kwargs) + return nn.Identity() + + with patch( + "tensorrt_llm._torch.models.modeling_nemotron_h.create_moe", + side_effect=fake_create_moe, + ): + with patch("torch.cuda.Event", side_effect=lambda: object()): + NemotronHMOE( + model_config=model_config, + layer_idx=1, + aux_stream_dict={AuxStreamType.MoeShared: None}, + ) + + assert captured["override_quant_config"] is layer_quant_config def test_nemotron_h_mtp_bf16_body_uses_cutlass_moe_backend(): diff --git a/tests/unittest/_torch/modules/moe/test_moe_backend.py b/tests/unittest/_torch/modules/moe/test_moe_backend.py index e8c8d9511e2c..72ad4870bd57 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_moe_backend.py @@ -31,7 +31,7 @@ import os from types import SimpleNamespace from typing import List, Optional -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest import torch @@ -64,9 +64,11 @@ from tensorrt_llm._torch.modules.fused_moe.mega_moe import MegaMoECuteDsl, MegaMoEDeepGemm from tensorrt_llm._torch.modules.fused_moe.quantization import ( FusedMoEMethodBase, + NVFP4FusedMoEMethod, NVFP4MarlinFusedMoEMethod, UnquantizedFusedMoEMethod, W4A8MXFP4MXFP8MegaMoEDeepGemmMethod, + W4A16NVFP4CutlassFusedMoEMethod, ) from tensorrt_llm._torch.utils import ActivationType, is_gated_activation from tensorrt_llm._utils import mpi_rank @@ -325,6 +327,75 @@ def test_configurable_moe_load_weights_invalidates_wrapper_transform_guard(): assert configurable_moe._weights_transformed is False +def test_configurable_moe_backend_uses_effective_layer_quant_config(): + from tensorrt_llm._torch.modules.fused_moe.configurable_moe import ( + _BACKEND_SYNC_ATTRS, + ConfigurableMoE, + ) + + global_quant_config = QuantConfig() + layer_quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16) + model_config = ModelConfig( + quant_config=global_quant_config, + quant_config_dict={ + "model.layers.3.mlp.experts": layer_quant_config, + }, + ) + + configurable_moe = ConfigurableMoE.__new__(ConfigurableMoE) + torch.nn.Module.__init__(configurable_moe) + configurable_moe.dtype = torch.bfloat16 + configurable_moe.layer_idx = 3 + configurable_moe.num_experts = 8 + configurable_moe.hidden_size = 32 + configurable_moe.intermediate_size = 64 + configurable_moe.reduce_results = False + configurable_moe.aux_stream_dict = None + configurable_moe.weight_loading_mode = MoEWeightLoadingMode.VANILLA + configurable_moe.apply_router_weight_on_input = False + configurable_moe.activation_type = int(ActivationType.Swiglu) + for attr in _BACKEND_SYNC_ATTRS: + if not hasattr(configurable_moe, attr): + setattr(configurable_moe, attr, None) + + backend = torch.nn.Module() + backend.create_weights = MagicMock() + captured = {} + + def create_backend(**kwargs): + captured.update(kwargs) + return backend + + with ( + patch( + "tensorrt_llm._torch.modules.fused_moe.create_moe.resolve_moe_cls", + return_value=object, + ) as resolve_moe_cls, + patch( + "tensorrt_llm._torch.modules.fused_moe.create_moe.create_moe_backend", + side_effect=create_backend, + ), + ): + configurable_moe._create_and_sync_backend( + model_config=model_config, + routing_method=MagicMock(), + effective_quant_config=layer_quant_config, + ) + + backend_model_config = captured["model_config"] + assert backend_model_config is not model_config + assert backend_model_config.quant_config is layer_quant_config + assert model_config.quant_config is global_quant_config + resolve_moe_cls.assert_called_once() + assert resolve_moe_cls.call_args.kwargs["override_quant_config"] is layer_quant_config + backend.create_weights.assert_called_once_with() + + +def test_moe_nvfp4_activation_quantization_capability(): + assert NVFP4FusedMoEMethod.quantizes_nvfp4_activations + assert not W4A16NVFP4CutlassFusedMoEMethod.quantizes_nvfp4_activations + + def test_marlin_moe_repack_is_transform_stage(): assert "transform_weights" in NVFP4MarlinFusedMoEMethod.__dict__ assert "post_load_weights" not in NVFP4MarlinFusedMoEMethod.__dict__ diff --git a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py index 61133312a5a7..18bcd4ceaca5 100644 --- a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py +++ b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py @@ -21,17 +21,24 @@ import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM +from tensorrt_llm._torch.modules.attention import Attention from tensorrt_llm._torch.modules.embedding import LMHead +from tensorrt_llm._torch.modules.gated_mlp import GatedMLP from tensorrt_llm._torch.modules.linear import ( Linear, MarlinNVFP4LinearMethod, NVFP4LinearMethod, + TensorParallelMode, W4A16NVFP4LinearMethod, get_quant_method, get_sm_version, + quant_config_has_nvfp4_activation_quantization, ) from tensorrt_llm._torch.modules.mlp import MLP -from tensorrt_llm._torch.utils import relu2 +from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.utils import gelu_tanh, is_nvfp4_marlin_enabled, model_extra_attrs, relu2 +from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig @@ -129,6 +136,147 @@ def test_get_quant_method_returns_w4a16_nvfp4_linear_method(): assert type(method) is W4A16NVFP4LinearMethod +def test_nvfp4_activation_quantization_excludes_w4a16(): + assert quant_config_has_nvfp4_activation_quantization(QuantConfig(quant_algo=QuantAlgo.NVFP4)) + assert not quant_config_has_nvfp4_activation_quantization( + QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) + ) + + +def test_nvfp4_marlin_utility_requires_explicit_opt_in(): + with ( + patch("tensorrt_llm._torch.utils.get_sm_version", return_value=90), + patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), + ): + assert not is_nvfp4_marlin_enabled() + with model_extra_attrs({"nvfp4_gemm_allowed_backends": ["cutlass", "marlin"]}): + assert is_nvfp4_marlin_enabled() + + +def test_nvfp4_rmsnorm_keeps_high_precision_output_for_hopper_marlin(): + with ( + patch("tensorrt_llm._torch.modules.rms_norm.get_sm_version", return_value=90), + patch("tensorrt_llm._torch.modules.rms_norm.is_nvfp4_marlin_enabled", return_value=True), + ): + norm = RMSNorm( + hidden_size=32, + eps=1e-5, + dtype=torch.bfloat16, + quantize_type="nvfp4", + return_hp_output=True, + ) + + assert not norm.is_nvfp4 + assert not norm.return_hp_output + assert norm.nvfp4_scale is None + + +def test_w4a16_attention_does_not_quantize_output_to_fp4(): + quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) + o_proj = Linear( + 32, + 32, + bias=False, + dtype=torch.bfloat16, + quant_config=quant_config, + reduce_output=False, + ) + attention = SimpleNamespace( + attn=SimpleNamespace(has_nvfp4=False), + o_proj=o_proj, + quant_config=quant_config, + has_quant_scale=True, + attn_output_gate=False, + is_marlin_enabled=False, + ) + + assert not Attention._use_quantize_output(attention) + + +def test_static_nvfp4_attention_can_quantize_output_to_fp4_on_blackwell(): + quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4) + with patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=100): + o_proj = Linear( + 32, + 32, + bias=False, + dtype=torch.bfloat16, + quant_config=quant_config, + reduce_output=False, + ) + attention = SimpleNamespace( + attn=SimpleNamespace(has_nvfp4=False), + o_proj=o_proj, + quant_config=quant_config, + has_quant_scale=True, + attn_output_gate=False, + is_marlin_enabled=False, + ) + + assert Attention._use_quantize_output(attention) + + +def test_nvfp4_attention_keeps_high_precision_output_for_hopper_marlin(): + quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4) + with ( + patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=90), + patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), + patch("torch.ops.trtllm.gptq_marlin_repack", create=True), + ): + o_proj = Linear( + 32, + 32, + bias=False, + dtype=torch.bfloat16, + quant_config=quant_config, + reduce_output=False, + ) + attention = SimpleNamespace( + attn=SimpleNamespace(has_nvfp4=False), + o_proj=o_proj, + quant_config=quant_config, + has_quant_scale=o_proj.has_nvfp4_activation_quantization, + attn_output_gate=False, + is_marlin_enabled=o_proj.uses_marlin_nvfp4, + ) + + assert o_proj.uses_marlin_nvfp4 + assert type(o_proj.quant_method) is NVFP4LinearMethod + assert o_proj.has_nvfp4_activation_quantization + assert not Attention._use_quantize_output(attention) + + +def test_w4a16_disables_fused_gemm_allreduce(monkeypatch): + monkeypatch.setenv("TRTLLM_GEMM_ALLREDUCE_FUSION_ENABLED", "1") + mapping = Mapping(world_size=2, rank=0, tp_size=2) + + with ( + patch("tensorrt_llm._torch.modules.linear.mpi_disabled", return_value=False), + patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=120), + patch("tensorrt_llm._torch.modules.linear.ipc_nvls_supported", return_value=True), + patch("tensorrt_llm._torch.distributed.AllReduce"), + ): + linear = Linear( + 256, + 64, + bias=False, + dtype=torch.bfloat16, + mapping=mapping, + tensor_parallel_mode=TensorParallelMode.ROW, + quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4), + reduce_output=True, + skip_create_weights_in_init=True, + ) + assert linear.use_fused_gemm_allreduce + + # Simulate apply_layerwise_quant_config rebinding a mixed-precision + # layer after Linear.__init__ but before deferred weight creation. + linear.quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) + linear.create_weights() + + assert not linear.use_fused_gemm_allreduce + + @pytest.mark.parametrize( ("sm_version", "dtype", "expected_backends"), [ @@ -170,6 +318,8 @@ def fake_nvfp4_gemm(*args, **kwargs): side_effect=fake_nvfp4_gemm, create=True, ), + patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), + patch("torch.ops.trtllm.gptq_marlin_repack", create=True), ): output = method.apply(module, input_tensor, bias=None) @@ -178,12 +328,16 @@ def fake_nvfp4_gemm(*args, **kwargs): @pytest.mark.parametrize("sm_version", [90, 120, 121]) -def test_nvfp4_linear_keeps_normal_quant_method(sm_version): +def test_nvfp4_linear_preserves_activation_quant_method(sm_version): quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4) - with patch( - "tensorrt_llm._torch.modules.linear.get_sm_version", - return_value=sm_version, + with ( + patch( + "tensorrt_llm._torch.modules.linear.get_sm_version", + return_value=sm_version, + ), + patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), + patch("torch.ops.trtllm.gptq_marlin_repack", create=True), ): linear = Linear( 32, @@ -193,6 +347,25 @@ def test_nvfp4_linear_keeps_normal_quant_method(sm_version): quant_config=quant_config, reduce_output=False, ) + assert type(linear.quant_method) is NVFP4LinearMethod + assert linear.has_nvfp4_activation_quantization + assert linear.uses_marlin_nvfp4 is (sm_version == 90) + + +def test_nvfp4_linear_hopper_fp16_keeps_normal_method(): + with ( + patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=90), + patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), + patch("torch.ops.trtllm.gptq_marlin_repack", create=True), + ): + linear = Linear( + 32, + 32, + bias=False, + dtype=torch.float16, + quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4), + reduce_output=False, + ) assert type(linear.quant_method) is NVFP4LinearMethod @@ -230,6 +403,8 @@ def fake_nvfp4_gemm(*args, **kwargs): side_effect=fake_nvfp4_gemm, create=True, ), + patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), + patch("torch.ops.trtllm.gptq_marlin_repack", create=True), ): output = method.apply(module, input_tensor, bias=bias) @@ -319,6 +494,96 @@ def test_w4a16_nvfp4_mlp_rechecks_relu2_fp4_fusion_before_forward(): torch.testing.assert_close(output, relu2(x_up)) +def test_w4a16_nvfp4_mlp_disables_cutedsl_gelu_fusion(): + model_config = ModelConfig(quant_config=QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4)) + + with ( + patch("tensorrt_llm._torch.modules.mlp.get_sm_version", return_value=100), + patch("torch.ops.trtllm.cute_dsl_nvfp4_dense_gemm_gelu_blackwell", create=True), + patch( + "torch.ops.trtllm.cute_dsl_nvfp4_dense_gemm_gelu_fp4out_blackwell", + create=True, + ), + ): + mlp = MLP( + hidden_size=32, + intermediate_size=64, + bias=False, + activation=gelu_tanh, + dtype=torch.bfloat16, + config=model_config, + reduce_output=False, + ) + mlp.create_weights() + + assert not mlp._use_fused_gelu + assert not mlp._use_fused_gelu_fp4out + + +def test_dynamic_nvfp4_mlp_keeps_bf16_cutedsl_gelu_fusion(): + model_config = ModelConfig( + quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4), + force_dynamic_quantization=True, + ) + + with ( + patch("tensorrt_llm._torch.modules.mlp.get_sm_version", return_value=100), + patch("torch.ops.trtllm.cute_dsl_nvfp4_dense_gemm_gelu_blackwell", create=True), + patch( + "torch.ops.trtllm.cute_dsl_nvfp4_dense_gemm_gelu_fp4out_blackwell", + create=True, + ), + ): + mlp = MLP( + hidden_size=32, + intermediate_size=64, + bias=False, + activation=gelu_tanh, + dtype=torch.bfloat16, + config=model_config, + reduce_output=False, + ) + mlp.create_weights() + + assert mlp._use_fused_gelu + assert not mlp._use_fused_gelu_fp4out + + +def test_w4a16_nvfp4_gated_mlp_disables_cutedsl_swiglu_fusion(): + model_config = ModelConfig(quant_config=QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4)) + mlp = GatedMLP( + hidden_size=32, + intermediate_size=64, + bias=False, + dtype=torch.bfloat16, + config=model_config, + reduce_output=False, + use_cute_dsl_blockscaling_mm=True, + ) + + assert not mlp._can_fuse_gate_up_swiglu() + assert not mlp._can_fuse_gate_up_swiglu_fp4out() + + +def test_dynamic_nvfp4_gated_mlp_keeps_bf16_cutedsl_swiglu_fusion(): + model_config = ModelConfig( + quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4), + force_dynamic_quantization=True, + ) + mlp = GatedMLP( + hidden_size=32, + intermediate_size=64, + bias=False, + dtype=torch.bfloat16, + config=model_config, + reduce_output=False, + use_cute_dsl_blockscaling_mm=True, + ) + + assert mlp._can_fuse_gate_up_swiglu() + assert not mlp._can_fuse_gate_up_swiglu_fp4out() + + def test_w4a16_nvfp4_linear_uses_high_precision_activation_without_fp4_quantize(): method = W4A16NVFP4LinearMethod() input_tensor = torch.ones((2, 32), dtype=torch.bfloat16) @@ -433,6 +698,29 @@ def fake_dequant(weight, weight_scale, weight_scale_2, **kwargs): torch.testing.assert_close(output, expected) +def test_w4a16_nvfp4_linear_scale_cache_is_nonpersistent_buffer(): + with patch.object(MarlinNVFP4LinearMethod, "is_supported", return_value=False): + linear = Linear( + 32, + 4, + bias=False, + dtype=torch.bfloat16, + quant_config=QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4), + reduce_output=False, + ) + + scale_linear = torch.arange(128 * 4, dtype=torch.int32).to(torch.uint8) + with patch( + "torch.ops.trtllm.block_scale_interleave_reverse", + return_value=scale_linear, + create=True, + ): + linear.quant_method.cache_derived_state(linear) + + assert linear._buffers["_w4a16_weight_scale_linear"].data_ptr() == scale_linear.data_ptr() + assert "_w4a16_weight_scale_linear" not in linear.state_dict() + + def test_w4a16_nvfp4_linear_uses_marlin_op_after_weight_transform(): method = MarlinNVFP4LinearMethod() input_tensor = torch.ones((1, 32), dtype=torch.bfloat16) @@ -653,6 +941,51 @@ def test_lm_head_uses_w4a16_nvfp4_quant_method_for_packed_lm_head(): assert lm_head.weight_scale_2.shape == (1,) +@pytest.mark.parametrize( + ("checkpoint_has_scale", "exclude_modules", "expected_width"), + [ + (True, None, 1344), + (True, ["lm_head"], 2688), + (False, None, 2688), + ], +) +def test_causal_lm_head_uses_global_w4a16_nvfp4_config( + checkpoint_has_scale, + exclude_modules, + expected_width, +): + quant_config = QuantConfig( + quant_algo=QuantAlgo.W4A16_NVFP4, + exclude_modules=exclude_modules, + ) + model_config = ModelConfig( + pretrained_config=SimpleNamespace( + torch_dtype=torch.float16, + tie_word_embeddings=False, + ), + quant_config=quant_config, + ) + + with patch.object( + DecoderModelForCausalLM, + "_checkpoint_has_lm_head_scale", + return_value=checkpoint_has_scale, + ): + causal_lm = DecoderModelForCausalLM( + torch.nn.Module(), + config=model_config, + hidden_size=2688, + vocab_size=32, + ) + + assert causal_lm.lm_head.weight.shape == (32, expected_width) + if checkpoint_has_scale and exclude_modules is None: + assert causal_lm.lm_head.quant_config is quant_config + assert isinstance(causal_lm.lm_head.quant_method, W4A16NVFP4LinearMethod) + else: + assert not causal_lm.lm_head.has_any_quant + + def test_lm_head_w4a16_nvfp4_forward_uses_triton_dequant(): quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) lm_head = LMHead( diff --git a/tests/unittest/_torch/test_model_config.py b/tests/unittest/_torch/test_model_config.py index 02d0767b05a3..60a1bfdc5a24 100644 --- a/tests/unittest/_torch/test_model_config.py +++ b/tests/unittest/_torch/test_model_config.py @@ -1,6 +1,7 @@ import json import struct import types +from contextlib import nullcontext import pytest import torch @@ -79,6 +80,65 @@ def test_qwen36_modelopt_mixed_precision_preserves_w4a16_layers(tmp_path): assert experts_config.group_size == 16 +def test_modelopt_mixed_precision_is_not_replaced_by_inline_config(tmp_path, monkeypatch): + (tmp_path / "hf_quant_config.json").write_text( + json.dumps( + { + "quantization": { + "quant_algo": "MIXED_PRECISION", + "kv_cache_quant_algo": "FP8", + "quantized_layers": { + "model.language_model.layers.0.mlp.experts": { + "quant_algo": "W4A16_NVFP4", + "group_size": 16, + }, + }, + }, + } + ), + encoding="utf-8", + ) + inline_quant_config = { + "quant_method": "compressed-tensors", + "format": "nvfp4-pack-quantized", + "config_groups": { + "group_0": { + "weights": { + "type": "float", + "num_bits": 4, + "strategy": "tensor_group", + "group_size": 16, + }, + }, + }, + } + pretrained_config = types.SimpleNamespace( + architectures=["Qwen3_5MoeForConditionalGeneration"], + dtype=torch.bfloat16, + quantization_config=inline_quant_config, + ) + + from tensorrt_llm._torch import model_config as model_config_module + + monkeypatch.setattr( + model_config_module, + "load_pretrained_config", + lambda *args, **kwargs: pretrained_config, + ) + monkeypatch.setattr(model_config_module, "config_file_lock", nullcontext) + monkeypatch.setattr(model_config_module, "get_sm_version", lambda: 121) + + model_config = ModelConfig.from_pretrained(str(tmp_path), moe_backend="CUTLASS") + + assert model_config.quant_config.quant_algo is None + assert model_config.quant_config.group_size == 16 + assert model_config.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 + assert model_config.quant_config_dict is not None + experts_config = model_config.quant_config_dict["model.language_model.layers.0.mlp.experts"] + assert experts_config.quant_algo == QuantAlgo.W4A16_NVFP4 + assert experts_config.group_size == 16 + + def test_qwen36_auto_moe_backend_uses_layer_w4a16_on_sm121(monkeypatch): layer_quant_config = { "model.language_model.layers.0.mlp.experts": QuantConfig( @@ -97,6 +157,58 @@ def test_qwen36_auto_moe_backend_uses_layer_w4a16_on_sm121(monkeypatch): assert moe_backend == "CUTEDSL" +def test_deepseek_v4_auto_backend_precedes_w4a16_sm121_override(monkeypatch): + monkeypatch.setattr("tensorrt_llm._torch.model_config.get_sm_version", lambda: 121) + + moe_backend = ModelConfig.resolve_moe_backend( + "AUTO", + "DeepseekV4ForCausalLM", + quant_config=QuantConfig( + quant_algo=QuantAlgo.W4A16_NVFP4, + group_size=16, + ), + ) + + assert moe_backend == "CUTLASS" + + +def test_modelopt_nvfp4_uses_inline_w4a16_activation_semantics(tmp_path): + hf_quant_config = { + "quant_method": "compressed-tensors", + "format": "nvfp4-pack-quantized", + "ignore": ["mtp.layers"], + "config_groups": { + "group_0": { + "targets": ["Linear", "lm_head"], + "weights": { + "type": "float", + "num_bits": 4, + "strategy": "tensor_group", + "group_size": 16, + }, + "input_activations": None, + }, + }, + } + + quant_config, layer_quant_config = ModelConfig._build_modelopt_quant_config( + { + "quant_algo": "NVFP4", + "kv_cache_quant_algo": None, + "group_size": 16, + "exclude_modules": ["mtp*"], + }, + str(tmp_path), + moe_backend="CUTLASS", + hf_quant_config=hf_quant_config, + ) + + assert quant_config.quant_algo == QuantAlgo.W4A16_NVFP4 + assert quant_config.group_size == 16 + assert quant_config.exclude_modules == ["mtp.layers"] + assert layer_quant_config is None + + @pytest.mark.parametrize( "num_key_value_heads", [ diff --git a/tests/unittest/_torch/thop/parallel/test_dense_gemm_act_fusion.py b/tests/unittest/_torch/thop/parallel/test_dense_gemm_act_fusion.py index a16dd771923c..586973673f87 100644 --- a/tests/unittest/_torch/thop/parallel/test_dense_gemm_act_fusion.py +++ b/tests/unittest/_torch/thop/parallel/test_dense_gemm_act_fusion.py @@ -210,9 +210,16 @@ def test_mlp_fp4out_min_m_switch(): fake_qm = types.SimpleNamespace(_input_prepare=lambda *a, **k: None) mlp._use_fused_gelu = True mlp._use_fused_gelu_fp4out = True - mlp.up_proj.quant_method = fake_qm + up = torch.nn.Identity() + up.quant_method = fake_qm + up.has_nvfp4_activation_quantization = True + mlp.up_proj = up down = torch.nn.Identity() # skip the real down GEMM down.quant_method = fake_qm + down.has_nvfp4_activation_quantization = True + down.force_dynamic_quantization = False + down.input_scale = torch.ones(1) + down.pre_quant_scale = None mlp.down_proj = down seen = {} diff --git a/tests/unittest/llmapi/test_kv_cache_dtype_override.py b/tests/unittest/llmapi/test_kv_cache_dtype_override.py index b5fed337c769..efdd768274e4 100644 --- a/tests/unittest/llmapi/test_kv_cache_dtype_override.py +++ b/tests/unittest/llmapi/test_kv_cache_dtype_override.py @@ -43,6 +43,15 @@ def _compressed_tensors_nvfp4_config(**overrides): return config +def _compressed_tensors_w4a16_nvfp4_config(): + config = _compressed_tensors_nvfp4_config( + format="nvfp4-pack-quantized", + ignore=["mtp.layers"], + ) + config["config_groups"]["group_0"]["input_activations"] = None + return config + + def test_get_llm_args_plumbs_kv_cache_dtype(): llm_args, _ = get_llm_args(model="dummy", kv_cache_dtype="nvfp4") assert llm_args["kv_cache_config"].dtype == "nvfp4" @@ -88,6 +97,36 @@ def test_update_from_hf_quant_config_explicit_dtype_overrides(tmp_path): assert llm_args.quant_config.kv_cache_quant_algo == QuantAlgo.NVFP4 +def test_update_from_hf_quant_config_uses_inline_w4a16_semantics(tmp_path): + with open(tmp_path / "hf_quant_config.json", "w") as f: + json.dump( + { + "quantization": { + "quant_algo": "NVFP4", + "kv_cache_quant_algo": None, + "group_size": 16, + "exclude_modules": ["mtp*"], + }, + }, + f, + ) + with open(tmp_path / "config.json", "w") as f: + json.dump( + { + "quantization_config": _compressed_tensors_w4a16_nvfp4_config(), + }, + f, + ) + + llm_args = TorchLlmArgs(model=str(tmp_path)) + model_loader = ModelLoader(llm_args) + + assert model_loader._update_from_hf_quant_config() is True + assert llm_args.quant_config.quant_algo == QuantAlgo.W4A16_NVFP4 + assert llm_args.quant_config.group_size == 16 + assert llm_args.quant_config.exclude_modules == ["mtp.layers"] + + def test_update_from_hf_quant_config_parses_compressed_tensors_model_kwargs(tmp_path): llm_args = TorchLlmArgs( model=str(tmp_path), diff --git a/tests/unittest/models/test_quant_config_utils.py b/tests/unittest/models/test_quant_config_utils.py index 8d58f0aacdbc..10598b5834a6 100644 --- a/tests/unittest/models/test_quant_config_utils.py +++ b/tests/unittest/models/test_quant_config_utils.py @@ -79,7 +79,6 @@ def test_update_quant_config_from_compressed_tensors_parses_w4a16_nvfp4(): "strategy": "tensor_group", "group_size": 16, }, - "input_activations": None, }, }, "ignore": ["lm_head"], From 137b96572819f79ad27a3c5d494f2e1e1128f337 Mon Sep 17 00:00:00 2001 From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> Date: Mon, 3 Aug 2026 16:12:55 +0000 Subject: [PATCH 18/25] [None][chore] Remove Qwen3.6 and decoder kernel changes Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> --- cpp/kernels/fmha_v2/setup.py | 4 +- cpp/kernels/xqa/mha.cu | 6 +- cpp/tensorrt_llm/common/attentionOp.cpp | 3 +- .../decoderXQARunner.cpp | 4 +- .../decoderXQARunnerUtils.h | 13 -- cpp/tests/unit_tests/kernels/CMakeLists.txt | 1 - .../kernels/decoderXQAImplCommonTest.cpp | 92 ------------ .../_torch/attention_backend/trtllm.py | 46 +++--- tensorrt_llm/_torch/model_config.py | 26 +--- .../_torch/models/modeling_qwen3_5.py | 28 ++-- .../_torch/models/modeling_qwen3_moe.py | 3 +- .../_torch/models/modeling_qwen3_next.py | 5 +- .../modules/fused_moe/configurable_moe.py | 19 +-- .../_torch/modules/fused_moe/create_moe.py | 19 +-- .../defs/accuracy/test_llm_api_pytorch.py | 15 +- .../test_lists/test-db/l0_gb10.yml | 1 - .../test_lists/test-db/l0_rtx_pro_6000.yml | 1 - .../modeling/test_modeling_qwen3_5_quant.py | 131 ------------------ .../hf/test_qwen3_5_weight_mapper.py | 120 ---------------- .../moe/test_cute_dsl_b12x_moe_backend.py | 16 --- .../_torch/modules/moe/test_moe_backend.py | 66 +-------- tests/unittest/_torch/test_model_config.py | 115 --------------- 22 files changed, 71 insertions(+), 663 deletions(-) delete mode 100644 cpp/tests/unit_tests/kernels/decoderXQAImplCommonTest.cpp delete mode 100644 tests/unittest/_torch/modeling/test_modeling_qwen3_5_quant.py delete mode 100644 tests/unittest/_torch/models/checkpoints/hf/test_qwen3_5_weight_mapper.py diff --git a/cpp/kernels/fmha_v2/setup.py b/cpp/kernels/fmha_v2/setup.py index 206fc2578cd7..cb854eab6205 100644 --- a/cpp/kernels/fmha_v2/setup.py +++ b/cpp/kernels/fmha_v2/setup.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2020-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2020-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -6810,7 +6810,7 @@ def enumerate_kernels(): enumerate_qmma_flash_kernels(specs, sm=120, dtype='e4m3_fp32', - head_sizes=[64, 128, 192, 256, 576], + head_sizes=[64, 128, 192, 576], output_dtype="bf16") if 'ENABLE_HMMA_FP32' in os.environ: diff --git a/cpp/kernels/xqa/mha.cu b/cpp/kernels/xqa/mha.cu index 47ed51be9a93..bc70bebe862a 100644 --- a/cpp/kernels/xqa/mha.cu +++ b/cpp/kernels/xqa/mha.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -2552,7 +2552,6 @@ CUBIN_EXPORT __global__ // merge results from different warp groups SharedMem::XSmemBuffer* smemOutTile = mergeAndSaveOutTile(outTile, inputElemSize == 2 && cacheElemSize == 1); - bool writeOutput = !isMultiBlock; if (isMultiBlock) { static_assert(ctaShapeInWarps.y == 1, "not implemented"); @@ -2627,7 +2626,6 @@ CUBIN_EXPORT __global__ bool const isLastCta = mbsmem.isLastCta; if (isLastCta) { - writeOutput = true; MultiBlockSMem::MBBuf& mbbuf = mbsmem.storage[warpIdx.y]; SMemWarpRowMax& smemRowMax = reinterpret_cast(smem); // get row max. @@ -2716,7 +2714,7 @@ CUBIN_EXPORT __global__ smemOutTile = mergeAndSaveOutTile(mergedOutTile, false); } } - if (warpGrpIdx == 0 && writeOutput) + if (warpGrpIdx == 0) { #if SPEC_DEC copyOutputToGlobalMem(warp, &output[reqSeqOffset * nbQHeads], nbQHeads, headGrpSize, diff --git a/cpp/tensorrt_llm/common/attentionOp.cpp b/cpp/tensorrt_llm/common/attentionOp.cpp index 95a2e4683659..bcf8f671939e 100644 --- a/cpp/tensorrt_llm/common/attentionOp.cpp +++ b/cpp/tensorrt_llm/common/attentionOp.cpp @@ -241,9 +241,8 @@ bool AttentionOp::convertMMHAParamsToXQAParams(tensorrt_llm::kernels::XQAParams& { xqaParams.kv_cache_data_type = xqaParams.data_type; } - // SM120/121 ship FP8 paged-KV XQA cubins with multi-block support. if (xqaParams.kv_cache_data_type == DATA_TYPE_INT8 - || (xqaParams.kv_cache_data_type == DATA_TYPE_E4M3 && mSM < kSM_90)) + || (xqaParams.kv_cache_data_type == DATA_TYPE_E4M3 && (mSM < kSM_90 || mSM >= kSM_120))) { xqaParams.multi_block_mode = false; } diff --git a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunner.cpp b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunner.cpp index caf532c77eab..e49dcf849ed1 100644 --- a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunner.cpp +++ b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunner.cpp @@ -423,8 +423,10 @@ void DecoderXQARunner::runImpl(XQAParams const& xqaParams, KVCacheBuffer const& // MultiQueryTokens (generation_input_length > 1) need extra parameters (like qSeqLen, headGrpSize, and // mask). Input parameters for MultiQueryTokens kernels. unsigned int headGrpSize = numQHeadsOverKV; + // Use mTileSize = 16 kernels when qSeqLen <= 16. unsigned int qSeqLen = static_cast(xqaParams.generation_input_length); - unsigned int nbTokenBlocksPerGrp = getSpecDecHmmaTokenBlocksPerGroup(headGrpSize, qSeqLen); + unsigned int mTileSize = qSeqLen <= 16 ? 16 : 32; + unsigned int nbTokenBlocksPerGrp = divUp(qSeqLen * headGrpSize, mTileSize); unsigned int maxQSeqLen = xqaParams.spec_decoding_is_generation_length_variable ? // true for ReDrafter xqaParams.spec_decoding_max_generation_length : qSeqLen; diff --git a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunnerUtils.h b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunnerUtils.h index 3e786c27c195..b57abf025583 100644 --- a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunnerUtils.h +++ b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunnerUtils.h @@ -106,19 +106,6 @@ inline uint32_t getKernelMTileSize( return gemmM < 16 ? 16 : 32; }; -inline uint32_t getSpecDecHmmaMTileSize(uint32_t headGrpSize, uint32_t qSeqLen) -{ - return getKernelMTileSize(headGrpSize, /*isSpecDec=*/true, qSeqLen, /*supportQGMMA=*/false, /*supportMLA=*/false); -} - -inline uint32_t getSpecDecHmmaTokenBlocksPerGroup(uint32_t headGrpSize, uint32_t qSeqLen) -{ - uint32_t const mTileSize = getSpecDecHmmaMTileSize(headGrpSize, qSeqLen); - TLLM_CHECK_WITH_INFO(mTileSize > 0U, "Spec-dec HMMA M tile size must be positive."); - uint32_t const headTokens = qSeqLen * headGrpSize; - return (headTokens + mTileSize - 1U) / mTileSize; -} - inline XQAKernelRuntimeHashKey getRuntimeHashKeyFromXQAParams(XQAParams const& xqaParams, int SM) { unsigned int headSize = xqaParams.head_size; diff --git a/cpp/tests/unit_tests/kernels/CMakeLists.txt b/cpp/tests/unit_tests/kernels/CMakeLists.txt index 6135dcdebf26..5f121d2406a9 100644 --- a/cpp/tests/unit_tests/kernels/CMakeLists.txt +++ b/cpp/tests/unit_tests/kernels/CMakeLists.txt @@ -15,7 +15,6 @@ add_gtest(banRepeatNGramsKernelsTest banRepeatNGramsKernelsTest.cpp) add_gtest(decodingKernelsTest decodingKernelTest.cpp) -add_gtest(decoderXQAImplCommonTest decoderXQAImplCommonTest.cpp) add_gtest(logitsBitmaskTest logitsBitmaskTest.cpp) add_gtest(cascadeAttentionKernelTest cascadeAttentionKernelTest.cpp) add_gtest(cascadeAttentionNumericsTest cascadeAttentionNumericsTest.cu) diff --git a/cpp/tests/unit_tests/kernels/decoderXQAImplCommonTest.cpp b/cpp/tests/unit_tests/kernels/decoderXQAImplCommonTest.cpp deleted file mode 100644 index eb2e2941301b..000000000000 --- a/cpp/tests/unit_tests/kernels/decoderXQAImplCommonTest.cpp +++ /dev/null @@ -1,92 +0,0 @@ -/* - * 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. - */ - -#include "tensorrt_llm/common/envUtils.h" -#include "tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunnerUtils.h" - -#include - -namespace tk = tensorrt_llm::kernels; -namespace tc = tensorrt_llm::common; - -namespace -{ - -tk::XQAParams makeSpecDecXqaParams(int numKvHeads, int maxPastKvLength) -{ - tk::XQAParams params{}; - params.num_kv_heads = numKvHeads; - params.max_past_kv_length = maxPastKvLength; - return params; -} - -void skipIfXqaBlocksPerSequenceIsForced() -{ - if (tc::getEnvXqaBlocksPerSequence().has_value()) - { - GTEST_SKIP() << "TRTLLM_XQA_BLOCKS_PER_SEQUENCE overrides multi-block tuning."; - } -} - -TEST(DecoderXQAImplCommonTest, computeMultiBlockCountSpecDecPopulatesLowCtaLaunch) -{ - skipIfXqaBlocksPerSequenceIsForced(); - - int constexpr kBatchSize = 1; - int constexpr kMultiprocessorCount = 128; - int constexpr kSpecDecBlocks = 1; - auto const params = makeSpecDecXqaParams(/*numKvHeads=*/1, /*maxPastKvLength=*/4096); - - EXPECT_EQ(tk::computeMultiBlockCountSpecDec(params, kBatchSize, kMultiprocessorCount, kSpecDecBlocks), 16); -} - -TEST(DecoderXQAImplCommonTest, computeMultiBlockCountSpecDecAccountsForTokenBlocks) -{ - skipIfXqaBlocksPerSequenceIsForced(); - - int constexpr kBatchSize = 1; - int constexpr kMultiprocessorCount = 128; - int constexpr kSpecDecBlocks = 8; - auto const params = makeSpecDecXqaParams(/*numKvHeads=*/1, /*maxPastKvLength=*/4096); - - EXPECT_EQ(tk::computeMultiBlockCountSpecDec(params, kBatchSize, kMultiprocessorCount, kSpecDecBlocks), 8); -} - -TEST(DecoderXQAImplCommonTest, computeMultiBlockCountSpecDecKeepsShortHistorySingleBlock) -{ - skipIfXqaBlocksPerSequenceIsForced(); - - int constexpr kBatchSize = 1; - int constexpr kMultiprocessorCount = 128; - int constexpr kSpecDecBlocks = 1; - auto const params = makeSpecDecXqaParams(/*numKvHeads=*/1, /*maxPastKvLength=*/1024); - - EXPECT_EQ(tk::computeMultiBlockCountSpecDec(params, kBatchSize, kMultiprocessorCount, kSpecDecBlocks), 1); -} - -TEST(DecoderXQAImplCommonTest, getSpecDecHmmaMTileSizeUsesHeadTokens) -{ - EXPECT_EQ(tk::getSpecDecHmmaMTileSize(/*headGrpSize=*/8, /*qSeqLen=*/4), 32U); -} - -TEST(DecoderXQAImplCommonTest, getSpecDecHmmaTokenBlocksPerGroupMatchesCompiledTile) -{ - EXPECT_EQ(tk::getSpecDecHmmaTokenBlocksPerGroup(/*headGrpSize=*/8, /*qSeqLen=*/4), 1U); - EXPECT_EQ(tk::getSpecDecHmmaTokenBlocksPerGroup(/*headGrpSize=*/16, /*qSeqLen=*/4), 2U); -} - -} // namespace diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 2bcdd5e5965f..cb2d5e308d61 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -588,31 +588,29 @@ def prepare(self) -> None: # kv block offsets assert self.request_ids is not None if self.kv_cache_manager is not None: + max_kv_len = int(self.kv_lens[:self.num_seqs].max()) + assert max_kv_len <= self.kv_cache_manager.max_seq_len, ( + f"The max KV cache length of input sequences ({max_kv_len}) " + f"exceeds the KV cache manager's maximum supported length " + f"({self.kv_cache_manager.max_seq_len}).") + + # On the non-speculative path the host kv_lens snapshot bounds + # every block-table access, so the staged/H2D width can be capped + # at the batch's maximum instead of max_seq_len's worth of + # columns. Speculative decoding must stage the full width: + # draft/tree sub-steps and the overlap scheduler advance + # kv_lens_cuda on device past the host snapshot, and their + # kernels dereference block columns a host-derived cap would + # leave unstaged (uninitialized in this buffer). + spec_active = (self.draft_kv_cache_manager is not None + or self.is_spec_decoding_enabled + or bool(self.kv_cache_params.num_extra_kv_tokens) or + (self.runtime_features is not None and + self.runtime_features.has_speculative_draft_tokens)) max_blocks = None - if self.num_seqs > 0: - max_kv_len = int(self.kv_lens[:self.num_seqs].max()) - assert max_kv_len <= self.kv_cache_manager.max_seq_len, ( - f"The max KV cache length of input sequences ({max_kv_len}) " - f"exceeds the KV cache manager's maximum supported length " - f"({self.kv_cache_manager.max_seq_len}).") - - # On the non-speculative path the host kv_lens snapshot bounds - # every block-table access, so the staged/H2D width can be capped - # at the batch's maximum instead of max_seq_len's worth of - # columns. Speculative decoding must stage the full width: - # draft/tree sub-steps and the overlap scheduler advance - # kv_lens_cuda on device past the host snapshot, and their - # kernels dereference block columns a host-derived cap would - # leave unstaged (uninitialized in this buffer). - spec_active = ( - self.draft_kv_cache_manager is not None - or self.is_spec_decoding_enabled - or bool(self.kv_cache_params.num_extra_kv_tokens) - or (self.runtime_features is not None - and self.runtime_features.has_speculative_draft_tokens)) - if not spec_active and self.kv_cache_manager.tokens_per_block: - max_blocks = ceil_div( - max_kv_len, self.kv_cache_manager.tokens_per_block) + if not spec_active and self.kv_cache_manager.tokens_per_block: + max_blocks = ceil_div(max_kv_len, + self.kv_cache_manager.tokens_per_block) self.kv_cache_manager.copy_batch_block_offsets( self.kv_cache_block_offsets, self.request_ids, diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 062d26f33083..7bcf3e0b0bf8 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -330,11 +330,9 @@ def is_generation_model(model_architectures: Optional[List[str]], # once ModelType is used in pytorch flow. @staticmethod - def resolve_moe_backend( - moe_backend: str, - architecture: str, - quant_config: Optional[QuantConfig] = None, - layer_quant_config: Optional[Dict[str, QuantConfig]] = None) -> str: + def resolve_moe_backend(moe_backend: str, + architecture: str, + quant_config: Optional[QuantConfig] = None) -> str: """Resolve AUTO moe_backend to a specific backend based on model architecture. Args: @@ -357,10 +355,6 @@ def resolve_moe_backend( is_w4a16_nvfp4 = (quant_config is not None and quant_config.quant_algo in (QuantAlgo.W4A16_NVFP4, "W4A16_NVFP4")) - if not is_w4a16_nvfp4 and layer_quant_config is not None: - is_w4a16_nvfp4 = any(config.quant_algo in (QuantAlgo.W4A16_NVFP4, - "W4A16_NVFP4") - for config in layer_quant_config.values()) if is_w4a16_nvfp4 and get_sm_version() in (120, 121): return "CUTEDSL" @@ -466,15 +460,12 @@ def _build_modelopt_quant_config(json_quant_configs, quant_config.exclude_modules = json_quant_configs.get( 'exclude_modules', quant_config.exclude_modules) - w4a16_nvfp4_group_size = None for layer in mixed_quant_configs: layer_cfg = mixed_quant_configs[layer] config = QuantConfig() config.kv_cache_quant_algo = kv_cache_quant_algo config.quant_algo = QuantAlgo(layer_cfg['quant_algo']) config.group_size = layer_cfg.get('group_size', None) - if config.quant_algo == QuantAlgo.W4A16_NVFP4: - w4a16_nvfp4_group_size = config.group_size or 16 # AWQ-specific extras emitted by modelopt per-layer. if 'has_zero_point' in layer_cfg: config.has_zero_point = layer_cfg['has_zero_point'] @@ -482,9 +473,6 @@ def _build_modelopt_quant_config(json_quant_configs, config.pre_quant_scale = layer_cfg['pre_quant_scale'] mixed_quant_configs[layer] = config layer_quant_config = mixed_quant_configs - if w4a16_nvfp4_group_size is not None: - quant_config.quant_algo = None - quant_config.group_size = w4a16_nvfp4_group_size elif quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES: if quant_config.group_size is None: quant_config.group_size = 128 @@ -1166,8 +1154,6 @@ def _recursive_update_config(config: transformers.PretrainedConfig, 'hf_quant_config.json'): with open(quant_config_file) as f: normalized = read_modelopt_quant_config(json.load(f)) - modelopt_declares_quant_algo = normalized.get( - "quant_algo") is not None # The file is authoritative; warn if the inline copy disagrees. # Done before _build_modelopt_quant_config since the builder may # mutate ``normalized`` via ``.update`` from quant_cfg.json. @@ -1188,9 +1174,7 @@ def _recursive_update_config(config: transformers.PretrainedConfig, "quantization_config", None)) hf_quant_config = getattr(pretrained_config, "quantization_config", None) - if (quant_config.quant_algo is None - and not modelopt_declares_quant_algo - and hf_quant_config is not None): + if quant_config.quant_algo is None and hf_quant_config is not None: hf_quant_config, hf_layer_quant_config = cls.load_hf_quant_config( hf_quant_config, moe_backend_hint, @@ -1225,7 +1209,7 @@ def _recursive_update_config(config: transformers.PretrainedConfig, requested_moe_backend, architecture, quant_config=quant_config, - layer_quant_config=layer_quant_config) + ) if architecture in _DEEPSEEK_V4_ARCHITECTURES: layer_quant_config = cls._set_deepseek_v4_routed_moe_quant_config( diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_5.py b/tensorrt_llm/_torch/models/modeling_qwen3_5.py index 619b13911ad1..6fbe3c08dfc8 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_5.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_5.py @@ -391,11 +391,12 @@ def _lm_head_nvfp4_enabled(model_config): """Whether the checkpoint's quantized lm_head should stay quantized. ModelOpt MIXED_PRECISION exports for Qwen3.5/3.6 quantize lm_head to - W4A16_NVFP4 (packed FP4 weight + per-group FP8 scales). On SM100/103 the - NVFP4 (W4A4) Linear path consumes it; on SM120/121 the W4A16 Linear path - consumes it directly. Both avoid the bf16 dequant fallback and cut lm_head - weight traffic 4x. Conditions mirror what the quantized LMHead supports - (see LMHead.__init__ guards) plus paths that bypass Linear entirely: + W4A16_NVFP4 (packed FP4 weight + per-group FP8 scales). On SM100/103 the + NVFP4 (W4A4) Linear path can consume it directly, cutting the lm_head + GEMM's weight traffic 4x vs the bf16 dequant fallback -- the decode + lm_head is purely weight-bandwidth-bound. Conditions mirror what the + quantized LMHead supports (see LMHead.__init__ guards) plus the paths + that bypass the Linear machinery entirely: - tie_word_embeddings shares the weight with the embedding lookup, which needs a dense bf16 weight; @@ -414,7 +415,7 @@ def _lm_head_nvfp4_enabled(model_config): return ( cfg is not None and cfg.quant_algo == QuantAlgo.W4A16_NVFP4 - and get_sm_version() in (100, 103, 120, 121) + and get_sm_version() in (100, 103) and not getattr(pretrained, "tie_word_embeddings", False) and not mapping.enable_attention_dp and getattr(pretrained, "vocab_size", 0) % mapping.tp_size == 0 @@ -506,11 +507,10 @@ def _normalize_qwen35_quant_config_dict(model_config, keep_lm_head_quant=False): shared scale (_requantize_linear_attn_fp8_qkvz). Incomplete or non-FP8 sets get no fused entry, and the mapper dequantizes them to bf16 instead. - When ``keep_lm_head_quant`` is true (see _lm_head_nvfp4_enabled), the - ``lm_head`` entry is promoted W4A16_NVFP4 -> NVFP4 on SM100/103 and remains - W4A16_NVFP4 on SM120/121. It is dropped otherwise: a leftover entry would - make DecoderModelForCausalLM build a quantized LMHead whose weights the - mapper had already dequantized to bf16. + The ``lm_head`` entry is promoted W4A16_NVFP4 -> NVFP4 when + ``keep_lm_head_quant`` (see _lm_head_nvfp4_enabled) and dropped otherwise: + a leftover entry would make DecoderModelForCausalLM build a quantized + LMHead whose weights the mapper had already dequantized to bf16. """ qcd = getattr(model_config, "quant_config_dict", None) if not qcd: @@ -528,9 +528,7 @@ def _normalize_qwen35_quant_config_dict(model_config, keep_lm_head_quant=False): continue if name == "lm_head": if keep_lm_head_quant: - if convert_to_nvfp4: - cfg = cfg.model_copy(update={"quant_algo": QuantAlgo.NVFP4}) - normalized[name] = cfg + normalized[name] = cfg.model_copy(update={"quant_algo": QuantAlgo.NVFP4}) else: # Make the fallback visible: the checkpoint quantizes lm_head # but this configuration can't keep it quantized (see @@ -538,7 +536,7 @@ def _normalize_qwen35_quant_config_dict(model_config, keep_lm_head_quant=False): logger.info( f"lm_head quant entry ({cfg.quant_algo}) dropped: " "unsupported configuration for quantized LMHead " - "(requires SM100/103/120/121, untied embeddings, no attention-DP, " + "(requires SM100/103, untied embeddings, no attention-DP, " "vocab divisible by tp_size); lm_head runs bf16" ) continue diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_moe.py b/tensorrt_llm/_torch/models/modeling_qwen3_moe.py index 601098acf402..571e3fe503c0 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_moe.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_moe.py @@ -263,8 +263,7 @@ def forward( self.fusion_config.POST_MOE_FUSION and hidden_states.shape[0] <= self.moe_allreduce.max_token and self.model_config.moe_backend == 'TRTLLM' - and self.mlp.experts.has_nvfp4_activation_quantization - and self.is_p2p_supported) + and self.mlp.experts.has_nvfp4 and self.is_p2p_supported) hidden_states = self.mlp( hidden_states, diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_next.py b/tensorrt_llm/_torch/models/modeling_qwen3_next.py index 0fea1ba422ca..15047b48a4d5 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_next.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_next.py @@ -379,8 +379,9 @@ def __init__(self, mlp: GatedMLP, mapping): self.mlp = mlp self.mapping = mapping self.enable_attention_dp = mapping.enable_attention_dp - # Match the activation-quantization capability exposed by MoE experts. - self.experts = SimpleNamespace(has_nvfp4_activation_quantization=False) + # Provide a dummy `experts` attribute so that + # `self.mlp.experts.has_nvfp4` checks in decoder forward don't crash. + self.experts = SimpleNamespace(has_nvfp4=False) def forward( self, diff --git a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py index 029de1d4302b..b2681db5a2a0 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py @@ -160,11 +160,6 @@ def __init__( override_quant_config: Optional["QuantConfig"] = None, **kwargs, ): - from tensorrt_llm._torch.modules.fused_moe.create_moe import _get_effective_moe_quant_config - - effective_quant_config = _get_effective_moe_quant_config( - model_config, override_quant_config, layer_idx - ) super().__init__( routing_method=routing_method, num_experts=num_experts, @@ -177,7 +172,8 @@ def __init__( layer_idx=layer_idx, # ConfigurableMoE needs correct layer_idx for EPLB initialization **kwargs, ) - self.quant_config = effective_quant_config + if override_quant_config is not None: + self.quant_config = override_quant_config # Store model_config and aux_stream_dict for later use (e.g., backend setter) self.model_config = model_config @@ -190,7 +186,7 @@ def __init__( self._create_and_sync_backend( model_config=model_config, routing_method=routing_method, - effective_quant_config=effective_quant_config, + override_quant_config=override_quant_config, **kwargs, ) @@ -273,7 +269,7 @@ def _create_and_sync_backend( *, model_config: ModelConfig, routing_method: BaseMoeRoutingMethod, - effective_quant_config: Optional["QuantConfig"], + override_quant_config: Optional["QuantConfig"], **kwargs, ) -> None: """Build the MoE backend, mirror EPLB attrs, then create weights. @@ -299,14 +295,13 @@ def _create_and_sync_backend( model_config, routing_method, self.dtype, - override_quant_config=effective_quant_config, - layer_idx=self.layer_idx, + override_quant_config=override_quant_config, ) backend_model_config = model_config - if effective_quant_config is not model_config.quant_config: + if override_quant_config is not None: backend_model_config = copy.deepcopy(model_config) - backend_model_config.quant_config = effective_quant_config + backend_model_config.quant_config = override_quant_config with self._temporarily_skip_weight_creation(backend_model_config): backend = create_moe_backend( diff --git a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py index 2876a4922066..cc6472636d03 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py @@ -45,16 +45,6 @@ def _get_layer_quant_config( return None -def _get_effective_moe_quant_config( - model_config: ModelConfig, - override_quant_config: Optional[QuantConfig] = None, - layer_idx: Optional[int] = None, -) -> Optional[QuantConfig]: - return (override_quant_config - or _get_layer_quant_config(model_config, layer_idx) - or model_config.quant_config) - - def _get_pretrained_megamoe_capability_args( model_config: ModelConfig) -> Dict[str, Optional[object]]: """Extract dtype / hidden / intermediate kwargs for MegaMoE @@ -86,9 +76,9 @@ def get_moe_cls( layer_idx: Optional[int] = None, ) -> Type[MoE]: moe_backend = model_config.moe_backend - quant_config = _get_effective_moe_quant_config(model_config, - override_quant_config, - layer_idx) + quant_config = model_config.quant_config + if override_quant_config is not None: + quant_config = override_quant_config layer_prefix = f"[layer_idx={layer_idx}] " if layer_idx is not None else "" if moe_backend.upper() == "MARLIN": # Marlin MoE is a Hopper-specific NVFP4 W4A16 backend. Layers without @@ -266,8 +256,7 @@ def resolve_moe_cls( ) -> Type[MoE]: moe_cls = get_moe_cls(model_config, override_quant_config, layer_idx) - effective_quant_config = _get_effective_moe_quant_config( - model_config, override_quant_config, layer_idx) + effective_quant_config = override_quant_config or model_config.quant_config has_quant = (effective_quant_config is not None and effective_quant_config.layer_quant_mode.has_any_quant( exclude_kv_cache=True)) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index e4bf32503d8f..df42b0615e65 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -6382,20 +6382,19 @@ class TestQwen3_6_35B_A3B(LlmapiAccuracyTestHarness): @pytest.mark.parametrize("moe_backend", ["TRTLLM", "CUTEDSL"]) def test_nvfp4(self, moe_backend, mocker): - sm_version = get_sm_version() - if moe_backend == "TRTLLM" and sm_version not in (100, 103): - pytest.skip("TRTLLM NVFP4 MoE supports SM100/103 only") - if moe_backend == "CUTEDSL" and sm_version not in (100, 103, 120, 121): - pytest.skip("CuteDSL NVFP4 MoE supports SM100/103/120/121 only") + # Qwen3.6-35B-A3B NVFP4 MoE checkpoint. The TRTLLM-Gen / CuteDSL NVFP4 + # MoE backends only support the SM100 family (B200/B300); RTX 6000 + # (SM120) uses a different MoE path, so restrict this test to SM100/103. + if get_sm_version() not in (100, 103): + pytest.skip("Qwen3.6-35B-A3B NVFP4 MoE test runs on SM100/103 only") if not os.path.exists(self.MODEL_PATH): pytest.skip(f"Model directory {self.MODEL_PATH} does not exist") - max_batch_size = 32 if sm_version in (120, 121) else 128 kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.8, enable_block_reuse=False) cuda_graph_config = CudaGraphConfig(enable_padding=True, - max_batch_size=max_batch_size) + max_batch_size=128) moe_config = MoeConfig(backend=moe_backend) with LLM(self.MODEL_PATH, @@ -6403,7 +6402,7 @@ def test_nvfp4(self, moe_backend, mocker): tensor_parallel_size=1, moe_expert_parallel_size=1, max_seq_len=4096, - max_batch_size=max_batch_size, + max_batch_size=128, kv_cache_config=kv_cache_config, cuda_graph_config=cuda_graph_config, moe_config=moe_config) as llm: diff --git a/tests/integration/test_lists/test-db/l0_gb10.yml b/tests/integration/test_lists/test-db/l0_gb10.yml index 3efa15b293ca..efc6a64530f9 100644 --- a/tests/integration/test_lists/test-db/l0_gb10.yml +++ b/tests/integration/test_lists/test-db/l0_gb10.yml @@ -40,4 +40,3 @@ l0_gb10: - unittest/_torch/modeling -k "modeling_out_of_tree" - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e8_k1_h512_i512-seq=8-dtype=torch.float16-backend=CUTLASS-quant=NVFP4-routing=Renormalize] - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=NVFP4-routing=Renormalize] - - accuracy/test_llm_api_pytorch.py::TestQwen3_6_35B_A3B::test_nvfp4[CUTEDSL] diff --git a/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml b/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml index 6b4793e191a1..f7372a8c4224 100644 --- a/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml +++ b/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml @@ -37,7 +37,6 @@ l0_rtx_pro_6000: - accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_bf16 - accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_fp8 - accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_nvfp4 - - accuracy/test_llm_api_pytorch.py::TestQwen3_6_35B_A3B::test_nvfp4[CUTEDSL] - condition: ranges: diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen3_5_quant.py b/tests/unittest/_torch/modeling/test_modeling_qwen3_5_quant.py deleted file mode 100644 index 404fa703b9d9..000000000000 --- a/tests/unittest/_torch/modeling/test_modeling_qwen3_5_quant.py +++ /dev/null @@ -1,131 +0,0 @@ -# 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 types import SimpleNamespace -from unittest.mock import patch - -import torch -from torch import nn - -from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.models.modeling_qwen3_5 import ( - _lm_head_nvfp4_enabled, - _normalize_qwen35_exclude_modules, - _normalize_qwen35_quant_config_dict, -) -from tensorrt_llm._torch.models.modeling_qwen3_next import Qwen3NextSparseMoeBlock -from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig - - -def test_qwen36_normalizes_modelopt_quantized_layer_paths_on_sm121(): - w4a16_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16) - fp8_config = QuantConfig(quant_algo=QuantAlgo.FP8) - model_config = ModelConfig( - pretrained_config=SimpleNamespace( - num_hidden_layers=40, - tie_word_embeddings=False, - vocab_size=248320, - ), - quant_config=QuantConfig( - exclude_modules=[ - "model.language_model.layers.0.linear_attn.in_proj_qkv", - "mtp.layers.0*", - ] - ), - quant_config_dict={ - "model.language_model.layers.0.linear_attn.in_proj_qkv": fp8_config, - "model.language_model.layers.0.linear_attn.in_proj_z": fp8_config, - "model.language_model.layers.0.mlp.experts": w4a16_config, - "model.language_model.layers.0.mlp.shared_expert.gate_proj": w4a16_config, - "model.language_model.layers.0.mlp.shared_expert.up_proj": w4a16_config, - "model.language_model.layers.0.mlp.shared_expert.down_proj": w4a16_config, - "model.visual.patch_embed": fp8_config, - "mtp.layers.0.mlp.experts": w4a16_config, - "lm_head": w4a16_config, - }, - ) - - with patch("tensorrt_llm._torch.models.modeling_qwen3_5.get_sm_version", return_value=121): - keep_lm_head_quant = _lm_head_nvfp4_enabled(model_config) - assert keep_lm_head_quant - _normalize_qwen35_exclude_modules(model_config, keep_lm_head_quant=keep_lm_head_quant) - _normalize_qwen35_quant_config_dict(model_config, keep_lm_head_quant=keep_lm_head_quant) - - assert model_config.quant_config.exclude_modules == [ - "*linear_attn.conv1d", - "model.layers.0.linear_attn.in_proj_qkvz*", - "model.layers.40*", - ] - assert set(model_config.quant_config_dict) == { - "model.layers.0.linear_attn.in_proj_qkvz", - "model.layers.0.mlp.experts", - "model.layers.0.mlp.shared_expert.down_proj", - "model.layers.0.mlp.shared_expert.gate_proj", - "model.layers.0.mlp.shared_expert.up_proj", - "model.layers.40.mlp.experts", - "lm_head", - } - for name, quant_config in model_config.quant_config_dict.items(): - if name.endswith("in_proj_qkvz"): - assert quant_config.quant_algo == QuantAlgo.FP8 - else: - assert quant_config.quant_algo == QuantAlgo.W4A16_NVFP4 - - -def test_qwen36_sparse_moe_uses_layer_w4a16_quant_config(): - experts_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16) - model_config = ModelConfig( - pretrained_config=SimpleNamespace( - hidden_size=16, - intermediate_size=32, - moe_intermediate_size=8, - num_experts=4, - num_experts_per_tok=1, - shared_expert_intermediate_size=8, - torch_dtype=torch.bfloat16, - model_type="qwen3_5_moe_text", - mlp_bias=False, - ), - moe_backend="CUTEDSL", - quant_config=QuantConfig(), - quant_config_dict={ - "model.layers.0.mlp.experts": experts_config, - }, - ) - captured = {} - - def fake_create_moe(**kwargs): - captured.update(kwargs) - return nn.Identity() - - with ( - patch( - "tensorrt_llm._torch.models.modeling_qwen3_next.create_moe", side_effect=fake_create_moe - ), - patch( - "tensorrt_llm._torch.models.modeling_qwen3_next.AllReduce", - side_effect=lambda **kwargs: nn.Identity(), - ), - patch( - "tensorrt_llm._torch.models.modeling_qwen3_next.GatedMLP", - side_effect=lambda **kwargs: nn.Identity(), - ), - patch("torch.cuda.Event", side_effect=lambda: object()), - ): - Qwen3NextSparseMoeBlock(model_config=model_config, aux_stream=None, layer_idx=0) - - assert captured["model_config"] is model_config - assert captured["override_quant_config"] is experts_config - assert captured["layer_idx"] == 0 diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_qwen3_5_weight_mapper.py b/tests/unittest/_torch/models/checkpoints/hf/test_qwen3_5_weight_mapper.py deleted file mode 100644 index e9a644845e05..000000000000 --- a/tests/unittest/_torch/models/checkpoints/hf/test_qwen3_5_weight_mapper.py +++ /dev/null @@ -1,120 +0,0 @@ -# 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 types import SimpleNamespace - -import torch - -from tensorrt_llm._torch.models.checkpoints.hf.qwen3_5_weight_mapper import Qwen3_5MoeHfWeightMapper -from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig - - -def _make_mapper(dtype=torch.bfloat16): - mapper = Qwen3_5MoeHfWeightMapper() - mapper._config = SimpleNamespace( - pretrained_config=SimpleNamespace(torch_dtype=dtype), - quant_config_dict={}, - ) - return mapper - - -def test_qwen35_modelopt_preprocess_preserves_scalar_fp8_scale_name(): - mapper = _make_mapper() - scale = torch.tensor(0.5, dtype=torch.float32) - - weights, is_modelopt_pb_wo = mapper._normalize_scale_names( - { - "model.layers.0.linear_attn.out_proj.weight_scale": scale, - }, - QuantAlgo.MIXED_PRECISION, - ) - - assert not is_modelopt_pb_wo - assert "model.layers.0.linear_attn.out_proj.weight_scale" in weights - assert weights["model.layers.0.linear_attn.out_proj.weight_scale"].shape == torch.Size([]) - - -def test_qwen36_preserves_quantized_modelopt_lm_head(): - mapper = _make_mapper() - mapper._config.quant_config_dict = { - "lm_head": QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16), - } - weights = { - "lm_head.weight": torch.empty((8, 4), dtype=torch.uint8), - "lm_head.weight_scale": torch.empty((8, 1), dtype=torch.float8_e4m3fn), - "lm_head.weight_scale_2": torch.tensor(0.00011189778888365254, dtype=torch.float32), - "lm_head.input_scale": torch.tensor(0.02771577425301075, dtype=torch.float32), - } - - mapped = mapper._dequantize_lm_head_nvfp4(weights) - - assert mapped is weights - - -def test_qwen35_rescales_per_tensor_fp8_linear_attention_qkvz_projection(): - mapper = _make_mapper(dtype=torch.bfloat16) - mapper._config.quant_config_dict = { - "model.layers.0.linear_attn.in_proj_qkvz": QuantConfig(quant_algo=QuantAlgo.FP8), - } - mapper._config.pretrained_config.linear_num_key_heads = 1 - mapper._config.pretrained_config.linear_num_value_heads = 1 - mapper._config.pretrained_config.linear_key_head_dim = 2 - mapper._config.pretrained_config.linear_value_head_dim = 2 - weight_name = "model.layers.0.linear_attn.in_proj_qkv.weight" - scale_name = "model.layers.0.linear_attn.in_proj_qkv.weight_scale" - input_scale_name = "model.layers.0.linear_attn.in_proj_qkv.input_scale" - z_weight_name = "model.layers.0.linear_attn.in_proj_z.weight" - z_scale_name = "model.layers.0.linear_attn.in_proj_z.weight_scale" - z_input_scale_name = "model.layers.0.linear_attn.in_proj_z.input_scale" - - weights = { - weight_name: torch.tensor( - [ - [1.0, 2.0], - [3.0, 4.0], - [1.0, 2.0], - [3.0, 4.0], - [1.0, 2.0], - [3.0, 4.0], - ], - dtype=torch.float8_e4m3fn, - ), - scale_name: torch.tensor(0.5, dtype=torch.float32), - input_scale_name: torch.tensor(2.0, dtype=torch.float32), - z_weight_name: torch.tensor([[5.0, 6.0], [7.0, 8.0]], dtype=torch.float8_e4m3fn), - z_scale_name: torch.tensor(1.0, dtype=torch.float32), - z_input_scale_name: torch.tensor(1.0, dtype=torch.float32), - } - - updated = mapper._requantize_linear_attn_fp8_qkvz(weights) - packed = mapper._pack_split_projections(updated) - - assert scale_name not in updated - assert input_scale_name not in updated - assert z_scale_name not in updated - assert z_input_scale_name not in updated - assert packed["model.layers.0.linear_attn.in_proj_qkvz.weight"].dtype == torch.float8_e4m3fn - torch.testing.assert_close( - packed["model.layers.0.linear_attn.in_proj_qkvz.weight_scale"], - torch.tensor(1.0, dtype=torch.float32), - ) - torch.testing.assert_close( - packed["model.layers.0.linear_attn.in_proj_qkvz.input_scale"], - torch.tensor(2.0, dtype=torch.float32), - ) - torch.testing.assert_close( - packed["model.layers.0.linear_attn.in_proj_qkvz.weight"][:6].to(torch.float32), - weights[weight_name].to(torch.float32) * 0.5, - ) diff --git a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py index 61c90629f071..fd849a9795ad 100644 --- a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py @@ -178,22 +178,6 @@ def test_get_moe_cls_cutedsl_selects_b12x_for_w4a16_nvfp4_on_supported_sm(sm_ver assert cls is CuteDslB12xFusedMoE -@pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) -def test_get_moe_cls_cutedsl_selects_b12x_for_layer_w4a16_nvfp4_on_supported_sm(sm_version): - """MIXED_PRECISION per-layer W4A16 must select the same backend that the - layer will use after apply_layerwise_quant_config(). - """ - cfg = ModelConfig() - cfg.moe_backend = "CUTEDSL" - cfg.quant_config = QuantConfig() - cfg.quant_config_dict = { - "model.layers.0.mlp.experts": QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16), - } - with patch("tensorrt_llm._utils.get_sm_version", return_value=sm_version): - cls = get_moe_cls(cfg, layer_idx=0) - assert cls is CuteDslB12xFusedMoE - - @pytest.mark.parametrize( "mapping", [ diff --git a/tests/unittest/_torch/modules/moe/test_moe_backend.py b/tests/unittest/_torch/modules/moe/test_moe_backend.py index 72ad4870bd57..2e6d27878d07 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_moe_backend.py @@ -31,7 +31,7 @@ import os from types import SimpleNamespace from typing import List, Optional -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock import pytest import torch @@ -327,70 +327,6 @@ def test_configurable_moe_load_weights_invalidates_wrapper_transform_guard(): assert configurable_moe._weights_transformed is False -def test_configurable_moe_backend_uses_effective_layer_quant_config(): - from tensorrt_llm._torch.modules.fused_moe.configurable_moe import ( - _BACKEND_SYNC_ATTRS, - ConfigurableMoE, - ) - - global_quant_config = QuantConfig() - layer_quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16) - model_config = ModelConfig( - quant_config=global_quant_config, - quant_config_dict={ - "model.layers.3.mlp.experts": layer_quant_config, - }, - ) - - configurable_moe = ConfigurableMoE.__new__(ConfigurableMoE) - torch.nn.Module.__init__(configurable_moe) - configurable_moe.dtype = torch.bfloat16 - configurable_moe.layer_idx = 3 - configurable_moe.num_experts = 8 - configurable_moe.hidden_size = 32 - configurable_moe.intermediate_size = 64 - configurable_moe.reduce_results = False - configurable_moe.aux_stream_dict = None - configurable_moe.weight_loading_mode = MoEWeightLoadingMode.VANILLA - configurable_moe.apply_router_weight_on_input = False - configurable_moe.activation_type = int(ActivationType.Swiglu) - for attr in _BACKEND_SYNC_ATTRS: - if not hasattr(configurable_moe, attr): - setattr(configurable_moe, attr, None) - - backend = torch.nn.Module() - backend.create_weights = MagicMock() - captured = {} - - def create_backend(**kwargs): - captured.update(kwargs) - return backend - - with ( - patch( - "tensorrt_llm._torch.modules.fused_moe.create_moe.resolve_moe_cls", - return_value=object, - ) as resolve_moe_cls, - patch( - "tensorrt_llm._torch.modules.fused_moe.create_moe.create_moe_backend", - side_effect=create_backend, - ), - ): - configurable_moe._create_and_sync_backend( - model_config=model_config, - routing_method=MagicMock(), - effective_quant_config=layer_quant_config, - ) - - backend_model_config = captured["model_config"] - assert backend_model_config is not model_config - assert backend_model_config.quant_config is layer_quant_config - assert model_config.quant_config is global_quant_config - resolve_moe_cls.assert_called_once() - assert resolve_moe_cls.call_args.kwargs["override_quant_config"] is layer_quant_config - backend.create_weights.assert_called_once_with() - - def test_moe_nvfp4_activation_quantization_capability(): assert NVFP4FusedMoEMethod.quantizes_nvfp4_activations assert not W4A16NVFP4CutlassFusedMoEMethod.quantizes_nvfp4_activations diff --git a/tests/unittest/_torch/test_model_config.py b/tests/unittest/_torch/test_model_config.py index 60a1bfdc5a24..bc8981501b01 100644 --- a/tests/unittest/_torch/test_model_config.py +++ b/tests/unittest/_torch/test_model_config.py @@ -1,7 +1,6 @@ import json import struct import types -from contextlib import nullcontext import pytest import torch @@ -43,120 +42,6 @@ def make_pretrained_config( ) -def test_qwen36_modelopt_mixed_precision_preserves_w4a16_layers(tmp_path): - quant_config_file = tmp_path / "hf_quant_config.json" - quant_config_file.write_text( - json.dumps( - { - "quantization": { - "quant_algo": "MIXED_PRECISION", - "kv_cache_quant_algo": "FP8", - "exclude_modules": ["mtp.layers.0*"], - "quantized_layers": { - "model.language_model.layers.0.linear_attn.out_proj": { - "quant_algo": "FP8", - }, - "model.language_model.layers.0.mlp.experts": { - "quant_algo": "W4A16_NVFP4", - "group_size": 16, - }, - }, - }, - } - ), - encoding="utf-8", - ) - - quant_config, layer_quant_config = ModelConfig.load_modelopt_quant_config( - str(quant_config_file), str(tmp_path), moe_backend="CUTLASS" - ) - - assert quant_config.quant_algo is None - assert quant_config.group_size == 16 - assert quant_config.kv_cache_quant_algo == QuantAlgo.FP8 - assert layer_quant_config is not None - experts_config = layer_quant_config["model.language_model.layers.0.mlp.experts"] - assert experts_config.quant_algo == QuantAlgo.W4A16_NVFP4 - assert experts_config.group_size == 16 - - -def test_modelopt_mixed_precision_is_not_replaced_by_inline_config(tmp_path, monkeypatch): - (tmp_path / "hf_quant_config.json").write_text( - json.dumps( - { - "quantization": { - "quant_algo": "MIXED_PRECISION", - "kv_cache_quant_algo": "FP8", - "quantized_layers": { - "model.language_model.layers.0.mlp.experts": { - "quant_algo": "W4A16_NVFP4", - "group_size": 16, - }, - }, - }, - } - ), - encoding="utf-8", - ) - inline_quant_config = { - "quant_method": "compressed-tensors", - "format": "nvfp4-pack-quantized", - "config_groups": { - "group_0": { - "weights": { - "type": "float", - "num_bits": 4, - "strategy": "tensor_group", - "group_size": 16, - }, - }, - }, - } - pretrained_config = types.SimpleNamespace( - architectures=["Qwen3_5MoeForConditionalGeneration"], - dtype=torch.bfloat16, - quantization_config=inline_quant_config, - ) - - from tensorrt_llm._torch import model_config as model_config_module - - monkeypatch.setattr( - model_config_module, - "load_pretrained_config", - lambda *args, **kwargs: pretrained_config, - ) - monkeypatch.setattr(model_config_module, "config_file_lock", nullcontext) - monkeypatch.setattr(model_config_module, "get_sm_version", lambda: 121) - - model_config = ModelConfig.from_pretrained(str(tmp_path), moe_backend="CUTLASS") - - assert model_config.quant_config.quant_algo is None - assert model_config.quant_config.group_size == 16 - assert model_config.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 - assert model_config.quant_config_dict is not None - experts_config = model_config.quant_config_dict["model.language_model.layers.0.mlp.experts"] - assert experts_config.quant_algo == QuantAlgo.W4A16_NVFP4 - assert experts_config.group_size == 16 - - -def test_qwen36_auto_moe_backend_uses_layer_w4a16_on_sm121(monkeypatch): - layer_quant_config = { - "model.language_model.layers.0.mlp.experts": QuantConfig( - quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16 - ), - } - monkeypatch.setattr("tensorrt_llm._torch.model_config.get_sm_version", lambda: 121) - - moe_backend = ModelConfig.resolve_moe_backend( - "AUTO", - "Qwen3_5MoeForConditionalGeneration", - quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8), - layer_quant_config=layer_quant_config, - ) - - assert moe_backend == "CUTEDSL" - - def test_deepseek_v4_auto_backend_precedes_w4a16_sm121_override(monkeypatch): monkeypatch.setattr("tensorrt_llm._torch.model_config.get_sm_version", lambda: 121) From e5eb6c999bb40263a614f0c56ccdecf0b82c9950 Mon Sep 17 00:00:00 2001 From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> Date: Tue, 4 Aug 2026 03:06:56 +0000 Subject: [PATCH 19/25] fix tests Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> --- tensorrt_llm/_torch/modules/linear.py | 123 ++++++------------ .../moe/test_cute_dsl_b12x_moe_backend.py | 31 +---- .../_torch/modules/test_w4a16_nvfp4_linear.py | 58 --------- .../_torch/thop/parallel/test_fp4_linear.py | 17 ++- 4 files changed, 62 insertions(+), 167 deletions(-) diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index dab6cac766a5..64568867a0bb 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -1694,12 +1694,16 @@ def _cleanup_nvfp4_tmp_attrs(self, def process_weights_after_loading_vanilla(self, module: Linear): input_scale, weight_scale_2, alpha = self._finalize_nvfp4_scales(module) - # For dynamic activation quantization, input_scale and alpha are computed at runtime - if input_scale is not None: + # For dynamic activation quantization, input_scale and alpha are computed at runtime. + # The destination checks cover weight-only subclasses (W4A16 / Marlin), which leave + # input_scale / inv_input_scale / alpha as None in create_weights: a checkpoint may + # still carry activation scales, but they are meaningless for a high-precision + # activation path, so drop them instead of copying into a missing Parameter. + if input_scale is not None and module.input_scale is not None: copy_weight(module.input_scale, input_scale) E2M1_MAX = 6.0 module.inv_input_scale.data = module.input_scale / E2M1_MAX - if alpha is not None: + if alpha is not None and module.alpha is not None: copy_weight(module.alpha, alpha) module.scalar_alpha = alpha.item() if weight_scale_2 is not None: @@ -1790,11 +1794,13 @@ def process_weights_after_loading_fused_qkv_linear(self, module: Linear): weight_scale = torch.ops.trtllm.block_scale_interleave(weight_scale) copy_weight(module.weight_scale, weight_scale) - # Finalize input_scale, weight_scale_2, alpha + # Finalize input_scale, weight_scale_2, alpha. The destination checks skip + # scales this linear method never allocated (see the weight-only note in + # process_weights_after_loading_vanilla). input_scale, weight_scale_2, alpha = self._finalize_nvfp4_scales(module) - if input_scale is not None: + if input_scale is not None and module.input_scale is not None: copy_weight(module.input_scale, input_scale) - if alpha is not None: + if alpha is not None and module.alpha is not None: copy_weight(module.alpha, alpha) module.scalar_alpha = alpha.item() if weight_scale_2 is not None: @@ -1876,11 +1882,13 @@ def process_weights_after_loading_fused_gate_up_linear( weight_scale = torch.ops.trtllm.block_scale_interleave(weight_scale) copy_weight(module.weight_scale, weight_scale) - # Finalize input_scale, weight_scale_2, alpha + # Finalize input_scale, weight_scale_2, alpha. The destination checks skip + # scales this linear method never allocated (see the weight-only note in + # process_weights_after_loading_vanilla). input_scale, weight_scale_2, alpha = self._finalize_nvfp4_scales(module) - if input_scale is not None: + if input_scale is not None and module.input_scale is not None: copy_weight(module.input_scale, input_scale) - if alpha is not None: + if alpha is not None and module.alpha is not None: copy_weight(module.alpha, alpha) module.scalar_alpha = alpha.item() if weight_scale_2 is not None: @@ -2004,85 +2012,18 @@ class W4A16NVFP4LinearMethod(NVFP4LinearMethod): def create_weights(self, module: Linear, in_features: int, out_features: int, bias: bool, dtype: torch.dtype): - module.scaling_vector_size = 16 - assert in_features % module.scaling_vector_size == 0, ( - f"in_features {in_features} must be divisible by scaling_vector_size " - f"{module.scaling_vector_size}") - - module.weight = Parameter(torch.empty([out_features, in_features // 2], - dtype=fp4_utils.float4_e2m1x2), - requires_grad=False) - - nrows = fp4_utils.pad_up(out_features, 128) - ncols = fp4_utils.pad_up(in_features // module.scaling_vector_size, 4) - module.weight_scale = Parameter(torch.empty( - [nrows * ncols], dtype=fp4_utils.float4_sf_dtype), - requires_grad=False) - module.weight_scale_2 = Parameter(torch.empty([1], dtype=torch.float32), - requires_grad=False) + super().create_weights(module, in_features, out_features, bias, dtype) + # W4A16 consumes high-precision activations, so there is no calibrated + # activation scale. These must be None from create_weights onwards (not + # merely left uninitialized): eligibility checks such as + # ``is_static_nvfp4_input_eligible`` run right after create_weights, + # before any weights are loaded, and read ``input_scale is not None``. + # The parent's process_weights_after_loading_* skip scales whose + # destination is None, so checkpoint-provided values are dropped here. module.input_scale = None module.inv_input_scale = None module.alpha = None - module.pre_quant_scale = None - module.kv_scales = Parameter(torch.ones(3, dtype=torch.float32), - requires_grad=False) - module.inv_kv_scales = Parameter(torch.ones(3, dtype=torch.float32), - requires_grad=False) - - if bias: - module.bias = Parameter(torch.empty((out_features), dtype=dtype), - requires_grad=False) - else: - module.register_parameter("bias", None) - - def _process_weights_without_static_activation_scale( - self, module: Linear, process_fn): - original_input_scale = module.input_scale - original_inv_input_scale = module.inv_input_scale - original_alpha = module.alpha - had_scalar_alpha = hasattr(module, "scalar_alpha") - original_scalar_alpha = getattr(module, "scalar_alpha", None) - - device = module.weight_scale_2.device - module.input_scale = Parameter(torch.empty([1], - dtype=torch.float32, - device=device), - requires_grad=False) - module.inv_input_scale = Parameter(torch.empty([1], - dtype=torch.float32, - device=device), - requires_grad=False) - module.alpha = Parameter(torch.empty([1], - dtype=torch.float32, - device=device), - requires_grad=False) - try: - process_fn(module) - finally: - module.input_scale = original_input_scale - module.inv_input_scale = original_inv_input_scale - module.alpha = original_alpha - if had_scalar_alpha: - module.scalar_alpha = original_scalar_alpha - elif hasattr(module, "scalar_alpha"): - delattr(module, "scalar_alpha") - - def process_weights_after_loading_vanilla(self, module: Linear): - self._process_weights_without_static_activation_scale( - module, - super().process_weights_after_loading_vanilla) - - def process_weights_after_loading_fused_qkv_linear(self, module: Linear): - self._process_weights_without_static_activation_scale( - module, - super().process_weights_after_loading_fused_qkv_linear) - - def process_weights_after_loading_fused_gate_up_linear( - self, module: Linear): - self._process_weights_without_static_activation_scale( - module, - super().process_weights_after_loading_fused_gate_up_linear) def transform_weights(self, module: Linear) -> None: # Materialize the smaller linear scale view once for Triton dequantization. @@ -2176,6 +2117,20 @@ def apply(self, module: Linear, input: torch.Tensor, class MarlinNVFP4LinearMethod(W4A16NVFP4LinearMethod): """W4A16 NVFP4 linear backed by Marlin.""" + # ``apply`` always allocates a plain output buffer (the Marlin GEMM has no + # NCCL-window output path) and ``apply_linear_allreduce`` is unsupported, so + # this must not inherit the True from NVFP4LinearMethod: Linear.forward reads + # this ClassVar to take a zero-copy branch that assumes the GEMM wrote into + # the symmetric-memory window. + supports_nccl_symmetric_memory_window_output: ClassVar[bool] = False + + def get_tp_alignment(self, tp_mode, quant_config=None): + # Same 32-element alignment as the parent NVFP4 path. The Marlin kernel + # itself wants K%64 and N%128, but ``transform_weights`` pads the weight + # and scales up to those bounds (and ``apply`` slices the N padding back + # off), so the sharding constraint stays unchanged. + return 32 + @staticmethod def is_supported(module: Linear) -> bool: sm_version = get_sm_version() diff --git a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py index fd849a9795ad..85dc479943d8 100644 --- a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py @@ -56,17 +56,10 @@ def test_can_implement_rejects_unsupported_sm(sm_version): @pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) -def test_can_implement_accepts_supported_sm_with_nvfp4(sm_version): +@pytest.mark.parametrize("quant_algo", [QuantAlgo.NVFP4, QuantAlgo.W4A16_NVFP4]) +def test_can_implement_accepts_supported_sm(sm_version, quant_algo): with patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=sm_version): - ok, reason = CuteDslB12xFusedMoE.can_implement(QuantAlgo.NVFP4) - assert ok - assert reason is None - - -@pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) -def test_can_implement_accepts_supported_sm_with_w4a16_nvfp4(sm_version): - with patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=sm_version): - ok, reason = CuteDslB12xFusedMoE.can_implement(QuantAlgo.W4A16_NVFP4) + ok, reason = CuteDslB12xFusedMoE.can_implement(quant_algo) assert ok assert reason is None @@ -157,22 +150,12 @@ def test_get_moe_cls_cutedsl_returns_cutlass_for_w4a16_nvfp4_on_unsupported_sm() @pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) -def test_get_moe_cls_cutedsl_selects_b12x_on_supported_sm(sm_version): - """CUTEDSL + NVFP4 + SM120/121 + flashinfer importable → CuteDslB12xFusedMoE.""" - cfg = ModelConfig() - cfg.moe_backend = "CUTEDSL" - cfg.quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4) - with patch("tensorrt_llm._utils.get_sm_version", return_value=sm_version): - cls = get_moe_cls(cfg) - assert cls is CuteDslB12xFusedMoE - - -@pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) -def test_get_moe_cls_cutedsl_selects_b12x_for_w4a16_nvfp4_on_supported_sm(sm_version): - """CUTEDSL + W4A16_NVFP4 + SM120/121 + flashinfer importable → CuteDslB12xFusedMoE.""" +@pytest.mark.parametrize("quant_algo", [QuantAlgo.NVFP4, QuantAlgo.W4A16_NVFP4]) +def test_get_moe_cls_cutedsl_selects_b12x_on_supported_sm(sm_version, quant_algo): + """CUTEDSL + NVFP4/W4A16_NVFP4 + SM120/121 + flashinfer importable → CuteDslB12xFusedMoE.""" cfg = ModelConfig() cfg.moe_backend = "CUTEDSL" - cfg.quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) + cfg.quant_config = QuantConfig(quant_algo=quant_algo) with patch("tensorrt_llm._utils.get_sm_version", return_value=sm_version): cls = get_moe_cls(cfg) assert cls is CuteDslB12xFusedMoE diff --git a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py index 18bcd4ceaca5..bc7e10b75d90 100644 --- a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py +++ b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py @@ -36,7 +36,6 @@ quant_config_has_nvfp4_activation_quantization, ) from tensorrt_llm._torch.modules.mlp import MLP -from tensorrt_llm._torch.modules.rms_norm import RMSNorm from tensorrt_llm._torch.utils import gelu_tanh, is_nvfp4_marlin_enabled, model_extra_attrs, relu2 from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig @@ -153,24 +152,6 @@ def test_nvfp4_marlin_utility_requires_explicit_opt_in(): assert is_nvfp4_marlin_enabled() -def test_nvfp4_rmsnorm_keeps_high_precision_output_for_hopper_marlin(): - with ( - patch("tensorrt_llm._torch.modules.rms_norm.get_sm_version", return_value=90), - patch("tensorrt_llm._torch.modules.rms_norm.is_nvfp4_marlin_enabled", return_value=True), - ): - norm = RMSNorm( - hidden_size=32, - eps=1e-5, - dtype=torch.bfloat16, - quantize_type="nvfp4", - return_hp_output=True, - ) - - assert not norm.is_nvfp4 - assert not norm.return_hp_output - assert norm.nvfp4_scale is None - - def test_w4a16_attention_does_not_quantize_output_to_fp4(): quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) o_proj = Linear( @@ -659,45 +640,6 @@ def fake_dequant(*args, **kwargs): assert output.shape == (2, 3, 8) -def test_w4a16_nvfp4_linear_uses_triton_dequant(): - method = W4A16NVFP4LinearMethod() - input_tensor = torch.ones((2, 32), dtype=torch.bfloat16) - bias = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.bfloat16) - module = SimpleNamespace( - weight=torch.empty((4, 16), dtype=torch.uint8), - weight_scale=torch.empty((128 * 4,), dtype=torch.uint8), - _w4a16_weight_scale_linear=torch.empty((128 * 4,), dtype=torch.uint8), - weight_scale_2=torch.tensor([0.5], dtype=torch.float32), - dtype=torch.bfloat16, - out_features=4, - scaling_vector_size=16, - pre_quant_scale=None, - use_custom_cublas_mm=False, - ) - captured = {} - - def fake_dequant(weight, weight_scale, weight_scale_2, **kwargs): - captured["weight"] = weight - captured["weight_scale"] = weight_scale - captured["weight_scale_2"] = weight_scale_2 - captured.update(kwargs) - return torch.ones((4, 32), dtype=torch.bfloat16) - - with patch( - "tensorrt_llm._torch.modules.fused_moe.triton_dequant_nvfp4.dequant_nvfp4_2d_triton", - side_effect=fake_dequant, - ): - output = method.apply(module, input_tensor, bias=bias) - - assert captured["weight"].data_ptr() == module.weight.data_ptr() - assert captured["weight_scale"] is module._w4a16_weight_scale_linear - assert captured["weight_scale_2"] is module.weight_scale_2 - assert captured["target_dtype"] is torch.bfloat16 - assert captured["sf_vec_size"] == 16 - expected = torch.tensor([33.0, 34.0, 35.0, 36.0], dtype=torch.bfloat16).expand(2, 4) - torch.testing.assert_close(output, expected) - - def test_w4a16_nvfp4_linear_scale_cache_is_nonpersistent_buffer(): with patch.object(MarlinNVFP4LinearMethod, "is_supported", return_value=False): linear = Linear( diff --git a/tests/unittest/_torch/thop/parallel/test_fp4_linear.py b/tests/unittest/_torch/thop/parallel/test_fp4_linear.py index 18888cab9c3d..e66eee2a8bd2 100644 --- a/tests/unittest/_torch/thop/parallel/test_fp4_linear.py +++ b/tests/unittest/_torch/thop/parallel/test_fp4_linear.py @@ -1,3 +1,18 @@ +# 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"); +# 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. + import sys import pytest @@ -797,7 +812,7 @@ def test_fp4_linear_marlin(dtype, mnk): out_features=OUTPUT_SIZE, bias=False, dtype=dtype, - quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4), + quant_config=QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4), nvfp4_allowed_backends=['marlin'], # key ) From 8d1a6018e53426ef3bb36b63efcd7e4879e4b21f Mon Sep 17 00:00:00 2001 From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> Date: Wed, 5 Aug 2026 00:10:32 +0000 Subject: [PATCH 20/25] fix comments Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> --- tensorrt_llm/_torch/model_config.py | 10 ++++- .../hf/nemotron_h_weight_mapper.py | 19 --------- .../_torch/models/modeling_nemotron_h.py | 31 ++++----------- .../_torch/models/modeling_qwen3_moe.py | 3 +- .../_torch/models/modeling_qwen3_next.py | 7 ++-- .../_torch/modules/fused_moe/quantization.py | 3 ++ tensorrt_llm/_torch/modules/linear.py | 14 ++++--- .../test_modeling_nemotron_h_moe_quant.py | 39 +++---------------- .../hf/test_nemotron_h_weight_mapper.py | 27 ++----------- 9 files changed, 43 insertions(+), 110 deletions(-) diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 7bcf3e0b0bf8..bfc96da6c8f6 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -355,8 +355,14 @@ def resolve_moe_backend(moe_backend: str, is_w4a16_nvfp4 = (quant_config is not None and quant_config.quant_algo in (QuantAlgo.W4A16_NVFP4, "W4A16_NVFP4")) - if is_w4a16_nvfp4 and get_sm_version() in (120, 121): - return "CUTEDSL" + if is_w4a16_nvfp4: + sm_version = get_sm_version() + # CuteDslB12xFusedMoE on SM120/121, MarlinFusedMoE on Hopper. Any + # other SM falls through to CUTLASS, which dequantizes on the fly. + if sm_version in (120, 121): + return "CUTEDSL" + if 90 <= sm_version < 100: + return "MARLIN" if architecture == "GptOssForCausalLM": sm_version = get_sm_version() diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py index 60c9730e766f..a8733a534217 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py @@ -53,23 +53,6 @@ def _split_mamba2_mixer_in_proj(w: torch.Tensor) -> torch.Tensor: # input_scale) do not and are left alone. d_in_proj = 2 * d_inner + 2 * n_groups * d_state + nheads - def _invert_compressed_tensors_scale(value) -> torch.Tensor: - value = value[...] if not isinstance(value, torch.Tensor) else value - value = value.to(torch.float32) - return torch.where(value > 0, value.reciprocal(), - torch.zeros_like(value)).contiguous() - - def _canonicalize_quant_weight(key: str, value): - if key.endswith(".weight_packed"): - return f"{key[:-len('.weight_packed')]}.weight", value - if key.endswith(".weight_global_scale"): - key = f"{key[:-len('.weight_global_scale')]}.weight_scale_2" - return key, _invert_compressed_tensors_scale(value) - if key.endswith(".input_global_scale"): - key = f"{key[:-len('.input_global_scale')]}.input_scale" - return key, _invert_compressed_tensors_scale(value) - return key, value - new_weights = {} for name, _ in weights.items(): key = name @@ -95,8 +78,6 @@ def _canonicalize_quant_weight(key: str, value): if "A_log" in key: key = key.replace("A_log", "A") - key, value = _canonicalize_quant_weight(key, value) - if "mixer.in_proj" in key and "_scale" in key: if self._num_rows(value) == d_in_proj: new_weights[key] = _split_mamba2_mixer_in_proj(value) diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index 94bfb22bcbcf..887ee9d42c96 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py @@ -34,7 +34,6 @@ from tensorrt_llm.logger import logger from tensorrt_llm.lora_helper import LoraConfig from tensorrt_llm.models.modeling_utils import QuantAlgo # noqa: E402 -from tensorrt_llm.models.modeling_utils import QuantConfig from ..attention_backend import AttentionMetadata from ..distributed import AllReduce, AllReduceFusionOp, AllReduceParams @@ -164,26 +163,6 @@ def forward( **kwargs) -def _get_nemotron_h_moe_quant_config( - model_config: ModelConfig[PretrainedConfig], - layer_idx: int) -> QuantConfig | None: - # Per-expert mixed precision config is more specific than the global config. - layer_quant_config = _get_layer_quant_config(model_config, layer_idx, - "mixer.experts") - if layer_quant_config is not None: - return layer_quant_config - - quant_config = model_config.quant_config - if (quant_config is not None and quant_config.quant_algo - in (QuantAlgo.W4A16_NVFP4, "W4A16_NVFP4") - and model_config.moe_backend.upper() != "CUTEDSL"): - values = quant_config.model_dump() - values["quant_algo"] = QuantAlgo.NVFP4 - return QuantConfig.model_validate(values) - - return None - - class NemotronHMOE(nn.Module): def __init__( @@ -268,8 +247,14 @@ def _moe(name): moe_backend=model_config.moe_backend, ) - moe_quant_config = _get_nemotron_h_moe_quant_config( - model_config, layer_idx) + # For MIXED_PRECISION models, the global quant_config has + # quant_algo=MIXED_PRECISION which maps to QuantMode(0) (no quant). This + # would cause the MoE backend to select UnquantizedFusedMoEMethod and + # allocate BF16 weight buffers, causing a shape mismatch when loading + # NVFP4/W4A8_NVFP4_FP8 quantized expert weights. The per-expert entry in + # quant_config_dict is more specific, so prefer it when present. + moe_quant_config = _get_layer_quant_config(model_config, layer_idx, + "mixer.experts") # Setup MoE experts. self.experts = create_moe( diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_moe.py b/tensorrt_llm/_torch/models/modeling_qwen3_moe.py index 571e3fe503c0..601098acf402 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_moe.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_moe.py @@ -263,7 +263,8 @@ def forward( self.fusion_config.POST_MOE_FUSION and hidden_states.shape[0] <= self.moe_allreduce.max_token and self.model_config.moe_backend == 'TRTLLM' - and self.mlp.experts.has_nvfp4 and self.is_p2p_supported) + and self.mlp.experts.has_nvfp4_activation_quantization + and self.is_p2p_supported) hidden_states = self.mlp( hidden_states, diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_next.py b/tensorrt_llm/_torch/models/modeling_qwen3_next.py index 15047b48a4d5..062e597e31cc 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_next.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_next.py @@ -379,9 +379,10 @@ def __init__(self, mlp: GatedMLP, mapping): self.mlp = mlp self.mapping = mapping self.enable_attention_dp = mapping.enable_attention_dp - # Provide a dummy `experts` attribute so that - # `self.mlp.experts.has_nvfp4` checks in decoder forward don't crash. - self.experts = SimpleNamespace(has_nvfp4=False) + # Provide a dummy `experts` attribute so that the + # `self.mlp.experts.has_nvfp4_activation_quantization` checks in decoder + # forward don't crash. + self.experts = SimpleNamespace(has_nvfp4_activation_quantization=False) def forward( self, diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index 52ddf2cf299e..1ed9771f00c0 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py @@ -2975,6 +2975,9 @@ class NVFP4MarlinFusedMoEMethod(NVFP4CutlassFusedMoEMethod): raw ``weight_scale_2`` values. """ + # BF16 activations in, so the NVFP4FusedMoEMethod default does not hold. + quantizes_nvfp4_activations = False + # Marlin's ``transform_weights`` repacks weights into Marlin tiled format # and rebuilds the module parameters, which is incompatible with dynamic # EPLB weight migration. diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 64568867a0bb..5b45ff2d014e 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -2051,11 +2051,13 @@ def _prepare_input(module: Linear, input: torch.Tensor): "W4A16NVFP4LinearMethod requires a high-precision input; " "disable upstream FP4 fusion") + # create_weights leaves inv_input_scale None on this path, so an FP8 + # activation has no scale to dequantize it with. Reject it up front the + # same way as a pre-quantized FP4 input. if input.dtype == torch.float8_e4m3fn: - assert module.inv_input_scale is not None, \ - "W4A16NVFP4LinearMethod: FP8 input requires static inv_input_scale" - input = (input.to(module.dtype) / module.inv_input_scale).to( - module.dtype) + raise RuntimeError( + "W4A16NVFP4LinearMethod requires a high-precision input; " + "disable upstream FP8 attention output") original_shape = None if input.dim() > 2: @@ -2124,7 +2126,9 @@ class MarlinNVFP4LinearMethod(W4A16NVFP4LinearMethod): # the symmetric-memory window. supports_nccl_symmetric_memory_window_output: ClassVar[bool] = False - def get_tp_alignment(self, tp_mode, quant_config=None): + def get_tp_alignment(self, + tp_mode: Optional[TensorParallelMode], + quant_config: Optional[QuantConfig] = None) -> int: # Same 32-element alignment as the parent NVFP4 path. The Marlin kernel # itself wants K%64 and N%128, but ``transform_weights`` pads the weight # and scales up to those bounds (and ``apply`` slices the N padding back diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py index 1e0adb064b3b..5a4614f6fe92 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py @@ -48,12 +48,12 @@ def _make_nemotron_h_moe_config( ) -def test_nemotron_h_moe_uses_w4a4_nvfp4_expert_config_for_w4a16_checkpoint(): +def test_nemotron_h_moe_passes_w4a16_config_through_unchanged(): + """Every MoE backend resolves W4A16_NVFP4 itself, so the layer must not + rewrite quant_algo on its way to create_moe.""" quant_config = QuantConfig( quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16, exclude_modules=["lm_head"] ) - quant_config.mamba_ssm_cache_dtype = "float32" - _ = quant_config.quant_mode model_config = _make_nemotron_h_moe_config(quant_config) captured = {} @@ -68,39 +68,12 @@ def fake_create_moe(**kwargs): aux_stream_dict = {AuxStreamType.MoeShared: None} NemotronHMOE(model_config=model_config, layer_idx=1, aux_stream_dict=aux_stream_dict) - moe_quant_config = captured["override_quant_config"] - assert moe_quant_config is not quant_config - assert moe_quant_config.quant_algo == QuantAlgo.NVFP4 - assert moe_quant_config.quant_mode.has_nvfp4() - assert moe_quant_config.group_size == 16 - assert moe_quant_config.exclude_modules == ["lm_head"] - assert moe_quant_config.mamba_ssm_cache_dtype == "float32" - assert captured["model_config"] is model_config + effective = captured["override_quant_config"] or captured["model_config"].quant_config + assert effective.quant_algo == QuantAlgo.W4A16_NVFP4 + assert effective.group_size == 16 assert model_config.quant_config.quant_algo == QuantAlgo.W4A16_NVFP4 -def test_nemotron_h_moe_preserves_w4a16_config_for_cutedsl_sm12x(): - quant_config = QuantConfig( - quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16, exclude_modules=["lm_head"] - ) - model_config = _make_nemotron_h_moe_config(quant_config, moe_backend="CUTEDSL") - captured = {} - - def fake_create_moe(**kwargs): - captured.update(kwargs) - return nn.Identity() - - with patch( - "tensorrt_llm._torch.models.modeling_nemotron_h.create_moe", side_effect=fake_create_moe - ): - with patch("torch.cuda.Event", side_effect=lambda: object()): - aux_stream_dict = {AuxStreamType.MoeShared: None} - NemotronHMOE(model_config=model_config, layer_idx=1, aux_stream_dict=aux_stream_dict) - - assert captured["override_quant_config"] is None - assert captured["model_config"] is model_config - - def test_nemotron_h_moe_uses_mixer_expert_layer_quant_config(): global_quant_config = QuantConfig() layer_quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16) diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_nemotron_h_weight_mapper.py b/tests/unittest/_torch/models/checkpoints/hf/test_nemotron_h_weight_mapper.py index 16525abe8d74..9859095264b1 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_nemotron_h_weight_mapper.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_nemotron_h_weight_mapper.py @@ -14,7 +14,6 @@ # limitations under the License. from types import SimpleNamespace -from typing import Optional import torch @@ -23,12 +22,10 @@ NemotronHHfWeightMapper, ) from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig +from tensorrt_llm.models.modeling_utils import QuantConfig -def _make_mapper( - quant_algo: Optional[QuantAlgo] = None, -) -> NemotronHHfWeightMapper: +def _make_mapper() -> NemotronHHfWeightMapper: mapper = NemotronHHfWeightMapper() mapper._config = ModelConfig( pretrained_config=SimpleNamespace( @@ -44,7 +41,7 @@ def _make_mapper( ), mapping=Mapping(), moe_backend="CUTLASS", - quant_config=QuantConfig(quant_algo=quant_algo), + quant_config=QuantConfig(), ) return mapper @@ -92,21 +89,3 @@ def test_nemotron_h_mapper_remaps_w4a16_moe_weights_without_input_scale(): assert mapped["model.layers.1.mixer.experts.0.w3.weight_scale_2"] is up_weight_scale_2 assert mapped["model.layers.1.mixer.experts.0.w2.weight_scale_2"] is down_weight_scale_2 assert not any(key.endswith(".input_scale") for key in mapped) - - -def test_nemotron_h_mapper_converts_compressed_tensors_global_scale(): - mapper = _make_mapper(QuantAlgo.W4A16_NVFP4) - global_scale = torch.tensor(9362.2861328125, dtype=torch.float32) - weights = { - "lm_head.weight_packed": torch.empty((8, 4), dtype=torch.uint8), - "lm_head.weight_scale": torch.empty((8, 1), dtype=torch.float8_e4m3fn), - "lm_head.weight_global_scale": global_scale, - } - - mapped = mapper.preprocess_weights(weights) - - assert mapped["lm_head.weight"] is weights["lm_head.weight_packed"] - assert "lm_head.weight_packed" not in mapped - assert "lm_head.weight_global_scale" not in mapped - assert "lm_head.input_scale" not in mapped - torch.testing.assert_close(mapped["lm_head.weight_scale_2"], global_scale.reciprocal()) From ec0c7b03f0d3a60f20da2946735bc56149ef5b2d Mon Sep 17 00:00:00 2001 From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> Date: Wed, 5 Aug 2026 03:41:34 +0000 Subject: [PATCH 21/25] remove unrelated changes Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> --- tensorrt_llm/_torch/model_config.py | 29 +++-------- tensorrt_llm/llmapi/llm_utils.py | 19 +------ tensorrt_llm/models/quant_config_utils.py | 22 ++------ tests/unittest/_torch/test_model_config.py | 52 ------------------- .../llmapi/test_kv_cache_dtype_override.py | 39 -------------- .../models/test_quant_config_utils.py | 26 ---------- 6 files changed, 10 insertions(+), 177 deletions(-) diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index bfc96da6c8f6..5a265c7e7395 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -383,21 +383,17 @@ def resolve_moe_backend(moe_backend: str, return "CUTLASS" @staticmethod - def load_modelopt_quant_config(quant_config_file, - checkpoint_dir, - moe_backend, - hf_quant_config=None): + def load_modelopt_quant_config(quant_config_file, checkpoint_dir, + moe_backend): with open(quant_config_file) as f: quant_config_dict = json.load(f) return ModelConfig._build_modelopt_quant_config( read_modelopt_quant_config(quant_config_dict), checkpoint_dir, - moe_backend, hf_quant_config) + moe_backend) @staticmethod - def _build_modelopt_quant_config(json_quant_configs, - checkpoint_dir, - moe_backend, - hf_quant_config=None): + def _build_modelopt_quant_config(json_quant_configs, checkpoint_dir, + moe_backend): """Build (quant_config, layer_quant_config) from a normalized modelopt 'quantization' inner dict. ``json_quant_configs`` should be a dict as produced by @@ -422,15 +418,6 @@ def _build_modelopt_quant_config(json_quant_configs, if 'pre_quant_scale' in json_quant_configs: quant_config.pre_quant_scale = json_quant_configs['pre_quant_scale'] - if (quant_config.quant_algo == QuantAlgo.NVFP4 - and hf_quant_config is not None and - hf_quant_config.get("quant_method") == "compressed-tensors"): - inline_quant_config = quant_config.model_copy(deep=True) - update_quant_config_from_compressed_tensors(inline_quant_config, - hf_quant_config) - if inline_quant_config.quant_algo == QuantAlgo.W4A16_NVFP4: - quant_config = inline_quant_config - if quant_config.quant_algo == QuantAlgo.MIXED_PRECISION: json_extended_quant_configs: dict = {} # See tests/unittest/llmapi/test_llm_quant.py @@ -1173,11 +1160,7 @@ def _recursive_update_config(config: transformers.PretrainedConfig, cls._has_deepseek_v4_layer_only_modelopt_quant_config( quant_config_file)) quant_config, layer_quant_config = cls._build_modelopt_quant_config( - normalized, - checkpoint_dir, - moe_backend_hint, - hf_quant_config=getattr(pretrained_config, - "quantization_config", None)) + normalized, checkpoint_dir, moe_backend_hint) hf_quant_config = getattr(pretrained_config, "quantization_config", None) if quant_config.quant_algo is None and hf_quant_config is not None: diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index 07fdd41acc22..63b360f5584b 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -160,14 +160,11 @@ def _update_from_hf_quant_config(self) -> bool: # Cross-check against inline config.json.quantization_config if any. # Done before _apply_modelopt_quant_config since the apply step # mutates ``normalized`` via ``.pop()``. - inline_quant_config = None try: with open(f"{self._model_dir}/config.json", "r") as f: - inline_quant_config = json.load(f).get( - "quantization_config") warn_if_inline_diverges( normalized, - inline_quant_config, + json.load(f).get("quantization_config"), source_file="hf_quant_config.json", ) except FileNotFoundError: @@ -183,22 +180,8 @@ def _update_from_hf_quant_config(self) -> bool: raise ValueError( "Pre-quantized checkpoint must have quant_algo.") else: - normalized_quant_algo = normalized.get("quant_algo") self._apply_modelopt_quant_config(normalized, explicit_kv_cache_quant_algo) - if (normalized_quant_algo == QuantAlgo.NVFP4 - and inline_quant_config is not None - and inline_quant_config.get("quant_method") - == "compressed-tensors"): - effective_quant_config = self.llm_args.quant_config.model_copy( - deep=True) - update_quant_config_from_compressed_tensors( - effective_quant_config, inline_quant_config) - if effective_quant_config.quant_algo == QuantAlgo.W4A16_NVFP4: - logger.info( - "Using W4A16_NVFP4 activation semantics from the " - "inline compressed-tensors config.") - self.llm_args.quant_config = effective_quant_config return True hf_config_path = f"{self._model_dir}/config.json" diff --git a/tensorrt_llm/models/quant_config_utils.py b/tensorrt_llm/models/quant_config_utils.py index fd638423629c..09d50a1727e3 100644 --- a/tensorrt_llm/models/quant_config_utils.py +++ b/tensorrt_llm/models/quant_config_utils.py @@ -38,26 +38,11 @@ def update_quant_config_from_compressed_tensors( ) group_config = next(iter(config_groups.values())) weights_quant_config = group_config["weights"] - inputs_quant_config = group_config.get("input_activations") + inputs_quant_config = group_config["input_activations"] weights_quant_strategy = weights_quant_config["strategy"] + inputs_quant_strategy = inputs_quant_config["strategy"] - if inputs_quant_config is None: - is_w4a16_nvfp4 = ( - hf_quant_config.get("format") == "nvfp4-pack-quantized" - and weights_quant_config["num_bits"] == 4 - and weights_quant_config.get("type") == "float" - and weights_quant_strategy == "tensor_group" - and weights_quant_config.get("group_size") == 16 - ) - if not is_w4a16_nvfp4: - raise ValueError( - "input_activations=None is only supported for W4A16 NVFP4 " - "with 4-bit float weights and group_size=16." - ) - quant_config.quant_algo = QuantAlgo.W4A16_NVFP4 - quant_config.group_size = 16 - elif weights_quant_config["num_bits"] == 8: - inputs_quant_strategy = inputs_quant_config["strategy"] + if weights_quant_config["num_bits"] == 8: if weights_quant_strategy == "channel": if inputs_quant_strategy != "token": raise ValueError(f"Unsupported inputs_quant_strategy: {inputs_quant_strategy}.") @@ -83,7 +68,6 @@ def update_quant_config_from_compressed_tensors( and weights_quant_config.get("type") == "float" and weights_quant_strategy == "tensor_group" ): - inputs_quant_strategy = inputs_quant_config["strategy"] # llm-compressor NVFP4: weights FP4 with FP8 per-group scales # (group_size=16), scaled by an FP32 global scale. if inputs_quant_strategy != "tensor_group": diff --git a/tests/unittest/_torch/test_model_config.py b/tests/unittest/_torch/test_model_config.py index bc8981501b01..28b2b0f9a96a 100644 --- a/tests/unittest/_torch/test_model_config.py +++ b/tests/unittest/_torch/test_model_config.py @@ -42,58 +42,6 @@ def make_pretrained_config( ) -def test_deepseek_v4_auto_backend_precedes_w4a16_sm121_override(monkeypatch): - monkeypatch.setattr("tensorrt_llm._torch.model_config.get_sm_version", lambda: 121) - - moe_backend = ModelConfig.resolve_moe_backend( - "AUTO", - "DeepseekV4ForCausalLM", - quant_config=QuantConfig( - quant_algo=QuantAlgo.W4A16_NVFP4, - group_size=16, - ), - ) - - assert moe_backend == "CUTLASS" - - -def test_modelopt_nvfp4_uses_inline_w4a16_activation_semantics(tmp_path): - hf_quant_config = { - "quant_method": "compressed-tensors", - "format": "nvfp4-pack-quantized", - "ignore": ["mtp.layers"], - "config_groups": { - "group_0": { - "targets": ["Linear", "lm_head"], - "weights": { - "type": "float", - "num_bits": 4, - "strategy": "tensor_group", - "group_size": 16, - }, - "input_activations": None, - }, - }, - } - - quant_config, layer_quant_config = ModelConfig._build_modelopt_quant_config( - { - "quant_algo": "NVFP4", - "kv_cache_quant_algo": None, - "group_size": 16, - "exclude_modules": ["mtp*"], - }, - str(tmp_path), - moe_backend="CUTLASS", - hf_quant_config=hf_quant_config, - ) - - assert quant_config.quant_algo == QuantAlgo.W4A16_NVFP4 - assert quant_config.group_size == 16 - assert quant_config.exclude_modules == ["mtp.layers"] - assert layer_quant_config is None - - @pytest.mark.parametrize( "num_key_value_heads", [ diff --git a/tests/unittest/llmapi/test_kv_cache_dtype_override.py b/tests/unittest/llmapi/test_kv_cache_dtype_override.py index ad58f0e93259..8049d034d568 100644 --- a/tests/unittest/llmapi/test_kv_cache_dtype_override.py +++ b/tests/unittest/llmapi/test_kv_cache_dtype_override.py @@ -45,15 +45,6 @@ def _compressed_tensors_nvfp4_config(**overrides): return config -def _compressed_tensors_w4a16_nvfp4_config(): - config = _compressed_tensors_nvfp4_config( - format="nvfp4-pack-quantized", - ignore=["mtp.layers"], - ) - config["config_groups"]["group_0"]["input_activations"] = None - return config - - def test_get_llm_args_plumbs_kv_cache_dtype(): llm_args, _ = get_llm_args(model="dummy", kv_cache_dtype="nvfp4", gpus_per_node=1) assert llm_args["kv_cache_config"].dtype == "nvfp4" @@ -99,36 +90,6 @@ def test_update_from_hf_quant_config_explicit_dtype_overrides(tmp_path): assert llm_args.quant_config.kv_cache_quant_algo == QuantAlgo.NVFP4 -def test_update_from_hf_quant_config_uses_inline_w4a16_semantics(tmp_path): - with open(tmp_path / "hf_quant_config.json", "w") as f: - json.dump( - { - "quantization": { - "quant_algo": "NVFP4", - "kv_cache_quant_algo": None, - "group_size": 16, - "exclude_modules": ["mtp*"], - }, - }, - f, - ) - with open(tmp_path / "config.json", "w") as f: - json.dump( - { - "quantization_config": _compressed_tensors_w4a16_nvfp4_config(), - }, - f, - ) - - llm_args = TorchLlmArgs(model=str(tmp_path)) - model_loader = ModelLoader(llm_args) - - assert model_loader._update_from_hf_quant_config() is True - assert llm_args.quant_config.quant_algo == QuantAlgo.W4A16_NVFP4 - assert llm_args.quant_config.group_size == 16 - assert llm_args.quant_config.exclude_modules == ["mtp.layers"] - - def test_update_from_hf_quant_config_parses_compressed_tensors_model_kwargs(tmp_path): llm_args = TorchLlmArgs( model=str(tmp_path), diff --git a/tests/unittest/models/test_quant_config_utils.py b/tests/unittest/models/test_quant_config_utils.py index 10598b5834a6..823e0ff740ff 100644 --- a/tests/unittest/models/test_quant_config_utils.py +++ b/tests/unittest/models/test_quant_config_utils.py @@ -64,32 +64,6 @@ def test_update_quant_config_from_compressed_tensors_parses_nvfp4(): assert set(quant_config.exclude_modules) == {gate_exclude, "lm_head"} -def test_update_quant_config_from_compressed_tensors_parses_w4a16_nvfp4(): - quant_config = QuantConfig() - update_quant_config_from_compressed_tensors( - quant_config, - { - "quant_method": "compressed-tensors", - "format": "nvfp4-pack-quantized", - "config_groups": { - "group_0": { - "weights": { - "num_bits": 4, - "type": "float", - "strategy": "tensor_group", - "group_size": 16, - }, - }, - }, - "ignore": ["lm_head"], - }, - ) - - assert quant_config.quant_algo == QuantAlgo.W4A16_NVFP4 - assert quant_config.group_size == 16 - assert quant_config.exclude_modules == ["lm_head"] - - def test_update_quant_config_from_compressed_tensors_parses_fp8_block_scales(): quant_config = QuantConfig() update_quant_config_from_compressed_tensors( From 6fa753269d480958f44f9c92d73cd25bfa2d3e66 Mon Sep 17 00:00:00 2001 From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> Date: Wed, 5 Aug 2026 21:19:10 +0000 Subject: [PATCH 22/25] address comments Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> --- tensorrt_llm/_torch/model_config.py | 1 - .../hf/nemotron_h_weight_mapper.py | 53 +++++------ .../_torch/models/modeling_deepseekv3.py | 5 +- .../_torch/models/modeling_deepseekv4.py | 2 +- .../_torch/models/modeling_exaone_moe.py | 2 +- tensorrt_llm/_torch/models/modeling_glm.py | 2 +- .../_torch/models/modeling_nemotron_h.py | 44 ++++----- .../_torch/models/modeling_qwen3_moe.py | 3 +- .../_torch/models/modeling_qwen3_next.py | 7 +- .../_torch/modules/fused_moe/create_moe.py | 20 +--- .../test_modeling_nemotron_h_moe_quant.py | 7 +- .../hf/test_nemotron_h_weight_mapper.py | 91 ------------------- .../_torch/thop/parallel/test_fp4_linear.py | 8 +- 13 files changed, 64 insertions(+), 181 deletions(-) delete mode 100644 tests/unittest/_torch/models/checkpoints/hf/test_nemotron_h_weight_mapper.py diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 5a265c7e7395..90721fc34cef 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -351,7 +351,6 @@ def resolve_moe_backend(moe_backend: str, sm_version = get_sm_version() if 100 <= sm_version < 120: return "TRTLLM" - return "CUTLASS" is_w4a16_nvfp4 = (quant_config is not None and quant_config.quant_algo in (QuantAlgo.W4A16_NVFP4, "W4A16_NVFP4")) diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py index a8733a534217..a1a38e6382f1 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py @@ -56,7 +56,6 @@ def _split_mamba2_mixer_in_proj(w: torch.Tensor) -> torch.Tensor: new_weights = {} for name, _ in weights.items(): key = name - value = weights[name] # change backbone root name to model if "backbone" in key: @@ -79,23 +78,24 @@ def _split_mamba2_mixer_in_proj(w: torch.Tensor) -> torch.Tensor: key = key.replace("A_log", "A") if "mixer.in_proj" in key and "_scale" in key: - if self._num_rows(value) == d_in_proj: - new_weights[key] = _split_mamba2_mixer_in_proj(value) + if self._num_rows(weights[name]) == d_in_proj: + new_weights[key] = _split_mamba2_mixer_in_proj( + weights[name]) else: - new_weights[key] = value + new_weights[key] = weights[name] elif "A" in key: - w = split(value, tp_size, tp_rank) + w = split(weights[name], tp_size, tp_rank) w = w.to(torch.float32) # Avoid extra temporaries: one fp32 cast, then in-place exp/neg. w.exp_() w.neg_() new_weights[key] = w elif "D" in key: - w = split(value, tp_size, tp_rank) + w = split(weights[name], tp_size, tp_rank) w = w.to(torch.float32) new_weights[key] = w elif "dt_bias" in key: - w = split(value, tp_size, tp_rank) + w = split(weights[name], tp_size, tp_rank) w = w.to(torch.float32) new_weights[key] = w elif "mixer.in_proj" in key: @@ -104,9 +104,9 @@ def _split_mamba2_mixer_in_proj(w: torch.Tensor) -> torch.Tensor: # ``weight_scale``, ``weight_scale_2``, …) under ``mixer.in_proj.*`` # — those are scalars / 1-D scales and must not go through the # Mamba2 split rearrangement. - new_weights[key] = _split_mamba2_mixer_in_proj(value) + new_weights[key] = _split_mamba2_mixer_in_proj(weights[name]) elif "conv1d" in key: - w = value + w = weights[name] # removing dim(1) because we are using Linear to store conv1d weights if "weight" in key: w = w.squeeze(1) @@ -124,12 +124,12 @@ def _split_mamba2_mixer_in_proj(w: torch.Tensor) -> torch.Tensor: w = torch.concat(w).contiguous() new_weights[key] = w elif "mixer.norm.weight" in key: - w = split(value, tp_size, tp_rank) + w = split(weights[name], tp_size, tp_rank) new_weights[key] = w # Remap MoE expert weights. elif "mixer.experts." in key: if self.config.moe_backend == 'VANILLA': - new_weights[key] = value + new_weights[key] = weights[name] else: # HF transformers 5.x exposes routed MoE experts as fused # tensors stacked along dim 0 ([num_experts, ...]) under keys @@ -138,7 +138,7 @@ def _split_mamba2_mixer_in_proj(w: torch.Tensor) -> torch.Tensor: # contrast, stores per-expert keys (``experts.{i}.up_proj``). # The VANILLA FusedMoE loader expects per-expert keys, so # unfuse the 3D HF format here before the standard rename. - val = value + val = weights[name] m = re.match(r"(.*\.mixer\.experts)\.(up_proj|down_proj)$", key) is_hf_fused = (m is not None @@ -161,38 +161,35 @@ def _split_mamba2_mixer_in_proj(w: torch.Tensor) -> torch.Tensor: elif "up_proj" in key: w1_key = key.replace("up_proj", "w1") w3_key = key.replace("up_proj", "w3") - # Per-tensor quantization parameters are shared by w1 - # and the empty w3 projection. - if ("input_scale" in key or "weight_scale_2" in key - or "input_quantizer" in key - or "weight_quantizer" in key): - new_weights[w3_key] = value - new_weights[w1_key] = value + # Don't need to handle with input_scale and weight_scale_2 since they are scalar for fp8 and nvfp4 models. + if "input_scale" in key or "weight_scale_2" in key or "input_quantizer" in key or "weight_quantizer" in key: + new_weights[w3_key] = weights[name] + new_weights[w1_key] = weights[name] elif "weight_scale" in key: # NVFP4 case. - if value.shape: + if weights[name].shape: # w3 weight (gate_proj) scale should be empty for Nemotron-H MoE model. # Use [:0] to keep the same input dimension as the other weights. # The w3 weight_scale shape should be [0, input_dim]. - new_weights[w3_key] = value[:0] - new_weights[w1_key] = value + new_weights[w3_key] = weights[name][:0] + new_weights[w1_key] = weights[name] # FP8 case. else: - new_weights[w3_key] = value - new_weights[w1_key] = value + new_weights[w3_key] = weights[name] + new_weights[w1_key] = weights[name] else: # w3 weight (gate_proj) should be empty for Nemotron-H MoE model. # Use [:0] to keep the same input dimension as the other weights. # The w3 weight shape should be [0, input_dim]. - new_weights[w3_key] = value[:0] - new_weights[w1_key] = value + new_weights[w3_key] = weights[name][:0] + new_weights[w1_key] = weights[name] elif "down_proj" in key: key = key.replace("down_proj", "w2") - new_weights[key] = value + new_weights[key] = weights[name] else: raise ValueError(f"Unknown MoE weight: {key}") else: - new_weights[key] = value + new_weights[key] = weights[name] return new_weights diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index e0d4cf9eb6da..ff44bf7d3b73 100755 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -1532,8 +1532,7 @@ def _run_MoE(hidden_states, hidden_states_fp4, do_finalize): not (self.fusion_config.POST_MOE_FUSION and hidden_states.shape[0] <= self.moe_allreduce.max_token and self.model_config.moe_backend == "TRTLLM" - and self.mlp.experts.has_nvfp4_activation_quantization - and self.is_p2p_supported)) + and self.mlp.experts.has_nvfp4 and self.is_p2p_supported)) hidden_states = _run_MoE(hidden_states, hidden_states_fp4=None, @@ -1605,7 +1604,7 @@ def forward_mlp( ) -> Tuple[torch.Tensor, torch.Tensor]: if self.fusion_config.PRE_MLP_FUSION: - if self.mlp.gate_up_proj.has_nvfp4_activation_quantization: + if self.mlp.gate_up_proj.has_nvfp4: act_fp4, act_sf, residual = self.allreduce( hidden_states, all_reduce_params=AllReduceParams( diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index ba9d6ca7c2c0..b988d1dd7e9a 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -2162,7 +2162,7 @@ def _run_MoE(hidden_states, hidden_states_fp4, do_finalize, input_ids): self.fusion_config.POST_MOE_FUSION and hidden_states.shape[0] <= self.moe_allreduce.max_token and self.model_config.moe_backend == "TRTLLM" - and self.mlp.experts.has_nvfp4_activation_quantization + and self.mlp.experts.has_nvfp4 and self.is_p2p_supported ) ) diff --git a/tensorrt_llm/_torch/models/modeling_exaone_moe.py b/tensorrt_llm/_torch/models/modeling_exaone_moe.py index d9307e36a85c..40ae3653d6e0 100644 --- a/tensorrt_llm/_torch/models/modeling_exaone_moe.py +++ b/tensorrt_llm/_torch/models/modeling_exaone_moe.py @@ -383,7 +383,7 @@ def _run_moe(hidden_states, hidden_states_fp4, do_finalize): self.fusion_config.POST_MOE_FUSION and hidden_states.shape[0] <= self.moe_allreduce.max_token and self.model_config.moe_backend == "TRTLLM" - and self.mlp.experts.has_nvfp4_activation_quantization + and self.mlp.experts.has_nvfp4 and self.is_p2p_supported ) ) diff --git a/tensorrt_llm/_torch/models/modeling_glm.py b/tensorrt_llm/_torch/models/modeling_glm.py index 0835319ba54c..2572ea548e48 100644 --- a/tensorrt_llm/_torch/models/modeling_glm.py +++ b/tensorrt_llm/_torch/models/modeling_glm.py @@ -720,7 +720,7 @@ def _run_MoE(hidden_states, hidden_states_fp4, do_finalize): hidden_states.shape[0] <= self.moe_allreduce.max_token and self.fusion_config.POST_MOE_FUSION and self.model_config.moe_backend == "TRTLLM" - and self.mlp.experts.has_nvfp4_activation_quantization + and self.mlp.experts.has_nvfp4 and self.is_p2p_supported ) ) diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index 887ee9d42c96..4484ee8315b4 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py @@ -42,7 +42,6 @@ from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding from ..modules.fused_moe import MoEWeightLoadingMode, create_moe -from ..modules.fused_moe.create_moe import _get_layer_quant_config from ..modules.fused_moe.fused_moe_cutlass import CutlassFusedMoE from ..modules.fused_moe.quantization import (NVFP4CutlassFusedMoEMethod, W4A16NVFP4CutlassFusedMoEMethod) @@ -247,14 +246,18 @@ def _moe(name): moe_backend=model_config.moe_backend, ) - # For MIXED_PRECISION models, the global quant_config has - # quant_algo=MIXED_PRECISION which maps to QuantMode(0) (no quant). This - # would cause the MoE backend to select UnquantizedFusedMoEMethod and - # allocate BF16 weight buffers, causing a shape mismatch when loading - # NVFP4/W4A8_NVFP4_FP8 quantized expert weights. The per-expert entry in - # quant_config_dict is more specific, so prefer it when present. - moe_quant_config = _get_layer_quant_config(model_config, layer_idx, - "mixer.experts") + # For MIXED_PRECISION models, the global quant_config has quant_algo=MIXED_PRECISION + # which maps to QuantMode(0) (no quant). This would cause the MoE backend to select + # UnquantizedFusedMoEMethod and allocate BF16 weight buffers, causing a shape mismatch + # when loading NVFP4/W4A8_NVFP4_FP8 quantized expert weights. + # Look up the per-expert quant config from quant_config_dict and use it for create_moe. + override_quant_config = None + if model_config.quant_config_dict is not None: + experts_prefix = f"model.layers.{layer_idx}.mixer.experts." + for key, cfg in model_config.quant_config_dict.items(): + if key.startswith(experts_prefix): + override_quant_config = cfg + break # Setup MoE experts. self.experts = create_moe( @@ -266,7 +269,7 @@ def _moe(name): dtype=config.torch_dtype, reduce_results=self.reduce_results, model_config=model_config, - override_quant_config=moe_quant_config, + override_quant_config=override_quant_config, layer_idx=self.layer_idx, weight_loading_mode=MoEWeightLoadingMode.VANILLA, bias=self.mlp_bias, @@ -1218,20 +1221,14 @@ def __init__( sublayer_quant_config = self._get_mtp_sublayer_quant_config( model_config, self.layer_idx) - sublayer_moe_backend = model_config.moe_backend - if (sublayer_quant_config is None - or sublayer_quant_config.quant_algo is None): - sublayer_moe_backend = "CUTLASS" # Create a model_config copy with quant_config overridden and # spec_config cleared. All other fields (use_cuda_graph, - # moe_max_num_tokens, etc.) must be inherited so MoE layers are - # configured correctly for CUDA graph capture and communication - # (e.g., DeepEP). BF16 MTP body layers cannot use NVFP4-only MoE - # backends, so route those sublayers to CUTLASS. + # moe_backend, moe_max_num_tokens, etc.) must be inherited + # so MoE layers are configured correctly for CUDA graph + # capture and communication (e.g., DeepEP). sublayer_model_config = replace(model_config, quant_config=sublayer_quant_config, - moe_backend=sublayer_moe_backend, spec_config=None) self.layers[str(step_rel_idx)] = NemotronHMTPDecoderLayer( @@ -1251,15 +1248,14 @@ def _get_mtp_sublayer_quant_config(self, model_config: NemotronHModelConfig, layer_idx: int): """ Get quantization config for MTP sublayer. - The MTP body tensors are stored in BF16. The shared MTP head still - receives the checkpoint-backed lm_head, so its logits path keeps the - lm_head precision instead of inheriting this sublayer override. + The MTP layer in the nvfp4 checkpoint is unquantized. Because the TRTLLM + moe_backend only supports fp8/fp4 quantization, we need to override + the quant_config for the MTP layer. """ from tensorrt_llm.models.modeling_utils import QuantConfig quant_config = model_config.quant_config - # This checkpoint's MTP body is unquantized, so force quant_algo=None - # only for the MTP sublayers constructed here. + # MTP layers are always unquantized, force quant_algo=None if quant_config is None: return None return QuantConfig( diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_moe.py b/tensorrt_llm/_torch/models/modeling_qwen3_moe.py index 601098acf402..571e3fe503c0 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_moe.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_moe.py @@ -263,8 +263,7 @@ def forward( self.fusion_config.POST_MOE_FUSION and hidden_states.shape[0] <= self.moe_allreduce.max_token and self.model_config.moe_backend == 'TRTLLM' - and self.mlp.experts.has_nvfp4_activation_quantization - and self.is_p2p_supported) + and self.mlp.experts.has_nvfp4 and self.is_p2p_supported) hidden_states = self.mlp( hidden_states, diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_next.py b/tensorrt_llm/_torch/models/modeling_qwen3_next.py index 062e597e31cc..15047b48a4d5 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_next.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_next.py @@ -379,10 +379,9 @@ def __init__(self, mlp: GatedMLP, mapping): self.mlp = mlp self.mapping = mapping self.enable_attention_dp = mapping.enable_attention_dp - # Provide a dummy `experts` attribute so that the - # `self.mlp.experts.has_nvfp4_activation_quantization` checks in decoder - # forward don't crash. - self.experts = SimpleNamespace(has_nvfp4_activation_quantization=False) + # Provide a dummy `experts` attribute so that + # `self.mlp.experts.has_nvfp4` checks in decoder forward don't crash. + self.experts = SimpleNamespace(has_nvfp4=False) def forward( self, diff --git a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py index efce45b75545..71fbfe5395a1 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py @@ -27,24 +27,6 @@ from .routing import BaseMoeRoutingMethod -def _is_same_or_child_module_path(lhs: str, rhs: str) -> bool: - return lhs == rhs or lhs.startswith(f"{rhs}.") or rhs.startswith(f"{lhs}.") - - -def _get_layer_quant_config( - model_config: ModelConfig, - layer_idx: Optional[int], - expert_module_suffix: str = "mlp.experts") -> Optional[QuantConfig]: - if layer_idx is None or model_config.quant_config_dict is None: - return None - - moe_module_name = f"model.layers.{layer_idx}.{expert_module_suffix}" - for name, quant_config in model_config.quant_config_dict.items(): - if _is_same_or_child_module_path(name, moe_module_name): - return quant_config - return None - - def _get_pretrained_megamoe_capability_args( model_config: ModelConfig) -> Dict[str, Optional[object]]: """Extract dtype / hidden / intermediate kwargs for MegaMoE @@ -113,7 +95,7 @@ def get_moe_cls( if sm_version in CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS: mapping = model_config.mapping if mapping.moe_ep_size > 1 or mapping.dp_size > 1: - logger.info( + logger.warning( "CuteDslB12xFusedMoE does not support expert " "parallelism or attention-DP/all-to-all; selecting " "CutlassFusedMoE.") diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py index 5a4614f6fe92..1a25b1e18a25 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py @@ -101,7 +101,7 @@ def fake_create_moe(**kwargs): assert captured["override_quant_config"] is layer_quant_config -def test_nemotron_h_mtp_bf16_body_uses_cutlass_moe_backend(): +def test_nemotron_h_mtp_overrides_quant_and_inherits_moe_backend(): quant_config = QuantConfig( quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16, exclude_modules=["lm_head"] ) @@ -141,6 +141,5 @@ def fake_decoder_layer(**kwargs): for layer_kwargs in captured: sublayer_model_config = layer_kwargs["model_config"] assert sublayer_model_config.quant_config.quant_algo is None - assert sublayer_model_config.moe_backend == "CUTLASS" - assert model_config.quant_config.quant_algo == QuantAlgo.W4A16_NVFP4 - assert model_config.moe_backend == "CUTEDSL" + assert sublayer_model_config.moe_backend == model_config.moe_backend + assert model_config.quant_config is quant_config diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_nemotron_h_weight_mapper.py b/tests/unittest/_torch/models/checkpoints/hf/test_nemotron_h_weight_mapper.py deleted file mode 100644 index 9859095264b1..000000000000 --- a/tests/unittest/_torch/models/checkpoints/hf/test_nemotron_h_weight_mapper.py +++ /dev/null @@ -1,91 +0,0 @@ -# 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 types import SimpleNamespace - -import torch - -from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.models.checkpoints.hf.nemotron_h_weight_mapper import ( - NemotronHHfWeightMapper, -) -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.modeling_utils import QuantConfig - - -def _make_mapper() -> NemotronHHfWeightMapper: - mapper = NemotronHHfWeightMapper() - mapper._config = ModelConfig( - pretrained_config=SimpleNamespace( - mamba_head_dim=1, - mamba_num_heads=1, - n_groups=1, - num_hidden_layers=52, - quantization_config={ - "producer": {"name": "modelopt", "version": "0.37.0"}, - "quant_method": "modelopt", - }, - ssm_state_size=1, - ), - mapping=Mapping(), - moe_backend="CUTLASS", - quant_config=QuantConfig(), - ) - return mapper - - -def test_nemotron_h_mapper_preserves_w4a16_lm_head_weights_without_input_scale(): - mapper = _make_mapper() - weight_scale_2 = torch.tensor(0.291 / (448 * 6), dtype=torch.float32) - weights = { - "lm_head.weight": torch.empty((8, 4), dtype=torch.uint8), - "lm_head.weight_scale": torch.empty((8, 1), dtype=torch.float8_e4m3fn), - "lm_head.weight_scale_2": weight_scale_2, - } - - mapped = mapper.preprocess_weights(weights) - - assert mapped["lm_head.weight"] is weights["lm_head.weight"] - assert mapped["lm_head.weight_scale"] is weights["lm_head.weight_scale"] - assert mapped["lm_head.weight_scale_2"] is weight_scale_2 - assert "lm_head.input_scale" not in mapped - - -def test_nemotron_h_mapper_remaps_w4a16_moe_weights_without_input_scale(): - mapper = _make_mapper() - up_prefix = "backbone.layers.1.mixer.experts.0.up_proj" - down_prefix = "backbone.layers.1.mixer.experts.0.down_proj" - up_weight_scale_2 = torch.tensor(0.134 / (448 * 6), dtype=torch.float32) - down_weight_scale_2 = torch.tensor(0.214 / (448 * 6), dtype=torch.float32) - weights = { - f"{up_prefix}.weight": torch.empty((8, 4), dtype=torch.uint8), - f"{up_prefix}.weight_scale": torch.empty((8, 1), dtype=torch.float8_e4m3fn), - f"{up_prefix}.weight_scale_2": up_weight_scale_2, - f"{down_prefix}.weight": torch.empty((8, 4), dtype=torch.uint8), - f"{down_prefix}.weight_scale": torch.empty((8, 1), dtype=torch.float8_e4m3fn), - f"{down_prefix}.weight_scale_2": down_weight_scale_2, - } - - mapped = mapper.preprocess_weights(weights) - - assert "model.layers.1.mixer.experts.0.w1.weight" in mapped - assert "model.layers.1.mixer.experts.0.w3.weight" in mapped - assert "model.layers.1.mixer.experts.0.w2.weight" in mapped - assert mapped["model.layers.1.mixer.experts.0.w3.weight"].shape == (0, 4) - assert mapped["model.layers.1.mixer.experts.0.w3.weight_scale_2"].shape == () - assert mapped["model.layers.1.mixer.experts.0.w1.weight_scale_2"] is up_weight_scale_2 - assert mapped["model.layers.1.mixer.experts.0.w3.weight_scale_2"] is up_weight_scale_2 - assert mapped["model.layers.1.mixer.experts.0.w2.weight_scale_2"] is down_weight_scale_2 - assert not any(key.endswith(".input_scale") for key in mapped) diff --git a/tests/unittest/_torch/thop/parallel/test_fp4_linear.py b/tests/unittest/_torch/thop/parallel/test_fp4_linear.py index 72d9bee4e3a9..0d829aeb1b54 100644 --- a/tests/unittest/_torch/thop/parallel/test_fp4_linear.py +++ b/tests/unittest/_torch/thop/parallel/test_fp4_linear.py @@ -769,6 +769,7 @@ def test_fp4_linear_cuda_core(dtype, mnk): not (89 <= get_sm_version() < 100 or get_sm_version() in (120, 121)), reason="Dense Marlin NVFP4 runs on SM89-99 and SM120/121", ) +@pytest.mark.parametrize("quant_algo", [QuantAlgo.NVFP4, QuantAlgo.W4A16_NVFP4]) @pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize( "mnk", @@ -790,7 +791,10 @@ def test_fp4_linear_cuda_core(dtype, mnk): (3, 176, 144), (128, 928, 1360), ]) -def test_fp4_linear_marlin(dtype, mnk): +def test_fp4_linear_marlin(quant_algo, dtype, mnk): + if quant_algo == QuantAlgo.NVFP4 and get_sm_version() in (120, 121): + pytest.skip( + "Marlin backend shouldn't be used for NVFP4 quant on SM120/121") SEQ_LEN, OUTPUT_SIZE, HIDDEN_SIZE = mnk torch.manual_seed(0) @@ -812,7 +816,7 @@ def test_fp4_linear_marlin(dtype, mnk): out_features=OUTPUT_SIZE, bias=False, dtype=dtype, - quant_config=QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4), + quant_config=QuantConfig(quant_algo=quant_algo), nvfp4_allowed_backends=['marlin'], # key ) From a10ffa56e26b0c43071095afe90bfd69ab32dd2e Mon Sep 17 00:00:00 2001 From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> Date: Wed, 5 Aug 2026 21:39:11 +0000 Subject: [PATCH 23/25] extend test list Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> --- tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py | 8 ++++++-- tests/integration/test_lists/test-db/l0_gb10.yml | 2 ++ tests/integration/test_lists/test-db/l0_h100.yml | 1 + 3 files changed, 9 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py index c01a99da51d2..5dc95306f55c 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py @@ -57,6 +57,10 @@ class MarlinFusedMoE(CutlassFusedMoE): "sm_constraint": ("in", set(range(89, 100))), "dtypes": {torch.bfloat16}, }, + QuantAlgo.W4A16_NVFP4: { + "sm_constraint": ("in", set(range(89, 100))), + "dtypes": {torch.bfloat16}, + }, } @classmethod @@ -68,9 +72,9 @@ def can_implement( ) -> Tuple[bool, Optional[str]]: sm_version = get_sm_version() - if quant_algo != QuantAlgo.NVFP4: + if quant_algo not in cls._QUANT_SUPPORT_TABLE: return _warn_and_return( - f"MarlinFusedMoE only supports NVFP4 (got quant_algo={quant_algo})" + f"MarlinFusedMoE only supports NVFP4 or W4A16_NVFP4 (got quant_algo={quant_algo})" ) if not is_nvfp4_marlin_supported_sm(sm_version): diff --git a/tests/integration/test_lists/test-db/l0_gb10.yml b/tests/integration/test_lists/test-db/l0_gb10.yml index efc6a64530f9..da77f6e22ace 100644 --- a/tests/integration/test_lists/test-db/l0_gb10.yml +++ b/tests/integration/test_lists/test-db/l0_gb10.yml @@ -17,6 +17,8 @@ l0_gb10: tests: # ------------- PyTorch tests --------------- - unittest/_torch/attention/test_attention_mla.py + - unittest/_torch/modules/test_w4a16_nvfp4_linear.py + - unittest/_torch/thop/parallel/test_fp4_linear.py - test_e2e.py::test_ptp_quickstart_bert[VANILLA-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] - test_e2e.py::test_ptp_quickstart_bert[TRTLLM-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_bf16[latency] diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 9c7d152e3fc2..19e5e32e7746 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -33,6 +33,7 @@ l0_h100: - unittest/_torch/modules/test_rotary_embedding.py - unittest/_torch/modules/mamba - unittest/_torch/modules/tests_lora_modules + - unittest/_torch/modules/test_w4a16_nvfp4_linear.py # ------------- MoE components tests --------------- - unittest/_torch/modules/test_moe_load_balancer.py - unittest/_torch/modules/test_moe_routing.py From 33712bf017cd0b1c9ef2089ffa2a8dd2db73b47c Mon Sep 17 00:00:00 2001 From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> Date: Thu, 6 Aug 2026 18:36:31 +0000 Subject: [PATCH 24/25] fix test Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> --- .../_torch/modules/test_w4a16_nvfp4_linear.py | 608 +++--------------- 1 file changed, 74 insertions(+), 534 deletions(-) diff --git a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py index 822b3cc38320..b3d6e5782bca 100644 --- a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py +++ b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py @@ -19,124 +19,20 @@ import pytest import torch -import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM from tensorrt_llm._torch.modules.attention import Attention from tensorrt_llm._torch.modules.embedding import LMHead -from tensorrt_llm._torch.modules.gated_mlp import GatedMLP from tensorrt_llm._torch.modules.linear import ( Linear, MarlinNVFP4LinearMethod, NVFP4LinearMethod, - TensorParallelMode, W4A16NVFP4LinearMethod, - get_quant_method, - get_sm_version, quant_config_has_nvfp4_activation_quantization, ) -from tensorrt_llm._torch.modules.mlp import MLP -from tensorrt_llm._torch.utils import gelu_tanh, is_nvfp4_marlin_enabled, model_extra_attrs, relu2 -from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig -def _run_w4a16_marlin_reference_case(m: int, n: int, k: int) -> None: - act, weight, weight_scale, weight_scale_2 = _make_w4a16_nvfp4_case(m, n, k, torch.bfloat16) - - linear = Linear( - k, - n, - bias=False, - dtype=torch.bfloat16, - quant_config=QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4), - reduce_output=False, - nvfp4_allowed_backends=["marlin"], - ).cuda() - assert isinstance(linear.quant_method, MarlinNVFP4LinearMethod) - linear.weight.data.copy_(weight) - linear.weight_scale.data.copy_(weight_scale) - linear.weight_scale_2.data.copy_(weight_scale_2) - linear.transform_weights() - - with patch.object(MarlinNVFP4LinearMethod, "is_supported", return_value=False): - reference = Linear( - k, - n, - bias=False, - dtype=torch.bfloat16, - quant_config=QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4), - reduce_output=False, - ).cuda() - assert type(reference.quant_method) is W4A16NVFP4LinearMethod - reference.weight.data.copy_(weight) - reference.weight_scale.data.copy_(weight_scale) - reference.weight_scale_2.data.copy_(weight_scale_2) - reference.transform_weights() - - expected = reference(act) - actual = linear(act) - torch.testing.assert_close(actual, expected, atol=0.75, rtol=0.02) - - -def _make_w4a16_nvfp4_case(m: int, n: int, k: int, dtype: torch.dtype): - torch.manual_seed(m + n + k) - act = torch.randn((m, k), device="cuda", dtype=dtype) - weight = torch.empty((n, k // 2), device="cuda", dtype=fp4_utils.float4_e2m1x2) - weight_u8 = torch.randint( - 0, - 256, - (n, k // 2), - device="cuda", - dtype=torch.uint8, - ) - weight.copy_(weight_u8.view(fp4_utils.float4_e2m1x2)) - - scale_cols = fp4_utils.pad_up(k // 16, 4) - scale_rows = fp4_utils.pad_up(n, 128) - # E4M3 bit patterns in [0x30, 0x40] represent scales from 0.5 to 2.0. - weight_scale_linear = torch.randint( - 0x30, - 0x41, - (scale_rows, scale_cols), - device="cuda", - dtype=torch.uint8, - ) - weight_scale = torch.ops.trtllm.block_scale_interleave(weight_scale_linear).view( - fp4_utils.float4_sf_dtype - ) - weight_scale_2 = torch.ones((1,), device="cuda", dtype=torch.float32) - return act, weight, weight_scale, weight_scale_2 - - -@pytest.mark.skipif( - not torch.cuda.is_available() - or not (89 <= get_sm_version() < 100 or get_sm_version() in (120, 121)), - reason="requires CUDA SM89-99 or SM120/121", -) -@pytest.mark.parametrize( - "shape", - [ - (1, 256, 256), - (4, 160, 288), - (4, 256, 256), - (32, 256, 256), - (128, 512, 1024), - ], -) -def test_w4a16_nvfp4_marlin_bf16_matches_triton(shape): - m, n, k = shape - _run_w4a16_marlin_reference_case(m, n, k) - - -def test_get_quant_method_returns_w4a16_nvfp4_linear_method(): - quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) - - method = get_quant_method(quant_config) - - assert type(method) is W4A16NVFP4LinearMethod - - def test_nvfp4_activation_quantization_excludes_w4a16(): assert quant_config_has_nvfp4_activation_quantization(QuantConfig(quant_algo=QuantAlgo.NVFP4)) assert not quant_config_has_nvfp4_activation_quantization( @@ -144,16 +40,6 @@ def test_nvfp4_activation_quantization_excludes_w4a16(): ) -def test_nvfp4_marlin_utility_requires_explicit_opt_in(): - with ( - patch("tensorrt_llm._torch.utils.get_sm_version", return_value=90), - patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), - ): - assert not is_nvfp4_marlin_enabled() - with model_extra_attrs({"nvfp4_gemm_allowed_backends": ["cutlass", "marlin"]}): - assert is_nvfp4_marlin_enabled() - - def test_w4a16_attention_does_not_quantize_output_to_fp4(): quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) o_proj = Linear( @@ -230,56 +116,25 @@ def test_nvfp4_attention_keeps_high_precision_output_for_hopper_marlin(): assert not Attention._use_quantize_output(attention) -def test_w4a16_disables_fused_gemm_allreduce(monkeypatch): - monkeypatch.setenv("TRTLLM_GEMM_ALLREDUCE_FUSION_ENABLED", "1") - mapping = Mapping(world_size=2, rank=0, tp_size=2) - - with ( - patch("tensorrt_llm._torch.modules.linear.mpi_disabled", return_value=False), - patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=120), - patch("tensorrt_llm._torch.modules.linear.ipc_nvls_supported", return_value=True), - patch("tensorrt_llm._torch.distributed.AllReduce"), - ): - linear = Linear( - 256, - 64, - bias=False, - dtype=torch.bfloat16, - mapping=mapping, - tensor_parallel_mode=TensorParallelMode.ROW, - quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4), - reduce_output=True, - skip_create_weights_in_init=True, - ) - assert linear.use_fused_gemm_allreduce - - # Simulate apply_layerwise_quant_config rebinding a mixed-precision - # layer after Linear.__init__ but before deferred weight creation. - linear.quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) - linear.create_weights() - - assert not linear.use_fused_gemm_allreduce - - @pytest.mark.parametrize( - ("sm_version", "dtype", "expected_backends"), + ("allowed_backends", "expected_backends"), [ - (89, torch.bfloat16, "cutlass,cublaslt,cuda_core"), - (90, torch.bfloat16, "cutlass,cublaslt,cuda_core"), - (90, torch.float16, "cutlass,cublaslt,cuda_core"), - (120, torch.bfloat16, "cutlass,cublaslt,cuda_core"), - (121, torch.bfloat16, "cutlass,cublaslt,cuda_core"), + (["cutlass", "cublaslt", "cuda_core"], "cutlass,cublaslt,cuda_core"), + (["marlin"], "marlin"), ], ) -def test_nvfp4_linear_uses_architecture_default_backend(sm_version, dtype, expected_backends): +def test_nvfp4_linear_forwards_allowed_backends_to_gemm(allowed_backends, expected_backends): + """The module's backend list reaches the unified GEMM op verbatim. Which + backends are eligible in the first place is covered by + ``test_nvfp4_linear_keeps_activation_quant_method``.""" method = NVFP4LinearMethod() - input_tensor = torch.ones((2, 32), dtype=dtype) + input_tensor = torch.ones((2, 32), dtype=torch.bfloat16) module = SimpleNamespace( weight=torch.empty((4, 16), dtype=torch.uint8), weight_scale=torch.empty((128 * 4,), dtype=torch.uint8), out_features=4, - dtype=dtype, - nvfp4_allowed_backends=["cutlass", "cublaslt", "cuda_core"], + dtype=torch.bfloat16, + nvfp4_allowed_backends=allowed_backends, all_reduce=None, mapping=None, ) @@ -290,19 +145,12 @@ def test_nvfp4_linear_uses_architecture_default_backend(sm_version, dtype, expec def fake_nvfp4_gemm(*args, **kwargs): captured["allowed_backends"] = kwargs["allowed_backends"] - return torch.ones((2, 4), dtype=dtype) + return torch.ones((2, 4), dtype=torch.bfloat16) with ( patch.object(method, "_input_prepare", return_value=(act_fp4, act_sf, alpha)), - patch( - "tensorrt_llm._torch.modules.linear.get_sm_version", - return_value=sm_version, - ), - patch( - "torch.ops.trtllm.nvfp4_gemm", - side_effect=fake_nvfp4_gemm, - create=True, - ), + patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=90), + patch("torch.ops.trtllm.nvfp4_gemm", side_effect=fake_nvfp4_gemm, create=True), patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), patch("torch.ops.trtllm.gptq_marlin_repack", create=True), ): @@ -312,58 +160,27 @@ def fake_nvfp4_gemm(*args, **kwargs): assert output.shape == (2, 4) -@pytest.mark.parametrize("sm_version", [89, 90, 120, 121]) -def test_nvfp4_linear_preserves_activation_quant_method(sm_version): - quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4) - - with ( - patch( - "tensorrt_llm._torch.modules.linear.get_sm_version", - return_value=sm_version, - ), - patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), - patch("torch.ops.trtllm.gptq_marlin_repack", create=True), - ): - linear = Linear( - 32, - 32, - bias=False, - dtype=torch.bfloat16, - quant_config=quant_config, - reduce_output=False, - ) - assert type(linear.quant_method) is NVFP4LinearMethod - assert linear.has_nvfp4_activation_quantization - assert not linear.uses_marlin_nvfp4 - - -@pytest.mark.parametrize("sm_version", [89, 90]) -def test_nvfp4_linear_uses_marlin_when_explicitly_enabled(sm_version): - with ( - patch( - "tensorrt_llm._torch.modules.linear.get_sm_version", - return_value=sm_version, - ), - patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), - patch("torch.ops.trtllm.gptq_marlin_repack", create=True), - ): - linear = Linear( - 32, - 32, - bias=False, - dtype=torch.bfloat16, - quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4), - reduce_output=False, - nvfp4_allowed_backends=["marlin"], - ) - assert type(linear.quant_method) is NVFP4LinearMethod - assert linear.has_nvfp4_activation_quantization - assert linear.uses_marlin_nvfp4 - +@pytest.mark.parametrize( + ("sm_version", "dtype", "allowed_backends", "expect_marlin"), + [ + (89, torch.bfloat16, None, False), + (90, torch.bfloat16, None, False), + (120, torch.bfloat16, None, False), + (121, torch.bfloat16, None, False), + (89, torch.bfloat16, ["marlin"], True), + (90, torch.bfloat16, ["marlin"], True), + (90, torch.float16, None, False), + ], +) +def test_nvfp4_linear_keeps_activation_quant_method( + sm_version, dtype, allowed_backends, expect_marlin +): + """Plain NVFP4 always keeps NVFP4LinearMethod -- Marlin is reached inside the + GEMM op, so only ``uses_marlin_nvfp4`` flips.""" + kwargs = {} if allowed_backends is None else {"nvfp4_allowed_backends": allowed_backends} -def test_nvfp4_linear_hopper_fp16_keeps_normal_method(): with ( - patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=90), + patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=sm_version), patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), patch("torch.ops.trtllm.gptq_marlin_repack", create=True), ): @@ -371,12 +188,15 @@ def test_nvfp4_linear_hopper_fp16_keeps_normal_method(): 32, 32, bias=False, - dtype=torch.float16, + dtype=dtype, quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4), reduce_output=False, + **kwargs, ) assert type(linear.quant_method) is NVFP4LinearMethod + assert linear.has_nvfp4_activation_quantization + assert linear.uses_marlin_nvfp4 is expect_marlin def test_nvfp4_linear_hopper_marlin_applies_bias_as_post_op(): @@ -422,177 +242,6 @@ def fake_nvfp4_gemm(*args, **kwargs): torch.testing.assert_close(output, torch.ones((2, 4), dtype=torch.bfloat16) + bias) -def test_w4a16_nvfp4_mlp_disables_relu2_fp4_fusion_without_input_scale(): - model_config = ModelConfig(quant_config=QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4)) - - with ( - patch("tensorrt_llm._torch.modules.mlp.get_sm_version", return_value=121), - patch("torch.ops.trtllm.fused_relu2_quantize", create=True), - ): - mlp = MLP( - hidden_size=32, - intermediate_size=64, - bias=False, - activation=relu2, - dtype=torch.bfloat16, - config=model_config, - reduce_output=False, - ) - mlp.create_weights() - - assert mlp.down_proj.has_nvfp4 - assert mlp.down_proj.input_scale is None - assert not mlp._use_fused_relu2_quant - - -def test_nvfp4_mlp_enables_relu2_fp4_fusion_with_static_input_scale(): - model_config = ModelConfig(quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4)) - - with ( - patch("tensorrt_llm._torch.modules.mlp.get_sm_version", return_value=121), - patch("torch.ops.trtllm.fused_relu2_quantize", create=True), - ): - mlp = MLP( - hidden_size=32, - intermediate_size=64, - bias=False, - activation=relu2, - dtype=torch.bfloat16, - config=model_config, - reduce_output=False, - ) - mlp.create_weights() - - assert mlp.down_proj.input_scale is not None - assert mlp._use_fused_relu2_quant - - -def test_w4a16_nvfp4_mlp_rechecks_relu2_fp4_fusion_before_forward(): - model_config = ModelConfig(quant_config=QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4)) - - with ( - patch("tensorrt_llm._torch.modules.mlp.get_sm_version", return_value=121), - patch("torch.ops.trtllm.fused_relu2_quantize", create=True), - ): - mlp = MLP( - hidden_size=32, - intermediate_size=64, - bias=False, - activation=relu2, - dtype=torch.bfloat16, - config=model_config, - reduce_output=False, - ) - mlp.create_weights() - - # Simulate eligibility cached before weight loading replaced the linear - # method with W4A16, whose high-precision activation has no input_scale. - mlp._use_fused_relu2_quant = True - x_up = torch.tensor([[-2.0, 3.0]], dtype=torch.bfloat16) - with ( - patch.object(mlp.up_proj, "forward", return_value=x_up), - patch.object(mlp.down_proj, "forward", side_effect=lambda x: x), - patch.object( - MLP, - "_fused_relu2_quant", - side_effect=AssertionError("missing input_scale must use unfused ReLU2"), - ), - ): - output = mlp(torch.empty((1, 32), dtype=torch.bfloat16)) - - torch.testing.assert_close(output, relu2(x_up)) - - -def test_w4a16_nvfp4_mlp_disables_cutedsl_gelu_fusion(): - model_config = ModelConfig(quant_config=QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4)) - - with ( - patch("tensorrt_llm._torch.modules.mlp.get_sm_version", return_value=100), - patch("torch.ops.trtllm.cute_dsl_nvfp4_dense_gemm_gelu_blackwell", create=True), - patch( - "torch.ops.trtllm.cute_dsl_nvfp4_dense_gemm_gelu_fp4out_blackwell", - create=True, - ), - ): - mlp = MLP( - hidden_size=32, - intermediate_size=64, - bias=False, - activation=gelu_tanh, - dtype=torch.bfloat16, - config=model_config, - reduce_output=False, - ) - mlp.create_weights() - - assert not mlp._use_fused_gelu - assert not mlp._use_fused_gelu_fp4out - - -def test_dynamic_nvfp4_mlp_keeps_bf16_cutedsl_gelu_fusion(): - model_config = ModelConfig( - quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4), - force_dynamic_quantization=True, - ) - - with ( - patch("tensorrt_llm._torch.modules.mlp.get_sm_version", return_value=100), - patch("torch.ops.trtllm.cute_dsl_nvfp4_dense_gemm_gelu_blackwell", create=True), - patch( - "torch.ops.trtllm.cute_dsl_nvfp4_dense_gemm_gelu_fp4out_blackwell", - create=True, - ), - ): - mlp = MLP( - hidden_size=32, - intermediate_size=64, - bias=False, - activation=gelu_tanh, - dtype=torch.bfloat16, - config=model_config, - reduce_output=False, - ) - mlp.create_weights() - - assert mlp._use_fused_gelu - assert not mlp._use_fused_gelu_fp4out - - -def test_w4a16_nvfp4_gated_mlp_disables_cutedsl_swiglu_fusion(): - model_config = ModelConfig(quant_config=QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4)) - mlp = GatedMLP( - hidden_size=32, - intermediate_size=64, - bias=False, - dtype=torch.bfloat16, - config=model_config, - reduce_output=False, - use_cute_dsl_blockscaling_mm=True, - ) - - assert not mlp._can_fuse_gate_up_swiglu() - assert not mlp._can_fuse_gate_up_swiglu_fp4out() - - -def test_dynamic_nvfp4_gated_mlp_keeps_bf16_cutedsl_swiglu_fusion(): - model_config = ModelConfig( - quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4), - force_dynamic_quantization=True, - ) - mlp = GatedMLP( - hidden_size=32, - intermediate_size=64, - bias=False, - dtype=torch.bfloat16, - config=model_config, - reduce_output=False, - use_cute_dsl_blockscaling_mm=True, - ) - - assert mlp._can_fuse_gate_up_swiglu() - assert not mlp._can_fuse_gate_up_swiglu_fp4out() - - def test_w4a16_nvfp4_linear_uses_high_precision_activation_without_fp4_quantize(): method = W4A16NVFP4LinearMethod() input_tensor = torch.ones((2, 32), dtype=torch.bfloat16) @@ -641,31 +290,32 @@ def fail_fp4_quantize(*args, **kwargs): torch.testing.assert_close(output, expected) -def test_w4a16_nvfp4_linear_restores_high_rank_input_shape(): +@pytest.mark.parametrize( + "bad_input", + [ + pytest.param((torch.empty(1), torch.empty(1)), id="prequantized_fp4_tuple"), + pytest.param(torch.ones((2, 32), dtype=torch.float8_e4m3fn), id="fp8_activation"), + ], +) +def test_w4a16_nvfp4_linear_rejects_quantized_input(bad_input): + """W4A16 has no activation scale, so an upstream FP4/FP8 fusion must fail + loudly here rather than silently produce garbage.""" method = W4A16NVFP4LinearMethod() - input_tensor = torch.ones((2, 3, 32), dtype=torch.float16) module = SimpleNamespace( - weight=torch.empty((8, 16), dtype=torch.uint8), + weight=torch.empty((4, 16), dtype=torch.uint8), weight_scale=torch.empty((128 * 4,), dtype=torch.uint8), _w4a16_weight_scale_linear=torch.empty((128 * 4,), dtype=torch.uint8), - weight_scale_2=torch.tensor([0.5], dtype=torch.float32), - dtype=torch.float16, - out_features=8, + weight_scale_2=torch.tensor([0.25], dtype=torch.float32), + dtype=torch.bfloat16, + out_features=4, scaling_vector_size=16, + inv_input_scale=None, pre_quant_scale=None, use_custom_cublas_mm=False, ) - def fake_dequant(*args, **kwargs): - return torch.ones((8, 32), dtype=torch.float16) - - with patch( - "tensorrt_llm._torch.modules.fused_moe.triton_dequant_nvfp4.dequant_nvfp4_2d_triton", - side_effect=fake_dequant, - ): - output = method.apply(module, input_tensor, bias=None) - - assert output.shape == (2, 3, 8) + with pytest.raises(RuntimeError, match="high-precision input"): + method.apply(module, bad_input, bias=None) def test_w4a16_nvfp4_linear_scale_cache_is_nonpersistent_buffer(): @@ -811,64 +461,25 @@ def test_w4a16_nvfp4_marlin_selection_requires_supported_module(dtype, use_fused assert not MarlinNVFP4LinearMethod.is_supported(module) -@pytest.mark.parametrize("sm_version", [120, 121]) -def test_w4a16_nvfp4_linear_selects_marlin_by_default(sm_version): - quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) - - with ( - patch( - "tensorrt_llm._torch.modules.linear.get_sm_version", - return_value=sm_version, - ), - patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), - patch("torch.ops.trtllm.gptq_marlin_repack", create=True), - ): - linear = Linear( - 32, - 32, - bias=False, - dtype=torch.bfloat16, - quant_config=quant_config, - reduce_output=False, - ) - - assert isinstance(linear.quant_method, MarlinNVFP4LinearMethod) - assert isinstance(linear.quant_method, W4A16NVFP4LinearMethod) - - -@pytest.mark.parametrize("sm_version", [89, 90]) -def test_w4a16_nvfp4_linear_selects_marlin_when_explicitly_enabled(sm_version): - quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) - - with ( - patch( - "tensorrt_llm._torch.modules.linear.get_sm_version", - return_value=sm_version, - ), - patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), - patch("torch.ops.trtllm.gptq_marlin_repack", create=True), - ): - linear = Linear( - 32, - 32, - bias=False, - dtype=torch.bfloat16, - quant_config=quant_config, - reduce_output=False, - nvfp4_allowed_backends=["marlin"], - ) - - assert isinstance(linear.quant_method, MarlinNVFP4LinearMethod) - assert isinstance(linear.quant_method, W4A16NVFP4LinearMethod) - +@pytest.mark.parametrize( + ("sm_version", "dtype", "allowed_backends", "expect_marlin"), + [ + # Marlin is the default on SM120/121, opt-in on SM89-99, never on SM100. + (120, torch.bfloat16, None, True), + (121, torch.bfloat16, None, True), + (89, torch.bfloat16, ["marlin"], True), + (90, torch.bfloat16, ["marlin"], True), + (89, torch.bfloat16, None, False), + (90, torch.bfloat16, None, False), + (100, torch.bfloat16, None, False), + (120, torch.float16, None, False), # Marlin is bf16-only + ], +) +def test_w4a16_nvfp4_linear_method_selection(sm_version, dtype, allowed_backends, expect_marlin): + kwargs = {} if allowed_backends is None else {"nvfp4_allowed_backends": allowed_backends} -@pytest.mark.parametrize("sm_version", [89, 90]) -def test_w4a16_nvfp4_linear_uses_fallback_without_opt_in(sm_version): with ( - patch( - "tensorrt_llm._torch.modules.linear.get_sm_version", - return_value=sm_version, - ), + patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=sm_version), patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), patch("torch.ops.trtllm.gptq_marlin_repack", create=True), ): @@ -876,52 +487,14 @@ def test_w4a16_nvfp4_linear_uses_fallback_without_opt_in(sm_version): 32, 32, bias=False, - dtype=torch.bfloat16, + dtype=dtype, quant_config=QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4), reduce_output=False, + **kwargs, ) - assert type(linear.quant_method) is W4A16NVFP4LinearMethod - - -def test_w4a16_nvfp4_linear_uses_default_method_on_sm100(): - quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) - - with ( - patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=100), - patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), - patch("torch.ops.trtllm.gptq_marlin_repack", create=True), - ): - linear = Linear( - 32, - 32, - bias=False, - dtype=torch.bfloat16, - quant_config=quant_config, - reduce_output=False, - ) - - assert type(linear.quant_method) is W4A16NVFP4LinearMethod - - -def test_w4a16_nvfp4_linear_keeps_default_method_for_fp16(): - quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) - - with ( - patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=120), - patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), - patch("torch.ops.trtllm.gptq_marlin_repack", create=True), - ): - linear = Linear( - 32, - 32, - bias=False, - dtype=torch.float16, - quant_config=quant_config, - reduce_output=False, - ) - - assert type(linear.quant_method) is W4A16NVFP4LinearMethod + assert isinstance(linear.quant_method, W4A16NVFP4LinearMethod) + assert isinstance(linear.quant_method, MarlinNVFP4LinearMethod) is expect_marlin def test_w4a16_nvfp4_post_load_preserves_checkpoint_weight_global_scale(): @@ -1002,36 +575,3 @@ def test_causal_lm_head_uses_global_w4a16_nvfp4_config( assert isinstance(causal_lm.lm_head.quant_method, W4A16NVFP4LinearMethod) else: assert not causal_lm.lm_head.has_any_quant - - -def test_lm_head_w4a16_nvfp4_forward_uses_triton_dequant(): - quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) - lm_head = LMHead( - num_embeddings=4, - embedding_dim=32, - dtype=torch.float16, - quant_config=quant_config, - ) - input_tensor = torch.ones((2, 32), dtype=torch.float16) - lm_head._w4a16_weight_scale_linear = torch.empty((128 * 4,), dtype=torch.uint8) - captured = {} - - def fake_dequant(weight, weight_scale, weight_scale_2, **kwargs): - captured["weight"] = weight - captured["weight_scale"] = weight_scale - captured["weight_scale_2"] = weight_scale_2 - captured.update(kwargs) - return torch.ones((4, 32), dtype=torch.float16) - - with patch( - "tensorrt_llm._torch.modules.fused_moe.triton_dequant_nvfp4.dequant_nvfp4_2d_triton", - side_effect=fake_dequant, - ): - output = lm_head(input_tensor) - - assert captured["weight"].data_ptr() == lm_head.weight.data_ptr() - assert captured["weight_scale"] is lm_head._w4a16_weight_scale_linear - assert captured["weight_scale_2"] is lm_head.weight_scale_2 - assert captured["target_dtype"] is torch.float16 - assert captured["sf_vec_size"] == 16 - assert output.shape == (2, 4) From 8316a9826e2f09481e8a1230962c8e6942d495f6 Mon Sep 17 00:00:00 2001 From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> Date: Fri, 7 Aug 2026 04:42:23 +0000 Subject: [PATCH 25/25] fix hopper nvfp4 Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> --- tensorrt_llm/_torch/modules/linear.py | 17 ++++-- .../_torch/modules/test_w4a16_nvfp4_linear.py | 52 ++++++++++++------- 2 files changed, 48 insertions(+), 21 deletions(-) diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index ec17b2a8d791..9524c4db2851 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -31,6 +31,7 @@ from ..._utils import get_sm_version, is_sm_100f from ...models.modeling_utils import QuantConfig from ..utils import (Fp4QuantizedTensor, get_model_extra_attrs, + is_nvfp4_marlin_supported_sm, replace_parameter_and_save_metadata, unswizzle_sf) @@ -3503,9 +3504,19 @@ def __init__( def get_quant_method(self, quant_config: Optional[QuantConfig] = None): quant_method = get_quant_method(quant_config) - use_marlin = type(quant_method) is W4A16NVFP4LinearMethod - if use_marlin and MarlinNVFP4LinearMethod.is_enabled(self): - return MarlinNVFP4LinearMethod() + method_type = type(quant_method) + if method_type is W4A16NVFP4LinearMethod: + # Marlin by default on SM120/121, explicit opt-in elsewhere. + if MarlinNVFP4LinearMethod.is_enabled(self): + return MarlinNVFP4LinearMethod() + elif method_type is NVFP4LinearMethod: + # The Marlin kernel is W4A16, so an explicit opt-in on a W4A4 + # checkpoint runs the weight-only method, which pads N/K to the + # tile sizes the kernel requires. + if ("marlin" in self.nvfp4_allowed_backends + and is_nvfp4_marlin_supported_sm() + and MarlinNVFP4LinearMethod.is_supported(self)): + return MarlinNVFP4LinearMethod() return quant_method @staticmethod diff --git a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py index b3d6e5782bca..f169105914f4 100644 --- a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py +++ b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py @@ -111,8 +111,8 @@ def test_nvfp4_attention_keeps_high_precision_output_for_hopper_marlin(): ) assert o_proj.uses_marlin_nvfp4 - assert type(o_proj.quant_method) is NVFP4LinearMethod - assert o_proj.has_nvfp4_activation_quantization + assert isinstance(o_proj.quant_method, MarlinNVFP4LinearMethod) + assert not o_proj.has_nvfp4_activation_quantization assert not Attention._use_quantize_output(attention) @@ -161,26 +161,37 @@ def fake_nvfp4_gemm(*args, **kwargs): @pytest.mark.parametrize( - ("sm_version", "dtype", "allowed_backends", "expect_marlin"), + ("sm_version", "allowed_backends", "expect_marlin"), [ - (89, torch.bfloat16, None, False), - (90, torch.bfloat16, None, False), - (120, torch.bfloat16, None, False), - (121, torch.bfloat16, None, False), - (89, torch.bfloat16, ["marlin"], True), - (90, torch.bfloat16, ["marlin"], True), - (90, torch.float16, None, False), + # Honoured only on Ada/Hopper, and only when opted in. + (89, ["marlin"], True), + (90, ["marlin"], True), + (89, None, False), + (90, None, False), + (120, None, False), + (121, None, False), + # Opt-in ignored off Ada/Hopper: Marlin is not the right NVFP4 backend + # on SM120/121, where the W4A4 kernels are faster. + (120, ["marlin"], False), + (121, ["marlin"], False), ], ) -def test_nvfp4_linear_keeps_activation_quant_method( - sm_version, dtype, allowed_backends, expect_marlin +def test_nvfp4_linear_marlin_opt_in_switches_to_weight_only_method( + sm_version, allowed_backends, expect_marlin ): - """Plain NVFP4 always keeps NVFP4LinearMethod -- Marlin is reached inside the - GEMM op, so only ``uses_marlin_nvfp4`` flips.""" + """The Marlin kernel is W4A16, so opting a W4A4 checkpoint into it converts + the layer to the weight-only method, which pads N/K for the kernel. Without + the opt-in the layer keeps NVFP4LinearMethod and its activation quantize.""" kwargs = {} if allowed_backends is None else {"nvfp4_allowed_backends": allowed_backends} with ( patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=sm_version), + # is_nvfp4_marlin_supported_sm reads its own get_sm_version binding in + # utils.py, so it has to be patched alongside linear's. + patch( + "tensorrt_llm._torch.modules.linear.is_nvfp4_marlin_supported_sm", + return_value=89 <= sm_version < 100, + ), patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), patch("torch.ops.trtllm.gptq_marlin_repack", create=True), ): @@ -188,15 +199,20 @@ def test_nvfp4_linear_keeps_activation_quant_method( 32, 32, bias=False, - dtype=dtype, + dtype=torch.bfloat16, quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4), reduce_output=False, **kwargs, ) - assert type(linear.quant_method) is NVFP4LinearMethod - assert linear.has_nvfp4_activation_quantization - assert linear.uses_marlin_nvfp4 is expect_marlin + # Asserted inside the patch: uses_marlin_nvfp4 re-reads the SM on every + # access rather than caching what create_weights decided. + assert isinstance(linear.quant_method, MarlinNVFP4LinearMethod) is expect_marlin + assert linear.uses_marlin_nvfp4 is expect_marlin + # Marlin consumes BF16 activations; the plain NVFP4 path quantizes them. + assert linear.has_nvfp4_activation_quantization is not expect_marlin + if not expect_marlin: + assert type(linear.quant_method) is NVFP4LinearMethod def test_nvfp4_linear_hopper_marlin_applies_bias_as_post_op():