Skip to content

[None][fix] Bypass FlashInfer SSD prefill to fix state dtype precision #14600

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
2ez4bz merged 1 commit into from
May 27, 2026
+8 −1
Merged

[None][fix] Bypass FlashInfer SSD prefill to fix state dtype precision #14600

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion tensorrt_llm/_torch/modules/mamba/ssd_combined.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
Loading