From 011fa68ca90c4b8745ac0376a4a44ec9cf2e8b23 Mon Sep 17 00:00:00 2001 From: Matthias Gehre Date: Thu, 15 Oct 2026 16:26:32 +0200 Subject: [PATCH] [ROCm][Perf] Pad fp16/bf16 GEMM weight row stride for L2 cache (#55090) Signed-off-by: Matthias Gehre --- .../layers/test_rocm_unquantized_gemm.py | 124 ++++++++++++++++++ vllm/model_executor/layers/linear.py | 3 + vllm/model_executor/layers/utils.py | 84 ++++++++++-- 3 files changed, 198 insertions(+), 13 deletions(-) diff --git a/tests/model_executor/layers/test_rocm_unquantized_gemm.py b/tests/model_executor/layers/test_rocm_unquantized_gemm.py index 491cfa5435..db736c7438 100644 --- a/tests/model_executor/layers/test_rocm_unquantized_gemm.py +++ b/tests/model_executor/layers/test_rocm_unquantized_gemm.py @@ -17,6 +17,15 @@ if current_platform.is_cuda(): from vllm.model_executor.layers import utils +def _on_gfx11() -> bool: + """gfx11 check that is safe to evaluate at collection time on any platform.""" + if not current_platform.is_rocm(): + return False + from vllm.platforms.rocm import on_gfx11 + + return on_gfx11() + + @pytest.mark.parametrize("m", [1, 128]) def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch, m): x = torch.randn(1, 64, dtype=torch.float16) @@ -145,6 +154,101 @@ def test_rocm_unquantized_gemm_rejects_unsupported_skinny_layouts( assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3) +def _row_stride_padded(weight: torch.Tensor, pad: int) -> torch.Tensor: + """A weight whose row stride is offset by ``pad`` elements, as produced by + ``maybe_pad_weight_avoid_cache_aliasing``.""" + padded = torch.nn.functional.pad(weight, (0, pad))[..., :-pad] + assert padded.shape == weight.shape + assert padded.stride() == (weight.shape[1] + pad, 1) + return padded + + +def test_rocm_unquantized_gemm_wvsplitk_accepts_row_stride_padded_weight(monkeypatch): + # wvSplitK reads the weight row stride, so a row-stride-padded weight must + # still reach it instead of falling back to torch.nn.functional.linear. + x = torch.randn(4, 1024, dtype=torch.float16) + weight = _row_stride_padded(torch.randn(128, 1024, dtype=torch.float16), 64) + + monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False) + monkeypatch.setattr(utils.rocm_aiter_ops, "is_tgemm_enabled", lambda: False) + monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx11", lambda: True) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: False) + monkeypatch.setattr(utils, "num_compute_units", lambda: 120) + + wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t()) + monkeypatch.setattr(utils.ops, "wvSplitK", wvsplitk_mock) + + out = utils.rocm_unquantized_gemm_impl(x, weight, None) + ref = torch.nn.functional.linear(x, weight, None) + + wvsplitk_mock.assert_called_once() + assert wvsplitk_mock.call_args.args[0].stride() == weight.stride() + assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3) + + +@pytest.mark.parametrize( + "k,element_size,expected", + [ + (4096, 2, 64), # 8192 B row stride, aliases + (12288, 2, 64), # 24576 B row stride, aliases + (1024, 2, 64), # smallest aliasing bf16/fp16 K + (2560, 2, 0), # 5120 B, no aliasing + (11008, 2, 0), # 22016 B, no aliasing + (2048, 1, 128), # fp8: 2048 B, aliases, twice the pad + (512, 4, 32), # fp32: 2048 B, aliases, half the pad + ], +) +def test_cache_aliasing_pad_elems(k, element_size, expected): + assert utils.cache_aliasing_pad_elems(k, element_size) == expected + + +def test_maybe_pad_weight_avoid_cache_aliasing_skips_non_gfx11(monkeypatch): + # Padding is gfx11-only, even for a K whose row stride would alias there. + monkeypatch.setattr(utils.current_platform, "is_rocm", lambda: True) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx11", lambda: False) + + weight = torch.randn(128, 4096, dtype=torch.bfloat16) + assert utils.maybe_pad_weight_avoid_cache_aliasing(weight) is weight + + +def test_maybe_pad_weight_avoid_cache_aliasing(monkeypatch): + monkeypatch.setattr(utils.current_platform, "is_rocm", lambda: True) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx11", lambda: True) + + weight = torch.randn(128, 4096, dtype=torch.bfloat16) + padded = utils.maybe_pad_weight_avoid_cache_aliasing(weight) + + assert padded.shape == weight.shape + assert padded.stride() == (4096 + 64, 1) + assert torch.equal(padded, weight) + # Idempotent: the padded stride no longer aliases, so a second call is a no-op + # and preserves the pointer captured by any graph. + repadded = utils.maybe_pad_weight_avoid_cache_aliasing(padded) + assert repadded.data_ptr() == padded.data_ptr() + + +@pytest.mark.parametrize( + "weight", + [ + torch.randn(128, 2560, dtype=torch.bfloat16), # no aliasing + torch.randn(2560, 128, dtype=torch.bfloat16).t(), # not contiguous + torch.randn(2, 128, 4096, dtype=torch.bfloat16), # not 2D + ], +) +def test_maybe_pad_weight_avoid_cache_aliasing_skips(monkeypatch, weight): + monkeypatch.setattr(utils.current_platform, "is_rocm", lambda: True) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx11", lambda: True) + + assert ( + utils.maybe_pad_weight_avoid_cache_aliasing(weight).data_ptr() + == weight.data_ptr() + ) + + @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-only kernel test") @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) def test_rocm_unquantized_gemm_noncontiguous_activation_real_kernel(monkeypatch, dtype): @@ -164,6 +268,26 @@ def test_rocm_unquantized_gemm_noncontiguous_activation_real_kernel(monkeypatch, torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2) +@pytest.mark.skipif( + not _on_gfx11(), reason="cache-aliasing padding only applies on gfx11" +) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_rocm_unquantized_gemm_padded_weight_real_kernel(monkeypatch, dtype): + x = torch.randn(4, 4096, device="cuda", dtype=dtype) + weight = _row_stride_padded(torch.randn(128, 4096, device="cuda", dtype=dtype), 64) + + monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True) + original_wvsplitk = utils.ops.wvSplitK + wvsplitk_mock = MagicMock(side_effect=original_wvsplitk) + monkeypatch.setattr(utils.ops, "wvSplitK", wvsplitk_mock) + + out = utils.rocm_unquantized_gemm_impl(x, weight, None) + ref = torch.nn.functional.linear(x, weight, None) + + wvsplitk_mock.assert_called_once() + torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2) + + @pytest.mark.parametrize("m", [1, 128]) def test_rocm_unquantized_gemm_gfx1x_n_gt_5_falls_back(monkeypatch, m): # wvSplitK skinny GEMM handles n in [1, 5] (see PR #40687); n > 5 must diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 5b77caa256..7a97321074 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -31,6 +31,7 @@ from vllm.model_executor.layers.quantization.base_config import ( ) from vllm.model_executor.layers.utils import ( dispatch_unquantized_gemm, + maybe_pad_weight_avoid_cache_aliasing, ) from vllm.model_executor.parameter import ( BasevLLMParameter, @@ -225,6 +226,8 @@ class UnquantizedLinearMethod(LinearMethodBase): and weight.stride(0) != 1 ): layer.weight.data = weight.t().contiguous().t() + elif current_platform.is_rocm(): + layer.weight.data = maybe_pad_weight_avoid_cache_aliasing(layer.weight.data) def apply( self, diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index 8d285dda75..b89a6443ad 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -285,6 +285,44 @@ def wvsplitkrc_dispatch(n: int, k: int, m: int, cu_count: int) -> tuple[int, boo return chunkk, fits +# On gfx11 a weight row stride that is a multiple of 2048 B makes the multi-row global +# loads of a GEMM collide on the L2/MALL channel hash. Offsetting the stride by one +# 128 B cache line dodges the collision and keeps each row cache-line aligned. +def cache_aliasing_pad_elems(k: int, element_size: int) -> int: + """Row-stride pad, in elements, that stops an ``[n, k]`` weight aliasing on gfx11, + or 0 when the stride does not alias. + + Args: + k: Number of elements per weight row. + element_size: Size of one element in bytes. + + Returns: + The number of elements to add to the row stride. + + """ + if (k * element_size) % 2048 != 0: + return 0 + return 128 // element_size + + +def maybe_pad_weight_avoid_cache_aliasing(weight: torch.Tensor) -> torch.Tensor: + """Offset a dense weight's row stride so it stops aliasing. + + Returns a view with the original shape whose row stride is one cache line + wider, or ``weight`` unchanged when the pad does not apply. + """ + from vllm.platforms.rocm import on_gfx11 + + if not (current_platform.is_rocm() and on_gfx11()): + return weight + if weight.dim() != 2 or not weight.is_contiguous(): + return weight + pad = cache_aliasing_pad_elems(weight.shape[1], weight.element_size()) + if pad == 0: + return weight + return torch.nn.functional.pad(weight, (0, pad))[..., :-pad] + + def rocm_unquantized_gemm_impl( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None ) -> torch.Tensor: @@ -295,7 +333,7 @@ def rocm_unquantized_gemm_impl( return linear_batch_invariant(x, weight, bias) - from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx950, on_gfx1250 + from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx11, on_gfx950, on_gfx1250 n = x.numel() // x.size(-1) m = weight.shape[0] @@ -337,29 +375,49 @@ def rocm_unquantized_gemm_impl( return gemm_a16w16(x, weight, bias) - use_skinny = ( + skinny_gemm_enabled = ( envs.VLLM_ROCM_USE_SKINNY_GEMM and (on_gfx9() or on_gfx1x()) # build (gfx9/gfx11 ISA); fall back to torch GEMM there. # TODO GFX1250: Include once skinny GEMM is supported on gfx1250 and x.dtype in [torch.float16, torch.bfloat16] and k % 8 == 0 - and skinny_operands_compatible ) - if use_skinny: + # wvSplitK allows row strides. Perf impact is only validated on gfx11. + weight_ok_for_wvsplitk = weight.is_contiguous() or ( + on_gfx11() and weight.stride(1) == 1 + ) + + use_wvsplitk = ( + skinny_gemm_enabled + and weight_ok_for_wvsplitk + and (bias is None or bias.is_contiguous()) + and (m == 1 or m > 8) + and 0 < n <= 5 + ) + + use_llmm1 = ( + skinny_gemm_enabled + and weight.is_contiguous() + and bias is None + and m % 4 == 0 + and n == 1 + and k <= 8192 + ) + + if use_wvsplitk: # The skinny kernels assume contiguous K elements. A shape-preserving # reshape can retain a transposed activation's non-contiguous strides. # Note: Only build that view inside the branches that consume it. - if (m == 1 or m > 8) and 0 < n <= 5: - x_view = x.reshape(-1, x.size(-1)).contiguous() - cu_count = num_compute_units() - out = ops.wvSplitK(weight, x_view, cu_count, bias) - return out.reshape(*x.shape[:-1], weight.shape[0]) - elif m % 4 == 0 and n == 1 and k <= 8192 and bias is None: - x_view = x.reshape(-1, x.size(-1)).contiguous() - out = ops.LLMM1(weight, x_view, 4) - return out.reshape(*x.shape[:-1], weight.shape[0]) + x_view = x.reshape(-1, x.size(-1)).contiguous() + out = ops.wvSplitK(weight, x_view, cu_count, bias) + return out.reshape(*x.shape[:-1], weight.shape[0]) + + if use_llmm1: + x_view = x.reshape(-1, x.size(-1)).contiguous() + out = ops.LLMM1(weight, x_view, 4) + return out.reshape(*x.shape[:-1], weight.shape[0]) if rocm_aiter_ops.is_tgemm_enabled(): from aiter.tuned_gemm import tgemm