iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm][Perf][GLM-5.3-Flash] Remove redundant copy after ragged sparse MLA (#58569)

Browse Source 
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Simon Danielsson
2026-10-02 11:45:51 +01:00
committed by GitHub
co-authored by mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent 5688a4dd4a
commit 751fa96a49
2 changed files with 130 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/attention/test_rocm_triton_attn_dsv4.py
+111 -2
View File
@@ -596,6 +596,114 @@ def test_sparse_attn_prefill_ragged_kernel() -> None:
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
def _sparse_prefill_ragged_inputs(nope_dim: int, rope_dim: int) -> dict:
device = torch.device("cuda")
set_random_seed(7)
head_dim = nope_dim + rope_dim
return dict(
q=torch.randn(3, 3, head_dim, dtype=torch.bfloat16, device=device) * 0.125,
kv=torch.randn(5, head_dim, dtype=torch.bfloat16, device=device) * 0.125,
indices=torch.tensor([0, 2, 1, 3, 4], dtype=torch.int32, device=device),
indptr=torch.tensor([0, 2, 5, 5], dtype=torch.int32, device=device),
scale=head_dim**-0.5,
attn_sink=torch.tensor([-0.25, 0.0, 0.25], dtype=torch.float32, device=device),
nope_head_dim=nope_dim,
rope_head_dim=rope_dim,
)
# 448+64 is DeepSeek V4/V4.1, 512+64 is V3.2, and 256+0 stands in for the
# NoPE-only layouts where the destination spans the whole head dim.
SPARSE_PREFILL_DIMS = [(NOPE_HEAD_DIM, ROPE_HEAD_DIM), (512, 64), (256, 0)]
@pytest.mark.parametrize(("nope_dim", "rope_dim"), SPARSE_PREFILL_DIMS)
@torch.inference_mode()
def test_sparse_attn_prefill_ragged_out_matches_allocated(
nope_dim: int, rope_dim: int
) -> None:
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
_rocm_sparse_attn_prefill_ragged_triton,
)
kwargs = _sparse_prefill_ragged_inputs(nope_dim, rope_dim)
allocated = _rocm_sparse_attn_prefill_ragged_triton(**kwargs)
q = kwargs["q"]
dest = torch.empty(q.shape[0], q.shape[1], nope_dim, dtype=q.dtype, device=q.device)
returned = _rocm_sparse_attn_prefill_ragged_triton(**kwargs, out=dest)
assert returned is dest
torch.testing.assert_close(dest, allocated[..., :nope_dim], atol=0.0, rtol=0.0)
@pytest.mark.parametrize(("nope_dim", "rope_dim"), SPARSE_PREFILL_DIMS)
@torch.inference_mode()
def test_sparse_attn_prefill_ragged_out_view_leaves_neighbors_untouched(
nope_dim: int, rope_dim: int
) -> None:
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
_rocm_sparse_attn_prefill_ragged_triton,
)
kwargs = _sparse_prefill_ragged_inputs(nope_dim, rope_dim)
q = kwargs["q"]
num_queries, num_heads = q.shape[0], q.shape[1]
buffer = torch.full(
(num_queries + 4, num_heads + 2, nope_dim),
-7.0,
dtype=q.dtype,
device=q.device,
)
dest = buffer[2 : 2 + num_queries, :num_heads]
assert not dest.is_contiguous()
_rocm_sparse_attn_prefill_ragged_triton(**kwargs, out=dest)
expected = _rocm_sparse_attn_prefill_ragged_triton(**kwargs)[..., :nope_dim]
torch.testing.assert_close(dest, expected, atol=0.0, rtol=0.0)
untouched = torch.ones_like(buffer, dtype=torch.bool)
untouched[2 : 2 + num_queries, :num_heads] = False
assert torch.all(buffer[untouched] == -7.0)
@pytest.mark.parametrize("out_dtype", [torch.float32, torch.float16])
@torch.inference_mode()
def test_sparse_attn_prefill_ragged_out_casts_to_dest_dtype(
out_dtype: torch.dtype,
) -> None:
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
_rocm_sparse_attn_prefill_ragged_triton,
)
kwargs = _sparse_prefill_ragged_inputs(512, 64)
q = kwargs["q"]
dest = torch.empty(q.shape[0], q.shape[1], 512, dtype=out_dtype, device=q.device)
_rocm_sparse_attn_prefill_ragged_triton(**kwargs, out=dest)
expected = _rocm_sparse_attn_prefill_ragged_triton(**kwargs)[..., :512]
assert dest.dtype == out_dtype
torch.testing.assert_close(dest.float(), expected.float(), atol=2e-2, rtol=2e-2)
@torch.inference_mode()
def test_sparse_attn_prefill_head_dim_wide_out_keeps_rope_rows() -> None:
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
_rocm_sparse_attn_prefill_ragged_triton,
)
kwargs = _sparse_prefill_ragged_inputs(NOPE_HEAD_DIM, ROPE_HEAD_DIM)
q = kwargs["q"]
dest = torch.empty_like(q)
_rocm_sparse_attn_prefill_ragged_triton(**kwargs, out=dest)
expected = _ref_sparse_prefill_ragged(
q, kwargs["kv"], [[0, 2], [1, 3, 4], []], kwargs["scale"], kwargs["attn_sink"]
)
torch.testing.assert_close(dest, expected, atol=2e-2, rtol=2e-2)
@pytest.mark.parametrize(
("num_queries", "on_gfx950", "expected"),
[(1023, True, False), (1024, True, True), (1024, False, False)],
@@ -718,10 +826,11 @@ def test_sparse_attn_prefill_preserves_dense_triton_fallback(monkeypatch) -> Non
output = torch.empty_like(q)
dense_fallback_calls = 0
def fake_dense_fallback(*args, **kwargs):
def fake_dense_fallback(*args, out, **kwargs):
nonlocal dense_fallback_calls
dense_fallback_calls += 1
return torch.zeros_like(q)
out.zero_()
return out
monkeypatch.setattr(mod, "_can_use_aiter_sparse_prefill_opus", lambda *args: True)
monkeypatch.setattr(mod, "_get_aiter_sparse_prefill_opus", lambda: None)
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+19 -5
View File
@@ -2190,6 +2190,7 @@ def _sparse_attn_prefill_ragged_kernel(
num_kv,
scale,
HAS_ATTN_SINK: tl.constexpr,
OUT_DV: tl.constexpr,
BLOCK_H: tl.constexpr,
BLOCK_D: tl.constexpr,
BLOCK_K: tl.constexpr,
@@ -2280,7 +2281,7 @@ def _sparse_attn_prefill_ragged_kernel(
+ head_offsets[:, None] * out_stride_h
+ dim_offsets[None, :] * out_stride_d,
out,
mask=head_mask[:, None] & dim_mask[None, :],
mask=head_mask[:, None] & (dim_offsets[None, :] < OUT_DV),
)
@@ -3442,6 +3443,7 @@ def _rocm_sparse_attn_prefill_ragged_triton(
attn_sink: torch.Tensor | None,
nope_head_dim: int,
rope_head_dim: int,
out: torch.Tensor | None = None,
) -> torch.Tensor:
assert q.ndim == 3, f"expected q=[sq,h,d], got {q.shape}"
assert kv.ndim == 2, f"expected kv=[skv,d], got {kv.shape}"
@@ -3472,7 +3474,15 @@ def _rocm_sparse_attn_prefill_ragged_triton(
block_d = triton.next_power_of_2(head_dim)
block_k = 16 if head_dim >= 256 else 32
num_warps = 4
out = torch.empty_like(q)
if out is None:
out = torch.empty_like(q)
else:
assert out.shape[:2] == q.shape[:2], (
f"expected out=[{num_queries},{num_heads},*], got {out.shape}"
)
assert out.shape[-1] in (nope_head_dim, head_dim), (
f"expected out width {nope_head_dim} or {head_dim}, got {out.shape[-1]}"
)
_sparse_attn_prefill_ragged_kernel[(num_queries, triton.cdiv(num_heads, block_h))](
q,
kv,
@@ -3493,6 +3503,7 @@ def _rocm_sparse_attn_prefill_ragged_triton(
kv.shape[0],
float(scale),
HAS_ATTN_SINK=has_attn_sink,
OUT_DV=out.shape[-1],
BLOCK_H=block_h,
BLOCK_D=block_d,
BLOCK_K=block_k,
@@ -3510,6 +3521,7 @@ def _rocm_sparse_attn_prefill_triton(
nope_head_dim: int,
rope_head_dim: int,
topk_length: torch.Tensor | None = None,
out: torch.Tensor | None = None,
) -> torch.Tensor:
ragged_indices, ragged_indptr = build_ragged_indices_from_dense(
indices,
@@ -3527,6 +3539,7 @@ def _rocm_sparse_attn_prefill_triton(
attn_sink=attn_sink,
nope_head_dim=nope_head_dim,
rope_head_dim=rope_head_dim,
out=out,
)
@@ -4180,7 +4193,7 @@ def rocm_sparse_attn_prefill(
return
if ragged_indices is not None and ragged_indptr is not None:
output_chunk = _rocm_sparse_attn_prefill_ragged_triton(
_rocm_sparse_attn_prefill_ragged_triton(
q=q,
kv=kv.squeeze(1),
indices=ragged_indices,
@@ -4189,11 +4202,12 @@ def rocm_sparse_attn_prefill(
attn_sink=None if attn_sink is None else attn_sink[: q.shape[1]],
nope_head_dim=nope_head_dim,
rope_head_dim=rope_head_dim,
out=output,
)
else:
assert indices is not None
indices_2d = indices.reshape(indices.shape[0], -1)
output_chunk = _rocm_sparse_attn_prefill_triton(
_rocm_sparse_attn_prefill_triton(
q=q,
kv=kv.squeeze(1),
indices=indices_2d,
@@ -4202,8 +4216,8 @@ def rocm_sparse_attn_prefill(
nope_head_dim=nope_head_dim,
rope_head_dim=rope_head_dim,
topk_length=topk_length,
out=output,
)
output.copy_(output_chunk[..., : output.shape[-1]].to(output.dtype))
def rocm_sparse_attn_decode(