[ROCm][Perf][GLM-5.3-Flash] Remove redundant copy after ragged sparse MLA (#58569)
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Simon Danielsson committed
Oct 2, 2026 at 12:45 UTC
751fa96a497fdba67ca594db832c02140478208f
2 files changed
+130
-7
tests/kernels/attention/test_rocm_triton_attn_dsv4.py
+111
-2
@@ -596,6 +596,114 @@ def test_sparse_attn_prefill_ragged_kernel() -> None:
596
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
597
598
599
+def _sparse_prefill_ragged_inputs(nope_dim: int, rope_dim: int) -> dict:
600
+ device = torch.device("cuda")
601
+ set_random_seed(7)
602
+ head_dim = nope_dim + rope_dim
603
+ return dict(
604
+ q=torch.randn(3, 3, head_dim, dtype=torch.bfloat16, device=device) * 0.125,
605
+ kv=torch.randn(5, head_dim, dtype=torch.bfloat16, device=device) * 0.125,
606
+ indices=torch.tensor([0, 2, 1, 3, 4], dtype=torch.int32, device=device),
607
+ indptr=torch.tensor([0, 2, 5, 5], dtype=torch.int32, device=device),
608
+ scale=head_dim**-0.5,
609
+ attn_sink=torch.tensor([-0.25, 0.0, 0.25], dtype=torch.float32, device=device),
610
+ nope_head_dim=nope_dim,
611
+ rope_head_dim=rope_dim,
612
+ )
613
+
614
+
615
+# 448+64 is DeepSeek V4/V4.1, 512+64 is V3.2, and 256+0 stands in for the
616
+# NoPE-only layouts where the destination spans the whole head dim.
617
+SPARSE_PREFILL_DIMS = [(NOPE_HEAD_DIM, ROPE_HEAD_DIM), (512, 64), (256, 0)]
618
+
619
+
620
+@pytest.mark.parametrize(("nope_dim", "rope_dim"), SPARSE_PREFILL_DIMS)
621
+@torch.inference_mode()
622
+def test_sparse_attn_prefill_ragged_out_matches_allocated(
623
+ nope_dim: int, rope_dim: int
624
+) -> None:
625
+ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
626
+ _rocm_sparse_attn_prefill_ragged_triton,
627
+ )
628
+
629
+ kwargs = _sparse_prefill_ragged_inputs(nope_dim, rope_dim)
630
+ allocated = _rocm_sparse_attn_prefill_ragged_triton(**kwargs)
631
+
632
+ q = kwargs["q"]
633
+ dest = torch.empty(q.shape[0], q.shape[1], nope_dim, dtype=q.dtype, device=q.device)
634
+ returned = _rocm_sparse_attn_prefill_ragged_triton(**kwargs, out=dest)
635
+
636
+ assert returned is dest
637
+ torch.testing.assert_close(dest, allocated[..., :nope_dim], atol=0.0, rtol=0.0)
638
+
639
+
640
+@pytest.mark.parametrize(("nope_dim", "rope_dim"), SPARSE_PREFILL_DIMS)
641
+@torch.inference_mode()
642
+def test_sparse_attn_prefill_ragged_out_view_leaves_neighbors_untouched(
643
+ nope_dim: int, rope_dim: int
644
+) -> None:
645
+ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
646
+ _rocm_sparse_attn_prefill_ragged_triton,
647
+ )
648
+
649
+ kwargs = _sparse_prefill_ragged_inputs(nope_dim, rope_dim)
650
+ q = kwargs["q"]
651
+ num_queries, num_heads = q.shape[0], q.shape[1]
652
+
653
+ buffer = torch.full(
654
+ (num_queries + 4, num_heads + 2, nope_dim),
655
+ -7.0,
656
+ dtype=q.dtype,
657
+ device=q.device,
658
+ )
659
+ dest = buffer[2 : 2 + num_queries, :num_heads]
660
+ assert not dest.is_contiguous()
661
+ _rocm_sparse_attn_prefill_ragged_triton(**kwargs, out=dest)
662
+
663
+ expected = _rocm_sparse_attn_prefill_ragged_triton(**kwargs)[..., :nope_dim]
664
+ torch.testing.assert_close(dest, expected, atol=0.0, rtol=0.0)
665
+
666
+ untouched = torch.ones_like(buffer, dtype=torch.bool)
667
+ untouched[2 : 2 + num_queries, :num_heads] = False
668
+ assert torch.all(buffer[untouched] == -7.0)
669
+
670
+
671
+@pytest.mark.parametrize("out_dtype", [torch.float32, torch.float16])
672
+@torch.inference_mode()
673
+def test_sparse_attn_prefill_ragged_out_casts_to_dest_dtype(
674
+ out_dtype: torch.dtype,
675
+) -> None:
676
+ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
677
+ _rocm_sparse_attn_prefill_ragged_triton,
678
+ )
679
+
680
+ kwargs = _sparse_prefill_ragged_inputs(512, 64)
681
+ q = kwargs["q"]
682
+ dest = torch.empty(q.shape[0], q.shape[1], 512, dtype=out_dtype, device=q.device)
683
+ _rocm_sparse_attn_prefill_ragged_triton(**kwargs, out=dest)
684
+
685
+ expected = _rocm_sparse_attn_prefill_ragged_triton(**kwargs)[..., :512]
686
+ assert dest.dtype == out_dtype
687
+ torch.testing.assert_close(dest.float(), expected.float(), atol=2e-2, rtol=2e-2)
688
+
689
+
690
+@torch.inference_mode()
691
+def test_sparse_attn_prefill_head_dim_wide_out_keeps_rope_rows() -> None:
692
+ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
693
+ _rocm_sparse_attn_prefill_ragged_triton,
694
+ )
695
+
696
+ kwargs = _sparse_prefill_ragged_inputs(NOPE_HEAD_DIM, ROPE_HEAD_DIM)
697
+ q = kwargs["q"]
698
+ dest = torch.empty_like(q)
699
+ _rocm_sparse_attn_prefill_ragged_triton(**kwargs, out=dest)
700
+
701
+ expected = _ref_sparse_prefill_ragged(
702
+ q, kwargs["kv"], [[0, 2], [1, 3, 4], []], kwargs["scale"], kwargs["attn_sink"]
703
+ )
704
+ torch.testing.assert_close(dest, expected, atol=2e-2, rtol=2e-2)
705
+
706
+
707
@pytest.mark.parametrize(
708
("num_queries", "on_gfx950", "expected"),
709
[(1023, True, False), (1024, True, True), (1024, False, False)],
@@ -718,10 +826,11 @@ def test_sparse_attn_prefill_preserves_dense_triton_fallback(monkeypatch) -> Non
826
output = torch.empty_like(q)
827
dense_fallback_calls = 0
828
721
- def fake_dense_fallback(*args, **kwargs):
829
+ def fake_dense_fallback(*args, out, **kwargs):
830
nonlocal dense_fallback_calls
831
dense_fallback_calls += 1
724
- return torch.zeros_like(q)
832
+ out.zero_()
833
+ return out
834
835
monkeypatch.setattr(mod, "_can_use_aiter_sparse_prefill_opus", lambda *args: True)
836
monkeypatch.setattr(mod, "_get_aiter_sparse_prefill_opus", lambda: None)
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+19
-5
@@ -2190,6 +2190,7 @@ def _sparse_attn_prefill_ragged_kernel(
2190
num_kv,
2191
scale,
2192
HAS_ATTN_SINK: tl.constexpr,
2193
+ OUT_DV: tl.constexpr,
2194
BLOCK_H: tl.constexpr,
2195
BLOCK_D: tl.constexpr,
2196
BLOCK_K: tl.constexpr,
@@ -2280,7 +2281,7 @@ def _sparse_attn_prefill_ragged_kernel(
2281
+ head_offsets[:, None] * out_stride_h
2282
+ dim_offsets[None, :] * out_stride_d,
2283
out,
2283
- mask=head_mask[:, None] & dim_mask[None, :],
2284
+ mask=head_mask[:, None] & (dim_offsets[None, :] < OUT_DV),
2285
)
2286
2287
@@ -3442,6 +3443,7 @@ def _rocm_sparse_attn_prefill_ragged_triton(
3443
attn_sink: torch.Tensor | None,
3444
nope_head_dim: int,
3445
rope_head_dim: int,
3446
+ out: torch.Tensor | None = None,
3447
) -> torch.Tensor:
3448
assert q.ndim == 3, f"expected q=[sq,h,d], got {q.shape}"
3449
assert kv.ndim == 2, f"expected kv=[skv,d], got {kv.shape}"
@@ -3472,7 +3474,15 @@ def _rocm_sparse_attn_prefill_ragged_triton(
3474
block_d = triton.next_power_of_2(head_dim)
3475
block_k = 16 if head_dim >= 256 else 32
3476
num_warps = 4
3475
- out = torch.empty_like(q)
3477
+ if out is None:
3478
+ out = torch.empty_like(q)
3479
+ else:
3480
+ assert out.shape[:2] == q.shape[:2], (
3481
+ f"expected out=[{num_queries},{num_heads},*], got {out.shape}"
3482
+ )
3483
+ assert out.shape[-1] in (nope_head_dim, head_dim), (
3484
+ f"expected out width {nope_head_dim} or {head_dim}, got {out.shape[-1]}"
3485
+ )
3486
_sparse_attn_prefill_ragged_kernel[(num_queries, triton.cdiv(num_heads, block_h))](
3487
q,
3488
kv,
@@ -3493,6 +3503,7 @@ def _rocm_sparse_attn_prefill_ragged_triton(
3503
kv.shape[0],
3504
float(scale),
3505
HAS_ATTN_SINK=has_attn_sink,
3506
+ OUT_DV=out.shape[-1],
3507
BLOCK_H=block_h,
3508
BLOCK_D=block_d,
3509
BLOCK_K=block_k,
@@ -3510,6 +3521,7 @@ def _rocm_sparse_attn_prefill_triton(
3521
nope_head_dim: int,
3522
rope_head_dim: int,
3523
topk_length: torch.Tensor | None = None,
3524
+ out: torch.Tensor | None = None,
3525
) -> torch.Tensor:
3526
ragged_indices, ragged_indptr = build_ragged_indices_from_dense(
3527
indices,
@@ -3527,6 +3539,7 @@ def _rocm_sparse_attn_prefill_triton(
3539
attn_sink=attn_sink,
3540
nope_head_dim=nope_head_dim,
3541
rope_head_dim=rope_head_dim,
3542
+ out=out,
3543
)
3544
3545
@@ -4180,7 +4193,7 @@ def rocm_sparse_attn_prefill(
4193
return
4194
4195
if ragged_indices is not None and ragged_indptr is not None:
4183
- output_chunk = _rocm_sparse_attn_prefill_ragged_triton(
4196
+ _rocm_sparse_attn_prefill_ragged_triton(
4197
q=q,
4198
kv=kv.squeeze(1),
4199
indices=ragged_indices,
@@ -4189,11 +4202,12 @@ def rocm_sparse_attn_prefill(
4202
attn_sink=None if attn_sink is None else attn_sink[: q.shape[1]],
4203
nope_head_dim=nope_head_dim,
4204
rope_head_dim=rope_head_dim,
4205
+ out=output,
4206
)
4207
else:
4208
assert indices is not None
4209
indices_2d = indices.reshape(indices.shape[0], -1)
4196
- output_chunk = _rocm_sparse_attn_prefill_triton(
4210
+ _rocm_sparse_attn_prefill_triton(
4211
q=q,
4212
kv=kv.squeeze(1),
4213
indices=indices_2d,
@@ -4202,8 +4216,8 @@ def rocm_sparse_attn_prefill(
4216
nope_head_dim=nope_head_dim,
4217
rope_head_dim=rope_head_dim,
4218
topk_length=topk_length,
4219
+ out=output,
4220
)
4206
- output.copy_(output_chunk[..., : output.shape[-1]].to(output.dtype))
4221
4222
4223
def rocm_sparse_attn_decode(