Skip to content

[ROCm][Perf] Fix wvSplitK tile choice for N>=5 on gfx950 and add N=6 #62846

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
jiacao-amd wants to merge 1 commit into from
Closed

[ROCm][Perf] Fix wvSplitK tile choice for N>=5 on gfx950 and add N=6 #62846

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 28 additions & 13 deletions csrc/rocm/skinny_gemms.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,15 @@ bool on_gfx12() {
return result;
}

bool on_gfx950() {
static const bool result = [] {
const auto* dprops = at::cuda::getCurrentDeviceProperties();
const std::string device_arch = dprops->gcnArchName;
return device_arch.find("gfx950") != std::string::npos;
}();
return result;
}

bool on_gfx1151() {
static const bool result = [] {
const auto* dprops = at::cuda::getCurrentDeviceProperties();
Expand Down Expand Up @@ -1239,19 +1248,22 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
CuCount); \
}

#define WVSPLIT_TILE_CFG(_THRDS, _WVPRGRP, _sYT, __N) \
{ \
bool fit_lds = (Kbp_in * N_in <= max_lds_len); \
if (_sYT <= 1) \
WVSPLITK_CFG(_THRDS, _WVPRGRP, 1, 4, __N) \
else if ((__N == 1) || (!fit_lds) || (_sYT <= 4 * 2)) \
WVSPLITK_CFG(_THRDS, _WVPRGRP, 2, 2, __N) \
else if (_sYT <= 4 * 3) \
WVSPLITK_CFG(_THRDS, _WVPRGRP, 3, 2, __N) \
else if (__N == 4) \
WVSPLITK_CFG(_THRDS, _WVPRGRP, 4, 1, __N) \
else \
WVSPLITK_CFG(_THRDS, _WVPRGRP, 4, 2, __N) \
// On gfx950, YTILE 3 and 4 are up to 4x slower than YTILE 2 for N >= 5 on
// wide outputs (e.g. 386 vs 95 us at M=50016, K=6144, N=5).
#define WVSPLIT_TILE_CFG(_THRDS, _WVPRGRP, _sYT, __N) \
{ \
bool fit_lds = (Kbp_in * N_in <= max_lds_len); \
if (_sYT <= 1) \
WVSPLITK_CFG(_THRDS, _WVPRGRP, 1, 4, __N) \
else if ((__N == 1) || (!fit_lds) || (_sYT <= 4 * 2) || \
(__N >= 5 && on_gfx950())) \
WVSPLITK_CFG(_THRDS, _WVPRGRP, 2, 2, __N) \
else if (_sYT <= 4 * 3) \
WVSPLITK_CFG(_THRDS, _WVPRGRP, 3, 2, __N) \
else if (__N == 4) \
WVSPLITK_CFG(_THRDS, _WVPRGRP, 4, 1, __N) \
else \
WVSPLITK_CFG(_THRDS, _WVPRGRP, 4, 2, __N) \
}

// WVSPLITK_CFG arguments are: (THRDS, WVPRGRP, YTILE, UNRL, N).
Expand Down Expand Up @@ -1323,6 +1335,9 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
case 5:
WVSPLIT_TILE(sYT, 5)
break;
case 6:
WVSPLIT_TILE(sYT, 6)
break;
default:
throw std::runtime_error(
"Unsupported N value: " + std::to_string(M_in) + "," +
Expand Down
41 changes: 41 additions & 0 deletions tests/model_executor/layers/test_rocm_unquantized_gemm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -316,6 +316,47 @@ def test_rocm_unquantized_gemm_gfx1x_n_gt_5_falls_back(monkeypatch, m):
assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3)


@pytest.mark.parametrize(
("on_gfx950", "n", "m", "k", "expect_wvsplitk"),
[
(True, 5, 8192, 64, True),
# gfx950 also takes N=6 on wide outputs with a mid-sized K ...
(True, 6, 2048, 2048, True),
# ... but not on narrow outputs, short or long K, or N=7.
(True, 6, 1024, 2048, False),
(True, 6, 2048, 1536, False),
(True, 6, 2048, 8200, False),
(True, 7, 2048, 2048, False),
# Other gfx9 parts stop at N=5.
(False, 5, 8192, 64, True),
(False, 6, 2048, 2048, False),
],
)
def test_rocm_unquantized_gemm_wvsplitk_batch_limit(
monkeypatch, on_gfx950, n, m, k, expect_wvsplitk
):
x = torch.randn(n, k, dtype=torch.float16)
weight = torch.randn(m, k, dtype=torch.float16)

monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False)
monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)
monkeypatch.setattr(utils.rocm_aiter_ops, "is_tgemm_enabled", lambda: False)
monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: False)
monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: True)
monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: on_gfx950)
monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: False)
monkeypatch.setattr(utils, "num_compute_units", lambda: 256)

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)

assert wvsplitk_mock.called == expect_wvsplitk
assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3)


def test_rocm_unquantized_gemm_gfx950_wvsplitkrc_path(monkeypatch):
x = torch.randn(1024, 16, dtype=torch.float16).t()
weight = torch.randn(256, 1024, dtype=torch.float16)
Expand Down
6 changes: 5 additions & 1 deletion vllm/model_executor/layers/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -395,7 +395,11 @@ def rocm_unquantized_gemm_impl(
and weight_ok_for_wvsplitk
and (bias is None or bias.is_contiguous())
and (m == 1 or m > 8)
and 0 < n <= 5
and (
0 < n <= 5
# On gfx950 N=6 also beats hipBLASLt, but only on these shapes.
or (n == 6 and on_gfx950() and m > 1024 and 2048 <= k <= 8192)
)
)

use_llmm1 = (
Expand Down
Loading