diff --git a/aiter/ops/gemm_op_a8w8.py b/aiter/ops/gemm_op_a8w8.py index 92832d8256..cbf913dfe8 100644 --- a/aiter/ops/gemm_op_a8w8.py +++ b/aiter/ops/gemm_op_a8w8.py @@ -988,6 +988,7 @@ def gemm_a8w8_blockscale_bpreshuffle( w_scale, dtype=dtype, config=_fallback_cfg, + is_x_scale_tranposed=x_scale.stride(0) != 1, ) config = get_CKGEMM_config( m, n, k, AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_BLOCKSCALE_BPRESHUFFLE_FILE @@ -1015,6 +1016,7 @@ def gemm_a8w8_blockscale_bpreshuffle( w_scale, dtype=dtype, backend=backend, + is_x_scale_tranposed=x_scale.stride(0) != 1, ) if config is not None: libtype = config["libtype"] diff --git a/aiter/ops/triton/gemm/basic/gemm_a8w8_blockscale.py b/aiter/ops/triton/gemm/basic/gemm_a8w8_blockscale.py index fead8a5ba5..7870b4791e 100644 --- a/aiter/ops/triton/gemm/basic/gemm_a8w8_blockscale.py +++ b/aiter/ops/triton/gemm/basic/gemm_a8w8_blockscale.py @@ -6,6 +6,7 @@ import torch import triton +from packaging.version import Version from aiter.ops.triton._triton_kernels.common.splitk_reduce import ( _gemm_splitk_reduce_kernel, @@ -25,6 +26,7 @@ _LOGGER = AiterTritonLogger() _FORCE_GFX1250_EX = os.environ.get("AITER_FORCE_GFX1250_EX", "0") == "1" +_TRITON_VERSION = Version(triton.__version__) _GLUON_SUPPORTED_ARCHS = ("gfx1250",) @@ -305,6 +307,16 @@ def gemm_a8w8_blockscale_preshuffle( if config is None: config, _ = _get_config(M, N, K, True, backend=backend) + # Triton 3.6 fails TritonAMDGPUConvertToBufferOps for gfx950 preshuffle + # configs with three pipeline stages. Keep the tuned tile and split-K. + if ( + backend == "triton" + and get_arch() == "gfx950" + and _TRITON_VERSION < Version("3.7.0") + and config.get("num_stages", 1) > 2 + ): + config["num_stages"] = 2 + kernel_type_from_config = config.pop("kernel_type", None) if kernel_type_from_config is not None: kernel_type = kernel_type_from_config diff --git a/op_tests/test_gemm_a8w8_blockscale.py b/op_tests/test_gemm_a8w8_blockscale.py index c99ce8f5ec..9281279bba 100644 --- a/op_tests/test_gemm_a8w8_blockscale.py +++ b/op_tests/test_gemm_a8w8_blockscale.py @@ -116,6 +116,17 @@ def test_gemm(dtype, m, n, k, ck_preshuffle=True, use_flydsl=False): b, avg_b = run_func(x, gemm_weight, gemm_x_scale, w_scale, dtype) err_ck = checkAllclose(a, b, msg="ck", catastrophic_check=True) + if ck_preshuffle: + x_scale_strided = x_scale.transpose(0, 1).contiguous().transpose(0, 1) + b_strided = aiter.gemm_a8w8_blockscale_bpreshuffle( + x, gemm_weight, x_scale_strided, w_scale, dtype + ) + checkAllclose( + a, + b_strided, + msg="ck strided x_scale", + catastrophic_check=True, + ) ret["ck us"] = avg_b ret["ck TFLOPS"] = m * n * k * 2 / avg_b / 1e6 ret["ck TB/s"] = (x.nbytes + weight.nbytes) / avg_b / 1e6