[ROCm][Perf][Bugfix] Probe the preshuffled FP4 GEMM config by logical K (#59406)

Signed-off-by: Mikko Tukiainen <Mikko.Tukiainen@amd.com>

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.

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