From a3b9e65a291e6bfbf957d8e160bc59782a7037bc Mon Sep 17 00:00:00 2001 From: Jipeng Li Date: Tue, 22 Sep 2026 11:13:33 +0800 Subject: [PATCH] [Perf][ROCm][Attention] Narrow the Triton prefill-attention KV tile on RDNA3/RDNA4 The Triton prefill attention kernel launches with BLOCK_N equal to BLOCK_M, which get_block_size() sets to 128 for bf16. RDNA3/RDNA4 are wave32 with a much smaller LDS budget than CDNA, which makes that KV tile a poor fit. I have not profiled the stall reason, so the mechanism is a hypothesis; what is measured is that the narrower tile is faster everywhere I tested. Narrowing only BLOCK_N to 32, leaving BLOCK_M, num_warps and num_stages at the shared defaults, is faster in every cell of a 9 head size x 4 mask mode matrix on gfx1100 and gfx1201, over four runs each. Matrix medians are 2.56x-2.67x and 2.25x-2.30x; the slowest cell seen anywhere is 1.08x. End to end, embedding models gain 1.08x-1.23x. Cross-encoder rerankers reach the kernel too but gain less: 1.016x for a 6-layer model, 1.0955x for a 12-layer one. Generation prefill and decode do not reach this kernel and measure unchanged. CDNA and NVIDIA keep the existing tile: the helper returns early unless vllm.platforms.rocm.on_gfx1x() is true, which excludes gfx10xx and gfx1250. min() leaves float32, whose default tile is already 32, untouched. Both tile widths are numerically correct, so the existing tests pass whichever one is selected, and CI has no RDNA agent to exercise the new branch. The tests added to test_triton_prefill_attention.py capture the launch configuration with the platform predicates mocked, so they need no RDNA part and cover the branch on the existing CDNA and NVIDIA agents. They fail if the gating regresses to a no-op, stops excluding other platforms, widens float32, or moves any launch parameter besides BLOCK_N. Addresses #58060. Signed-off-by: Jipeng Li Co-authored-by: GitHub Copilot CLI --- .../test_triton_prefill_attention.py | 85 +++++++++++++++++++ .../attention/ops/triton_prefill_attention.py | 16 +++- 2 files changed, 100 insertions(+), 1 deletion(-) diff --git a/tests/kernels/attention/test_triton_prefill_attention.py b/tests/kernels/attention/test_triton_prefill_attention.py index 6316a926ae33..e9b90251a492 100644 --- a/tests/kernels/attention/test_triton_prefill_attention.py +++ b/tests/kernels/attention/test_triton_prefill_attention.py @@ -5,6 +5,7 @@ import torch import torch.nn.functional as F +import vllm.v1.attention.ops.triton_prefill_attention as prefill_ops from vllm.platforms import current_platform from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd @@ -230,3 +231,87 @@ def test_context_attention_sliding_window( # Compare outputs torch.testing.assert_close(o, o_ref, rtol=2e-2, atol=2e-2) + + +class _LaunchCapture: + """Records the launch configuration in place of ``_fwd_kernel``.""" + + grid: tuple + kwargs: dict + + def __getitem__(self, grid): + self.grid = grid + return self._record + + def _record(self, *args, **kwargs) -> None: + self.kwargs = kwargs + + +def _capture_tile_config( + monkeypatch, *, is_rocm: bool, on_gfx1x: bool, dtype=torch.bfloat16 +) -> _LaunchCapture: + """Capture the tile configuration with the platform predicates mocked. + + Both tile widths are numerically correct, so the tests above pass whichever + one is selected. This needs no RDNA part and allocates no device memory, so + it covers the RDNA branch on the CDNA and NVIDIA agents CI actually runs. + """ + platform = prefill_ops.current_platform + monkeypatch.setattr(platform, "is_rocm", lambda: is_rocm) + # Every device this kernel targets is cuda-alike at capability 80 or better, + # so the stock tile is 128 for 16-bit dtypes. + monkeypatch.setattr(platform, "is_cuda_alike", lambda: True) + monkeypatch.setattr(platform, "has_device_capability", lambda *a, **k: True) + if is_rocm: + import vllm.platforms.rocm as rocm_platform + + monkeypatch.setattr(rocm_platform, "on_gfx1x", lambda: on_gfx1x) + + capture = _LaunchCapture() + monkeypatch.setattr(prefill_ops, "_fwd_kernel", capture) + + def meta(*shape): + return torch.empty(shape, dtype=dtype, device="meta") + + seq_lens = torch.empty(2, dtype=torch.int32, device="meta") + context_attention_fwd( + meta(256, 8, 128), + meta(256, 2, 128), + meta(256, 2, 128), + meta(256, 8, 128), + seq_lens, + seq_lens, + 128, + ) + return capture + + +@pytest.mark.parametrize( + ("is_rocm", "on_gfx1x", "dtype", "expected_block_n"), + [ + pytest.param(True, True, torch.bfloat16, 32, id="rdna"), + # on_gfx1x() is what excludes gfx10xx, CDNA and gfx1250. + pytest.param(True, False, torch.bfloat16, 128, id="rocm-not-rdna"), + pytest.param(False, False, torch.bfloat16, 128, id="not-rocm"), + # get_block_size already returns 32 here, so min() must not widen it. + pytest.param(True, True, torch.float32, 32, id="rdna-float32"), + ], +) +def test_kv_tile_width_is_gated_by_platform( + monkeypatch, is_rocm: bool, on_gfx1x: bool, dtype, expected_block_n: int +) -> None: + capture = _capture_tile_config( + monkeypatch, is_rocm=is_rocm, on_gfx1x=on_gfx1x, dtype=dtype + ) + assert capture.kwargs["BLOCK_N"] == expected_block_n + + +def test_rdna_narrows_the_kv_tile_and_nothing_else(monkeypatch) -> None: + with monkeypatch.context() as m: + tuned = _capture_tile_config(m, is_rocm=True, on_gfx1x=True) + with monkeypatch.context() as m: + stock = _capture_tile_config(m, is_rocm=True, on_gfx1x=False) + + assert tuned.kwargs.pop("BLOCK_N") != stock.kwargs.pop("BLOCK_N") + assert tuned.kwargs == stock.kwargs + assert tuned.grid == stock.grid diff --git a/vllm/v1/attention/ops/triton_prefill_attention.py b/vllm/v1/attention/ops/triton_prefill_attention.py index 00d630edcdd5..1ed55c3e28b5 100644 --- a/vllm/v1/attention/ops/triton_prefill_attention.py +++ b/vllm/v1/attention/ops/triton_prefill_attention.py @@ -32,6 +32,16 @@ from vllm.utils.math_utils import RCP_LN2 +def _prefer_narrow_kv_tile() -> bool: + """RDNA3/RDNA4 prefer a narrower KV tile than the shared default.""" + if not current_platform.is_rocm(): + return False + from vllm.platforms.rocm import on_gfx1x + + # RDNA3/RDNA4 only: on_gfx1x() excludes gfx10xx and gfx1250. + return on_gfx1x() + + @triton.jit def _fwd_kernel( Q, @@ -249,6 +259,10 @@ def context_attention_fwd( grid = (batch, head, triton.cdiv(max_input_len, BLOCK)) num_warps = 4 if Lk <= 64 else 8 + # BLOCK_M, num_warps and num_stages stay at the shared defaults; min() + # leaves dtypes whose default tile is already 32, such as float32, alone. + BLOCK_N = min(BLOCK, 32) if _prefer_narrow_kv_tile() else BLOCK + sliding_window_q = sliding_window_q if sliding_window_q is not None else 0 sliding_window_k = sliding_window_k if sliding_window_k is not None else 0 @@ -272,7 +286,7 @@ def context_attention_fwd( kv_group_num=kv_group_num, BLOCK_M=BLOCK, BLOCK_DMODEL=triton.next_power_of_2(Lk), - BLOCK_N=BLOCK, + BLOCK_N=BLOCK_N, IS_CAUSAL=is_causal, SLIDING_WINDOW_Q=sliding_window_q, SLIDING_WINDOW_K=sliding_window_k,