[ROCm] Refactor tuned gemms (#55001)
Signed-off-by: Andy Friedrich <afriedri@amd.com> Signed-off-by: afriedri <afriedri@amd.com> Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> Co-authored-by: Shanshan Shen <467638484@qq.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
afriedri committed
Sep 21, 2026 at 11:30 UTC
874a6b5a06f19c9d565d14ebe64131d6aaab91cd
2 files changed
+52
-99
tests/kernels/quantization/test_rocm_mxfp4.py
+3
-11
@@ -50,16 +50,7 @@ SKINNY_GEMM_PASS_RATES = {
50
}
51
PRESHUFFLED_SHAPES = [
52
(64, 4096, 8192),
53
- # aiter 0.1.20 Triton preshuffled fp4 GEMM OOBs (GPU memory access fault ->
54
- # SIGABRT, crashing the process) for (M=32, N=8192, K=8192) on gfx950, despite
55
- # advertising the (N, K) as tuned. A hard GPU fault can't be xfail'd, so skip it
56
- # until the aiter kernel is fixed (ROCm/aiter#4867).
57
- pytest.param(
58
- (32, 8192, 8192),
59
- marks=pytest.mark.skip(
60
- reason="aiter 0.1.20 preshuffled fp4 GEMM OOB (M32,N8192,K8192, gfx950)"
61
- ),
62
- ),
53
+ (32, 8192, 16384),
54
]
55
56
@@ -430,7 +421,8 @@ def test_aiter_fp4_gemm_preshuffled_tuned_shapes(shape):
421
M, K, N = shape
422
423
assert M <= 64
433
- assert rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K)
424
+ # The predicate takes K as packed bytes, matching `weight.shape[1]`.
425
+ assert rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K // 2)
426
427
A = torch.randn(M, K, dtype=torch.bfloat16)
428
B = torch.randn(N, K, dtype=torch.bfloat16)
vllm/_aiter_ops.py
+49
-88
@@ -13,7 +13,6 @@ from torch._ops import OpOverload
13
import vllm.envs as envs
14
from vllm.logger import init_logger
15
from vllm.platforms import current_platform
16
-from vllm.utils.import_utils import PlaceholderModule
16
from vllm.utils.torch_utils import direct_register_custom_op
17
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
18
rocm_aiter_sparse_attn_indexer,
@@ -25,11 +24,6 @@ if TYPE_CHECKING:
24
25
logger = init_logger(__name__)
26
28
-try:
29
- import pandas as pd
30
-except ImportError:
31
- pd = PlaceholderModule("pandas")
32
-
27
# fp8_dtype is not cached.
28
# on ROCm the fp8_dtype always calls is_fp8_fnuz
29
# which is a host op, so we cache it once here.
@@ -243,27 +237,42 @@ def is_aiter_found_and_supported_on_rdna4() -> bool:
237
return False
238
239
246
-@functools.cache
247
-def _load_gemm_tuned_configs(
248
- q_dtype_w: torch.dtype, csv_path: str
249
-) -> set[tuple[int, int, int]]:
250
- try:
251
- df = pd.read_csv(csv_path).drop_duplicates()
252
- df = df[df["q_dtype_w"] == str(q_dtype_w)]
253
- return set(zip(df["N"].astype(int), df["K"].astype(int), df["M"].astype(int)))
254
- except Exception:
255
- return set()
240
+def _triton_gemm_config_is_tuned(config_name: str, N: int, K: int) -> bool:
241
+ from aiter.ops.triton.utils.gemm_config_utils import get_gemm_config
242
+
243
+ # any M shape under 1024 should be tuned
244
+ M_TUNE_PROBE = 128
245
246
+ return get_gemm_config(config_name, M_TUNE_PROBE, N, K)[1]
247
258
-def _check_kernel_tuned(N: int, K: int, q_dtype_w: torch.dtype, csv_path: str) -> bool:
259
- configs = _load_gemm_tuned_configs(q_dtype_w, csv_path)
248
+
249
+def _ck_gemm_shape_is_tuned(
250
+ N: int, K: int, q_dtype_w: torch.dtype, csv_attr: str
251
+) -> bool:
252
l_m = (
253
[1, 2, 4]
254
+ list(range(8, 513, 8))
255
+ [1024, 1536]
256
+ [2**i for i in range(11, 19)]
257
)
266
- return any((N, K, M) in configs for M in l_m)
258
+ try:
259
+ from aiter.ops.gemm_op_a8w8 import (
260
+ AITER_CONFIGS,
261
+ get_GEMM_config_with_quant_type,
262
+ )
263
+
264
+ csv_path = getattr(AITER_CONFIGS, csv_attr)
265
+ return any(
266
+ get_GEMM_config_with_quant_type(M, N, K, q_dtype_w, csv_path) is not None
267
+ for M in l_m
268
+ )
269
+ except (AttributeError, ImportError, OSError):
270
+ logger.warning_once(
271
+ "Could not read aiter CK GEMM configs from AITER_CONFIGS.%s; "
272
+ "treating all shapes as untuned.",
273
+ csv_attr,
274
+ )
275
+ return False
276
277
278
def if_aiter_supported(func: Callable) -> Callable:
@@ -3458,89 +3467,41 @@ class rocm_aiter_ops:
3467
)
3468
3469
@staticmethod
3470
+ @functools.cache
3471
def is_triton_gemm_w8a8_tuned(n: int, k: int) -> bool:
3472
if not current_platform.is_rocm():
3473
return False
3464
- from vllm.platforms.rocm import on_gfx950, on_rdna4
3465
-
3466
- gfx950_tuned = {
3467
- (1024, 8192),
3468
- (2112, 7168),
3469
- (3072, 1536),
3470
- (32768, 8192),
3471
- (4096, 7168),
3472
- (4608, 7168),
3473
- (512, 7168),
3474
- (7168, 2048),
3475
- (7168, 256),
3476
- (8192, 1024),
3477
- (8192, 32768),
3478
- }
3479
- rdna4_tuned = gfx950_tuned | {
3480
- (2048, 2048),
3481
- (2624, 6144),
3482
- (3072, 6144),
3483
- (3584, 512),
3484
- (4096, 512),
3485
- (6144, 1536),
3486
- (6144, 2048),
3487
- (7168, 2304),
3488
- (7168, 16384),
3489
- (7168, 18432),
3490
- (8192, 8192),
3491
- (16384, 1536),
3492
- (24576, 1536),
3493
- (32768, 512),
3494
- (36864, 7168),
3495
- }
3496
- if on_rdna4():
3497
- return (n, k) in rdna4_tuned
3498
- if on_gfx950():
3499
- return (n, k) in gfx950_tuned
3500
- return False
3474
+ try:
3475
+ return _triton_gemm_config_is_tuned("GEMM-A8W8_BLOCKSCALE", n, k)
3476
+ except (AssertionError, ImportError):
3477
+ return False
3478
3479
@staticmethod
3503
- def is_triton_gemm_afp4wfp4_presh_ws_tuned(n: int, k: int) -> bool:
3504
- return (n, k) in [
3505
- (8192, 4096),
3506
- (1280, 8192),
3507
- (16384, 53248),
3508
- (106496, 16384),
3509
- (57344, 8192),
3510
- (8192, 2048),
3511
- (2560, 8192),
3512
- (10240, 8192),
3513
- (16384, 16384),
3514
- (8192, 28672),
3515
- (28672, 8192),
3516
- (18432, 16384),
3517
- (8192, 1024),
3518
- (7168, 8192),
3519
- (5120, 8192),
3520
- (8192, 8192),
3521
- (8192, 7168),
3522
- (14336, 8192),
3523
- (8192, 14336),
3524
- (8192, 3584),
3525
- ]
3480
+ @functools.cache
3481
+ def is_triton_gemm_afp4wfp4_presh_ws_tuned(n: int, k_bytes: int) -> bool:
3482
+ if not current_platform.is_rocm():
3483
+ return False
3484
+ # Input = weight.shape[1] (bytes); *4 for fp4 and aiter config 2*k
3485
+ try:
3486
+ return _triton_gemm_config_is_tuned(
3487
+ "GEMM-AFP4WFP4_PRESHUFFLED", n, 4 * k_bytes
3488
+ )
3489
+ except (AssertionError, ImportError):
3490
+ return False
3491
3492
@staticmethod
3493
+ @functools.cache
3494
def is_shuffled_per_token_w8a8_gemm_tuned(
3495
N: int, K: int, q_dtype_w: torch.dtype
3496
) -> bool:
3531
- import aiter.ops.gemm_op_a8w8 as aiter_gemm_a8w8_ops
3532
-
3533
- csv_path = (
3534
- aiter_gemm_a8w8_ops.AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_BPRESHUFFLE_FILE
3497
+ return _ck_gemm_shape_is_tuned(
3498
+ N, K, q_dtype_w, "AITER_CONFIG_GEMM_A8W8_BPRESHUFFLE_FILE"
3499
)
3536
- return _check_kernel_tuned(N, K, q_dtype_w, csv_path)
3500
3501
@staticmethod
3502
+ @functools.cache
3503
def is_per_token_w8a8_gemm_tuned(N: int, K: int, q_dtype_w: torch.dtype) -> bool:
3540
- import aiter.ops.gemm_op_a8w8 as aiter_gemm_a8w8_ops
3541
-
3542
- csv_path = aiter_gemm_a8w8_ops.AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_FILE
3543
- return _check_kernel_tuned(N, K, q_dtype_w, csv_path)
3504
+ return _ck_gemm_shape_is_tuned(N, K, q_dtype_w, "AITER_CONFIG_GEMM_A8W8_FILE")
3505
3506
@staticmethod
3507
def shuffle_weight(