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

[ROCm][Perf] Pad fp16/bf16 GEMM weight row stride for L2 cache (#55090)

Browse Source 
Signed-off-by: Matthias Gehre <matthias.gehre@amd.com>
This commit is contained in:
Matthias Gehre
2026-10-15 14:26:32 +00:00
committed by GitHub
parent d12d5460a5
commit 011fa68ca9
3 changed files with 198 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/model_executor/layers/test_rocm_unquantized_gemm.py
+124
View File
@@ -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
vllm/model_executor/layers/linear.py
+3
View File
@@ -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,
vllm/model_executor/layers/utils.py
+71 -13
View File
@@ -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