Skip to content

Commit e2d019f

Browse files
authored
[Triton/Gluon] [JIT] [Feature] Add tuned native group32 A8W8 Triton GEMM for gfx950 (#5750)
* feat(gemm): add native group32 A8W8 Triton backend Integrate unshuffled E4M3/E8M0 inputs into gemm_a8w8_blockscale. Keep microscaling and packed-K kernels in Triton, with gfx950 tuning tables loaded by get_gemm_config. Extend shared loading with JSON M bounds and reduce config-copy overhead while retaining defensive nested copies. Reuse the shared non-atomic split-K reduction and preserve legacy routes. Validated with 98 AITER tests, 73,728 matching config selections, and 285 bitwise-equal BF16 migration cases; repeated 40 GPU cases six times. Signed-off-by: Lingpeng Jin <103567126+valarLip@users.noreply.github.com> * Tune native group32 GEMM configurations for V4.1 projections Sweep 18 gfx950 geometries across TP4, TP2, and attention-DP local projections. Store tile, split-K, launch, and M-bound selections in the existing JSON family; retain baseline intervals without stable gains. Keep the existing large-M fallbacks and merge equal adjacent configs. Exercise N-first masking independently of the selected traversal. Validation: 49,304 candidates screened, 603 regression cases, 98 AITER tests and 85 ATOM tests passed. Paired regression geometric-mean speedup is 1.179x; changed configurations have no sampled slowdown above 1%. Signed-off-by: Lingpeng Jin <103567126+valarLip@users.noreply.github.com> * Respect tuned backend selection before group32 GEMM fallback Signed-off-by: Lingpeng Jin <103567126+valarLip@users.noreply.github.com> * Configure V4.1 group32 GEMMs to use the Triton backend Signed-off-by: Lingpeng Jin <103567126+valarLip@users.noreply.github.com> * Preserve native split-K controls through the public A8W8 GEMM interface Signed-off-by: Lingpeng Jin <103567126+valarLip@users.noreply.github.com> * fix(gemm): complete group32 dispatch and validation coverage * perf(gemm): fuse group32 split-K reduction within one XCD Small-M group32 GEMMs on narrow or deep projections are starved for CTAs: N=512 gives the packed kernel 32 CTAs on 256 CUs, and splitting K costs a second reduce launch. Both kernels now take FUSED_SPLITS, which sums the split partials in the same launch. A 1D grid places every split of a tile on one XCD (CTAs go to XCDs by pid % 8, a multiple of every gfx950 XCD count). That XCD's L2 is then the splits' coherence point, so the last CTA to arrive can read the other partials after a workgroup-scope acq_rel counter add: no agent-scope release, and no buffer_wbl2 of the whole L2. With agent scope the same fused kernels are 2.3x slower, which is why a plain fused split-K lost to the two-launch path. Partials are summed in split order whatever the arrival order, so results are deterministic, and the last CTA re-zeroes its counter, so graph replay is safe. Counters follow the a16w16 split-K semaphores: one zeroed buffer per (device, stream) from persistent_alloc, so concurrent streams never share one. Where a stream's first use falls inside a capture that cannot allocate outside the graph pool (torch < 2.10), the launch takes the unfused path instead of failing the capture. The fused path is selected only by the tuning table (FUSED_SPLITK on a tile config, NUM_KSPLIT on a packed config). An explicit split_k keeps the existing reduce-kernel path. With the original configs, all 567 sampled outputs across the 18 V4.1 geometries, tails, FP32/BF16 and explicit split_k are bitwise identical to before. Tests cover the reference match for both variants and output dtypes, 1x32 weight scales, 50-call determinism with counters re-zeroed, graph replay on a warmed stream, first use inside a capture, and interleaved launches on two streams. * perf(gemm): tune fused split-K intervals and time every V4.1 group32 row 57 M intervals of 11 V4.1 geometries now use the fused split-K path, 42 at M <= 64 and 15 above. Candidates were the best fused tile and packed configs from full sweeps; an interval changed only if a candidate was at least 3% faster at both of its bounds. Every changed interval was then re-measured against its previous config in the same process (run_perftest, rotated inputs, median of 5): at each CSV M inside it, or at its lower bound, midpoint and upper bound where no CSV M falls inside. Three intervals that did not hold up were left unchanged. No measured point is slower than the previous config. Speedup over the previous config, geometric mean of the rows re-measured: o_b.dp M <= 64 1.26x draft_kvf M 129..320 1.22-1.26x draft_kv M <= 64 1.22x sh_up.tp4 M 321..448 1.26x sh_up.tp4 M <= 64 1.19x sh_up.tp2 M 129..192 1.20x draft_kvf M <= 64 1.13x o_b.dp M 129..192 1.20x o_b.tp2 M <= 64 1.12x draft_kv M 513..576 1.14x The V4.1 model CSV previously left 188 us fields and every tflops, bw and errRatio field empty. All 450 rows are now timed in one run with the final configs, through gemm_a8w8_blockscale, with run_perftest and rotated inputs, so the us column comes from a single method. tflops and bw use the a8w8 blockscale tuner formula (bw excludes scales). errRatio is checkAllclose at its default tolerance against an FP64 reference, as in the tuner: at most 8.7e-4, and 0 rows above 0.05. * perf(gemm): constexpr K loop, weight cache modifier and 68-bucket retune for group32 GEMM Retune all 18 V4.1 geometries on the repo's 68 DSv4 M buckets (fused split-K, packed K_PACK=1, .cg weight loads); drop unused GROUP_SIZE_M; re-time the 1224-row CSV. --------- Signed-off-by: Lingpeng Jin <103567126+valarLip@users.noreply.github.com>
1 parent 37b1425 commit e2d019f

31 files changed

Lines changed: 23468 additions & 54 deletions
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
gfx,cu_num,M,N,K,libtype,kernelId,splitK,us,kernelName,tflops,bw,errRatio
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
M,N,K

‎aiter/configs/model_configs/dsv41_a8w8_blockscale_group32_tuned_gemm.csv‎

Lines changed: 1225 additions & 0 deletions
Large diffs are not rendered by default.

‎aiter/jit/core.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,6 +135,13 @@ def mp_lock(
135135
f"{AITER_ROOT_DIR}/aiter/configs/a8w8_blockscale_tuned_gemm.csv",
136136
)
137137

138+
# Native E8M0 group32 scales have a different operand contract from the
139+
# FP32 128x128 blockscale family, so shape-identical rows must stay separate.
140+
AITER_CONFIG_GEMM_A8W8_BLOCKSCALE_GROUP32 = os.getenv(
141+
"AITER_CONFIG_GEMM_A8W8_BLOCKSCALE_GROUP32",
142+
f"{AITER_ROOT_DIR}/aiter/configs/a8w8_blockscale_group32_tuned_gemm.csv",
143+
)
144+
138145
AITER_CONFIG_FMOE = os.getenv(
139146
"AITER_CONFIG_FMOE",
140147
f"{AITER_ROOT_DIR}/aiter/configs/tuned_fmoe.csv",
@@ -286,6 +293,14 @@ def AITER_CONFIG_GEMM_A8W8_BLOCKSCALE_FILE(self):
286293
"a8w8_blockscale_tuned_gemm",
287294
)
288295

296+
@property
297+
def AITER_CONFIG_GEMM_A8W8_BLOCKSCALE_GROUP32_FILE(self):
298+
return self.get_config_file(
299+
"AITER_CONFIG_GEMM_A8W8_BLOCKSCALE_GROUP32",
300+
AITER_CONFIG_GEMM_A8W8_BLOCKSCALE_GROUP32,
301+
"a8w8_blockscale_group32_tuned_gemm",
302+
)
303+
289304
@property
290305
def AITER_CONFIG_FMOE_FILE(self):
291306
return self.get_config_file(

‎aiter/ops/gemm_op_a8w8.py‎

Lines changed: 114 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -812,8 +812,28 @@ def _blockscale_triton(
812812
x_scale: Tensor,
813813
w_scale: Tensor,
814814
dtype: torch.dtype,
815+
*,
816+
group32: bool = False,
817+
split_k: int | None = None,
815818
) -> Tensor:
816-
"""Run the triton kernel on CK's inputs: row-major x_scale, (N, K) weight, JIT per arch."""
819+
"""Run Triton on unshuffled weights with the selected scale format."""
820+
if group32:
821+
from aiter.ops.triton.gemm.basic.gemm_a8w8_blockscale_group32 import (
822+
gemm_a8w8_blockscale_group32,
823+
)
824+
825+
assert WQ.ndim == w_scale.ndim == 2, "Expected matrix weights and scales"
826+
group_n = 32 if w_scale.shape[0] == -(-WQ.shape[0] // 32) else 1
827+
return gemm_a8w8_blockscale_group32(
828+
XQ,
829+
WQ,
830+
x_scale,
831+
w_scale,
832+
dtype=dtype,
833+
weight_group_rows=group_n,
834+
split_k=split_k,
835+
)
836+
817837
from aiter.ops.triton.gemm.basic.gemm_a8w8_blockscale import (
818838
gemm_a8w8_blockscale as _gemm_a8w8_blockscale_triton,
819839
)
@@ -830,77 +850,123 @@ def gemm_a8w8_blockscale_fake(
830850
w_scale: Tensor,
831851
dtype: torch.dtype = dtypes.bf16,
832852
isBpreshuffled=False,
853+
split_k: int | None = None,
833854
) -> torch.Tensor:
834855
m = XQ.shape[0]
835856
n = WQ.shape[0]
836857
Y = torch.empty(m, n, dtype=dtype, device=XQ.device)
837858
return Y
838859

839860

840-
@torch_compile_guard(gen_fake=gemm_a8w8_blockscale_fake)
861+
@torch_compile_guard(mutates_args=[], gen_fake=gemm_a8w8_blockscale_fake)
841862
def gemm_a8w8_blockscale(
842863
XQ: Tensor,
843864
WQ: Tensor,
844865
x_scale: Tensor,
845866
w_scale: Tensor,
846867
dtype: torch.dtype = dtypes.bf16,
847868
isBpreshuffled: bool = False,
869+
split_k: int | None = None,
848870
) -> torch.Tensor:
849-
assert dtype in [
850-
dtypes.bf16,
851-
dtypes.fp16,
852-
], f"Output {dtype=} is currently not supported in gemm_a8w8"
871+
"""Blockscaled A8W8 GEMM with configuration-first backend dispatch.
872+
873+
Native E8M0 group32 scales (typed tensors or uint8 views) use
874+
AITER_CONFIG_GEMM_A8W8_BLOCKSCALE_GROUP32;
875+
FP32 128x128 scales use AITER_CONFIG_GEMM_A8W8_BLOCKSCALE. Both tables are
876+
queried through get_CKGEMM_config before choosing a fallback. Group32
877+
currently supports libtype="triton", also its default on a config miss.
878+
Triton tile and split-K parameters come from its own tuning tables.
879+
For native group32 operands, split_k optionally overrides the positive
880+
partition count without changing configured backend selection.
881+
"""
882+
is_group32 = (
883+
x_scale.dtype in (dtypes.fp8_e8m0, torch.uint8)
884+
and w_scale.dtype in (dtypes.fp8_e8m0, torch.uint8)
885+
and XQ.ndim == WQ.ndim == 2
886+
and x_scale.shape == (XQ.shape[0], XQ.shape[1] // 32)
887+
and w_scale.shape
888+
in (
889+
(WQ.shape[0], WQ.shape[1] // 32),
890+
(-(-WQ.shape[0] // 32), WQ.shape[1] // 32),
891+
)
892+
)
893+
# A malformed byte-scale layout must not fall through to FP32 CK dispatch.
894+
assert (
895+
is_group32 or x_scale.dtype == w_scale.dtype == dtypes.fp32
896+
), "Expected E8M0 group32 scale shapes (typed or uint8), or FP32 128x128 scales"
897+
assert (
898+
split_k is None or is_group32
899+
), "split_k override requires native group32 operands"
900+
assert dtype in (dtypes.bf16, dtypes.fp16) or (
901+
is_group32 and dtype == dtypes.fp32
902+
), f"Output {dtype=} is currently not supported in gemm_a8w8"
853903
m = XQ.shape[0]
854904
n = WQ.shape[0]
855905
k = XQ.shape[1]
856-
Y = torch.empty(m, n, dtype=dtype, device=XQ.device)
857906
if isBpreshuffled:
907+
assert not is_group32, "Group32 FP8 GEMM requires native weights"
858908
if get_gfx() in ["gfx950"] and m >= 16 and k >= 512 and dtype == dtypes.bf16:
909+
Y = torch.empty(m, n, dtype=dtype, device=XQ.device)
859910
return gfx950_a8w8_blockscale_ASM(XQ, WQ, x_scale, w_scale, Y)
860911
else:
861912
assert 0, "asm kernel only support B preshuffle and m >= 16"
862-
else:
863-
if not _hip_blockscale_supported():
864-
return _blockscale_triton(XQ, WQ, x_scale, w_scale, dtype)
865-
config = get_CKGEMM_config(
866-
m, n, k, AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_BLOCKSCALE_FILE
913+
914+
tuned_file = (
915+
AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_BLOCKSCALE_GROUP32_FILE
916+
if is_group32
917+
else AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_BLOCKSCALE_FILE
918+
)
919+
config = get_CKGEMM_config(m, n, k, tuned_file)
920+
if config is not None:
921+
libtype = config["libtype"]
922+
if libtype == "triton":
923+
return _blockscale_triton(
924+
XQ, WQ, x_scale, w_scale, dtype, group32=is_group32, split_k=split_k
925+
)
926+
# CK/CKTile currently consume FP32 128x128 scales. A misconfigured
927+
# group32 row must fail instead of reinterpreting its E8M0 bytes.
928+
assert not is_group32, f"Unsupported libtype {libtype} for group32 GEMM"
929+
splitK = int(config.get("splitK", 0))
930+
kernelName = str(config.get("kernelName", ""))
931+
Y = torch.empty(m, n, dtype=dtype, device=XQ.device)
932+
if libtype == "ck":
933+
return gemm_a8w8_blockscale_ck(
934+
XQ,
935+
WQ,
936+
x_scale,
937+
w_scale,
938+
Y,
939+
splitK=splitK,
940+
kernelName=kernelName,
941+
)
942+
elif libtype == "cktile":
943+
return gemm_a8w8_blockscale_cktile(
944+
XQ,
945+
WQ,
946+
x_scale,
947+
w_scale,
948+
Y,
949+
splitK=splitK,
950+
kernelName=kernelName,
951+
)
952+
else:
953+
assert 0, f"Unsupported libtype {libtype} for gemm_a8w8_blockscale"
954+
955+
if is_group32 or not _hip_blockscale_supported():
956+
return _blockscale_triton(
957+
XQ, WQ, x_scale, w_scale, dtype, group32=is_group32, split_k=split_k
867958
)
868-
if config is not None:
869-
libtype = config["libtype"]
870-
splitK = int(config.get("splitK", 0))
871-
kernelName = str(config.get("kernelName", ""))
872-
if libtype == "ck":
873-
return gemm_a8w8_blockscale_ck(
874-
XQ,
875-
WQ,
876-
x_scale,
877-
w_scale,
878-
Y,
879-
splitK=splitK,
880-
kernelName=kernelName,
881-
)
882-
elif libtype == "cktile":
883-
return gemm_a8w8_blockscale_cktile(
884-
XQ,
885-
WQ,
886-
x_scale,
887-
w_scale,
888-
Y,
889-
splitK=splitK,
890-
kernelName=kernelName,
891-
)
892-
else:
893-
assert 0, f"Unsupported libtype {libtype} for gemm_a8w8_blockscale"
894-
min_m = _BLOCKSCALE_TRITON_FALLBACK_MIN_M.get(get_gfx())
895-
if min_m is not None and m >= min_m:
896-
return _blockscale_triton(XQ, WQ, x_scale, w_scale, dtype)
897-
try:
898-
return gemm_a8w8_blockscale_ck(XQ, WQ, x_scale, w_scale, Y)
899-
except RuntimeError as e:
900-
raise RuntimeError(
901-
f"gemm_a8w8_blockscale failed for shape M={m}, N={n}, K={k}, "
902-
f"{dtype=}, config={config}: {e}"
903-
) from e
959+
min_m = _BLOCKSCALE_TRITON_FALLBACK_MIN_M.get(get_gfx())
960+
if min_m is not None and m >= min_m:
961+
return _blockscale_triton(XQ, WQ, x_scale, w_scale, dtype)
962+
Y = torch.empty(m, n, dtype=dtype, device=XQ.device)
963+
try:
964+
return gemm_a8w8_blockscale_ck(XQ, WQ, x_scale, w_scale, Y)
965+
except RuntimeError as e:
966+
raise RuntimeError(
967+
f"gemm_a8w8_blockscale failed for shape M={m}, N={n}, K={k}, "
968+
f"{dtype=}, config={config}: {e}"
969+
) from e
904970

905971

906972
def flatmm_a8w8_blockscale_ASM(

0 commit comments

Comments
 (0)