[XPU] Add tuned Mamba SSU configs for Intel Arc Pro B70 (#50534)
Signed-off-by: pmanczak <pawel.manczak@intel.com>
pmanczak committed
Aug 13, 2026 at 14:36 UTC
c4e969294ecab9ffefb995b734303f13f62b723f
5 files changed
+247
-12
benchmarks/kernels/benchmark_selective_state_update.py
+19
-12
@@ -34,6 +34,7 @@ from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
34
override_ssm_config,
35
selective_state_update,
36
)
37
+from vllm.platforms import current_platform
38
from vllm.triton_utils import triton
39
40
# bf16 shares configs with fp16 - same bit width.
@@ -93,10 +94,10 @@ def _make_inputs(
94
ngroups: int,
95
dtype: torch.dtype,
96
state_dtype: torch.dtype | None = None,
96
- device: str = "cuda",
97
):
98
if state_dtype is None:
99
state_dtype = dtype
100
+ device = current_platform.device_type
101
state = torch.randn(batch, nheads, dim, dstate, dtype=state_dtype, device=device)
102
x = torch.randn(batch, nheads, dim, dtype=dtype, device=device)
103
dt = torch.randn(batch, nheads, dim, dtype=dtype, device=device)
@@ -127,9 +128,9 @@ def benchmark_config(
128
Time one (BLOCK_SIZE_M, num_warps) config for selective_state_update.
129
Returns elapsed time in microseconds, or None on error.
130
130
- Uses CUDA graph capture-and-replay to isolate kernel time from Python
131
- eager-mode dispatch / kwarg-resolution overhead, mirroring the timing
132
- methodology in benchmarks/kernels/benchmark_moe.py.
131
+ Uses accelerator graph capture-and-replay to isolate kernel time from
132
+ Python eager-mode dispatch / kwarg-resolution overhead, mirroring the
133
+ timing methodology in benchmarks/kernels/benchmark_moe.py.
134
"""
135
state, x, dt, A, B, C, D, dt_bias, out = _make_inputs(
136
batch, nheads, dim, dstate, ngroups, dtype, state_dtype=state_dtype
@@ -157,10 +158,15 @@ def benchmark_config(
158
_call_kernel()
159
torch.accelerator.synchronize()
160
160
- # Capture graph_batch_size invocations into a CUDA graph so the
161
+ # Capture graph_batch_size invocations into a device graph so the
162
# timed region runs without Python dispatch overhead per call.
162
- graph = torch.cuda.CUDAGraph()
163
- with torch.cuda.graph(graph):
163
+ # Capture via graph(), not the graph object: CUDA needs a side stream.
164
+ graph = (
165
+ torch.cuda.CUDAGraph()
166
+ if current_platform.is_cuda_alike()
167
+ else torch.xpu.XPUGraph()
168
+ )
169
+ with current_platform.graph(graph):
170
for _ in range(graph_batch_size):
171
_call_kernel()
172
torch.accelerator.synchronize()
@@ -170,8 +176,8 @@ def benchmark_config(
176
graph.replay()
177
torch.accelerator.synchronize()
178
173
- start = torch.cuda.Event(enable_timing=True)
174
- end = torch.cuda.Event(enable_timing=True)
179
+ start = torch.Event(enable_timing=True)
180
+ end = torch.Event(enable_timing=True)
181
latencies: list[float] = []
182
for _ in range(num_iters):
183
start.record()
@@ -671,8 +677,8 @@ def main():
677
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
678
state_dtype = _SSM_CACHE_DTYPE_MAP[args.mamba_ssm_cache_dtype]
679
device_name = get_ssm_device_name()
674
- cap = torch.cuda.get_device_capability()
675
- is_blackwell = cap[0] >= 10
680
+ cap = current_platform.get_device_capability()
681
+ is_blackwell = cap is not None and cap[0] >= 10
682
683
# Mirror all output to a results file (like Unix tee).
684
buf = StringIO()
@@ -690,7 +696,8 @@ def main():
696
sys.stdout = _Tee() # type: ignore[assignment]
697
698
try:
693
- print(f"Device : {device_name} (sm_{cap[0]}{cap[1]})")
699
+ cap_str = f"sm_{cap[0]}{cap[1]}" if cap is not None else "n/a"
700
+ print(f"Device : {device_name} ({cap_str})")
701
print(f"Blackwell: {is_blackwell}")
702
print(f"dtype : {args.dtype}")
703
print(f"ssm_cache_dtype: {args.mamba_ssm_cache_dtype}")
vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json
new
+51
@@ -0,0 +1,51 @@
1
+{
2
+ "triton_version": "3.7.1",
3
+ "8": {
4
+ "BLOCK_SIZE_M": 8,
5
+ "num_warps": 1
6
+ },
7
+ "16": {
8
+ "BLOCK_SIZE_M": 4,
9
+ "num_warps": 1
10
+ },
11
+ "32": {
12
+ "BLOCK_SIZE_M": 4,
13
+ "num_warps": 1
14
+ },
15
+ "64": {
16
+ "BLOCK_SIZE_M": 4,
17
+ "num_warps": 1
18
+ },
19
+ "128": {
20
+ "BLOCK_SIZE_M": 4,
21
+ "num_warps": 1
22
+ },
23
+ "256": {
24
+ "BLOCK_SIZE_M": 4,
25
+ "num_warps": 1
26
+ },
27
+ "512": {
28
+ "BLOCK_SIZE_M": 4,
29
+ "num_warps": 1
30
+ },
31
+ "1024": {
32
+ "BLOCK_SIZE_M": 4,
33
+ "num_warps": 1
34
+ },
35
+ "2048": {
36
+ "BLOCK_SIZE_M": 4,
37
+ "num_warps": 1
38
+ },
39
+ "4096": {
40
+ "BLOCK_SIZE_M": 4,
41
+ "num_warps": 1
42
+ },
43
+ "8192": {
44
+ "BLOCK_SIZE_M": 4,
45
+ "num_warps": 1
46
+ },
47
+ "16384": {
48
+ "BLOCK_SIZE_M": 4,
49
+ "num_warps": 1
50
+ }
51
+}
\ No newline at end of file
vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json
new
+51
@@ -0,0 +1,51 @@
1
+{
2
+ "triton_version": "3.7.1",
3
+ "8": {
4
+ "BLOCK_SIZE_M": 8,
5
+ "num_warps": 2
6
+ },
7
+ "16": {
8
+ "BLOCK_SIZE_M": 8,
9
+ "num_warps": 2
10
+ },
11
+ "32": {
12
+ "BLOCK_SIZE_M": 4,
13
+ "num_warps": 2
14
+ },
15
+ "64": {
16
+ "BLOCK_SIZE_M": 4,
17
+ "num_warps": 2
18
+ },
19
+ "128": {
20
+ "BLOCK_SIZE_M": 4,
21
+ "num_warps": 2
22
+ },
23
+ "256": {
24
+ "BLOCK_SIZE_M": 8,
25
+ "num_warps": 2
26
+ },
27
+ "512": {
28
+ "BLOCK_SIZE_M": 8,
29
+ "num_warps": 2
30
+ },
31
+ "1024": {
32
+ "BLOCK_SIZE_M": 8,
33
+ "num_warps": 2
34
+ },
35
+ "2048": {
36
+ "BLOCK_SIZE_M": 4,
37
+ "num_warps": 1
38
+ },
39
+ "4096": {
40
+ "BLOCK_SIZE_M": 8,
41
+ "num_warps": 2
42
+ },
43
+ "8192": {
44
+ "BLOCK_SIZE_M": 4,
45
+ "num_warps": 1
46
+ },
47
+ "16384": {
48
+ "BLOCK_SIZE_M": 8,
49
+ "num_warps": 2
50
+ }
51
+}
\ No newline at end of file
vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json
new
+63
@@ -0,0 +1,63 @@
1
+{
2
+ "triton_version": "3.7.1",
3
+ "8": {
4
+ "BLOCK_SIZE_M": 4,
5
+ "num_warps": 1
6
+ },
7
+ "16": {
8
+ "BLOCK_SIZE_M": 16,
9
+ "num_warps": 8
10
+ },
11
+ "32": {
12
+ "BLOCK_SIZE_M": 4,
13
+ "num_warps": 1
14
+ },
15
+ "64": {
16
+ "BLOCK_SIZE_M": 4,
17
+ "num_warps": 1
18
+ },
19
+ "128": {
20
+ "BLOCK_SIZE_M": 4,
21
+ "num_warps": 1
22
+ },
23
+ "256": {
24
+ "BLOCK_SIZE_M": 64,
25
+ "num_warps": 8
26
+ },
27
+ "512": {
28
+ "BLOCK_SIZE_M": 4,
29
+ "num_warps": 1
30
+ },
31
+ "1024": {
32
+ "BLOCK_SIZE_M": 4,
33
+ "num_warps": 1
34
+ },
35
+ "2048": {
36
+ "BLOCK_SIZE_M": 4,
37
+ "num_warps": 1
38
+ },
39
+ "4096": {
40
+ "BLOCK_SIZE_M": 4,
41
+ "num_warps": 1
42
+ },
43
+ "8192": {
44
+ "BLOCK_SIZE_M": 8,
45
+ "num_warps": 2
46
+ },
47
+ "16384": {
48
+ "BLOCK_SIZE_M": 8,
49
+ "num_warps": 2
50
+ },
51
+ "32768": {
52
+ "BLOCK_SIZE_M": 16,
53
+ "num_warps": 8
54
+ },
55
+ "65536": {
56
+ "BLOCK_SIZE_M": 16,
57
+ "num_warps": 8
58
+ },
59
+ "131072": {
60
+ "BLOCK_SIZE_M": 16,
61
+ "num_warps": 8
62
+ }
63
+}
\ No newline at end of file
vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json
new
+63
@@ -0,0 +1,63 @@
1
+{
2
+ "triton_version": "3.7.1",
3
+ "8": {
4
+ "BLOCK_SIZE_M": 4,
5
+ "num_warps": 1
6
+ },
7
+ "16": {
8
+ "BLOCK_SIZE_M": 8,
9
+ "num_warps": 1
10
+ },
11
+ "32": {
12
+ "BLOCK_SIZE_M": 4,
13
+ "num_warps": 1
14
+ },
15
+ "64": {
16
+ "BLOCK_SIZE_M": 4,
17
+ "num_warps": 1
18
+ },
19
+ "128": {
20
+ "BLOCK_SIZE_M": 4,
21
+ "num_warps": 1
22
+ },
23
+ "256": {
24
+ "BLOCK_SIZE_M": 4,
25
+ "num_warps": 1
26
+ },
27
+ "512": {
28
+ "BLOCK_SIZE_M": 4,
29
+ "num_warps": 1
30
+ },
31
+ "1024": {
32
+ "BLOCK_SIZE_M": 4,
33
+ "num_warps": 1
34
+ },
35
+ "2048": {
36
+ "BLOCK_SIZE_M": 4,
37
+ "num_warps": 1
38
+ },
39
+ "4096": {
40
+ "BLOCK_SIZE_M": 4,
41
+ "num_warps": 1
42
+ },
43
+ "8192": {
44
+ "BLOCK_SIZE_M": 8,
45
+ "num_warps": 1
46
+ },
47
+ "16384": {
48
+ "BLOCK_SIZE_M": 8,
49
+ "num_warps": 1
50
+ },
51
+ "32768": {
52
+ "BLOCK_SIZE_M": 16,
53
+ "num_warps": 8
54
+ },
55
+ "65536": {
56
+ "BLOCK_SIZE_M": 16,
57
+ "num_warps": 8
58
+ },
59
+ "131072": {
60
+ "BLOCK_SIZE_M": 16,
61
+ "num_warps": 8
62
+ }
63
+}
\ No newline at end of file