From 783380af397ceeae91a40f4f759439eb8c1ab690 Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Date: Sat, 17 Oct 2026 04:24:51 +0400 Subject: [PATCH] [Bugfix] Fix NaN output of the Mamba2 chunk scan kernel on sm90 (#61718) Signed-off-by: Vadim Gimpelson Co-authored-by: Claude --- tests/kernels/mamba/test_mamba_ssm_ssd.py | 47 +++ .../layers/mamba/ops/ssd_chunk_scan.py | 277 ++++++++++-------- 2 files changed, 197 insertions(+), 127 deletions(-) diff --git a/tests/kernels/mamba/test_mamba_ssm_ssd.py b/tests/kernels/mamba/test_mamba_ssm_ssd.py index 22abeab1ed..8522477045 100644 --- a/tests/kernels/mamba/test_mamba_ssm_ssd.py +++ b/tests/kernels/mamba/test_mamba_ssm_ssd.py @@ -6,6 +6,10 @@ import torch import torch.nn.functional as F from einops import rearrange, repeat +from vllm.model_executor.layers.mamba.ops.ssd_chunk_scan import ( + _chunk_scan_fwd, + _chunk_scan_fwd_kernel, +) from vllm.model_executor.layers.mamba.ops.ssd_combined import ( mamba_chunk_scan_combined_varlen, ) @@ -578,3 +582,46 @@ def test_mamba_chunk_scan_cont_batch_prefill_chunking(chunk_size, seqlens): rtol=rtol, msg=lambda x, i=i: f"seq{i} state " + x, ) + + +def test_chunk_scan_fwd_every_config_matches(monkeypatch): + """The autotuner picks a config by speed only, so every config in the list + must give the same finite output. On sm90, configs that pipeline a wgmma + with `cb` in registers wrote NaN into whole rows (#61603).""" + from triton.runtime.errors import OutOfResources + + set_random_seed(0) + nheads, headdim, ngroups, dstate, chunk_size = 64, 64, 8, 128, 128 + cu_chunk_seqlens = torch.tensor([0, 128, 181], dtype=torch.int32, device=DEVICE) + seqlen, nchunks = 181, 2 + x = torch.randn(seqlen, nheads, headdim, dtype=torch.bfloat16, device=DEVICE) + C = torch.randn(seqlen, ngroups, dstate, dtype=torch.bfloat16, device=DEVICE) + cb = torch.randn(nchunks, ngroups, chunk_size, chunk_size, device=DEVICE) + dt = torch.rand(nheads, nchunks, chunk_size, device=DEVICE) * 0.1 + dA_cumsum = -torch.rand(nheads, nchunks, chunk_size, device=DEVICE).cumsum(-1) + states = torch.randn(nchunks, nheads, headdim, dstate, device=DEVICE) * 0.1 + seq_idx = torch.zeros(nchunks, dtype=torch.int32, device=DEVICE) + D = torch.randn(nheads, headdim, device=DEVICE) * 0.1 + + ref = None + for config in list(_chunk_scan_fwd_kernel.configs): + monkeypatch.setattr(_chunk_scan_fwd_kernel, "configs", [config]) + for _ in range(10): + out = torch.empty_like(x) + try: + _chunk_scan_fwd( + cb, x, dt, dA_cumsum, C, states, cu_chunk_seqlens, out, seq_idx, D + ) + except OutOfResources: + # The autotuner skips configs that do not fit this device. + break + assert torch.isfinite(out).all(), f"non-finite output with {config}" + if ref is None: + ref = out + torch.testing.assert_close( + out, + ref, + atol=2e-2, + rtol=2e-2, + msg=lambda m, config=config: f"{config}: {m}", + ) diff --git a/vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py b/vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py index bd554ede5a..237dd33b52 100644 --- a/vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py +++ b/vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py @@ -9,139 +9,162 @@ from packaging import version from vllm.model_executor.layers.mamba.ops.triton_helpers import fast_exp +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton TRITON_22 = version.parse(triton.__version__) >= version.parse("2.2.0") +def _drop_racy_wgmma_configs(configs: list[triton.Config]) -> list[triton.Config]: + """Drop the configs that Triton miscompiles on sm90. + + With whole warpgroups (num_warps % 4 == 0) and BLOCK_SIZE_M % 64 == 0, + `tl.dot(cb, x)` lowers to wgmma with `cb` in registers. With + num_stages > 1, Triton leaves that wgmma running across the loop + back-edge while the next iteration overwrites the registers of `cb`, + which writes NaN/Inf into whole output rows. + """ + if not (current_platform.is_cuda() and current_platform.is_device_capability(90)): + return configs + return [ + c + for c in configs + if c.num_stages == 1 + or c.num_warps % 4 != 0 + or c.kwargs["BLOCK_SIZE_M"] % 64 != 0 + ] + + @triton.autotune( - configs=[ - # ================================================================= - # Higher warp count configs for better latency hiding - # More warps = more instructions in flight = better memory latency hiding - # ================================================================= - triton.Config( - {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32}, - num_stages=2, - num_warps=8, - ), - triton.Config( - {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, - num_stages=2, - num_warps=8, - ), - triton.Config( - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32}, - num_stages=2, - num_warps=8, - ), - # Smaller tiles with more stages for software pipelining - triton.Config( - {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32}, - num_stages=3, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 64}, - num_stages=2, - num_warps=4, - ), - # ================================================================= - # Low register pressure configs (num_stages=1) for large dstate - # ================================================================= - triton.Config( - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64}, - num_stages=1, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, - num_stages=1, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, - num_stages=1, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, - num_stages=1, - num_warps=4, - ), - # num_stages=2 configs - moderate register pressure - triton.Config( - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64}, - num_stages=2, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, - num_stages=2, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, - num_stages=2, - num_warps=4, - ), - # Original configs for larger dstate values - triton.Config( - {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 64}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 32}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 64}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32}, - num_stages=5, - num_warps=2, - ), - triton.Config( - {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, - num_stages=5, - num_warps=2, - ), - triton.Config( - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, - num_stages=4, - num_warps=2, - ), - ], + configs=_drop_racy_wgmma_configs( + [ + # ================================================================= + # Higher warp count configs for better latency hiding + # More warps = more instructions in flight = better memory latency hiding + # ================================================================= + triton.Config( + {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32}, + num_stages=2, + num_warps=8, + ), + triton.Config( + {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, + num_stages=2, + num_warps=8, + ), + triton.Config( + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32}, + num_stages=2, + num_warps=8, + ), + # Smaller tiles with more stages for software pipelining + triton.Config( + {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32}, + num_stages=3, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 64}, + num_stages=2, + num_warps=4, + ), + # ================================================================= + # Low register pressure configs (num_stages=1) for large dstate + # ================================================================= + triton.Config( + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64}, + num_stages=1, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, + num_stages=1, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, + num_stages=1, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, + num_stages=1, + num_warps=4, + ), + # num_stages=2 configs - moderate register pressure + triton.Config( + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64}, + num_stages=2, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, + num_stages=2, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, + num_stages=2, + num_warps=4, + ), + # Original configs for larger dstate values + triton.Config( + {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 64}, + num_stages=3, + num_warps=8, + ), + triton.Config( + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 32}, + num_stages=4, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32}, + num_stages=4, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, + num_stages=4, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32}, + num_stages=4, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64}, + num_stages=4, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 64}, + num_stages=4, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32}, + num_stages=4, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32}, + num_stages=5, + num_warps=2, + ), + triton.Config( + {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, + num_stages=5, + num_warps=2, + ), + triton.Config( + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, + num_stages=4, + num_warps=2, + ), + ] + ), key=["chunk_size", "hdim", "dstate", "IS_CAUSAL"], ) @triton.jit