A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 03:30 UTC
[ROCm][CI] Fix stale wvSplitK GEMM fallback test for N=5 (#44368)
Signed-off-by: JartX <sagformas@epdcenter.es>
JartX committed
Jun 3, 2026 at 05:00 UTC
4454a18695287a8b0ef5701e5390636d692e3aa5
1 file changed
+4
-2
tests/model_executor/layers/test_rocm_unquantized_gemm.py
+4
-2
@@ -41,8 +41,10 @@ def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch):
41
assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3)
42
43
44
-def test_rocm_unquantized_gemm_gfx1x_n_gt_4_falls_back(monkeypatch):
45
- x = torch.randn(5, 64, dtype=torch.float16)
44
+def test_rocm_unquantized_gemm_gfx1x_n_gt_5_falls_back(monkeypatch):
45
+ # wvSplitK skinny GEMM handles n in [1, 5] (see PR #40687); n > 5 must
46
+ # fall back to torch.nn.functional.linear.
47
+ x = torch.randn(6, 64, dtype=torch.float16)
48
weight = torch.randn(128, 64, dtype=torch.float16)
49
50
monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False)