[ROCm][Perf][DSV4] Enable split sparse decode on gfx942 (#46275)
Signed-off-by: Tuukka Sarvi <tuukka.sarvi@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Tuukka Sarvi committed
Jul 16, 2026 at 15:28 UTC
b8168e33e000df5c60d63d7ca244ba2ca2c34079
2 files changed
+45
-42
tests/kernels/attention/test_rocm_triton_attn_dsv4.py
+44
-41
@@ -11,22 +11,22 @@ pytestmark = pytest.mark.skipif(
11
)
12
13
14
-def _on_gfx950() -> bool:
14
+def _on_split_decode_arch() -> bool:
15
if not current_platform.is_rocm():
16
return False
17
try:
18
- from vllm.platforms.rocm import _ON_GFX950
18
+ from vllm.platforms.rocm import _ON_GFX942, _ON_GFX950
19
20
- return bool(_ON_GFX950)
20
+ return bool(_ON_GFX942 or _ON_GFX950)
21
except Exception:
22
return False
23
24
25
-# The flash-decode split-K decode path is only tuned for AMD gfx950; other
25
+# The flash-decode split-K decode path is only tuned for AMD gfx942/gfx950; other
26
# architectures take the fallback decode kernel, so its tests are skipped there.
27
-requires_gfx950 = pytest.mark.skipif(
28
- not _on_gfx950(),
29
- reason="split-K decode kernel is only tuned for AMD gfx950",
27
+requires_split_decode_arch = pytest.mark.skipif(
28
+ not _on_split_decode_arch(),
29
+ reason="split-K decode kernel is only tuned for AMD gfx942/gfx950",
30
)
31
32
NOPE_HEAD_DIM = 448
@@ -91,9 +91,13 @@ def _ref_sparse_prefill_ragged(
91
92
93
def _pack_fp8_ds_mla_cache(
94
- kv: torch.Tensor, block_size: int, is_extra: bool = False
94
+ kv: torch.Tensor, block_size: int, use_fnuz: bool
95
) -> torch.Tensor:
96
assert kv.shape[-1] == HEAD_DIM
97
+ from vllm.models.deepseek_v4.common.ops.cache_utils import (
98
+ quantize_and_insert_k_cache,
99
+ )
100
+
101
num_tokens = kv.shape[0]
102
num_blocks = (num_tokens + block_size - 1) // block_size
103
cache = torch.zeros(
@@ -101,41 +105,34 @@ def _pack_fp8_ds_mla_cache(
105
dtype=torch.uint8,
106
device=kv.device,
107
)
104
- cache_flat = cache.view(torch.uint8).flatten()
105
- kv_nope_fp8 = (
106
- kv[:, :NOPE_HEAD_DIM]
107
- .to(torch.float8_e4m3fn if is_extra else current_platform.fp8_dtype())
108
- .view(torch.uint8)
109
- )
110
- kv_rope_u8 = kv[:, NOPE_HEAD_DIM:].contiguous().view(torch.uint8)
111
-
112
- for slot in range(num_tokens):
113
- block_idx = slot // block_size
114
- pos = slot % block_size
115
- block_base = block_idx * cache.stride(0)
116
- token_base = block_base + pos * 576
117
- scale_base = block_base + block_size * 576 + pos * 8
118
- cache_flat[token_base : token_base + NOPE_HEAD_DIM].copy_(kv_nope_fp8[slot])
119
- cache_flat[
120
- token_base + NOPE_HEAD_DIM : token_base + NOPE_HEAD_DIM + ROPE_HEAD_DIM * 2
121
- ].copy_(kv_rope_u8[slot])
122
- cache_flat[scale_base : scale_base + 7].fill_(127)
108
+ slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=kv.device)
109
+ quantize_and_insert_k_cache(
110
+ kv,
111
+ cache,
112
+ slot_mapping,
113
+ block_size=block_size,
114
+ use_fnuz=use_fnuz,
115
+ )
116
return cache
117
118
119
def _read_fp8_ds_mla_cache(
127
- cache: torch.Tensor, slot: int, block_size: int, is_extra: bool = False
120
+ cache: torch.Tensor, slot: int, block_size: int, use_fnuz: bool
121
) -> torch.Tensor:
122
cache_flat = cache.view(torch.uint8).flatten()
123
block_idx = slot // block_size
124
pos = slot % block_size
125
block_base = block_idx * cache.stride(0)
126
token_base = block_base + pos * 576
127
+ scale_base = block_base + block_size * 576 + pos * 8
128
129
+ fp8_dtype = torch.float8_e4m3fnuz if use_fnuz else torch.float8_e4m3fn
130
nope_u8 = cache_flat[token_base : token_base + NOPE_HEAD_DIM]
136
- nope = nope_u8.view(
137
- torch.float8_e4m3fn if is_extra else current_platform.fp8_dtype()
138
- ).to(torch.float32)
131
+ nope = nope_u8.view(fp8_dtype).to(torch.float32)
132
+ scales = torch.exp2(
133
+ cache_flat[scale_base : scale_base + 7].to(torch.float32) - 127.0
134
+ )
135
+ nope = nope * scales.repeat_interleave(64)
136
rope_u8 = cache_flat[
137
token_base + NOPE_HEAD_DIM : token_base + NOPE_HEAD_DIM + ROPE_HEAD_DIM * 2
138
]
@@ -152,19 +149,21 @@ def _ref_sparse_decode_ragged(
149
block_size: int,
150
extra_cache: torch.Tensor | None = None,
151
extra_rows: list[list[int]] | None = None,
152
+ main_use_fnuz: bool = False,
153
+ extra_use_fnuz: bool = False,
154
) -> torch.Tensor:
155
q_f32 = q.float()
156
out = torch.empty_like(q_f32)
157
158
for query_idx in range(q.shape[0]):
159
row_kv = [
161
- _read_fp8_ds_mla_cache(main_cache, int(slot), block_size)
160
+ _read_fp8_ds_mla_cache(main_cache, int(slot), block_size, main_use_fnuz)
161
for slot in main_rows[query_idx]
162
]
163
if extra_cache is not None and extra_rows is not None:
164
row_kv.extend(
165
_read_fp8_ds_mla_cache(
167
- extra_cache, int(slot), block_size, is_extra=True
166
+ extra_cache, int(slot), block_size, extra_use_fnuz
167
)
168
for slot in extra_rows[query_idx]
169
)
@@ -290,11 +289,12 @@ def test_sparse_attn_decode_ragged_kernel() -> None:
289
device = torch.device("cuda")
290
torch.manual_seed(1)
291
block_size = 4
292
+ main_use_fnuz = current_platform.is_fp8_fnuz()
293
q = torch.randn(2, 3, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
294
main_kv = torch.randn(6, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
295
extra_kv = torch.randn(5, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
296
- main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size)
297
- extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, is_extra=True)
296
+ main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size, use_fnuz=main_use_fnuz)
297
+ extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, use_fnuz=False)
298
main_indices = torch.tensor([0, 2, 4, 1], dtype=torch.int32, device=device)
299
main_indptr = torch.tensor([0, 2, 4], dtype=torch.int32, device=device)
300
extra_indices = torch.tensor([1, 3, 0], dtype=torch.int32, device=device)
@@ -324,12 +324,13 @@ def test_sparse_attn_decode_ragged_kernel() -> None:
324
block_size=block_size,
325
extra_cache=extra_cache,
326
extra_rows=[[1], [3, 0]],
327
+ main_use_fnuz=main_use_fnuz,
328
)
329
330
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
331
332
332
-@requires_gfx950
333
+@requires_split_decode_arch
334
@torch.inference_mode()
335
def test_decode_num_splits_heuristic(monkeypatch) -> None:
336
"""Split-count heuristic added with the flash-decode split-K decode path."""
@@ -353,7 +354,7 @@ def test_decode_num_splits_heuristic(monkeypatch) -> None:
354
assert mod._decode_num_splits(2, 1, avg_main_len=0.0, avg_extra_len=0.0) >= 1
355
356
356
-@requires_gfx950
357
+@requires_split_decode_arch
358
@pytest.mark.parametrize("num_splits", [1, 2, 3, 4, 8])
359
@pytest.mark.parametrize("with_extra", [True, False])
360
@pytest.mark.parametrize("with_sink", [True, False])
@@ -363,8 +364,8 @@ def test_sparse_attn_decode_split_k_kernel(
364
) -> None:
365
"""Flash-decode split-K decode path (partial + reduce kernels).
366
366
- This path is the gfx950 production path (``_ON_GFX950``), so the test only
367
- runs on gfx950. The split count is pinned so the partial/reduce kernels are
367
+ This path is the gfx942/gfx950 production path, so the test only runs on
368
+ those architectures. The split count is pinned so the partial/reduce kernels are
369
exercised across split counts. ``num_splits=8`` drives splits past the
370
shortest segment length, covering the empty-split edge case handled by the
371
reduce kernel.
@@ -375,6 +376,7 @@ def test_sparse_attn_decode_split_k_kernel(
376
torch.manual_seed(7)
377
block_size = 4
378
num_heads = 3
379
+ main_use_fnuz = current_platform.is_fp8_fnuz()
380
381
main_rows = [[0, 2, 4, 6, 1, 3, 7, 5], [4, 1, 6, 0, 2]]
382
num_queries = len(main_rows)
@@ -385,7 +387,7 @@ def test_sparse_attn_decode_split_k_kernel(
387
* 0.125
388
)
389
main_kv = torch.randn(8, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
388
- main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size)
390
+ main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size, use_fnuz=main_use_fnuz)
391
main_indices, main_indptr = _ragged_from_rows(main_rows, device)
392
393
extra_rows: list[list[int]] | None = None
@@ -396,7 +398,7 @@ def test_sparse_attn_decode_split_k_kernel(
398
rows = [[1, 3, 0, 5, 2, 4], [3, 0, 6]]
399
extra_kv = torch.randn(7, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
400
extra_rows = rows
399
- extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, is_extra=True)
401
+ extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, use_fnuz=False)
402
extra_indices, extra_indptr = _ragged_from_rows(rows, device)
403
404
attn_sink = (
@@ -431,6 +433,7 @@ def test_sparse_attn_decode_split_k_kernel(
433
block_size=block_size,
434
extra_cache=extra_cache,
435
extra_rows=extra_rows,
436
+ main_use_fnuz=main_use_fnuz,
437
)
438
439
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+1
-1
@@ -2104,7 +2104,7 @@ def _rocm_sparse_attn_decode_ragged_triton(
2104
comb_dim = nope_head_dim + rope_head_dim
2105
is_fnuz = current_platform.is_fp8_fnuz()
2106
2107
- if not _ON_GFX950: # Fallback path for un-tuned architectures.
2107
+ if not (_ON_GFX942 or _ON_GFX950): # Fallback path for un-tuned architectures.
2108
block_k = 16 if head_dim >= 256 else 32
2109
_sparse_attn_decode_ragged_kernel[(num_queries, heads_blocks)](
2110
q,