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

System Info

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).

Start in tensorrt_llm/_torch/modules/linear.py at NVFP4LinearMethod._input_prepare and compare the static and dynamic branches. Reproduce the RTX 5080 warmup failure with the provided LLM configuration, then inspect the subsequent unquantized kv_b_proj BF16 call. Done means static NVFP4 execution no longer causes the illegal-memory-access or following CUBLAS error while retaining the working dynamic-path behavior.