[Bugfix] Fix NaN output of the Mamba2 chunk scan kernel on sm90 (#61718)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
Vadim Gimpelson
co-authored by
Claude
Claude
parent
be25977532
commit
783380af39
@@ -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}",
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user