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
[CI Perf]Prune Tests in kernel/mamba (#26538)
Signed-off-by: Fardin Hoque <kfhfar@amazon.com> Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Fardin Hoque committed
Oct 13, 2025 at 15:22 UTC
577c72a22743bb3ea9976d7029731237d8346e73
4 files changed
+21
-36
tests/kernels/mamba/test_causal_conv1d.py
+2
-2
@@ -183,7 +183,7 @@ def test_causal_conv1d_update(dim, width, seqlen, has_bias, silu_activation, ity
183
assert torch.allclose(out, out_ref, rtol=rtol, atol=atol)
184
185
186
-@pytest.mark.parametrize("itype", [torch.float32, torch.float16, torch.bfloat16])
186
+@pytest.mark.parametrize("itype", [torch.float32, torch.bfloat16])
187
@pytest.mark.parametrize("silu_activation", [False, True])
188
@pytest.mark.parametrize("has_bias", [False, True])
189
@pytest.mark.parametrize("seqlen", [1, 3])
@@ -265,7 +265,7 @@ def test_causal_conv1d_update_with_batch_gather(
265
@pytest.mark.parametrize("silu_activation", [True])
266
@pytest.mark.parametrize("has_bias", [True])
267
@pytest.mark.parametrize("width", [4])
268
-@pytest.mark.parametrize("seqlen", [8, 30, 249, 2049, 4096])
268
+@pytest.mark.parametrize("seqlen", [8, 249, 4096])
269
@pytest.mark.parametrize("dim", [64, 4096])
270
@pytest.mark.parametrize("with_padding", [True, False])
271
@pytest.mark.parametrize("batch", [4, 10])
tests/kernels/mamba/test_mamba_mixer2.py
-1
@@ -25,7 +25,6 @@ from vllm.utils import update_environment_variables
25
(64, 1),
26
(64, 2),
27
(64, 4), # hidden_size be divisible by num_gpus
28
- (100, 5), # and n_groups must divide hidden_size
28
],
29
)
30
@pytest.mark.parametrize("dtype", [torch.float16])
tests/kernels/mamba/test_mamba_ssm.py
+11
-11
@@ -229,8 +229,8 @@ def selective_scan_opcheck_fn(
229
230
231
@pytest.mark.parametrize("wtype", [torch.float32])
232
-@pytest.mark.parametrize("itype", [torch.float32, torch.float16, torch.bfloat16])
233
-@pytest.mark.parametrize("seqlen", [128, 256, 512, 1024, 2048, 4096])
232
+@pytest.mark.parametrize("itype", [torch.float32, torch.bfloat16])
233
+@pytest.mark.parametrize("seqlen", [128, 1024, 4096])
234
@pytest.mark.parametrize("has_delta_bias", [True])
235
@pytest.mark.parametrize("delta_softplus", [True])
236
@pytest.mark.parametrize("has_z", [True])
@@ -238,7 +238,7 @@ def selective_scan_opcheck_fn(
238
@pytest.mark.parametrize("varBC_groups", [1, 2])
239
@pytest.mark.parametrize("is_variable_C", [True])
240
@pytest.mark.parametrize("is_variable_B", [True])
241
-@pytest.mark.parametrize("scan_chunks", [1, 2, 3])
241
+@pytest.mark.parametrize("scan_chunks", [1, 3])
242
def test_selective_scan(
243
is_variable_B,
244
is_variable_C,
@@ -375,9 +375,9 @@ def test_selective_scan(
375
)
376
377
378
-@pytest.mark.parametrize("itype", [torch.float32, torch.float16, torch.bfloat16])
378
+@pytest.mark.parametrize("itype", [torch.float32, torch.bfloat16])
379
@pytest.mark.parametrize("has_z", [False, True])
380
-@pytest.mark.parametrize("dstate", [16, 32, 64])
380
+@pytest.mark.parametrize("dstate", [16, 64])
381
@pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
382
def test_selective_state_update(dim, dstate, has_z, itype):
383
device = "cuda"
@@ -413,7 +413,7 @@ def test_selective_state_update(dim, dstate, has_z, itype):
413
414
@pytest.mark.parametrize("wtype", [torch.float32])
415
@pytest.mark.parametrize("itype", [torch.float32])
416
-@pytest.mark.parametrize("seqlen", [1, 128, 129, 256, 512, 1024, 2048, 4096])
416
+@pytest.mark.parametrize("seqlen", [1, 256, 1024, 4096])
417
@pytest.mark.parametrize("return_last_state", [True])
418
@pytest.mark.parametrize("has_delta_bias", [True])
419
@pytest.mark.parametrize("delta_softplus", [True])
@@ -589,9 +589,9 @@ def test_selective_scan_varlen(
589
)
590
591
592
-@pytest.mark.parametrize("itype", [torch.float32, torch.float16, torch.bfloat16])
592
+@pytest.mark.parametrize("itype", [torch.float32, torch.bfloat16])
593
@pytest.mark.parametrize("has_z", [True])
594
-@pytest.mark.parametrize("dstate", [16, 32, 64])
594
+@pytest.mark.parametrize("dstate", [16, 64])
595
@pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
596
# tests correctness in case subset of the sequences are padded
597
@pytest.mark.parametrize("with_padding", [True, False])
@@ -679,11 +679,11 @@ def test_selective_state_update_with_batch_indices(
679
assert torch.allclose(out[:batch_size], out_ref, rtol=rtol, atol=atol)
680
681
682
-@pytest.mark.parametrize("itype", [torch.float32, torch.float16, torch.bfloat16])
682
+@pytest.mark.parametrize("itype", [torch.float32, torch.bfloat16])
683
@pytest.mark.parametrize("has_z", [False, True])
684
@pytest.mark.parametrize("tie_hdim", [False, True])
685
-@pytest.mark.parametrize("ngroups", [1, 2, 4])
686
-@pytest.mark.parametrize("dstate", [16, 32, 64])
685
+@pytest.mark.parametrize("ngroups", [1, 4])
686
+@pytest.mark.parametrize("dstate", [16, 64])
687
@pytest.mark.parametrize("dim", [2048, 4096])
688
def test_selective_state_update_with_heads_with_batch_indices(
689
dim, dstate, ngroups, has_z, tie_hdim, itype
tests/kernels/mamba/test_mamba_ssm_ssd.py
+8
-22
@@ -188,9 +188,9 @@ def generate_continuous_batched_examples(
188
)
189
190
191
-@pytest.mark.parametrize("itype", [torch.float32, torch.float16, torch.bfloat16])
192
-@pytest.mark.parametrize("n_heads", [3, 4, 11, 16, 32])
193
-@pytest.mark.parametrize("d_head", [5, 8, 19, 32, 128])
191
+@pytest.mark.parametrize("itype", [torch.float32, torch.bfloat16])
192
+@pytest.mark.parametrize("n_heads", [4, 16, 32])
193
+@pytest.mark.parametrize("d_head", [5, 8, 32, 128])
194
@pytest.mark.parametrize("seq_len_chunk_size", [(112, 16), (128, 32)])
195
def test_mamba_chunk_scan_single_example(d_head, n_heads, seq_len_chunk_size, itype):
196
# this tests the kernels on a single example (bs=1)
@@ -254,15 +254,14 @@ def test_mamba_chunk_scan_single_example(d_head, n_heads, seq_len_chunk_size, it
254
)
255
256
257
-@pytest.mark.parametrize("itype", [torch.float32, torch.float16])
258
-@pytest.mark.parametrize("n_heads", [4, 8, 13])
259
-@pytest.mark.parametrize("d_head", [5, 16, 21, 32])
257
+@pytest.mark.parametrize("itype", [torch.float32])
258
+@pytest.mark.parametrize("n_heads", [4, 8])
259
+@pytest.mark.parametrize("d_head", [5, 16, 32])
260
@pytest.mark.parametrize(
261
"seq_len_chunk_size_cases",
262
[
263
# small-ish chunk_size (8)
264
(64, 8, 2, [(64, 32), (64, 32)]),
265
- (64, 8, 2, [(32, 32), (32, 32), (32, 32)]),
265
(64, 8, 2, [(8, 8), (8, 8), (8, 8)]), # chunk size boundary
266
(
267
64,
@@ -270,16 +269,7 @@ def test_mamba_chunk_scan_single_example(d_head, n_heads, seq_len_chunk_size, it
269
2,
270
[(4, 4), (4, 4), (4, 4), (4, 4)],
271
), # chunk_size larger than cont batches
273
- (
274
- 64,
275
- 8,
276
- 5,
277
- [
278
- (64, 32, 16, 8, 8),
279
- (8, 16, 32, 16, 8),
280
- (8, 8, 16, 32, 16),
281
- ],
282
- ), # mode examples with varied lengths
272
+ (64, 8, 5, [(64, 32, 16, 8, 8)]),
273
# large-ish chunk_size (256)
274
(64, 256, 1, [(5,), (1,), (1,), (1,)]), # irregular sizes with small sequences
275
(
@@ -359,11 +349,7 @@ def test_mamba_chunk_scan_cont_batch(d_head, n_heads, seq_len_chunk_size_cases,
349
@pytest.mark.parametrize("chunk_size", [8, 256])
350
@pytest.mark.parametrize(
351
"seqlens",
362
- [
363
- (16, 2, 8, 13),
364
- (270, 88, 212, 203),
365
- (16, 20),
366
- ],
352
+ [(16, 20), (270, 88, 212, 203)],
353
)
354
def test_mamba_chunk_scan_cont_batch_prefill_chunking(chunk_size, seqlens):
355
# This test verifies the correctness of the chunked prefill implementation