From 1a0b80b1e50fecefab9a8622b124ae9dc8ff2886 Mon Sep 17 00:00:00 2001 From: Ahmed Date: Sun, 27 Sep 2026 03:57:51 +0000 Subject: [PATCH] replacing the fp32 mfma with two bf16 mfma. fn is split inside the kernel into hi = bf16(fn) and lo = bf16(fn -hi). a simple bf16 downcast for fn inside the kernel drops accuracy more and performs worse for large M. also change the concatenation of the four streams from column major to stream major (k = stream * TILE_K + column). retuned the gfx950 configs. --- .../mhc_fused_post_pre_delayed_rmsnorm.py | 33 ++++--- .../DEFAULT.json | 94 +++++++++++++++---- .../DEFAULT.json | 94 +++++++++++++++---- 3 files changed, 169 insertions(+), 52 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/fusions/mhc_fused_post_pre_delayed_rmsnorm.py b/aiter/ops/triton/_triton_kernels/fusions/mhc_fused_post_pre_delayed_rmsnorm.py index d76594b9d4c..73ca419ee20 100644 --- a/aiter/ops/triton/_triton_kernels/fusions/mhc_fused_post_pre_delayed_rmsnorm.py +++ b/aiter/ops/triton/_triton_kernels/fusions/mhc_fused_post_pre_delayed_rmsnorm.py @@ -18,12 +18,13 @@ comb' are consumed by the next one. Two launches. The main kernel is a split-K GEMM over the hidden dimension: a program -owns 16 tokens and a slice of columns of all four streams, post-mixes them in fp32, +owns BLOCK_M tokens and a slice of columns of all four streams, post-mixes them in fp32, stores the bf16 residual, writes the collapse un-normalised while the tile is in -registers, and accumulates its slice of the projection (fp32 MFMA) and of the two -square sums into one partial row per (token, slice). The reduce kernel finishes one -token per program: it sums the partial rows, computes the gates (Sinkhorn with the -hardware reciprocal) and applies the RMSNorm to the staged collapse in place. +registers, and accumulates its slice of the projection (bf16 MFMA against fn split into +bf16 hi + lo) and of the two square sums into one partial row per (token, slice). The +reduce kernel finishes one token per program: it sums the partial rows, computes the +gates (Sinkhorn with the hardware reciprocal) and applies the RMSNorm to the staged +collapse in place. """ import triton @@ -70,8 +71,8 @@ def _mhc_fused_post_pre_delayed_rmsnorm_main_kernel( H / NUM_KSPLIT hidden columns of all four streams, walked in k-steps of TILE_K columns. Each k-step loads the four stream tiles and the x tile, post-mixes them in fp32, - stores the bf16 residual and the collapse, and feeds one fp32 MFMA dot with the - streams concatenated along K (k = 4 * column + stream). + stores the bf16 residual and the collapse, and feeds two bf16 MFMA dots (fn hi, fn lo) + with the streams concatenated along K (k = stream * TILE_K + column). """ tl.static_assert(BLOCK_M % 16 == 0, "BLOCK_M is the M dimension of the MFMA dot") tl.static_assert(H % (NUM_KSPLIT * TILE_K) == 0, "the k-loop has no column mask") @@ -185,9 +186,11 @@ def _mhc_fused_post_pre_delayed_rmsnorm_main_kernel( tl.store(layer_input_ptr + rm[:, None] * H + rc[None, :], x1, mask=m2) x1f = x1.to(tl.float32) sq_x += x1f * x1f - # projection: A = the four stream tiles interleaved along K (k = 4*c + stream), - # B = the matching fn columns; Triton loads fn straight into the MFMA B layout - a = tl.reshape(tl.join(tl.join(g0, g2), tl.join(g1, g3)), (BLOCK_M, 4 * TILE_K)) + # concatanate the fours streams along K, k = stream * TILE_K + column, same as fn_t + a = tl.reshape( + tl.permute(tl.join(tl.join(r0, r2), tl.join(r1, r3)), (0, 2, 3, 1)), + (BLOCK_M, 4 * TILE_K), + ) fn_t = tl.load( fn_ptr + rn[:, None, None] * (4 * H) @@ -196,8 +199,14 @@ def _mhc_fused_post_pre_delayed_rmsnorm_main_kernel( mask=n_mask[:, None, None], other=0.0, ) # (32, 4, TILE_K) fp32 - b = tl.reshape(tl.permute(fn_t, (0, 2, 1)), (32, 4 * TILE_K)) - acc = tl.dot(a, tl.trans(b), acc=acc, input_precision="ieee") + b = tl.reshape(fn_t, (32, 4 * TILE_K)) + + # split fn into two bf16s: hi = bf16(fn), lo = bf16(fn - hi) (what hi lost) + # both products accumulate into acc + b_hi = b.to(tl.bfloat16) + b_lo = (b - b_hi.to(tl.float32)).to(tl.bfloat16) + acc = tl.dot(a, tl.trans(b_hi), acc=acc) + acc = tl.dot(a, tl.trans(b_lo), acc=acc) vals = tl.where( rn[None, :] == 24, diff --git a/aiter/ops/triton/configs/gfx942/triton/mhc/mhc_fused_post_pre_delayed_rmsnorm/DEFAULT.json b/aiter/ops/triton/configs/gfx942/triton/mhc/mhc_fused_post_pre_delayed_rmsnorm/DEFAULT.json index 46ec1bba5b3..56e49a02455 100644 --- a/aiter/ops/triton/configs/gfx942/triton/mhc/mhc_fused_post_pre_delayed_rmsnorm/DEFAULT.json +++ b/aiter/ops/triton/configs/gfx942/triton/mhc/mhc_fused_post_pre_delayed_rmsnorm/DEFAULT.json @@ -1,13 +1,31 @@ { - "M_LEQ_144": { + "M_LEQ_16": { "NUM_KSPLIT": 160, "BLOCK_M": 16, "TILE_K": 32, "num_warps": 2, + "num_stages": 1, + "BLOCK_C": 1024, + "reduce_num_warps": 1 + }, + "M_LEQ_64": { + "NUM_KSPLIT": 160, + "BLOCK_M": 16, + "TILE_K": 32, + "num_warps": 1, "num_stages": 3, "BLOCK_C": 1024, "reduce_num_warps": 1 }, + "M_LEQ_96": { + "NUM_KSPLIT": 40, + "BLOCK_M": 16, + "TILE_K": 32, + "num_warps": 2, + "num_stages": 2, + "BLOCK_C": 1024, + "reduce_num_warps": 1 + }, "M_LEQ_192": { "NUM_KSPLIT": 40, "BLOCK_M": 16, @@ -17,58 +35,94 @@ "BLOCK_C": 1024, "reduce_num_warps": 1 }, - "M_LEQ_1024": { - "NUM_KSPLIT": 160, + "M_LEQ_352": { + "NUM_KSPLIT": 40, "BLOCK_M": 16, "TILE_K": 32, "num_warps": 2, - "num_stages": 3, + "num_stages": 2, "BLOCK_C": 1024, "reduce_num_warps": 1 }, - "M_LEQ_2048": { - "NUM_KSPLIT": 80, + "M_LEQ_512": { + "NUM_KSPLIT": 32, "BLOCK_M": 16, "TILE_K": 32, "num_warps": 2, - "num_stages": 3, + "num_stages": 2, "BLOCK_C": 1024, "reduce_num_warps": 1 }, - "M_LEQ_4096": { - "NUM_KSPLIT": 10, + "M_LEQ_768": { + "NUM_KSPLIT": 20, "BLOCK_M": 16, "TILE_K": 32, "num_warps": 2, - "num_stages": 3, + "num_stages": 2, "BLOCK_C": 1024, "reduce_num_warps": 1 }, - "M_LEQ_8192": { - "NUM_KSPLIT": 8, + "M_LEQ_1024": { + "NUM_KSPLIT": 16, "BLOCK_M": 16, "TILE_K": 32, "num_warps": 2, - "num_stages": 3, + "num_stages": 2, "BLOCK_C": 1024, "reduce_num_warps": 1 }, - "M_LEQ_16384": { - "NUM_KSPLIT": 4, + "M_LEQ_1280": { + "NUM_KSPLIT": 40, "BLOCK_M": 16, "TILE_K": 32, "num_warps": 2, - "num_stages": 3, + "num_stages": 2, "BLOCK_C": 1024, "reduce_num_warps": 1 }, - "any": { - "NUM_KSPLIT": 4, + "M_LEQ_2048": { + "NUM_KSPLIT": 8, "BLOCK_M": 16, "TILE_K": 32, "num_warps": 2, - "num_stages": 3, + "num_stages": 2, + "BLOCK_C": 1024, + "reduce_num_warps": 1 + }, + "M_LEQ_2304": { + "NUM_KSPLIT": 16, + "BLOCK_M": 16, + "TILE_K": 32, + "num_warps": 2, + "num_stages": 2, + "BLOCK_C": 1024, + "reduce_num_warps": 1 + }, + "M_LEQ_14336": { + "NUM_KSPLIT": 16, + "BLOCK_M": 64, + "TILE_K": 64, + "num_warps": 4, + "num_stages": 2, + "BLOCK_C": 1024, + "reduce_num_warps": 1 + }, + "M_LEQ_16384": { + "NUM_KSPLIT": 1, + "BLOCK_M": 64, + "TILE_K": 64, + "num_warps": 4, + "num_stages": 2, + "BLOCK_C": 1024, + "reduce_num_warps": 1 + }, + "any": { + "NUM_KSPLIT": 4, + "BLOCK_M": 64, + "TILE_K": 64, + "num_warps": 4, + "num_stages": 2, "BLOCK_C": 1024, "reduce_num_warps": 1 } -} +} \ No newline at end of file diff --git a/aiter/ops/triton/configs/gfx950/triton/mhc/mhc_fused_post_pre_delayed_rmsnorm/DEFAULT.json b/aiter/ops/triton/configs/gfx950/triton/mhc/mhc_fused_post_pre_delayed_rmsnorm/DEFAULT.json index 46ec1bba5b3..56e49a02455 100644 --- a/aiter/ops/triton/configs/gfx950/triton/mhc/mhc_fused_post_pre_delayed_rmsnorm/DEFAULT.json +++ b/aiter/ops/triton/configs/gfx950/triton/mhc/mhc_fused_post_pre_delayed_rmsnorm/DEFAULT.json @@ -1,13 +1,31 @@ { - "M_LEQ_144": { + "M_LEQ_16": { "NUM_KSPLIT": 160, "BLOCK_M": 16, "TILE_K": 32, "num_warps": 2, + "num_stages": 1, + "BLOCK_C": 1024, + "reduce_num_warps": 1 + }, + "M_LEQ_64": { + "NUM_KSPLIT": 160, + "BLOCK_M": 16, + "TILE_K": 32, + "num_warps": 1, "num_stages": 3, "BLOCK_C": 1024, "reduce_num_warps": 1 }, + "M_LEQ_96": { + "NUM_KSPLIT": 40, + "BLOCK_M": 16, + "TILE_K": 32, + "num_warps": 2, + "num_stages": 2, + "BLOCK_C": 1024, + "reduce_num_warps": 1 + }, "M_LEQ_192": { "NUM_KSPLIT": 40, "BLOCK_M": 16, @@ -17,58 +35,94 @@ "BLOCK_C": 1024, "reduce_num_warps": 1 }, - "M_LEQ_1024": { - "NUM_KSPLIT": 160, + "M_LEQ_352": { + "NUM_KSPLIT": 40, "BLOCK_M": 16, "TILE_K": 32, "num_warps": 2, - "num_stages": 3, + "num_stages": 2, "BLOCK_C": 1024, "reduce_num_warps": 1 }, - "M_LEQ_2048": { - "NUM_KSPLIT": 80, + "M_LEQ_512": { + "NUM_KSPLIT": 32, "BLOCK_M": 16, "TILE_K": 32, "num_warps": 2, - "num_stages": 3, + "num_stages": 2, "BLOCK_C": 1024, "reduce_num_warps": 1 }, - "M_LEQ_4096": { - "NUM_KSPLIT": 10, + "M_LEQ_768": { + "NUM_KSPLIT": 20, "BLOCK_M": 16, "TILE_K": 32, "num_warps": 2, - "num_stages": 3, + "num_stages": 2, "BLOCK_C": 1024, "reduce_num_warps": 1 }, - "M_LEQ_8192": { - "NUM_KSPLIT": 8, + "M_LEQ_1024": { + "NUM_KSPLIT": 16, "BLOCK_M": 16, "TILE_K": 32, "num_warps": 2, - "num_stages": 3, + "num_stages": 2, "BLOCK_C": 1024, "reduce_num_warps": 1 }, - "M_LEQ_16384": { - "NUM_KSPLIT": 4, + "M_LEQ_1280": { + "NUM_KSPLIT": 40, "BLOCK_M": 16, "TILE_K": 32, "num_warps": 2, - "num_stages": 3, + "num_stages": 2, "BLOCK_C": 1024, "reduce_num_warps": 1 }, - "any": { - "NUM_KSPLIT": 4, + "M_LEQ_2048": { + "NUM_KSPLIT": 8, "BLOCK_M": 16, "TILE_K": 32, "num_warps": 2, - "num_stages": 3, + "num_stages": 2, + "BLOCK_C": 1024, + "reduce_num_warps": 1 + }, + "M_LEQ_2304": { + "NUM_KSPLIT": 16, + "BLOCK_M": 16, + "TILE_K": 32, + "num_warps": 2, + "num_stages": 2, + "BLOCK_C": 1024, + "reduce_num_warps": 1 + }, + "M_LEQ_14336": { + "NUM_KSPLIT": 16, + "BLOCK_M": 64, + "TILE_K": 64, + "num_warps": 4, + "num_stages": 2, + "BLOCK_C": 1024, + "reduce_num_warps": 1 + }, + "M_LEQ_16384": { + "NUM_KSPLIT": 1, + "BLOCK_M": 64, + "TILE_K": 64, + "num_warps": 4, + "num_stages": 2, + "BLOCK_C": 1024, + "reduce_num_warps": 1 + }, + "any": { + "NUM_KSPLIT": 4, + "BLOCK_M": 64, + "TILE_K": 64, + "num_warps": 4, + "num_stages": 2, "BLOCK_C": 1024, "reduce_num_warps": 1 } -} +} \ No newline at end of file