[TRTLLM-12720][feat] Support nvfp4 w4a16 on sm120 - #16511
[TRTLLM-12720][feat] Support nvfp4 w4a16 on sm120#16511pamelap-nvidia wants to merge 23 commits into
Conversation
Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
|
PR_Github #60774 [ run ] completed with state
|
Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> # Conflicts: # tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py # tensorrt_llm/_torch/modules/mlp.py
Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #63362 [ run ] triggered by Bot. Commit: |
|
PR_Github #63362 [ run ] completed with state |
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/models/modeling_deepseekv3.py (1)
755-756: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winAdd the NVIDIA copyright header to
tensorrt_llm/_torch/models/modeling_deepseekv3.py. Keep the DeepSeek 2023 notice and add the 2026 NVIDIA SPDX header used intensorrt_llm/_torch/models/modeling_deepseekv4.py.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/models/modeling_deepseekv3.py` around lines 755 - 756, Add the 2026 NVIDIA SPDX copyright header to modeling_deepseekv3.py while preserving its existing DeepSeek 2023 notice; use modeling_deepseekv4.py as the reference for the exact header format. The referenced line in modeling_deepseekv4.py requires no direct change.Source: Coding guidelines
🧹 Nitpick comments (1)
tensorrt_llm/_torch/models/modeling_deepseekv3.py (1)
755-756: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse the Python 3.10 annotation form for all new
aux_stream_dictparameters.The same new parameter uses
Optional[Dict[...]]in both files. Applydict[...] | Noneat each site.
tensorrt_llm/_torch/models/modeling_deepseekv3.py#L755-L756: updateDeepseekV3Attention.aux_stream_dict.tensorrt_llm/_torch/models/modeling_deepseekv3.py#L805-L806: updateDeepseekV32Attention.aux_stream_dict.tensorrt_llm/_torch/models/modeling_deepseekv4.py#L1319-L1319: updateDeepseekV4Attention.aux_stream_dict.As per coding guidelines, prefer built-in generic types and
|for Python 3.10+ annotations.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/models/modeling_deepseekv3.py` around lines 755 - 756, Replace the Optional[Dict[...]] annotation with the Python 3.10 dict[...] | None form for aux_stream_dict in DeepseekV3Attention and DeepseekV32Attention in tensorrt_llm/_torch/models/modeling_deepseekv3.py at lines 755-756 and 805-806, and in DeepseekV4Attention in tensorrt_llm/_torch/models/modeling_deepseekv4.py at line 1319. Use built-in generic types consistently at all three sites.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@tensorrt_llm/_torch/models/modeling_deepseekv3.py`:
- Around line 755-756: Add the 2026 NVIDIA SPDX copyright header to
modeling_deepseekv3.py while preserving its existing DeepSeek 2023 notice; use
modeling_deepseekv4.py as the reference for the exact header format. The
referenced line in modeling_deepseekv4.py requires no direct change.
---
Nitpick comments:
In `@tensorrt_llm/_torch/models/modeling_deepseekv3.py`:
- Around line 755-756: Replace the Optional[Dict[...]] annotation with the
Python 3.10 dict[...] | None form for aux_stream_dict in DeepseekV3Attention and
DeepseekV32Attention in tensorrt_llm/_torch/models/modeling_deepseekv3.py at
lines 755-756 and 805-806, and in DeepseekV4Attention in
tensorrt_llm/_torch/models/modeling_deepseekv4.py at line 1319. Use built-in
generic types consistently at all three sites.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 60d3f155-b180-4cb0-a174-45b6cd94b449
📒 Files selected for processing (4)
cpp/tensorrt_llm/common/attentionOp.cpptensorrt_llm/_torch/models/modeling_deepseekv3.pytensorrt_llm/_torch/models/modeling_deepseekv4.pytests/integration/defs/accuracy/test_llm_api_pytorch.py
🚧 Files skipped from review as they are similar to previous changes (2)
- cpp/tensorrt_llm/common/attentionOp.cpp
- tests/integration/defs/accuracy/test_llm_api_pytorch.py
Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
|
/bot run |
|
PR_Github #63492 [ run ] triggered by Bot. Commit: |
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tensorrt_llm/_torch/modules/fused_moe/create_moe.py (2)
133-150: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winReturn
CutlassFusedMoEwhen FlashInfer is unavailable on SM120/SM121.For standard NVFP4, this branch returns
CuteDslFusedMoE, butCuteDslFusedMoE.can_implement()rejects SM120/SM121.CutlassFusedMoEsupports NVFP4 on both architectures.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/modules/fused_moe/create_moe.py` around lines 133 - 150, Update the standard NVFP4 fallback in the surrounding MoE selection logic to return CutlassFusedMoE when FlashInfer is unavailable on SM120 or SM121. Preserve the existing CuteDslFusedMoE selection for architectures it supports, and align this path with the can_implement behavior of CuteDslFusedMoE and CutlassFusedMoE.
79-81: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winResolve the per-layer quantization configuration before backend selection.
Qwen3MoEpasseslayer_idxbut nooverride_quant_config, so both functions use the globalmodel_config.quant_configand ignorequant_config_dict. This selects the wrong backend for mixed-precision experts. Resolvemodel.layers.{layer_idx}.mlp.expertsin both functions, or pass it from every caller.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/modules/fused_moe/create_moe.py` around lines 79 - 81, The quantization setup currently ignores layer-specific entries in quant_config_dict when override_quant_config is absent. Update both affected functions to resolve model.layers.{layer_idx}.mlp.experts before backend selection, while preserving override_quant_config precedence; alternatively ensure every caller, including Qwen3MoE, passes the resolved per-layer configuration.
🧹 Nitpick comments (1)
tests/unittest/_torch/modules/moe/test_moe_backend.py (1)
330-334: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd CI coverage for the new capability test and fallback paths.
test_moe_nvfp4_activation_quantization_capabilityis not selected by the existingtest_moe_backendentries in the applicabletest-dbfiles. Add an explicit entry and cover W4A16 NVFP4 selection, missing FlashInfer, and unsupported SM fallback.Test coverage summary — needs follow-up. Added:
test_moe_nvfp4_activation_quantization_capability. Removed test: function name not determinable.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/modules/moe/test_moe_backend.py` around lines 330 - 334, Add explicit test-db entries for test_moe_nvfp4_activation_quantization_capability and the related W4A16 NVFP4 selection, missing FlashInfer, and unsupported-SM fallback tests, ensuring the applicable test_moe_backend suites select them in CI. Preserve existing test-db conventions and include every relevant fallback-path test.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@tensorrt_llm/_torch/modules/fused_moe/create_moe.py`:
- Around line 133-150: Update the standard NVFP4 fallback in the surrounding MoE
selection logic to return CutlassFusedMoE when FlashInfer is unavailable on
SM120 or SM121. Preserve the existing CuteDslFusedMoE selection for
architectures it supports, and align this path with the can_implement behavior
of CuteDslFusedMoE and CutlassFusedMoE.
- Around line 79-81: The quantization setup currently ignores layer-specific
entries in quant_config_dict when override_quant_config is absent. Update both
affected functions to resolve model.layers.{layer_idx}.mlp.experts before
backend selection, while preserving override_quant_config precedence;
alternatively ensure every caller, including Qwen3MoE, passes the resolved
per-layer configuration.
---
Nitpick comments:
In `@tests/unittest/_torch/modules/moe/test_moe_backend.py`:
- Around line 330-334: Add explicit test-db entries for
test_moe_nvfp4_activation_quantization_capability and the related W4A16 NVFP4
selection, missing FlashInfer, and unsupported-SM fallback tests, ensuring the
applicable test_moe_backend suites select them in CI. Preserve existing test-db
conventions and include every relevant fallback-path test.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 0f87bc57-198a-4222-bf72-1654d8ab2cfe
📒 Files selected for processing (5)
tensorrt_llm/_torch/model_config.pytensorrt_llm/_torch/modules/fused_moe/create_moe.pytests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.pytests/unittest/_torch/modules/moe/test_moe_backend.pytests/unittest/_torch/test_model_config.py
💤 Files with no reviewable changes (2)
- tests/unittest/_torch/test_model_config.py
- tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/_torch/model_config.py
|
PR_Github #63492 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63514 [ run ] triggered by Bot. Commit: |
|
PR_Github #63514 [ run ] completed with state
|
|
/bot run |
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
tensorrt_llm/_torch/modules/linear.py (1)
2127-2132: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd type annotations to the override.
get_tp_alignment()omits parameter and return annotations. Add precise types fortp_mode,quant_config, and the integer return value.Based on coding guidelines, “Annotate every function.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/modules/linear.py` around lines 2127 - 2132, Update the get_tp_alignment override to add precise type annotations for tp_mode and quant_config, and annotate its return value as int; preserve the existing parameters, default value, and return of 32.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/modules/linear.py`:
- Around line 2024-2026: Update W4A16NVFP4LinearMethod._prepare_input() to
explicitly reject FP8 input with a RuntimeError rather than an assertion,
consistent with the cleared module.inv_input_scale state. Add a regression test
verifying FP8 input raises RuntimeError while preserving supported W4A16
behavior.
In `@tests/unittest/_torch/thop/parallel/test_fp4_linear.py`:
- Line 815: Remove the obsolete FP4 activation quantization setup in the test
around the W4A16_NVFP4 QuantConfig, including the unused x_fp4 and x_sf_block
assignments; pass the high-precision x directly to l_marlin and retain only
values still consumed by the test.
---
Nitpick comments:
In `@tensorrt_llm/_torch/modules/linear.py`:
- Around line 2127-2132: Update the get_tp_alignment override to add precise
type annotations for tp_mode and quant_config, and annotate its return value as
int; preserve the existing parameters, default value, and return of 32.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 7036ddf7-7d41-457e-8579-d319b1695e0e
📒 Files selected for processing (4)
tensorrt_llm/_torch/modules/linear.pytests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.pytests/unittest/_torch/modules/test_w4a16_nvfp4_linear.pytests/unittest/_torch/thop/parallel/test_fp4_linear.py
💤 Files with no reviewable changes (1)
- tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py
| module.input_scale = None | ||
| module.inv_input_scale = None | ||
| module.alpha = None |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ast-grep outline tensorrt_llm/_torch/modules/linear.py \
--items all \
--match 'W4A16NVFP4LinearMethod|MarlinNVFP4LinearMethod' \
--view expanded
rg -n -C 3 --type py \
'torch\.float8_e4m3fn|Fp4QuantizedTensor|W4A16NVFP4LinearMethod|is_static_nvfp4_input_eligible' \
tensorrt_llm testsRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- linear.py target methods ---'
sed -n '1980,2155p' tensorrt_llm/_torch/modules/linear.py
printf '%s\n' '--- W4A16NVFP4LinearMethod references ---'
rg -n -C 5 --type py 'W4A16NVFP4LinearMethod|_prepare_input\(' tensorrt_llm/_torch tests/unittest/_torch
printf '%s\n' '--- FP8 input production and invocation paths ---'
rg -n -C 5 --type py \
'to\(torch\.float8_e4m3fn\)|dtype=torch\.float8_e4m3fn|input_scale|inv_input_scale' \
tensorrt_llm/_torch/modules tensorrt_llm/_torch/models tests/unittest/_torch/modulesRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- target implementation ---'
sed -n '1995,2145p' tensorrt_llm/_torch/modules/linear.py
printf '%s\n' '--- direct references in linear.py and nearby tests ---'
rg -n -C 4 \
'W4A16NVFP4LinearMethod|torch\.float8_e4m3fn|inv_input_scale|input_scale' \
tensorrt_llm/_torch/modules/linear.py \
tests/unittest/_torch/modules/linear.py \
tests/unittest/_torch/modules/test_linear.py \
tests/unittest/_torch/modules 2>/dev/null | head -n 1000
printf '%s\n' '--- callers of Linear with FP8 activation configuration ---'
rg -n -C 4 --type py \
'dtype_activation|activation_dtype|torch\.float8_e4m3fn|w4a16|W4A16' \
tensorrt_llm/_torch/modules tensorrt_llm/_torch/models tests/unittest/_torch/modules \
| head -n 1500Repository: NVIDIA/TensorRT-LLM
Length of output: 50375
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- W4A16 and activation-quantization symbols ---'
rg -n -C 6 --type py \
'W4A16NVFP4LinearMethod|MarlinNVFP4LinearMethod|quantizes_nvfp4_activations|is_static_nvfp4_input_eligible' \
tensorrt_llm/_torch/modules/linear.py \
tensorrt_llm/_torch/modules/mla.py \
tensorrt_llm/_torch/modules/rms_norm.py \
tests/unittest/_torch/modules 2>/dev/null | head -n 1600
printf '%s\n' '--- upstream FP8-producing paths near Linear calls ---'
rg -n -C 5 --type py \
'torch\.float8_e4m3fn|fp8_quant|quantize.*fp8|return.*fp8|Linear\(' \
tensorrt_llm/_torch/modules/mla.py \
tensorrt_llm/_torch/modules/rms_norm.py \
tensorrt_llm/_torch/modules/linear.py \
tests/unittest/_torch/modules 2>/dev/null | head -n 1600
printf '%s\n' '--- tests mentioning W4A16/Marlin/NVFP4 linear ---'
rg -l --type py \
'W4A16NVFP4LinearMethod|MarlinNVFP4LinearMethod|W4A16|Marlin.*NVFP4|NVFP4.*Marlin' \
tests tensorrt_llm/_torch | sortRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- current change for W4A16 input handling ---'
git diff --unified=12 -- tensorrt_llm/_torch/modules/linear.py \
| sed -n '/W4A16NVFP4LinearMethod/,/MarlinNVFP4LinearMethod/p'
printf '%s\n' '--- canonical eligibility predicate ---'
sed -n '3890,3935p' tensorrt_llm/_torch/modules/linear.py
printf '%s\n' '--- relevant existing tests ---'
sed -n '540,640p' tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py
sed -n '780,860p' tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py
printf '%s\n' '--- Linear input dispatch ---'
sed -n '3600,3760p' tensorrt_llm/_torch/modules/linear.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 15225
Reject FP8 input explicitly in W4A16NVFP4LinearMethod._prepare_input().
W4A16 is a weight-only path and clears module.inv_input_scale. Replace the FP8 assertion with a RuntimeError and add a regression test.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/modules/linear.py` around lines 2024 - 2026, Update
W4A16NVFP4LinearMethod._prepare_input() to explicitly reject FP8 input with a
RuntimeError rather than an assertion, consistent with the cleared
module.inv_input_scale state. Add a regression test verifying FP8 input raises
RuntimeError while preserving supported W4A16 behavior.
| bias=False, | ||
| dtype=dtype, | ||
| quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4), | ||
| quant_config=QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4), |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Remove unused FP4 activation quantization.
After this change, the test passes high-precision x to l_marlin. x_fp4 and x_sf_block at lines 839-840 have no use. Ruff F841 rejects these assignments. Remove the obsolete quantization setup.
Proposed fix
- x_sf_global = (448 * 6) / x.abs().max().float()
- x_fp4, x_sf_block = torch.ops.trtllm.fp4_quantize(
- x, x_sf_global, scaling_vector_size, False)
-
with torch.inference_mode():
output = l_marlin(x)Based on learnings, Ruff enables F rules in this repository.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/_torch/thop/parallel/test_fp4_linear.py` at line 815, Remove
the obsolete FP4 activation quantization setup in the test around the
W4A16_NVFP4 QuantConfig, including the unused x_fp4 and x_sf_block assignments;
pass the high-precision x directly to l_marlin and retain only values still
consumed by the test.
Source: Learnings
|
PR_Github #63638 [ run ] triggered by Bot. Commit: |
|
PR_Github #63638 [ run ] completed with state
|
|
|
||
| def _canonicalize_quant_weight(key: str, value): |
There was a problem hiding this comment.
please add a comment about why and which quantization config will use the if condition.
| return layer_quant_config | ||
|
|
||
| quant_config = model_config.quant_config | ||
| if (quant_config is not None and quant_config.quant_algo |
There was a problem hiding this comment.
it is quite confusing to me.
what if we pass MARLIN moe backend on Hopper, would it fallback to NVFP4 will cause issue?
Dev Engineer Review
The changes add SM120/121 PyTorch support for W4A16 NVFP4 inference across dense linear, LM-head, and MoE paths.
The reported local validation passes 74 merge-focused tests and 69 W4A16/model regression tests. Six CI runs failed or had L0 pipeline failures. Failure analysis and a new NVIDIA-triggered run are required.
QA Engineer Review
The PR adds test functions and coverage for:
No files under
tests/integration/test_lists/,test-db/,qa/, orwaives.txtwere modified. Coverage registration for the added test functions is not shown.Verdict: needs follow-up.
Summary
Adds end-to-end support for W4A16 NVFP4 checkpoints (packed FP4 weights + per-group FP8
scales, high-precision activations) in the PyTorch backend, targeting Qwen3.5/3.6 and
Nemotron-H NVFP4 exports on SM120/121 (RTX 5090/6000, GB10) and Hopper.
The core problem this solves: TRT-LLM previously conflated "has NVFP4 weights" with
"quantizes activations to NVFP4". W4A16 breaks that assumption — it has FP4 weights but
consumes BF16/FP16 activations and has no calibrated
input_scale. Every FP4 fusion sitekeyed on
has_nvfp4would therefore feed anFp4QuantizedTensorinto a layer that can'tconsume one, or read a
Noneactivation scale.Changes
Checkpoint parsing
quant_config_utils.py: compressed-tensors configs withinput_activations: null+format: nvfp4-pack-quantized+ 4-bit float weights atgroup_size=16now resolve toQuantAlgo.W4A16_NVFP4instead of raising.model_config.py/llm_utils.py: whenhf_quant_config.jsondeclaresNVFP4but theinline
config.jsoncompressed-tensors block says weight-only, the inline activationsemantics win. Guarded so it doesn't re-enter the config.json global-quant fallback.
nemotron_h_weight_mapper.py: canonicalizes compressed-tensors key names(
weight_packed->weight,weight_global_scale->weight_scale_2with reciprocalconversion).
Activation-quantization capability, made explicit
LinearMethodBase.quantizes_nvfp4_activationsClassVar, surfaced asLinear.has_nvfp4_activation_quantizationandMoE.has_nvfp4_activation_quantization.has_nvfp4: RMSNorm folding(
is_static_nvfp4_input_eligible), attention FP4 output, fused ReLU2/GELU/SwiGLUepilogues, GEMM-allreduce fusion, and the
POST_MOE_FUSIONpaths in DeepSeek V3/V4, GLM,EXAONE and Qwen3-MoE.
Linear
W4A16NVFP4LinearMethod: on-the-fly Triton weight dequant. Leavesinput_scale/inv_input_scale/alphaasNonefromcreate_weightsonward; the parentprocess_weights_after_loading_*now skip scales whose destination was never allocated.MarlinNVFP4LinearMethod: Marlin-backed W4A16, selected by default for BF16 onSM90/120/121. Handles its own K->64 / N->128 padding.
buffer, which also makes the GMS
cache_derived_staterefresh idempotent (the previousin-place unswizzle would double-apply on re-materialization).
MoE
CuteDslB12xFusedMoEacceptsW4A16_NVFP4and stays on the b12x path for both prefilland decode (
quant_mode="w4a16"); falls back to CUTLASS for EP / attention-DP ornon-SM120/121.
W4A16NVFP4CutlassFusedMoEMethodon the CUTLASS path; shared per-layer quant-configlookup so mixed-precision checkpoints resolve expert configs correctly.
body sublayers to CUTLASS.
C++
120fin addition to90; relax theruntime SM asserts to
9.x || 12.x. No kernel logic changes.No LLM API surface change —
llm_args.pyis untouched, so no golden-manifestregeneration.
Behavior changes for existing users
via
nvfp4_gemm_config.allowed_backends: [..., marlin]). Attention FP4 output and thetorch.compile inplace op are disabled on those layers as a consequence.
AUTOresolves toCUTEDSLfor W4A16 NVFP4 on SM120/121.Test coverage
test_w4a16_nvfp4_linear.py(new) covers method selection across SM90/100/120/121 and dtype,the Marlin-vs-Triton numerical match on GPU (including K/N padding shapes), scale-cache buffer
semantics, the mixed-precision
use_fused_gemm_allreducerebind, and every fusion gate thatchanged. Plus new checkpoint-parsing tests (
test_model_config.py,test_quant_config_utils.py,test_kv_cache_dtype_override.py), Nemotron-H MoE/MTP configtests, the weight-mapper tests, and W4A16 cases added to the b12x MoE backend suite.
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.