[SM120] NVFP4 static _input_prepare causes illegal memory access on RTX 5080

[lucaspirola]

System Info

Component	Version
GPU	NVIDIA GeForce RTX 5080 (SM120, compute_cap 12.0)
TRT-LLM	1.3.0rc9 (built from source, SM120 target)
PyTorch	2.10.0a0+b4e4ee81d3.nv25.12
CUDA	13.1
Driver	590.48.01
Container	nvcr.io/nvidia/pytorch:25.12-py3

Bug Description

The static branch of NVFP4LinearMethod._input_prepare (in tensorrt_llm/_torch/modules/linear.py, line ~1287) causes cudaErrorIllegalAddress on SM120 during the TRT-LLM model warmup forward pass. The dynamic branch (line ~1280, triggered by force_dynamic_quantization=True) works correctly.

Both branches call the same fp4_quantize + nvfp4_gemm kernels — the only difference is whether input_scale / alpha come from module Parameters (static) or freshly computed tensors (dynamic).

Affected code

# tensorrt_llm/_torch/modules/linear.py, NVFP4LinearMethod._input_prepare:

if module.input_scale is None or module.force_dynamic_quantization:
    # Dynamic mode — WORKS on SM120
    amax_input = torch.amax(torch.abs(input)).float()
    input_scale = FP8_MAX * E2M1_MAX / amax_input
    alpha = (amax_input / (FP8_MAX * E2M1_MAX)) * module.weight_scale_2
else:
    # Static mode — CRASHES on SM120
    input_scale = module.input_scale
    alpha = module.alpha

# Same call in both branches:
act_fp4, act_sf = torch.ops.trtllm.fp4_quantize(
    input, input_scale, module.scaling_vector_size, False)
return act_fp4, act_sf, alpha

Key observations

• The crash is NOT value-dependent: input_scale=1.0 (checkpoint placeholder) and input_scale=270.0 (reasonable estimate) both crash
• fp4_quantize and nvfp4_gemm work correctly in isolation with both 0D and [1]-shaped scale tensors
• The crash only manifests during the full model forward through PyExecutor warmup
• The error surfaces as CUBLAS_STATUS_INTERNAL_ERROR in a subsequent unquantized BF16 cublasGemmEx call, suggesting corrupted CUDA state from the static NVFP4 path
• Disabling CUDA graphs (cuda_graph_config=None) does not fix it — the static path is broken regardless

Reproduction

from tensorrt_llm import LLM
from tensorrt_llm.llmapi import KvCacheConfig

kv = KvCacheConfig(max_tokens=512, free_gpu_memory_fraction=0.05,
                   enable_block_reuse=False, dtype="fp8")

# ✅ This works (dynamic path):
llm = LLM(model=NVFP4_CHECKPOINT, force_dynamic_quantization=True,
          tensor_parallel_size=1, max_num_tokens=128, max_batch_size=1,
          kv_cache_config=kv, skip_tokenizer_init=True)

# ❌ This crashes with "illegal memory access" (static path):
llm = LLM(model=NVFP4_CHECKPOINT,
          tensor_parallel_size=1, max_num_tokens=128, max_batch_size=1,
          kv_cache_config=kv, skip_tokenizer_init=True)

Tested with a Glm4MoeLiteForCausalLM (GLM-4.7-Flash) NVFP4 checkpoint (235 NVFP4 linear modules). The model uses MLA attention with BF16 kv_b_proj excluded from NVFP4 — the cuBLAS error occurs in this unquantized layer after the static NVFP4 path corrupts CUDA state.

Error trace

RuntimeError: CUDA error: CUBLAS_STATUS_INTERNAL_ERROR when calling
cublasGemmEx(..., CUDA_R_16BF, ..., CUBLAS_GEMM_DEFAULT_TENSOR_OP)

# Originates from warmup forward → layer 0 → MLA attention →
# kv_b_proj (unquantized BF16 Linear) → F.linear()

Workaround

Set force_dynamic_quantization=True either via the LLM API or per-module in post_load_weights:

from tensorrt_llm._torch.modules.linear import NVFP4LinearMethod

# In model's post_load_weights():
for name, module in self.named_modules():
    if hasattr(module, 'quant_method') and isinstance(
        module.quant_method, NVFP4LinearMethod
    ):
        module.force_dynamic_quantization = True

This costs ~18% decode throughput (235 extra torch.amax() reductions per decode step for dynamic activation scaling).

[pengbowang-nv]

Hi @lucaspirola , could you please provide the checkpoint you used so we can take investigate it? Thanks!

[waynehacking8]

Cannot reproduce on RTX PRO 6000 (SM120) — trtllm 1.2.1

Tested the static NVFP4 _input_prepare path on an RTX PRO 6000 Blackwell (SM120, 96 GB) with tensorrt_llm 1.2.1 (pip). Three independent test axes, all clean:

1. Dense model (static input_scale):
Qwen3-0.6B-FP4 — warmup + generate clean, EXIT=0, coherent output.

2. MLA model (static input_scale + BF16 kv_b_proj):
DeepSeek-V2-Lite (MLA, kv_lora_rank=512, 64 experts) quantized to NVFP4 with modelopt 0.37 (NVFP4_DEFAULT_CFG, kv_b_proj excluded / kept BF16). This exercises the exact code path described: static NVFP4 _input_prepare interleaved with BF16 cublasGemmEx in the unquantized kv_b_proj during warmup.

torch 2.9.1+cu128 cc (12, 0) NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition
[TensorRT-LLM] TensorRT LLM version: 1.2.1
>>> LLM constructed + warmed up OK (static NVFP4 MLA path did NOT crash)
>>> GENERATE OK: ' the most beautiful city in the world.\nThe Eiffel Tower, ...'
>>> RESULT: #12411 does NOT reproduce on RTX PRO 6000 (MLA, trtllm 1.2.1)

3. Isolated NVFP4 kernel sweep:
All input_scale provenance variants (host/device, scalar/per-tensor, contiguous/non-contiguous) × all backends (cutlass, cublaslt, cuda_core) — all numerically correct.

4. compute-sanitizer memcheck (dense path):
34 hits, all benign CUDA_ERROR_INVALID_VALUE on cuGetProcAddress_v2 at import time. Zero illegal memory access / invalid read / invalid write.

Loader bugs found during MLA triage (separate from #12411)

While loading DeepSeek-V2-Lite NVFP4 through the DeepseekV3 weight loader, I hit three crashes in modeling_deepseekv3.py — all caused by V3-specific assumptions that break on V2-Lite (q_lora_rank=None, no e_score_correction_bias):

1. q_a_proj KeyError — nvfp4_fused_a check reads q_a_proj.weight unconditionally, but V2-Lite has no q_a_proj. (Already fixed in main with a cleaner if not is_lite guard.)
2. e_score_correction_bias KeyError — DeepseekV3Gate.load_weights accesses this unconditionally, but V2-Lite doesn't have the correction bias. (Still unfixed in main — will open a PR.)
3. kv_b_proj reshape — NVFP4-packed kv_b_proj has shape [out, in//2] but MLA weight absorption needs full BF16 [out, in]. (Workaround: exclude kv_b_proj from quantization.)

Remaining gaps

• Version gap: tested 1.2.1 (pip max); report is on 1.3.0rc9 (source-build only).
• Checkpoint gap: GLM-4.7-Flash NVFP4 is not publicly available. DeepSeek-V2-Lite is the closest public MLA model.
• Card gap: RTX PRO 6000 (workstation) vs RTX 5080 (consumer). Both SM120 but different memory subsystems.

@pengbowang-nv — I can share the DeepSeek-V2-Lite NVFP4 checkpoint (bf16 kv_b_proj) and repro scripts if useful.