[ROCm][CI] Cover the AITER MQA logits dispatch on gfx950 (#58724)
Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com> Co-authored-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
stefankoncarevic committed
Sep 25, 2026 at 19:43 UTC
b21757555ac5f853a8eff02a687d2777437e055b
1 file changed
+59
tests/kernels/attention/test_rocm_triton_attn_dsv4.py
+59
@@ -52,6 +52,65 @@ ROPE_HEAD_DIM = 64
52
HEAD_DIM = NOPE_HEAD_DIM + ROPE_HEAD_DIM
53
54
55
+@pytest.fixture
56
+def enable_aiter_mqa(monkeypatch):
57
+ from vllm._aiter_ops import rocm_aiter_ops
58
+
59
+ with monkeypatch.context() as context:
60
+ context.setenv("VLLM_ROCM_USE_AITER", "1")
61
+ rocm_aiter_ops.refresh_env_variables()
62
+ yield
63
+ rocm_aiter_ops.refresh_env_variables()
64
+
65
+
66
+@requires_gfx950
67
+@pytest.mark.parametrize("column_scales", [False, True], ids=["vector", "column"])
68
+@torch.inference_mode()
69
+def test_unpaged_mqa_logits_preserves_intervals_and_scale_layout(
70
+ column_scales, monkeypatch, enable_aiter_mqa
71
+) -> None:
72
+ """The vLLM AITER dispatch preserves ragged masking and per-token scales."""
73
+ from vllm._aiter_ops import rocm_aiter_ops
74
+ from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
75
+
76
+ assert rocm_aiter_ops.is_enabled(), "AITER is required on gfx950"
77
+ assert mod.mqa_logits_module() is not None, "AITER MQA logits are required"
78
+ torch.manual_seed(42)
79
+ device = torch.device("cuda")
80
+ dtype = current_platform.fp8_dtype()
81
+ q = torch.randn(7, 32, 128, device=device).to(dtype)
82
+ k = torch.randn(259, 128, device=device).to(dtype)
83
+ scales = torch.linspace(0.25, 1.25, k.shape[0], device=device)
84
+ if column_scales:
85
+ scales = scales[:, None]
86
+ weights = torch.rand(q.shape[:2], device=device) / q.shape[1]
87
+ starts = torch.tensor([0, 3, 128, 5, 0, 257, 259], device=device, dtype=torch.int32)
88
+ ends = torch.tensor(
89
+ [1, 129, 259, 5, 259, 259, 259], device=device, dtype=torch.int32
90
+ )
91
+
92
+ module = mod.mqa_logits_module()
93
+ implementation = module.fp8_mqa_logits
94
+ calls = 0
95
+
96
+ def traced_implementation(*args, **kwargs):
97
+ nonlocal calls
98
+ calls += 1
99
+ return implementation(*args, **kwargs)
100
+
101
+ monkeypatch.setattr(module, "fp8_mqa_logits", traced_implementation)
102
+ actual = mod.rocm_fp8_mqa_logits(q, (k, scales), weights, starts, ends)
103
+ assert calls == 1
104
+ score = torch.einsum("mhd,nd->mhn", q.float(), k.float()) * scales.reshape(-1)
105
+ expected = (score.relu() * weights[..., None]).sum(dim=1)
106
+ positions = torch.arange(k.shape[0], device=device)
107
+ valid = (positions >= starts[:, None]) & (positions < ends[:, None])
108
+ assert actual.shape == expected.shape
109
+ assert actual.dtype == torch.float32
110
+ assert torch.equal(torch.isneginf(actual), ~valid)
111
+ torch.testing.assert_close(actual[valid], expected[valid], rtol=1e-3, atol=1e-3)
112
+
113
+
114
def _ref_global_topk_ragged(
115
topk_indices: torch.Tensor,
116
token_to_req_indices: torch.Tensor,