From 0c2584929202353a1b2e33934611125c604b7eed Mon Sep 17 00:00:00 2001 From: Mikko Tukiainen Date: Fri, 16 Oct 2026 00:11:33 +0300 Subject: [PATCH] [ROCm][Perf][Bugfix] Probe the preshuffled FP4 GEMM config by logical K (#59406) Signed-off-by: Mikko Tukiainen --- tests/kernels/quantization/test_rocm_mxfp4.py | 20 +++++++++++++++++-- vllm/_aiter_ops.py | 6 ++++-- 2 files changed, 22 insertions(+), 4 deletions(-) diff --git a/tests/kernels/quantization/test_rocm_mxfp4.py b/tests/kernels/quantization/test_rocm_mxfp4.py index fa39c641ba..f83445b39c 100644 --- a/tests/kernels/quantization/test_rocm_mxfp4.py +++ b/tests/kernels/quantization/test_rocm_mxfp4.py @@ -50,7 +50,7 @@ SKINNY_GEMM_PASS_RATES = { } PRESHUFFLED_SHAPES = [ (64, 4096, 8192), - (32, 8192, 16384), + (32, 16384, 16384), ] @@ -421,8 +421,12 @@ def test_aiter_fp4_gemm_preshuffled_tuned_shapes(shape): M, K, N = shape assert M <= 64 - # The predicate takes K as packed bytes, matching `weight.shape[1]`. + # The predicate takes K as packed bytes, matching `weight.shape[1]`, and + # must agree with the config the kernel itself resolves. assert rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K // 2) + from aiter.ops.triton.utils.gemm_config_utils import get_gemm_config + + assert get_gemm_config("GEMM-AFP4WFP4_PRESHUFFLED", M, N, K)[1] A = torch.randn(M, K, dtype=torch.bfloat16) B = torch.randn(N, K, dtype=torch.bfloat16) @@ -464,6 +468,18 @@ def test_aiter_fp4_gemm_preshuffled_tuned_shapes(shape): _assert_deterministic(run_preshuffled, n_runs=3) +@pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only") +def test_aiter_fp4_preshuffled_tuned_gate_uses_logical_k(): + """AITER ships N=16384-K=16384 but not N=16384-K=8192, so a K=8192 layer + has no tuned preshuffled config even though twice its K does.""" + _assert_aiter_supported() + from vllm._aiter_ops import rocm_aiter_ops + + rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned.cache_clear() + assert rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(16384, 16384 // 2) + assert not rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(16384, 8192 // 2) + + @pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only") def test_aiter_fp4_gemm_a4w4_determinism(): """AITER FP4 A4W4 GEMM is bitwise deterministic across repeated runs. diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index ba25cc8b5a..41a76bfc60 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -4079,10 +4079,12 @@ class rocm_aiter_ops: def is_triton_gemm_afp4wfp4_presh_ws_tuned(n: int, k_bytes: int) -> bool: if not current_platform.is_rocm(): return False - # Input = weight.shape[1] (bytes); *4 for fp4 and aiter config 2*k + # k_bytes = weight.shape[1] (two fp4 per byte). AITER names these + # config files by the logical K, which is what gemm_afp4wfp4_preshuffle + # itself looks up, so probe the same file. try: return _triton_gemm_config_is_tuned( - "GEMM-AFP4WFP4_PRESHUFFLED", n, 4 * k_bytes + "GEMM-AFP4WFP4_PRESHUFFLED", n, 2 * k_bytes ) except (AssertionError, ImportError): return False