[Perf][ROCm][Attention] Narrow the Triton prefill-attention KV tile on RDNA3/RDNA4 (#58225)
Signed-off-by: Jipeng Li <jipengli@amd.com> Co-authored-by: GitHub Copilot CLI <noreply@github.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
lijipeng committed
Sep 23, 2026 at 16:39 UTC
1904ed9402ce81e7db4391841c2205bf785380d9
2 files changed
+100
-1
tests/kernels/attention/test_triton_prefill_attention.py
+85
@@ -5,6 +5,7 @@ import pytest
5
import torch
6
import torch.nn.functional as F
7
8
+import vllm.v1.attention.ops.triton_prefill_attention as prefill_ops
9
from vllm.platforms import current_platform
10
from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd
11
@@ -230,3 +231,87 @@ def test_context_attention_sliding_window(
231
232
# Compare outputs
233
torch.testing.assert_close(o, o_ref, rtol=2e-2, atol=2e-2)
234
+
235
+
236
+class _LaunchCapture:
237
+ """Records the launch configuration in place of ``_fwd_kernel``."""
238
+
239
+ grid: tuple
240
+ kwargs: dict
241
+
242
+ def __getitem__(self, grid):
243
+ self.grid = grid
244
+ return self._record
245
+
246
+ def _record(self, *args, **kwargs) -> None:
247
+ self.kwargs = kwargs
248
+
249
+
250
+def _capture_tile_config(
251
+ monkeypatch, *, is_rocm: bool, on_gfx1x: bool, dtype=torch.bfloat16
252
+) -> _LaunchCapture:
253
+ """Capture the tile configuration with the platform predicates mocked.
254
+
255
+ Both tile widths are numerically correct, so the tests above pass whichever
256
+ one is selected. This needs no RDNA part and allocates no device memory, so
257
+ it covers the RDNA branch on the CDNA and NVIDIA agents CI actually runs.
258
+ """
259
+ platform = prefill_ops.current_platform
260
+ monkeypatch.setattr(platform, "is_rocm", lambda: is_rocm)
261
+ # Every device this kernel targets is cuda-alike at capability 80 or better,
262
+ # so the stock tile is 128 for 16-bit dtypes.
263
+ monkeypatch.setattr(platform, "is_cuda_alike", lambda: True)
264
+ monkeypatch.setattr(platform, "has_device_capability", lambda *a, **k: True)
265
+ if is_rocm:
266
+ import vllm.platforms.rocm as rocm_platform
267
+
268
+ monkeypatch.setattr(rocm_platform, "on_gfx1x", lambda: on_gfx1x)
269
+
270
+ capture = _LaunchCapture()
271
+ monkeypatch.setattr(prefill_ops, "_fwd_kernel", capture)
272
+
273
+ def meta(*shape):
274
+ return torch.empty(shape, dtype=dtype, device="meta")
275
+
276
+ seq_lens = torch.empty(2, dtype=torch.int32, device="meta")
277
+ context_attention_fwd(
278
+ meta(256, 8, 128),
279
+ meta(256, 2, 128),
280
+ meta(256, 2, 128),
281
+ meta(256, 8, 128),
282
+ seq_lens,
283
+ seq_lens,
284
+ 128,
285
+ )
286
+ return capture
287
+
288
+
289
+@pytest.mark.parametrize(
290
+ ("is_rocm", "on_gfx1x", "dtype", "expected_block_n"),
291
+ [
292
+ pytest.param(True, True, torch.bfloat16, 32, id="rdna"),
293
+ # on_gfx1x() is what excludes gfx10xx, CDNA and gfx1250.
294
+ pytest.param(True, False, torch.bfloat16, 128, id="rocm-not-rdna"),
295
+ pytest.param(False, False, torch.bfloat16, 128, id="not-rocm"),
296
+ # get_block_size already returns 32 here, so min() must not widen it.
297
+ pytest.param(True, True, torch.float32, 32, id="rdna-float32"),
298
+ ],
299
+)
300
+def test_kv_tile_width_is_gated_by_platform(
301
+ monkeypatch, is_rocm: bool, on_gfx1x: bool, dtype, expected_block_n: int
302
+) -> None:
303
+ capture = _capture_tile_config(
304
+ monkeypatch, is_rocm=is_rocm, on_gfx1x=on_gfx1x, dtype=dtype
305
+ )
306
+ assert capture.kwargs["BLOCK_N"] == expected_block_n
307
+
308
+
309
+def test_rdna_narrows_the_kv_tile_and_nothing_else(monkeypatch) -> None:
310
+ with monkeypatch.context() as m:
311
+ tuned = _capture_tile_config(m, is_rocm=True, on_gfx1x=True)
312
+ with monkeypatch.context() as m:
313
+ stock = _capture_tile_config(m, is_rocm=True, on_gfx1x=False)
314
+
315
+ assert tuned.kwargs.pop("BLOCK_N") != stock.kwargs.pop("BLOCK_N")
316
+ assert tuned.kwargs == stock.kwargs
317
+ assert tuned.grid == stock.grid
vllm/v1/attention/ops/triton_prefill_attention.py
+15
-1
@@ -32,6 +32,16 @@ from vllm.triton_utils import tl, triton
32
from vllm.utils.math_utils import RCP_LN2
33
34
35
+def _prefer_narrow_kv_tile() -> bool:
36
+ """RDNA3/RDNA4 prefer a narrower KV tile than the shared default."""
37
+ if not current_platform.is_rocm():
38
+ return False
39
+ from vllm.platforms.rocm import on_gfx1x
40
+
41
+ # RDNA3/RDNA4 only: on_gfx1x() excludes gfx10xx and gfx1250.
42
+ return on_gfx1x()
43
+
44
+
45
@triton.jit
46
def _fwd_kernel(
47
Q,
@@ -249,6 +259,10 @@ def context_attention_fwd(
259
grid = (batch, head, triton.cdiv(max_input_len, BLOCK))
260
num_warps = 4 if Lk <= 64 else 8
261
262
+ # BLOCK_M, num_warps and num_stages stay at the shared defaults; min()
263
+ # leaves dtypes whose default tile is already 32, such as float32, alone.
264
+ BLOCK_N = min(BLOCK, 32) if _prefer_narrow_kv_tile() else BLOCK
265
+
266
sliding_window_q = sliding_window_q if sliding_window_q is not None else 0
267
sliding_window_k = sliding_window_k if sliding_window_k is not None else 0
268
@@ -272,7 +286,7 @@ def context_attention_fwd(
286
kv_group_num=kv_group_num,
287
BLOCK_M=BLOCK,
288
BLOCK_DMODEL=triton.next_power_of_2(Lk),
275
- BLOCK_N=BLOCK,
289
+ BLOCK_N=BLOCK_N,
290
IS_CAUSAL=is_causal,
291
SLIDING_WINDOW_Q=sliding_window_q,
292
SLIDING_WINDOW_K=sliding_window_k,