A high-throughput and memory-efficient inference and serving engine for LLMs
[Kernel][Test] Make kernel tests for mamba dual-HW (CUDA + XPU) (#42736)
Signed-off-by: Dobrzyniewicz, Agata <agata.dobrzyniewicz@intel.com>
Agata Dobrzyniewicz committed
Jun 8, 2026 at 02:22 UTC
2ed0a9627b50d9d1d26eddb84abb7e7afc2984c1
3 files changed
+49
-14
tests/kernels/mamba/test_causal_conv1d.py
+12
-4
@@ -11,9 +11,17 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
11
causal_conv1d_fn,
12
causal_conv1d_update,
13
)
14
+from vllm.platforms import current_platform
15
from vllm.utils.torch_utils import set_random_seed
16
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
17
18
+DEVICE = current_platform.device_type
19
+
20
+pytestmark = pytest.mark.skipif(
21
+ not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
22
+ reason="causal_conv1d Triton kernels require CUDA-alike or XPU",
23
+)
24
+
25
26
def causal_conv1d_ref(
27
x: torch.Tensor,
@@ -149,7 +157,7 @@ def causal_conv1d_opcheck_fn(
157
@pytest.mark.parametrize("width", [4])
158
@pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
159
def test_causal_conv1d_update(dim, width, seqlen, has_bias, silu_activation, itype):
152
- device = "cuda"
160
+ device = DEVICE
161
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3)
162
if itype == torch.bfloat16:
163
rtol, atol = 1e-2, 5e-2
@@ -196,7 +204,7 @@ def test_causal_conv1d_update(dim, width, seqlen, has_bias, silu_activation, ity
204
def test_causal_conv1d_update_with_batch_gather(
205
batch_size, with_padding, dim, width, seqlen, has_bias, silu_activation, itype
206
):
199
- device = "cuda"
207
+ device = DEVICE
208
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3)
209
if itype == torch.bfloat16:
210
rtol, atol = 1e-2, 5e-2
@@ -275,7 +283,7 @@ def test_causal_conv1d_update_with_batch_gather(
283
def test_causal_conv1d_varlen(
284
batch, with_padding, dim, seqlen, width, has_bias, silu_activation, itype
285
):
278
- device = "cuda"
286
+ device = DEVICE
287
torch.accelerator.empty_cache()
288
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3)
289
if itype == torch.bfloat16:
@@ -341,7 +349,7 @@ def test_causal_conv1d_varlen(
349
weight,
350
bias=bias,
351
conv_states=final_states,
344
- query_start_loc=cumsum.cuda(),
352
+ query_start_loc=cumsum.to(device),
353
cache_indices=padded_state_indices,
354
has_initial_state=has_initial_states,
355
activation=activation,
tests/kernels/mamba/test_mamba_ssm.py
+24
-7
@@ -17,6 +17,21 @@ from vllm.platforms import current_platform
17
from vllm.utils.torch_utils import set_random_seed
18
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
19
20
+DEVICE = current_platform.device_type
21
+
22
+pytestmark = pytest.mark.skipif(
23
+ not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
24
+ reason="mamba_ssm kernels require CUDA-alike or XPU",
25
+)
26
+
27
+# selective_scan_fn is backed by the CUDA-only `ops.selective_scan_fwd` C++ op,
28
+# so tests exercising it must be skipped on XPU. selective_state_update is
29
+# pure Triton and runs on both CUDA-alike and XPU.
30
+skip_unless_cuda_alike = pytest.mark.skipif(
31
+ not current_platform.is_cuda_alike(),
32
+ reason="selective_scan_fn uses CUDA-only custom op",
33
+)
34
+
35
36
def selective_scan_ref(
37
u,
@@ -181,6 +196,7 @@ def selective_scan_opcheck_fn(
196
@pytest.mark.parametrize("is_variable_C", [True])
197
@pytest.mark.parametrize("is_variable_B", [True])
198
@pytest.mark.parametrize("scan_chunks", [1, 3])
199
+@skip_unless_cuda_alike
200
def test_selective_scan(
201
is_variable_B,
202
is_variable_C,
@@ -327,7 +343,7 @@ def test_selective_scan(
343
@pytest.mark.parametrize("dstate", [16, 64])
344
@pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
345
def test_selective_state_update(dim, dstate, has_z, itype):
330
- device = "cuda"
346
+ device = DEVICE
347
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
348
if itype == torch.bfloat16:
349
rtol, atol = 1e-2, 5e-2
@@ -370,7 +386,7 @@ def test_selective_state_update(dim, dstate, has_z, itype):
386
" on compute capability 10.0 CUDA devices.",
387
)
388
def test_selective_state_update_stochastic_rounding(dim, dstate, has_z, philox_rounds):
373
- device = "cuda"
389
+ device = DEVICE
390
rtol, atol = 5e-3, 1e-1
391
# set seed
392
set_random_seed(0)
@@ -417,7 +433,7 @@ def test_selective_state_update_stochastic_rounding(dim, dstate, has_z, philox_r
433
@pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
434
@pytest.mark.parametrize("max_seq_len", [1, 2, 4])
435
def test_selective_state_update_varlen(dim, dstate, has_z, itype, max_seq_len):
420
- device = "cuda"
436
+ device = DEVICE
437
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
438
if itype == torch.bfloat16:
439
rtol, atol = 5e-2, 1.5e-1
@@ -498,6 +514,7 @@ def test_selective_state_update_varlen(dim, dstate, has_z, itype, max_seq_len):
514
@pytest.mark.parametrize("is_variable_B", [True])
515
# tests correctness in case subset of the sequences are padded
516
@pytest.mark.parametrize("with_padding", [False, True])
517
+@skip_unless_cuda_alike
518
def test_selective_scan_varlen(
519
with_padding,
520
is_variable_B,
@@ -679,7 +696,7 @@ def test_selective_scan_varlen(
696
def test_selective_state_update_with_batch_indices(
697
with_padding, dim, dstate, has_z, itype
698
):
682
- device = "cuda"
699
+ device = DEVICE
700
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
701
if itype == torch.bfloat16:
702
rtol, atol = 1e-1, 1e-1
@@ -771,7 +788,7 @@ def test_selective_state_update_with_batch_indices(
788
def test_selective_state_update_with_heads_with_batch_indices(
789
dim, dstate, ngroups, has_z, tie_hdim, itype
790
):
774
- device = "cuda"
791
+ device = DEVICE
792
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 3e-2)
793
if itype == torch.bfloat16:
794
rtol, atol = 1e-1, 1e-1
@@ -844,7 +861,7 @@ def test_selective_state_update_with_heads_with_batch_indices(
861
def test_selective_state_update_with_num_accepted_tokens(
862
dim, dstate, has_z, itype, max_seq_len
863
):
847
- device = "cuda"
864
+ device = DEVICE
865
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
866
if itype == torch.bfloat16:
867
rtol, atol = 5e-2, 1.5e-1
@@ -970,7 +987,7 @@ def test_selective_state_update_with_num_accepted_tokens(
987
def test_selective_state_update_varlen_with_num_accepted(
988
dim, dstate, has_z, itype, max_seq_len
989
):
973
- device = "cuda"
990
+ device = DEVICE
991
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
992
if itype == torch.bfloat16:
993
rtol, atol = 5e-2, 1.5e-1
tests/kernels/mamba/test_mamba_ssm_ssd.py
+13
-3
@@ -9,9 +9,19 @@ from einops import rearrange, repeat
9
from vllm.model_executor.layers.mamba.ops.ssd_combined import (
10
mamba_chunk_scan_combined_varlen,
11
)
12
+from vllm.platforms import current_platform
13
from vllm.utils.torch_utils import set_random_seed
14
from vllm.v1.attention.backends.mamba2_attn import compute_varlen_chunk_metadata
15
16
+# All kernels exercised here are pure Triton, so they run on any backend
17
+# that the vLLM platform layer treats as a CUDA-alike device or as XPU.
18
+DEVICE = current_platform.device_type
19
+
20
+pytestmark = pytest.mark.skipif(
21
+ not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
22
+ reason="Mamba2 SSD Triton kernels require a CUDA-alike or XPU device.",
23
+)
24
+
25
# Added by the IBM Team, 2024
26
27
# Adapted from https://github.com/state-spaces/mamba/blob/v2.2.4/mamba_ssm/modules/ssd_minimal.py
@@ -81,7 +91,7 @@ def ssd_minimal_discrete(X, A, B, C, block_len, initial_states=None):
91
return Y, final_state
92
93
84
-def generate_random_inputs(batch_size, seqlen, n_heads, d_head, itype, device="cuda"):
94
+def generate_random_inputs(batch_size, seqlen, n_heads, d_head, itype, device=DEVICE):
95
set_random_seed(0)
96
A = -torch.exp(torch.rand(n_heads, dtype=itype, device=device))
97
dt = F.softplus(
@@ -103,7 +113,7 @@ def generate_continuous_batched_examples(
113
n_heads,
114
d_head,
115
itype,
106
- device="cuda",
116
+ device=DEVICE,
117
return_naive_ref=True,
118
):
119
# this function generates a random examples of certain length
@@ -215,7 +225,7 @@ def test_mamba_chunk_scan_single_example(d_head, n_heads, seq_len_chunk_size, it
225
X * dt.unsqueeze(-1), A * dt, B, C, chunk_size
226
)
227
218
- cu_seqlens = torch.tensor((0, seqlen), device="cuda").cumsum(dim=0)
228
+ cu_seqlens = torch.tensor((0, seqlen), device=DEVICE).cumsum(dim=0)
229
cu_chunk_seqlens, last_chunk_indices, seq_idx_chunks = (
230
compute_varlen_chunk_metadata(cu_seqlens, chunk_size)
231
)