A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Bugfix] Mamba2 SSD varlen bug fix initstates decay, improve test, assert chunk pwr 2 (#21783)
Signed-off-by: Rishi Astra <40644327+RishiAstra@users.noreply.github.com>
RishiAstra committed
Aug 12, 2025 at 05:04 UTC
46ae7f666699496f45c0349b87f08d5119720951
3 files changed
+15
-13
tests/kernels/mamba/test_mamba_ssm_ssd.py
+8
-9
@@ -187,7 +187,7 @@ def generate_continuous_batched_examples(example_lens_by_batch,
187
[torch.float32, torch.float16, torch.bfloat16])
188
@pytest.mark.parametrize("n_heads", [3, 4, 11, 16, 32])
189
@pytest.mark.parametrize("d_head", [5, 8, 19, 32, 128])
190
-@pytest.mark.parametrize("seq_len_chunk_size", [(119, 17), (128, 32)])
190
+@pytest.mark.parametrize("seq_len_chunk_size", [(112, 16), (128, 32)])
191
def test_mamba_chunk_scan_single_example(d_head, n_heads, seq_len_chunk_size,
192
itype):
193
@@ -253,15 +253,15 @@ def test_mamba_chunk_scan_single_example(d_head, n_heads, seq_len_chunk_size,
253
(8, 8, 16, 32, 16),
254
]), # mode examples with varied lengths
255
256
- # odd chunk_size
257
- (64, 29, 2, [(11, 4), (13, 23), (19, 22),
258
- (21, 15)]), # irregular sizes
259
-
256
# large-ish chunk_size (256)
257
(64, 256, 1, [(5, ), (1, ), (1, ),
258
(1, )]), # irregular sizes with small sequences
259
(64, 256, 2, [(5, 30), (1, 2), (1, 2),
260
(1, 2)]), # irregular sizes with small sequences
261
+
262
+ # we also need to test some large seqlen
263
+ # to catch errors with init states decay
264
+ (768, 128, 2, [(138, 225), (138, 225)]),
265
])
266
def test_mamba_chunk_scan_cont_batch(d_head, n_heads, seq_len_chunk_size_cases,
267
itype):
@@ -271,10 +271,9 @@ def test_mamba_chunk_scan_cont_batch(d_head, n_heads, seq_len_chunk_size_cases,
271
272
seqlen, chunk_size, num_examples, cases = seq_len_chunk_size_cases
273
274
- # TODO: the irregular chunk size cases have some issues and require higher
275
- # tolerance. This is to be invesigated
276
- if chunk_size not in {8, 256}:
277
- atol, rtol = 5e-1, 5e-1
274
+ # This test can have larger error for longer sequences
275
+ if seqlen > 256:
276
+ atol, rtol = 1e-2, 5e-3
277
else:
278
atol, rtol = 5e-3, 5e-3
279
vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py
+2
-4
@@ -290,10 +290,8 @@ def _chunk_scan_fwd_kernel(
290
# get the cs at the offset boundary
291
# - c_off == 0 is a passthrough
292
dA_cs_m_boundary = tl.load(
293
- dA_cumsum_ptr +
294
- (pid_m * BLOCK_SIZE_M + c_off - 1) * stride_dA_cs_csize,
295
- mask=(((pid_m * BLOCK_SIZE_M + c_off - 1) > -1)
296
- and ((pid_m * BLOCK_SIZE_M + c_off) < chunk_size)),
293
+ dA_cumsum_ptr + (c_off - 1) * stride_dA_cs_csize,
294
+ mask=(((c_off - 1) > -1) and ((c_off) < chunk_size)),
295
other=0.0).to(tl.float32)
296
297
if HAS_SEQ_IDX:
vllm/model_executor/layers/mamba/ops/ssd_combined.py
+5
@@ -21,6 +21,10 @@ from .ssd_state_passing import _state_passing_fwd
21
TRITON_22 = version.parse(triton.__version__) >= version.parse('2.2.0')
22
23
24
+def is_int_pow_2(n):
25
+ return isinstance(n, int) and n > 0 and (n & (n - 1)) == 0
26
+
27
+
28
def _mamba_chunk_scan_combined_fwd(x,
29
dt,
30
A,
@@ -38,6 +42,7 @@ def _mamba_chunk_scan_combined_fwd(x,
42
dt_softplus=False,
43
dt_limit=(0.0, float("inf")),
44
out=None):
45
+ assert is_int_pow_2(chunk_size), "chunk_size must be integer power of 2"
46
batch, seqlen, nheads, headdim = x.shape
47
_, _, ngroups, dstate = B.shape
48
assert nheads % ngroups == 0