From 32b20db5e7d4ec2ccdfd205539373b08976d9017 Mon Sep 17 00:00:00 2001 From: tijyojwad <1127155+tijyojwad@users.noreply.github.com> Date: Tue, 26 May 2026 16:22:18 -0700 Subject: [PATCH] [None][fix] Bypass FlashInfer SSD prefill to fix state dtype precision The FlashInfer SSD prefill path (added in PR #12731) hardcodes state_dtype=bfloat16 in the SSDCombined kernel, while the Triton path accumulates states in float32 (states_in_fp32=True). This silently drops state accumulation precision on Blackwell (SM100+), which can cause accuracy drift for long sequences. Disable the FlashInfer dispatch so mamba_chunk_scan_combined always uses the Triton SSD kernels for prefill. The FlashInfer code is preserved for re-enablement once state_dtype is fixed to float32. Signed-off-by: tijyojwad <1127155+tijyojwad@users.noreply.github.com> --- tensorrt_llm/_torch/modules/mamba/ssd_combined.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/modules/mamba/ssd_combined.py b/tensorrt_llm/_torch/modules/mamba/ssd_combined.py index 35a836c872ff..efb02364d291 100644 --- a/tensorrt_llm/_torch/modules/mamba/ssd_combined.py +++ b/tensorrt_llm/_torch/modules/mamba/ssd_combined.py @@ -437,9 +437,16 @@ def mamba_chunk_scan_combined( # FlashInfer fused CUTLASS kernel on Blackwell (SM100+); both varlen and # non-varlen route here based on cu_seqlens. Falls back to Triton when the # MMA tile constraints on (chunk_size, dstate, headdim) aren't met. + # + # DISABLED: The FlashInfer SSD kernel hardcodes state_dtype=bf16 while the + # Triton path accumulates in fp32 (states_in_fp32=True), causing silent + # precision loss during prefill. + # TODO: Re-enable once _get_flashinfer_ssd_kernel uses state_dtype=float32. + _USE_FLASHINFER_SSD_PREFILL = False dstate = B.shape[-1] headdim = x.shape[-1] - flashinfer_eligible = (z is None and is_sm_100f()) + flashinfer_eligible = (_USE_FLASHINFER_SSD_PREFILL and z is None + and is_sm_100f()) if flashinfer_eligible and _flashinfer_ssd_supported( chunk_size, dstate, headdim): return _mamba_chunk_scan_flashinfer_fwd(