From 287c1d1fd4377e680dcc59144aac4e587c29f0c2 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 29 Sep 2026 21:35:56 +0000 Subject: [PATCH 01/22] Add FlyDSL persistent work-stealing MXFP8 grouped GEMM (gfx950) Integrate a persistent, work-stealing grouped GEMM around the existing FlyDSL MXFP8 kernel into GroupedLinear, dispatched inline at the forward GEMM point (like the Triton grouped path) so it rides off TE's existing MXFP8 quantization rather than quantizing itself. The grouped kernel is a compile-time specialization (group=True) of the same mxfp8_gemm source as the dense kernel; the hot MFMA pipeline is unchanged and runs in a process_tile closure. Grouping only shifts B/B-scale addressing to the packed [G*N, K] layout (expert*N column, G*N stride) under const_expr(group), so the dense binary is byte-identical. A fixed num-CU persistent grid claims tiles from a precomputed (pid_m, pid_n, expert) descriptor queue via an atomic counter with LDS broadcast; a static grid-stride schedule is also provided, selected by a host heuristic (NVTE_FLYDSL_GROUPED_SCHED override). Opt in with NVTE_USE_GROUPED_GEMM_FLYDSL=1 (HIP, gfx950, MXFP8 recipe, 128-aligned in/out features). Bitwise-equal to TE's default grouped MXFP8 GEMM and to a per-expert sequential FlyDSL dense reference, forward and backward. Co-Authored-By: Claude --- .../pytorch/flydsl_kernels/gemm/__init__.py | 3 +- .../flydsl_kernels/gemm/gemm_wrappers.py | 134 ++ .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 1943 +++++++++-------- .../pytorch/module/grouped_linear.py | 162 +- 4 files changed, 1353 insertions(+), 889 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py b/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py index 9ee82e83b0..60c6592c8f 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py @@ -41,9 +41,10 @@ def _check_flydsl_version() -> None: _check_flydsl_version() from .exceptions import FlyDSLUnsupportedError -from .gemm_wrappers import te_generic_gemm_flydsl +from .gemm_wrappers import te_generic_gemm_flydsl, _run_mxfp8_grouped __all__ = [ "FlyDSLUnsupportedError", "te_generic_gemm_flydsl", + "_run_mxfp8_grouped", ] diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 4a1dbb09a4..8f581c6f9d 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -1008,6 +1008,140 @@ def _run_mxfp8( return D, aux +# Above this many persistent-loop rounds, dynamic atomic work-stealing pays for +# itself; below it a single grid-stride sweep already balances the tiles, so the +# atomic-free static schedule is preferred. Matches the HK reference's SCHED_ROUNDS. +_GROUPED_SCHED_ROUNDS = 4 + + +def _select_grouped_sched(num_tiles: int, device) -> str: + """Pick the grouped-GEMM schedule: 'worksteal' (atomic) or 'static' (grid-stride). + + Honors the NVTE_FLYDSL_GROUPED_SCHED override ('auto', 'static', 'worksteal') + for testing; otherwise applies the round-count heuristic against the CU grid. + """ + override = os.getenv("NVTE_FLYDSL_GROUPED_SCHED", "auto").lower() + if override in ("static", "worksteal"): + return override + if override not in ("auto", ""): + raise ValueError( + f"NVTE_FLYDSL_GROUPED_SCHED must be auto|static|worksteal, got {override!r}" + ) + num_cus = torch.cuda.get_device_properties(device).multi_processor_count + return "static" if num_tiles <= _GROUPED_SCHED_ROUNDS * num_cus else "worksteal" + + +def _run_mxfp8_grouped( + a_data, + a_scale, + b_data, + b_scale, + D, + *, + group_offs, + m_total, + n, + k, + output_dtype: torch.dtype, + bias=None, +): + """Dispatch a packed MXFP8 grouped GEMM through the persistent kernel (TN only). + + Operand contract (TN, single packed launch spanning all groups): + + a = packed activations [M_total, K] scales [M_total, K/32] + b = packed weights [G*N, K] scales [G*N, K/32] + group_offs [G+1] exclusive scan of per-group row counts + + Each group ``g`` owns rows ``[group_offs[g], group_offs[g+1])`` of ``a`` and the + weight block ``b[g*N:(g+1)*N]``. ``group_offs`` is consumed on the host to build + the persistent kernel's per-tile ``(pid_m, pid_n, expert)`` descriptor queue, so + the kernel itself does no group lookup. This iteration assumes the operands arrive + already MXFP8-quantized and packed -- it takes the pre-selected rowwise payloads + directly and performs no quantization, packing, or BLAS operand-ownership swap + (unlike :func:`_run_mxfp8`). + """ + if k % 32 != 0: + raise FlyDSLUnsupportedError(f"K={k} must be divisible by MXFP8 scale group size 32") + if n % 128 != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP8 grouped GEMM requires N divisible by 128, got N={n}" + ) + + a_fp8_dtype = getattr(a_data, "_fp8_dtype", None) + b_fp8_dtype = getattr(b_data, "_fp8_dtype", None) + if a_data.dtype == torch.uint8: + a_data = reinterpret_as_fp8_tensor(a_data, a_fp8_dtype) + if b_data.dtype == torch.uint8: + b_data = reinterpret_as_fp8_tensor(b_data, b_fp8_dtype) + + if a_scale.dtype != torch.uint8 or b_scale.dtype != torch.uint8: + raise TypeError("FlyDSL MXFP8 grouped GEMM expects raw E8M0 scales as torch.uint8") + + D = _validate_or_allocate_output( + D, + shape=torch.Size((m_total, n)), + dtype=output_dtype, + device=a_data.device, + backend_name="MXFP8 grouped TN", + ) + + # Persistent-kernel work distribution: + # work_queue [num_tiles, 3] int32 -- one precomputed (pid_m, pid_n, expert) + # descriptor per output tile. The host bakes the tile->expert + # mapping here so the kernel needs no in-kernel search, and may + # permute the rows (XCD-aware, group-balanced) freely. + # work_counter [1] int32 -- shared atomic slot dispenser, zeroed per launch. + # A workgroup claims slot = atomicAdd(work_counter, 1), reads work_queue[slot], + # and repeats until slot >= num_tiles. + from .mxfp8_gemm import BLOCK_M, BLOCK_N + + num_groups = int(group_offs.numel()) - 1 + num_pid_n = n // BLOCK_N + group_offs_cpu = group_offs.to(device="cpu", dtype=torch.int64) + + # Per-expert M-block ranges over the packed [M_total, N] output. Each expert's + # token count must be BLOCK_M-aligned (enforced upstream by the tiling check), + # so tiles never straddle a group boundary. + descriptors = [] + for g in range(num_groups): + m_start_blk = int(group_offs_cpu[g]) // BLOCK_M + m_end_blk = int(group_offs_cpu[g + 1]) // BLOCK_M + for pid_m in range(m_start_blk, m_end_blk): + for pid_n in range(num_pid_n): + descriptors.append((pid_m, pid_n, g)) + + num_tiles = len(descriptors) + work_queue = torch.tensor(descriptors, dtype=torch.int32, device=a_data.device).reshape(-1) + work_counter = torch.zeros(1, dtype=torch.int32, device=a_data.device) + + # Schedule selection (compile-time specialization on the kernel side): + # "worksteal": persistent atomic tile-claiming, best for imbalanced experts. + # "static": grid-stride slot ownership, no atomic; deadlock-proof. + # Heuristic mirrors the HK reference: when the tile count is small relative to + # the launch grid, one grid-stride sweep already balances well, so the static + # path avoids atomic contention. NVTE_FLYDSL_GROUPED_SCHED overrides for tests. + sched = _select_grouped_sched(num_tiles, a_data.device) + + epilogue, bias_arg = _resolve_epilogue(bias, n) + mxfp8_matmul( + a_data, + a_scale, + b_data, + b_scale, + D.view(m_total, n), + layout="TN", + epilogue=epilogue, + bias=bias_arg, + group=True, + num_groups=num_groups, + work_queue=work_queue, + work_counter=work_counter, + sched=sched, + ) + return D + + def _select_fp8_storage_for_layout(A, transa, B, transb): """Select existing TE FP8 backings and normalize to one core contract. diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index 375f11d5a1..b06fc1443a 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -28,6 +28,8 @@ import flydsl.compiler as flyc import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import fly as _fly from flydsl._mlir.dialects import llvm from flydsl.expr import arith, const_expr, gpu, math, range_constexpr, rocdl from flydsl.expr.typing import T @@ -59,6 +61,10 @@ SCALE_GROUP_SIZE = 32 +def get_num_cus() -> int: + return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + + def _debug_enabled() -> bool: value = os.getenv("DEBUG_FLYDSL_MXFP8_GEMM", "") return value.lower() not in ("", "0", "false", "no", "off") @@ -252,6 +258,8 @@ def _compile_kernel( output_dtype: torch.dtype, layout: str, epilogue: str = "DEFAULT", + group: bool = False, + sched: str = "worksteal", ): """Build one compile-time-specialized TN, NN, or NT kernel. @@ -491,6 +499,9 @@ class SharedStorage: b0_1: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] b1_0: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] b1_1: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + # Persistent grouped path: one-int32 scratch to broadcast the atomically + # claimed queue slot from thread 0 to the whole workgroup. + claim_slot: fx.Array[fx.Int32, 1, 16] @flyc.kernel(known_block_size=[NUM_THREADS, 1, 1]) def kernel_gemm( @@ -503,7 +514,15 @@ def kernel_gemm( Aux: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32, + num_groups: fx.Int32 = 0, + work_queue: fx.Tensor = None, + work_counter: fx.Tensor = None, ): + # Total output tiles over the packed [M_total, N] result. Dense uses this + # as the grid extent; grouped uses it as the drain bound for the persistent + # work-stealing loop (see the dispatch at the end of the kernel). + num_tiles = (c_m // BLOCK_M) * (c_n // BLOCK_N) + lds = fx.SharedAllocator().allocate(SharedStorage).peek() lds_a0 = (lds.a0_0, lds.a0_1) lds_a1 = (lds.a1_0, lds.a1_1) @@ -522,972 +541,1084 @@ def kernel_gemm( bs_div = fx.logical_divide(bs_rsrc, fx.make_layout(1, 1)) scale_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32) tx = gpu.thread_id("x") + tx_i32 = fx.Int32(tx) num_blocks_m = c_m // BLOCK_M num_blocks_n = c_n // BLOCK_N - pid_m, pid_n = xcd_swizzle(num_blocks_m, num_blocks_n) - - bx_m = pid_m * BLOCK_M - by_n = pid_n * BLOCK_N - - # The flattened/XCD-swizzled block coordinates are i32, while global - # address arithmetic below is expressed in MLIR index type. Convert - # once here and use these index-typed tile bases for every address. - bx_m_idx = fx.Index(bx_m) - by_n_idx = fx.Index(by_n) - - tx_i32 = fx.Int32(tx) - wave_id = tx_i32 // fx.Int32(WARP_SIZE) - lane = tx_i32 % fx.Int32(WARP_SIZE) - - # Compile-time global leading dimensions: - # normal source [X,K] -> leading dimension K - # transpose source [K,X] -> leading dimension X - a_leading_dim = _a_leading_dim(c_m) - b_leading_dim = _b_leading_dim(c_n) - - # gl_off_a/gl_off_b: per-lane lists of LOAD_PASSES_HALF swizzled flat - # global offsets, indexed by DMA pass. gl_off_a[step] is this lane's - # static 16-byte source for pass `step`; G2SLoader adds the dynamic - # K-tile base as soffset. Offsets are pre-swizzled so bytes land in the - # bank-conflict-free LDS slots the MFMA read (S2RLoader) expects. - gl_off_a = compute_global_swizzle( - lane, - wave_id, - a_leading_dim, - LOAD_PASSES_HALF, - preshuffled=False, - ) - gl_off_b = compute_global_swizzle( - lane, - wave_id, - b_leading_dim, - LOAD_PASSES_HALF, - preshuffled=False, - ) - a_g2s = G2SLoader( - a_div, - gl_off_a, - LOAD_PASSES_HALF, - a_f8_ir_t, - wave_id, - ) - b_g2s = G2SLoader( - b_div, - gl_off_b, - LOAD_PASSES_HALF, - b_f8_ir_t, - wave_id, - ) - s2r = S2RLoader(fx.Int32(0), 1) - - layout_lane16 = fx.make_layout((4, 16), (16, 1)) - coord_lane16 = fx.idx2crd(fx.Int32(lane), layout_lane16) - lane_div_16 = fx.get(coord_lane16, 0) - lane_mod_16 = fx.get(coord_lane16, 1) - - # Per-CTA tile base in C elements, folded into each store's linear - # coordinate below (matching the scale-load addressing on this build; - # add_offset on a dynamic Index is unsupported here). - c_tile_base_elems = bx_m_idx * fx.Index(c_n) + by_n_idx - gC = fx.rocdl.make_buffer_tensor(C, max_size=True) - c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) - c_store_atom = fx.make_copy_atom( - fx.rocdl.BufferCopy32b() if output_element_bytes == 4 else fx.rocdl.BufferCopy16b(), - output_fx_dtype, - ) - - # Bias is a length-N fp32 vector indexed by the output-feature (N) - # coordinate and broadcast across the M/token rows. Only staged when - # the epilogue requests it; DEFAULT receives a dummy 1-element tensor. - # const_expr folds this compile-time flag at trace time so the setup - # (and load_bias) are inlined into the kernel scope with no runtime - # dispatch branch. - if const_expr(has_bias): - gBias = fx.rocdl.make_buffer_tensor(Bias, max_size=True) - bias_div = fx.logical_divide(gBias, fx.make_layout(1, 1)) - bias_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) - - def load_bias(col): - reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) - fx.copy(bias_ld_atom, fx.slice(bias_div, (None, fx.Int32(col))), reg) + # Persistent work-stealing driver. + # + # Dense path (group is False): the XCD-swizzled block id owns exactly one + # (pid_m, pid_n) output tile; the body below runs once. + # + # Grouped path (group is True, packed TN, const_expr-folded so the dense + # binary is untouched): the launch is a fixed num-CU grid of persistent + # workgroups. Each workgroup claims a queue slot by atomically fetch-adding + # the shared counter (``work_counter[0]``), reads the precomputed + # ``(pid_m, pid_n, expert)`` descriptor at that slot from ``work_queue``, and + # keeps claiming until the queue drains -- so faster CUs process more tiles. + # The host bakes the tile->expert mapping into the descriptor list (no + # in-kernel search), and may permute it (XCD-aware, group-balanced, ...) + # without any kernel change. B is packed [G*N, K] and B-scale [K128, G*N], + # so B reads use the expert-shifted column ``expert*N + n`` with stride + # ``G*N`` while A / A-scale / C / bias keep absolute-row / local-N coords. + if const_expr(group): + wq_rsrc = fx.rocdl.make_buffer_tensor(work_queue, max_size=True) + wq_div = fx.logical_divide(wq_rsrc, fx.make_layout(1, 1)) + + # LDS broadcast scratch (one int32), following the persistent-kernel + # idiom in mega_moe_stage1: recast the LDS slot to an int32 iterator and + # a 1-element view. Thread 0 writes via ptr_store; every thread reads via + # the view after a workgroup barrier. + claim_iter = fx.recast_iter(fx.Int32, lds.claim_slot.ptr) + claim_view = fx.make_view(claim_iter, fx.make_layout(1, 1)) + + def _global_i32_ptr(tensor, index): + # Build an address-space-1 !llvm.ptr to tensor[index] (int32). + # llvm.atomicrmw requires an !llvm.ptr (not !fly.ptr). + ptr_ty = ir.Type.parse("!llvm.ptr<1>") + base = _fly.extract_aligned_pointer_as_index(ptr_ty, tensor) + base = llvm.PtrToIntOp(T.i64, base).result + byte_off = fx.as_ir_value(fx.Int64(fx.Index(index) * fx.Index(4))) + addr = llvm.AddOp(base, byte_off, llvm.IntegerOverflowFlags(0)).result + return llvm.IntToPtrOp(ptr_ty, addr).result + + def _claim_next_slot(): + # One atomic per workgroup: thread 0 fetch-adds the shared counter + # and stores the pre-update value into the LDS scratch; a barrier + # then broadcasts it so every thread returns the same slot. + if tx_i32 == fx.Int32(0): + old = llvm.AtomicRMWOp( + llvm.AtomicBinOp.add, + _global_i32_ptr(work_counter, fx.Int32(0)), + fx.as_ir_value(fx.Int32(1)), + llvm.AtomicOrdering.monotonic, + syncscope="agent", + alignment=4, + ).result + fx.ptr_store(fx.Vector.from_elements([fx.Int32(old)], fx.Int32), claim_iter) + fx.barrier() + return fx.Vector(claim_view.load())[0] + + def _desc_field(slot, field): + # Read one int32 of the packed (pid_m, pid_n, expert) descriptor. + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Int32) + fx.copy(scale_ld_atom, fx.slice(wq_div, (None, slot * fx.Int32(3) + field)), reg) return fx.memref_load_vec(reg)[0] - # GELU_AUX saves the pre-activation value (A@B[+bias]) to a second M x N - # output so the backward pass can recompute the GELU gradient. Same tile - # base / store atom shape as C; DEFAULT gets a dummy 1-element tensor. - if const_expr(has_gelu): - gAux = fx.rocdl.make_buffer_tensor(Aux, max_size=True) - aux_div = fx.logical_divide(gAux, fx.make_layout(1, 1)) - aux_store_atom = fx.make_copy_atom( + def process_tile(pid_m, pid_n, expert): + bx_m = pid_m * BLOCK_M + by_n = pid_n * BLOCK_N + + # The flattened/XCD-swizzled block coordinates are i32, while global + # address arithmetic below is expressed in MLIR index type. Convert + # once here and use these index-typed tile bases for every address. + bx_m_idx = fx.Index(bx_m) + by_n_idx = fx.Index(by_n) + + # B-operand column base and B-scale row stride. In the dense path these + # collapse to the C column base / c_n; in the grouped path they carry the + # expert offset (from the descriptor) into the packed [G*N, K] weight and + # [K128, G*N] scale. + if const_expr(group): + by_n_b_idx = fx.Index(expert * fx.Int32(c_n) + by_n) + c_n_bs = num_groups * c_n + else: + by_n_b_idx = by_n_idx + c_n_bs = c_n + + wave_id = tx_i32 // fx.Int32(WARP_SIZE) + lane = tx_i32 % fx.Int32(WARP_SIZE) + + # Compile-time global leading dimensions: + # normal source [X,K] -> leading dimension K + # transpose source [K,X] -> leading dimension X + a_leading_dim = _a_leading_dim(c_m) + b_leading_dim = _b_leading_dim(c_n) + + # gl_off_a/gl_off_b: per-lane lists of LOAD_PASSES_HALF swizzled flat + # global offsets, indexed by DMA pass. gl_off_a[step] is this lane's + # static 16-byte source for pass `step`; G2SLoader adds the dynamic + # K-tile base as soffset. Offsets are pre-swizzled so bytes land in the + # bank-conflict-free LDS slots the MFMA read (S2RLoader) expects. + gl_off_a = compute_global_swizzle( + lane, + wave_id, + a_leading_dim, + LOAD_PASSES_HALF, + preshuffled=False, + ) + gl_off_b = compute_global_swizzle( + lane, + wave_id, + b_leading_dim, + LOAD_PASSES_HALF, + preshuffled=False, + ) + a_g2s = G2SLoader( + a_div, + gl_off_a, + LOAD_PASSES_HALF, + a_f8_ir_t, + wave_id, + ) + b_g2s = G2SLoader( + b_div, + gl_off_b, + LOAD_PASSES_HALF, + b_f8_ir_t, + wave_id, + ) + s2r = S2RLoader(fx.Int32(0), 1) + + layout_lane16 = fx.make_layout((4, 16), (16, 1)) + coord_lane16 = fx.idx2crd(fx.Int32(lane), layout_lane16) + lane_div_16 = fx.get(coord_lane16, 0) + lane_mod_16 = fx.get(coord_lane16, 1) + + # Per-CTA tile base in C elements, folded into each store's linear + # coordinate below (matching the scale-load addressing on this build; + # add_offset on a dynamic Index is unsupported here). + c_tile_base_elems = bx_m_idx * fx.Index(c_n) + by_n_idx + gC = fx.rocdl.make_buffer_tensor(C, max_size=True) + c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) + c_store_atom = fx.make_copy_atom( fx.rocdl.BufferCopy32b() if output_element_bytes == 4 else fx.rocdl.BufferCopy16b(), output_fx_dtype, ) - def gelu_tanh(x): - # tanh-approx GELU (matches PyTorch approximate='tanh' and the - # FlyDSL preshuffle reference), expressed through a non-positive - # exponent so exp() cannot overflow: - # 0.5*x*(1 + tanh(y)), y = sqrt(2/pi)*(x + 0.044715*x^3) - half_f32 = fx.Float32(0.5) - one_f32 = fx.Float32(1.0) - zero_f32 = fx.Float32(0.0) - two_f32 = fx.Float32(2.0) - x3 = x * x * x - y = fx.Float32(0.7978845608) * (x + fx.Float32(0.044715) * x3) - abs_y = fx.Float32(y).maximumf(zero_f32 - y) - e_neg2abs = math.exp(fx.Float32(-2.0) * abs_y) - denom = one_f32 + e_neg2abs - numerator = (y > zero_f32).select(two_f32, two_f32 * e_neg2abs) - return half_f32 * x * (numerator * (one_f32 / denom)) - - PIN_ACC_BASE = 0 - - def _reg_list(prefix, start, end): - return ",".join(f"~{{{prefix}{r}}}" for r in range(start, end + 1)) - - def reserve_pinned_accumulators(): - # Reserve a fixed physical AGPR bank for all accumulators. In the - # SSA-lowered path, the compiler generated heavy AGPR <-> VGPR traffic, - # including v_accvgpr_mov/read sequences, s_nop stalls, and accumulator - # spills. Pinning each f32x4 accumulator to a stable AGPR range keeps the - # scaled MFMA accumulation in place and avoids those transfers and spills. - # - # ACCS_PER_WAVE = 64 accumulator objects and each object is f32x4, - # so the physical bank is exactly 64 * 4 = 256 AGPRs: a[0:255]. - clobbers = _reg_list("a", PIN_ACC_BASE, PIN_ACC_BASE + ACCS_PER_WAVE * 4 - 1) - llvm.InlineAsmOp( - None, - [], - "", - clobbers, - has_side_effects=True, - ) + # Bias is a length-N fp32 vector indexed by the output-feature (N) + # coordinate and broadcast across the M/token rows. Only staged when + # the epilogue requests it; DEFAULT receives a dummy 1-element tensor. + # const_expr folds this compile-time flag at trace time so the setup + # (and load_bias) are inlined into the kernel scope with no runtime + # dispatch branch. + if const_expr(has_bias): + gBias = fx.rocdl.make_buffer_tensor(Bias, max_size=True) + bias_div = fx.logical_divide(gBias, fx.make_layout(1, 1)) + bias_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) + + def load_bias(col): + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) + fx.copy(bias_ld_atom, fx.slice(bias_div, (None, fx.Int32(col))), reg) + return fx.memref_load_vec(reg)[0] + + # GELU_AUX saves the pre-activation value (A@B[+bias]) to a second M x N + # output so the backward pass can recompute the GELU gradient. Same tile + # base / store atom shape as C; DEFAULT gets a dummy 1-element tensor. + if const_expr(has_gelu): + gAux = fx.rocdl.make_buffer_tensor(Aux, max_size=True) + aux_div = fx.logical_divide(gAux, fx.make_layout(1, 1)) + aux_store_atom = fx.make_copy_atom( + fx.rocdl.BufferCopy32b() if output_element_bytes == 4 else fx.rocdl.BufferCopy16b(), + output_fx_dtype, + ) - def zero_pinned_accumulators(): - for ai in range_constexpr(ACCS_PER_WAVE * 4): + def gelu_tanh(x): + # tanh-approx GELU (matches PyTorch approximate='tanh' and the + # FlyDSL preshuffle reference), expressed through a non-positive + # exponent so exp() cannot overflow: + # 0.5*x*(1 + tanh(y)), y = sqrt(2/pi)*(x + 0.044715*x^3) + half_f32 = fx.Float32(0.5) + one_f32 = fx.Float32(1.0) + zero_f32 = fx.Float32(0.0) + two_f32 = fx.Float32(2.0) + x3 = x * x * x + y = fx.Float32(0.7978845608) * (x + fx.Float32(0.044715) * x3) + abs_y = fx.Float32(y).maximumf(zero_f32 - y) + e_neg2abs = math.exp(fx.Float32(-2.0) * abs_y) + denom = one_f32 + e_neg2abs + numerator = (y > zero_f32).select(two_f32, two_f32 * e_neg2abs) + return half_f32 * x * (numerator * (one_f32 / denom)) + + PIN_ACC_BASE = 0 + + def _reg_list(prefix, start, end): + return ",".join(f"~{{{prefix}{r}}}" for r in range(start, end + 1)) + + def reserve_pinned_accumulators(): + # Reserve a fixed physical AGPR bank for all accumulators. In the + # SSA-lowered path, the compiler generated heavy AGPR <-> VGPR traffic, + # including v_accvgpr_mov/read sequences, s_nop stalls, and accumulator + # spills. Pinning each f32x4 accumulator to a stable AGPR range keeps the + # scaled MFMA accumulation in place and avoids those transfers and spills. + # + # ACCS_PER_WAVE = 64 accumulator objects and each object is f32x4, + # so the physical bank is exactly 64 * 4 = 256 AGPRs: a[0:255]. + clobbers = _reg_list("a", PIN_ACC_BASE, PIN_ACC_BASE + ACCS_PER_WAVE * 4 - 1) llvm.InlineAsmOp( None, [], - f"v_accvgpr_write_b32 a[{PIN_ACC_BASE + ai}], 0", - f"~{{a{PIN_ACC_BASE + ai}}}", + "", + clobbers, has_side_effects=True, ) - def _inline_asm_i32(asm_string, constraints, operands=None): - op = llvm.InlineAsmOp( - T.i32, - operands or [], - asm_string, - constraints, - has_side_effects=True, - ) - return _one_i32_result(op) - - def _one_i32_result(op): - # Accept the result attribute names exposed by the supported MLIR Python bindings. - return getattr(op, "result", getattr(op, "res", op.results[0])) - - def _to_raw_inline_asm_operand(value): - # TODO: Replace arith._to_raw once FlyDSL exposes a supported public - # API for passing wrapped values to llvm.InlineAsmOp. _to_raw is - # deprecated, but remains heavily used internally by FlyDSL. - return arith._to_raw(value) - - def read_physical_accumulator_slot(slot_idx): - acc_pin = PIN_ACC_BASE + slot_idx * 4 - r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") - r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") - r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") - r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") - return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) - - # As/Bs are MFMA-ready packed scale words: [K128, row] uint32. - # Each loaded dword already contains the four 16-row/16-col MFMA scale - # bytes for this lane's 64-row A/B half. The MFMA instruction selects - # the byte via op_sel/op_sel_hi, so there is intentionally no hot-loop - # byte extraction and no 0x01010101 broadcast here. - c_m_idx = fx.Index(c_m) - c_n_idx = fx.Index(c_n) - - def hot_loop_scheduler_q_refill_2n(): - # Steady-state Q1 schedule: eight chunks of one K+2 VMEM/LDS - # refill pass followed by two MFMAs. - for _ in range_constexpr(8): - rocdl.sched_vmem(1) - rocdl.sched_mfma(2) - - rocdl.sched_barrier(0) - - def hot_loop_scheduler_q0_refill_a1_2n(): - # TN/NN: one normal A-bottom LDS read per chunk. - # NT: one transpose-read A half plus the matching transpose-read - # scheduling pressure retained from the passing NT specialization. - for _ in range_constexpr(8): - rocdl.sched_vmem(1) - rocdl.sched_dsrd(Q0_SCHED_DSRD) - rocdl.sched_mfma(2) - - rocdl.sched_barrier(0) - - def hot_loop_scheduler_q_prefetch_4n(): - # TN/NN retain two scheduled DS reads per chunk. NT retains four - # because both carried operands use two DS_READ_TR instructions. - for _ in range_constexpr(8): - rocdl.sched_dsrd(PREFETCH_SCHED_DSRD) - rocdl.sched_mfma(4) - - rocdl.sched_barrier(0) - - def load_a_scale_row(k128, row): - off = k128 * c_m_idx + bx_m_idx + row - reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Int32) - fx.copy(scale_ld_atom, fx.slice(as_div, (None, fx.Int32(off))), reg) - return fx.memref_load_vec(reg)[0] - - def load_b_scale_row(k128, row): - off = k128 * c_n_idx + by_n_idx + row - reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Int32) - fx.copy(scale_ld_atom, fx.slice(bs_div, (None, fx.Int32(off))), reg) - return fx.memref_load_vec(reg)[0] - - def load_a_scale_subtile(k128, sm): - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - a_row = subtile_m_idx * fx.Index(SUBTILE_M) + fx.Index(lane) - a_scale = load_a_scale_row(k128, a_row) - return (a_scale, a_scale, a_scale, a_scale) + def zero_pinned_accumulators(): + for ai in range_constexpr(ACCS_PER_WAVE * 4): + llvm.InlineAsmOp( + None, + [], + f"v_accvgpr_write_b32 a[{PIN_ACC_BASE + ai}], 0", + f"~{{a{PIN_ACC_BASE + ai}}}", + has_side_effects=True, + ) + + def _inline_asm_i32(asm_string, constraints, operands=None): + op = llvm.InlineAsmOp( + T.i32, + operands or [], + asm_string, + constraints, + has_side_effects=True, + ) + return _one_i32_result(op) + + def _one_i32_result(op): + # Accept the result attribute names exposed by the supported MLIR Python bindings. + return getattr(op, "result", getattr(op, "res", op.results[0])) + + def _to_raw_inline_asm_operand(value): + # TODO: Replace arith._to_raw once FlyDSL exposes a supported public + # API for passing wrapped values to llvm.InlineAsmOp. _to_raw is + # deprecated, but remains heavily used internally by FlyDSL. + return arith._to_raw(value) + + def read_physical_accumulator_slot(slot_idx): + acc_pin = PIN_ACC_BASE + slot_idx * 4 + r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") + r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") + r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") + r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") + return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) + + # As/Bs are MFMA-ready packed scale words: [K128, row] uint32. + # Each loaded dword already contains the four 16-row/16-col MFMA scale + # bytes for this lane's 64-row A/B half. The MFMA instruction selects + # the byte via op_sel/op_sel_hi, so there is intentionally no hot-loop + # byte extraction and no 0x01010101 broadcast here. + c_m_idx = fx.Index(c_m) + c_n_idx = fx.Index(c_n) + # B-scale row stride: c_n dense, packed G*N grouped (const_expr-folded). + c_n_bs_idx = fx.Index(c_n_bs) + + def hot_loop_scheduler_q_refill_2n(): + # Steady-state Q1 schedule: eight chunks of one K+2 VMEM/LDS + # refill pass followed by two MFMAs. + for _ in range_constexpr(8): + rocdl.sched_vmem(1) + rocdl.sched_mfma(2) + + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q0_refill_a1_2n(): + # TN/NN: one normal A-bottom LDS read per chunk. + # NT: one transpose-read A half plus the matching transpose-read + # scheduling pressure retained from the passing NT specialization. + for _ in range_constexpr(8): + rocdl.sched_vmem(1) + rocdl.sched_dsrd(Q0_SCHED_DSRD) + rocdl.sched_mfma(2) + + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q_prefetch_4n(): + # TN/NN retain two scheduled DS reads per chunk. NT retains four + # because both carried operands use two DS_READ_TR instructions. + for _ in range_constexpr(8): + rocdl.sched_dsrd(PREFETCH_SCHED_DSRD) + rocdl.sched_mfma(4) + + rocdl.sched_barrier(0) + + def load_a_scale_row(k128, row): + off = k128 * c_m_idx + bx_m_idx + row + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Int32) + fx.copy(scale_ld_atom, fx.slice(as_div, (None, fx.Int32(off))), reg) + return fx.memref_load_vec(reg)[0] - def load_b_scale_subtile(k128, sn): - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - b_row = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(lane) - b_scale = load_b_scale_row(k128, b_row) - return (b_scale, b_scale, b_scale, b_scale) - - def load_scale_tile(k128): - # Load all scale VGPRs needed by this wave for this K128 tile once. - # Return order: A-top, A-bottom, B-left, B-right. - return ( - load_a_scale_subtile(k128, 0), - load_a_scale_subtile(k128, 1), - load_b_scale_subtile(k128, 0), - load_b_scale_subtile(k128, 1), - ) + def load_b_scale_row(k128, row): + off = k128 * c_n_bs_idx + by_n_b_idx + row + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Int32) + fx.copy(scale_ld_atom, fx.slice(bs_div, (None, fx.Int32(off))), reg) + return fx.memref_load_vec(reg)[0] - def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): - a_g2s.load_one( - lds_a[subtile], - fx.Int32(_a_global_base(k_base, subtile, c_m, bx_m_idx)), - pass_in_subtile, - ) + def load_a_scale_subtile(k128, sm): + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + a_row = subtile_m_idx * fx.Index(SUBTILE_M) + fx.Index(lane) + a_scale = load_a_scale_row(k128, a_row) + return (a_scale, a_scale, a_scale, a_scale) + + def load_b_scale_subtile(k128, sn): + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + b_row = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(lane) + b_scale = load_b_scale_row(k128, b_row) + return (b_scale, b_scale, b_scale, b_scale) + + def load_scale_tile(k128): + # Load all scale VGPRs needed by this wave for this K128 tile once. + # Return order: A-top, A-bottom, B-left, B-right. + return ( + load_a_scale_subtile(k128, 0), + load_a_scale_subtile(k128, 1), + load_b_scale_subtile(k128, 0), + load_b_scale_subtile(k128, 1), + ) - def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): - b_g2s.load_one( - lds_b[subtile], - fx.Int32(_b_global_base(k_base, subtile, c_n, by_n_idx)), - pass_in_subtile, - ) + def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): + a_g2s.load_one( + lds_a[subtile], + fx.Int32(_a_global_base(k_base, subtile, c_m, bx_m_idx)), + pass_in_subtile, + ) - def stage_a_subtile(k_base, subtile, lds_a): - for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): - stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a) - - def stage_b_subtile(k_base, subtile, lds_b): - for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): - stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b) - - def load_frag_half_at_byte_base(lds_page, row_byte_base, half): - # Issue exactly one 16-byte LDS read for one K64 half of an MFMA operand. - # Keeping the halves separate allows steady-state Q0 to schedule one - # A-bottom ds_read_b128 in each refill/MFMA chunk. - k_col = reg_lds_k_col0 if half == 0 else reg_lds_k_col1 - return s2r.load_one(lds_page, fx.Int32(row_byte_base + k_col)) - - def pack_frag_halves(x0, x1): - return pack_i32x4_i32x8(x0, x1) - - def load_frag_at_byte_base(lds_page, row_byte_base): - # Default complete-fragment path used outside the dedicated Q0 schedule. - x0 = load_frag_half_at_byte_base(lds_page, row_byte_base, 0) - x1 = load_frag_half_at_byte_base(lds_page, row_byte_base, 1) - return pack_frag_halves(x0, x1) - - def load_normal_b_frag(lds_b, local_row, half): - # Physical [N,K] page, ordinary TN-style fixed-row read. - half_row = local_row - fx.Index(half * (BLOCK_N // 2)) - return load_frag_at_byte_base( - lds_b[half], - half_row * fx.Index(BLOCK_K), - ) + def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): + b_g2s.load_one( + lds_b[subtile], + fx.Int32(_b_global_base(k_base, subtile, c_n, by_n_b_idx)), + pass_in_subtile, + ) - def load_transposed_frag_half(lds_page, local_x_tile, half): - # Exact inverse mapping validated by the MXFP8 NN fragment probe. - lane_div16_i32 = fx.Int32(lane_div_16) - lane_in16_i32 = fx.Int32(lane_mod_16) - source_k = lane_div16_i32 * fx.Int32(16) + lane_in16_i32 // fx.Int32(2) - source_x = fx.Int32(local_x_tile) + (lane_in16_i32 % fx.Int32(2)) * fx.Int32(8) - - physical_k, physical_x = swizzle_128(source_k, source_x) - base = physical_k * fx.Int32(128) + physical_x - other = base ^ fx.Int32(0x440) - immediate_offset = 0 if half == 0 else 0x2000 - - return s2r.load_one_transpose( - lds_page, - base, - other, - immediate_offset=immediate_offset, - ) + def stage_a_subtile(k_base, subtile, lds_a): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a) + + def stage_b_subtile(k_base, subtile, lds_b): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b) + + def load_frag_half_at_byte_base(lds_page, row_byte_base, half): + # Issue exactly one 16-byte LDS read for one K64 half of an MFMA operand. + # Keeping the halves separate allows steady-state Q0 to schedule one + # A-bottom ds_read_b128 in each refill/MFMA chunk. + k_col = reg_lds_k_col0 if half == 0 else reg_lds_k_col1 + return s2r.load_one(lds_page, fx.Int32(row_byte_base + k_col)) + + def pack_frag_halves(x0, x1): + return pack_i32x4_i32x8(x0, x1) + + def load_frag_at_byte_base(lds_page, row_byte_base): + # Default complete-fragment path used outside the dedicated Q0 schedule. + x0 = load_frag_half_at_byte_base(lds_page, row_byte_base, 0) + x1 = load_frag_half_at_byte_base(lds_page, row_byte_base, 1) + return pack_frag_halves(x0, x1) + + def load_normal_b_frag(lds_b, local_row, half): + # Physical [N,K] page, ordinary TN-style fixed-row read. + half_row = local_row - fx.Index(half * (BLOCK_N // 2)) + return load_frag_at_byte_base( + lds_b[half], + half_row * fx.Index(BLOCK_K), + ) - def load_transposed_frag(lds_page, local_x_tile): - x0 = load_transposed_frag_half(lds_page, local_x_tile, 0) - x1 = load_transposed_frag_half(lds_page, local_x_tile, 1) - return pack_frag_halves(x0, x1) + def load_transposed_frag_half(lds_page, local_x_tile, half): + # Exact inverse mapping validated by the MXFP8 NN fragment probe. + lane_div16_i32 = fx.Int32(lane_div_16) + lane_in16_i32 = fx.Int32(lane_mod_16) + source_k = lane_div16_i32 * fx.Int32(16) + lane_in16_i32 // fx.Int32(2) + source_x = fx.Int32(local_x_tile) + (lane_in16_i32 % fx.Int32(2)) * fx.Int32(8) + + physical_k, physical_x = swizzle_128(source_k, source_x) + base = physical_k * fx.Int32(128) + physical_x + other = base ^ fx.Int32(0x440) + immediate_offset = 0 if half == 0 else 0x2000 + + return s2r.load_one_transpose( + lds_page, + base, + other, + immediate_offset=immediate_offset, + ) - def _acc_idx(subtile_id, mi, ni): - return ( - subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + mi * MFMA_N_PER_SUBTILE + ni - ) + def load_transposed_frag(lds_page, local_x_tile): + x0 = load_transposed_frag_half(lds_page, local_x_tile, 0) + x1 = load_transposed_frag_half(lds_page, local_x_tile, 1) + return pack_frag_halves(x0, x1) - def pinned_mfma(acc_idx, a_frag, b_frag, a_scale, b_scale, mi, ni): - # Fixed physical accumulator bank, visible SSA A/B/scale operands. - # acc_idx maps directly to a[PIN_ACC_BASE + 4*acc_idx : +3]. - # The scale operands are MFMA-ready packed dwords. mi/ni choose - # which of the four bytes inside the A/B scale dword the MFMA uses. - acc_pin = PIN_ACC_BASE + acc_idx * 4 - llvm.InlineAsmOp( - None, - [ - _to_raw_inline_asm_operand(a_frag), - _to_raw_inline_asm_operand(b_frag), - _to_raw_inline_asm_operand(a_scale), - _to_raw_inline_asm_operand(b_scale), - ], - "v_mfma_scale_f32_16x16x128_f8f6f4 " - f"a[{acc_pin}:{acc_pin + 3}], " - "$0, $1, " - f"a[{acc_pin}:{acc_pin + 3}], " - "$2, $3 " - f"op_sel:[{mi & 1},{ni & 1},0] " - f"op_sel_hi:[{mi >> 1},{ni >> 1},0] " - f"cbsz:{a_matrix_format} blgp:{b_matrix_format}", - f"v,v,v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}},~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}", - has_side_effects=True, - ) + def _acc_idx(subtile_id, mi, ni): + return ( + subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + mi * MFMA_N_PER_SUBTILE + ni + ) - def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag, a_scale, b_scale, mi, ni): - # Final-page form used by HK: destination and previous partial sum - # may be different AGPR ranges. Once old_acc_idx is consumed, its - # physical slot is dead and can be reused as a later destination. - dst_pin = PIN_ACC_BASE + dst_slot * 4 - old_pin = PIN_ACC_BASE + old_acc_idx * 4 - llvm.InlineAsmOp( - None, - [ - _to_raw_inline_asm_operand(a_frag), - _to_raw_inline_asm_operand(b_frag), - _to_raw_inline_asm_operand(a_scale), - _to_raw_inline_asm_operand(b_scale), - ], - "v_mfma_scale_f32_16x16x128_f8f6f4 " - f"a[{dst_pin}:{dst_pin + 3}], " - "$0, $1, " - f"a[{old_pin}:{old_pin + 3}], " - "$2, $3 " - f"op_sel:[{mi & 1},{ni & 1},0] " - f"op_sel_hi:[{mi >> 1},{ni >> 1},0] " - f"cbsz:{a_matrix_format} blgp:{b_matrix_format}", - f"v,v,v,v,~{{a{dst_pin}}},~{{a{dst_pin + 1}}},~{{a{dst_pin + 2}}},~{{a{dst_pin + 3}}}", - has_side_effects=True, - ) + def pinned_mfma(acc_idx, a_frag, b_frag, a_scale, b_scale, mi, ni): + # Fixed physical accumulator bank, visible SSA A/B/scale operands. + # acc_idx maps directly to a[PIN_ACC_BASE + 4*acc_idx : +3]. + # The scale operands are MFMA-ready packed dwords. mi/ni choose + # which of the four bytes inside the A/B scale dword the MFMA uses. + acc_pin = PIN_ACC_BASE + acc_idx * 4 + llvm.InlineAsmOp( + None, + [ + _to_raw_inline_asm_operand(a_frag), + _to_raw_inline_asm_operand(b_frag), + _to_raw_inline_asm_operand(a_scale), + _to_raw_inline_asm_operand(b_scale), + ], + "v_mfma_scale_f32_16x16x128_f8f6f4 " + f"a[{acc_pin}:{acc_pin + 3}], " + "$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}], " + "$2, $3 " + f"op_sel:[{mi & 1},{ni & 1},0] " + f"op_sel_hi:[{mi >> 1},{ni >> 1},0] " + f"cbsz:{a_matrix_format} blgp:{b_matrix_format}", + f"v,v,v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}},~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}", + has_side_effects=True, + ) - def mfma_4n(acc_base, a_frag, a_scale, b0, b1, b2, b3, bs0, bs1, bs2, bs3): - """Emit four N-direction scaled MFMAs into fixed physical AGPR accumulators.""" - mi = (acc_base // MFMA_N_PER_SUBTILE) % MFMA_M_PER_SUBTILE - pinned_mfma(acc_base + 0, a_frag, b0, a_scale, bs0, mi, 0) - pinned_mfma(acc_base + 1, a_frag, b1, a_scale, bs1, mi, 1) - pinned_mfma(acc_base + 2, a_frag, b2, a_scale, bs2, mi, 2) - pinned_mfma(acc_base + 3, a_frag, b3, a_scale, bs3, mi, 3) - - def mfma_2n(acc_base, a_frag, a_scale, b0, b1, bs0, bs1, ni_base): - mi = (acc_base // MFMA_N_PER_SUBTILE) % MFMA_M_PER_SUBTILE - pinned_mfma(acc_base + 0, a_frag, b0, a_scale, bs0, mi, ni_base + 0) - pinned_mfma(acc_base + 1, a_frag, b1, a_scale, bs1, mi, ni_base + 1) - - def store_acc_vector_for_logical_idx(logical_acc_idx, acc): - subtile_id = logical_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) - local_idx = logical_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) - sm = subtile_id // 2 - sn = subtile_id % 2 - mi = local_idx // MFMA_N_PER_SUBTILE - ni = local_idx % MFMA_N_PER_SUBTILE + def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag, a_scale, b_scale, mi, ni): + # Final-page form used by HK: destination and previous partial sum + # may be different AGPR ranges. Once old_acc_idx is consumed, its + # physical slot is dead and can be reused as a later destination. + dst_pin = PIN_ACC_BASE + dst_slot * 4 + old_pin = PIN_ACC_BASE + old_acc_idx * 4 + llvm.InlineAsmOp( + None, + [ + _to_raw_inline_asm_operand(a_frag), + _to_raw_inline_asm_operand(b_frag), + _to_raw_inline_asm_operand(a_scale), + _to_raw_inline_asm_operand(b_scale), + ], + "v_mfma_scale_f32_16x16x128_f8f6f4 " + f"a[{dst_pin}:{dst_pin + 3}], " + "$0, $1, " + f"a[{old_pin}:{old_pin + 3}], " + "$2, $3 " + f"op_sel:[{mi & 1},{ni & 1},0] " + f"op_sel_hi:[{mi >> 1},{ni >> 1},0] " + f"cbsz:{a_matrix_format} blgp:{b_matrix_format}", + f"v,v,v,v,~{{a{dst_pin}}},~{{a{dst_pin + 1}}},~{{a{dst_pin + 2}}},~{{a{dst_pin + 3}}}", + has_side_effects=True, + ) - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - row_base = subtile_m_idx * SUBTILE_M + fx.Index(mi * MFMA_M) + lane_div_16 * 4 - col = subtile_n_idx * SUBTILE_N + fx.Index(ni * MFMA_N) + lane_mod_16 - - # Bias depends only on the output-feature (N) coordinate, so read it - # once per column and reuse across the four M rows below. Index by - # the GLOBAL feature (by_n_idx + col), matching the C store's - # c_tile_base_elems fold; the tile-local col alone would reread - # bias[0..BLOCK_N) for every N-block. - if const_expr(has_bias): - bias_value = load_bias(by_n_idx + col) + def mfma_4n(acc_base, a_frag, a_scale, b0, b1, b2, b3, bs0, bs1, bs2, bs3): + """Emit four N-direction scaled MFMAs into fixed physical AGPR accumulators.""" + mi = (acc_base // MFMA_N_PER_SUBTILE) % MFMA_M_PER_SUBTILE + pinned_mfma(acc_base + 0, a_frag, b0, a_scale, bs0, mi, 0) + pinned_mfma(acc_base + 1, a_frag, b1, a_scale, bs1, mi, 1) + pinned_mfma(acc_base + 2, a_frag, b2, a_scale, bs2, mi, 2) + pinned_mfma(acc_base + 3, a_frag, b3, a_scale, bs3, mi, 3) + + def mfma_2n(acc_base, a_frag, a_scale, b0, b1, bs0, bs1, ni_base): + mi = (acc_base // MFMA_N_PER_SUBTILE) % MFMA_M_PER_SUBTILE + pinned_mfma(acc_base + 0, a_frag, b0, a_scale, bs0, mi, ni_base + 0) + pinned_mfma(acc_base + 1, a_frag, b1, a_scale, bs1, mi, ni_base + 1) + + def store_acc_vector_for_logical_idx(logical_acc_idx, acc): + subtile_id = logical_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = logical_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE - for ii in range_constexpr(4): - row = row_base + fx.Index(ii) - c_idx = c_tile_base_elems + row * fx.Index(c_n) + col + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + row_base = subtile_m_idx * SUBTILE_M + fx.Index(mi * MFMA_M) + lane_div_16 * 4 + col = subtile_n_idx * SUBTILE_N + fx.Index(ni * MFMA_N) + lane_mod_16 - # Epilogue stages run on the fp32 accumulator, in order, before - # the output-dtype narrowing: - # value = acc [+ bias] (pre-activation) - # GELU_AUX: save pre-activation to Aux, then value = gelu(value) - value = Vec(acc)[ii] + # Bias depends only on the output-feature (N) coordinate, so read it + # once per column and reuse across the four M rows below. Index by + # the GLOBAL feature (by_n_idx + col), matching the C store's + # c_tile_base_elems fold; the tile-local col alone would reread + # bias[0..BLOCK_N) for every N-block. if const_expr(has_bias): - value = value + bias_value + bias_value = load_bias(by_n_idx + col) + + for ii in range_constexpr(4): + row = row_base + fx.Index(ii) + c_idx = c_tile_base_elems + row * fx.Index(c_n) + col + + # Epilogue stages run on the fp32 accumulator, in order, before + # the output-dtype narrowing: + # value = acc [+ bias] (pre-activation) + # GELU_AUX: save pre-activation to Aux, then value = gelu(value) + value = Vec(acc)[ii] + if const_expr(has_bias): + value = value + bias_value + + if const_expr(has_gelu): + # Save the pre-activation (post-bias) value for backward, + # then apply GELU to the C output. + aux_val = value + if output_dtype != torch.float32: + aux_val = aux_val.to(output_fx_dtype) + aux_reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) + fx.memref_store_vec(Vec.filled(1, aux_val, output_fx_dtype), aux_reg) + fx.copy(aux_store_atom, aux_reg, fx.slice(aux_div, (None, fx.Int32(c_idx)))) + value = gelu_tanh(value) - if const_expr(has_gelu): - # Save the pre-activation (post-bias) value for backward, - # then apply GELU to the C output. - aux_val = value if output_dtype != torch.float32: - aux_val = aux_val.to(output_fx_dtype) - aux_reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) - fx.memref_store_vec(Vec.filled(1, aux_val, output_fx_dtype), aux_reg) - fx.copy(aux_store_atom, aux_reg, fx.slice(aux_div, (None, fx.Int32(c_idx)))) - value = gelu_tanh(value) - - if output_dtype != torch.float32: - value = value.to(output_fx_dtype) - reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) - fx.memref_store_vec(Vec.filled(1, value, output_fx_dtype), reg) - fx.copy(c_store_atom, reg, fx.slice(c_div, (None, fx.Int32(c_idx)))) - - # Explicit register coordinates for HK-style four-quadrant mapping. - # BLOCK_M/BLOCK_N are 256x256. Four waves map to warp positions - # inside each 128x128 quadrant: - # cA: (warp_m, warp_n) - # cB: (warp_m, warp_n + 2) - # cC: (warp_m + 2, warp_n) - # cD: (warp_m + 2, warp_n + 2) - reg_lds_k_col0, reg_lds_k_col1 = _normal_read_columns( - lane_div_16, - lane_mod_16, - ) - - reg_subtile_m_idx0 = wave_id // 2 - reg_subtile_n_idx0 = wave_id % 2 - - reserve_pinned_accumulators() - zero_pinned_accumulators() - - def load_b_subtile_ni_regs(lds_b, scale_tile, sn, ni): - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - b_scales = scale_tile[2] if sn == 0 else scale_tile[3] - - b_ni = _load_b_ni( - load_transposed_frag, - load_normal_b_frag, - lds_b, - sn, - ni, - reg_subtile_n_idx0, - lane_mod_16, - ) - return b_ni, b_scales[ni] - - def load_b_subtile_regs(lds_b, scale_tile, sn): - b0, bs0 = load_b_subtile_ni_regs(lds_b, scale_tile, sn, 0) - b1, bs1 = load_b_subtile_ni_regs(lds_b, scale_tile, sn, 1) - b2, bs2 = load_b_subtile_ni_regs(lds_b, scale_tile, sn, 2) - b3, bs3 = load_b_subtile_ni_regs(lds_b, scale_tile, sn, 3) - return b0, b1, b2, b3, bs0, bs1, bs2, bs3 - - def load_a_subtile_mi_half(lds_a, sm, mi, half): - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - - return _load_a_half( - load_transposed_frag_half, - load_frag_half_at_byte_base, - lds_a, - sm, - mi, - half, - reg_subtile_m_idx0, + value = value.to(output_fx_dtype) + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) + fx.memref_store_vec(Vec.filled(1, value, output_fx_dtype), reg) + fx.copy(c_store_atom, reg, fx.slice(c_div, (None, fx.Int32(c_idx)))) + + # Explicit register coordinates for HK-style four-quadrant mapping. + # BLOCK_M/BLOCK_N are 256x256. Four waves map to warp positions + # inside each 128x128 quadrant: + # cA: (warp_m, warp_n) + # cB: (warp_m, warp_n + 2) + # cC: (warp_m + 2, warp_n) + # cD: (warp_m + 2, warp_n + 2) + reg_lds_k_col0, reg_lds_k_col1 = _normal_read_columns( + lane_div_16, lane_mod_16, ) - def load_a_subtile_mi_regs(lds_a, scale_tile, sm, mi): - # Fine-grained A register load for one 16-row M-direction MFMA slice. - a_scales = scale_tile[0] if sm == 0 else scale_tile[1] - x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) - x1 = load_a_subtile_mi_half(lds_a, sm, mi, 1) - a_mi = pack_frag_halves(x0, x1) - a_scale_mi = a_scales[mi] - return a_mi, a_scale_mi - - def load_a_subtile_regs(lds_a, scale_tile, sm): - a0, as0 = load_a_subtile_mi_regs(lds_a, scale_tile, sm, 0) - a1, as1 = load_a_subtile_mi_regs(lds_a, scale_tile, sm, 1) - a2, as2 = load_a_subtile_mi_regs(lds_a, scale_tile, sm, 2) - a3, as3 = load_a_subtile_mi_regs(lds_a, scale_tile, sm, 3) - return a0, a1, a2, a3, as0, as1, as2, as3 - - def hk_one_k_with_refill( - k128, - cur_a, - cur_b, - next_a, - next_b, - refill_a, - refill_b, - a0_regs, - b0_regs, - cur_scales, - prev_refill_scales, - ): - # Scale invariant: - # cur_scales is HK MFMA-ready for K. - # prev_refill_scales is HK MFMA-ready for K+1. - # This iteration issues K+2 scale loads and returns them for the - # next steady iteration or final tail. - - # Wait only far enough for the current page; the next-page refill may remain in flight. - barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) - rocdl.sched_barrier(0) + reg_subtile_m_idx0 = wave_id // 2 + reg_subtile_n_idx0 = wave_id % 2 - # Immediately issue MFMA-ready K+2 scale loads. - # They are returned for the next iteration without any in-kernel - # byte extraction or broadcast. - refill_scales = load_scale_tile(fx.Index(k128 + 2)) - next_scales_ready = prev_refill_scales - # A-top and B-left are both carried as complete 64-row register tiles, - # so their LDS half-pages can be refilled immediately. - a00, a01, a02, a03, as00, as01, as02, as03 = a0_regs - b00, b01, b02, b03, bs00, bs01, bs02, bs03 = b0_regs - - b10, bs10 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 0) - b11, bs11 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 1) - b12, bs12 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 2) - b13, bs13 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 3) - - # Refill the current ping-pong page with K+2, alternating A and B passes. - k_refill = fx.Index((k128 + 2) * BLOCK_K) - - # Q0: interleave the current tile's A-bottom LDS reads with K+2 - # refills and Q0 compute. Each complete A-bottom fragment is assembled - # from two independently scheduled K64 halves. - rocdl.sched_barrier(0) - a10_x0 = load_a_subtile_mi_half(cur_a, 1, 0, 0) - stage_a_subtile_pass(k_refill, 0, 0, refill_a) - mfma_2n(_acc_idx(0, 0, 0), a00, as00, b00, b01, bs00, bs01, 0) + reserve_pinned_accumulators() + zero_pinned_accumulators() - a10_x1 = load_a_subtile_mi_half(cur_a, 1, 0, 1) - stage_b_subtile_pass(k_refill, 0, 0, refill_b) - mfma_2n(_acc_idx(0, 0, 2), a00, as00, b02, b03, bs02, bs03, 2) + def load_b_subtile_ni_regs(lds_b, scale_tile, sn, ni): + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + b_scales = scale_tile[2] if sn == 0 else scale_tile[3] - a11_x0 = load_a_subtile_mi_half(cur_a, 1, 1, 0) - stage_a_subtile_pass(k_refill, 0, 1, refill_a) - mfma_2n(_acc_idx(0, 1, 0), a01, as01, b00, b01, bs00, bs01, 0) + b_ni = _load_b_ni( + load_transposed_frag, + load_normal_b_frag, + lds_b, + sn, + ni, + reg_subtile_n_idx0, + lane_mod_16, + ) + return b_ni, b_scales[ni] + + def load_b_subtile_regs(lds_b, scale_tile, sn): + b0, bs0 = load_b_subtile_ni_regs(lds_b, scale_tile, sn, 0) + b1, bs1 = load_b_subtile_ni_regs(lds_b, scale_tile, sn, 1) + b2, bs2 = load_b_subtile_ni_regs(lds_b, scale_tile, sn, 2) + b3, bs3 = load_b_subtile_ni_regs(lds_b, scale_tile, sn, 3) + return b0, b1, b2, b3, bs0, bs1, bs2, bs3 + + def load_a_subtile_mi_half(lds_a, sm, mi, half): + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + + return _load_a_half( + load_transposed_frag_half, + load_frag_half_at_byte_base, + lds_a, + sm, + mi, + half, + reg_subtile_m_idx0, + lane_mod_16, + ) - a11_x1 = load_a_subtile_mi_half(cur_a, 1, 1, 1) - stage_b_subtile_pass(k_refill, 0, 1, refill_b) - mfma_2n(_acc_idx(0, 1, 2), a01, as01, b02, b03, bs02, bs03, 2) + def load_a_subtile_mi_regs(lds_a, scale_tile, sm, mi): + # Fine-grained A register load for one 16-row M-direction MFMA slice. + a_scales = scale_tile[0] if sm == 0 else scale_tile[1] + x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) + x1 = load_a_subtile_mi_half(lds_a, sm, mi, 1) + a_mi = pack_frag_halves(x0, x1) + a_scale_mi = a_scales[mi] + return a_mi, a_scale_mi + + def load_a_subtile_regs(lds_a, scale_tile, sm): + a0, as0 = load_a_subtile_mi_regs(lds_a, scale_tile, sm, 0) + a1, as1 = load_a_subtile_mi_regs(lds_a, scale_tile, sm, 1) + a2, as2 = load_a_subtile_mi_regs(lds_a, scale_tile, sm, 2) + a3, as3 = load_a_subtile_mi_regs(lds_a, scale_tile, sm, 3) + return a0, a1, a2, a3, as0, as1, as2, as3 + + def hk_one_k_with_refill( + k128, + cur_a, + cur_b, + next_a, + next_b, + refill_a, + refill_b, + a0_regs, + b0_regs, + cur_scales, + prev_refill_scales, + ): + # Scale invariant: + # cur_scales is HK MFMA-ready for K. + # prev_refill_scales is HK MFMA-ready for K+1. + # This iteration issues K+2 scale loads and returns them for the + # next steady iteration or final tail. + + # Wait only far enough for the current page; the next-page refill may remain in flight. + barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + # Immediately issue MFMA-ready K+2 scale loads. + # They are returned for the next iteration without any in-kernel + # byte extraction or broadcast. + refill_scales = load_scale_tile(fx.Index(k128 + 2)) + next_scales_ready = prev_refill_scales + # A-top and B-left are both carried as complete 64-row register tiles, + # so their LDS half-pages can be refilled immediately. + a00, a01, a02, a03, as00, as01, as02, as03 = a0_regs + b00, b01, b02, b03, bs00, bs01, bs02, bs03 = b0_regs + + b10, bs10 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 0) + b11, bs11 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 1) + b12, bs12 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 2) + b13, bs13 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 3) + + # Refill the current ping-pong page with K+2, alternating A and B passes. + k_refill = fx.Index((k128 + 2) * BLOCK_K) + + # Q0: interleave the current tile's A-bottom LDS reads with K+2 + # refills and Q0 compute. Each complete A-bottom fragment is assembled + # from two independently scheduled K64 halves. + rocdl.sched_barrier(0) + a10_x0 = load_a_subtile_mi_half(cur_a, 1, 0, 0) + stage_a_subtile_pass(k_refill, 0, 0, refill_a) + mfma_2n(_acc_idx(0, 0, 0), a00, as00, b00, b01, bs00, bs01, 0) + + a10_x1 = load_a_subtile_mi_half(cur_a, 1, 0, 1) + stage_b_subtile_pass(k_refill, 0, 0, refill_b) + mfma_2n(_acc_idx(0, 0, 2), a00, as00, b02, b03, bs02, bs03, 2) + + a11_x0 = load_a_subtile_mi_half(cur_a, 1, 1, 0) + stage_a_subtile_pass(k_refill, 0, 1, refill_a) + mfma_2n(_acc_idx(0, 1, 0), a01, as01, b00, b01, bs00, bs01, 0) + + a11_x1 = load_a_subtile_mi_half(cur_a, 1, 1, 1) + stage_b_subtile_pass(k_refill, 0, 1, refill_b) + mfma_2n(_acc_idx(0, 1, 2), a01, as01, b02, b03, bs02, bs03, 2) + + a12_x0 = load_a_subtile_mi_half(cur_a, 1, 2, 0) + stage_a_subtile_pass(k_refill, 0, 2, refill_a) + mfma_2n(_acc_idx(0, 2, 0), a02, as02, b00, b01, bs00, bs01, 0) + + a12_x1 = load_a_subtile_mi_half(cur_a, 1, 2, 1) + stage_b_subtile_pass(k_refill, 0, 2, refill_b) + mfma_2n(_acc_idx(0, 2, 2), a02, as02, b02, b03, bs02, bs03, 2) + + a13_x0 = load_a_subtile_mi_half(cur_a, 1, 3, 0) + stage_a_subtile_pass(k_refill, 0, 3, refill_a) + mfma_2n(_acc_idx(0, 3, 0), a03, as03, b00, b01, bs00, bs01, 0) + + a13_x1 = load_a_subtile_mi_half(cur_a, 1, 3, 1) + stage_b_subtile_pass(k_refill, 0, 3, refill_b) + mfma_2n(_acc_idx(0, 3, 2), a03, as03, b02, b03, bs02, bs03, 2) + + hot_loop_scheduler_q0_refill_a1_2n() + + # Retire the eight distributed A-bottom LDS reads before K+2 refills + # overwrite the current page's A-bottom half-page. Keep this wait as + # late as possible to maximize read/compute overlap. + rocdl.sched_barrier(0) + barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = pack_frag_halves(a10_x0, a10_x1) + a11 = pack_frag_halves(a11_x0, a11_x1) + a12 = pack_frag_halves(a12_x0, a12_x1) + a13 = pack_frag_halves(a13_x0, a13_x1) + as10 = cur_scales[1][0] + as11 = cur_scales[1][1] + as12 = cur_scales[1][2] + as13 = cur_scales[1][3] + + rocdl.sched_barrier(0) + stage_b_subtile_pass(k_refill, 1, 0, refill_b) + mfma_2n(_acc_idx(1, 0, 0), a00, as00, b10, b11, bs10, bs11, 0) + + stage_a_subtile_pass(k_refill, 1, 0, refill_a) + mfma_2n(_acc_idx(1, 0, 2), a00, as00, b12, b13, bs12, bs13, 2) + + stage_b_subtile_pass(k_refill, 1, 1, refill_b) + mfma_2n(_acc_idx(1, 1, 0), a01, as01, b10, b11, bs10, bs11, 0) + + stage_a_subtile_pass(k_refill, 1, 1, refill_a) + mfma_2n(_acc_idx(1, 1, 2), a01, as01, b12, b13, bs12, bs13, 2) + + stage_b_subtile_pass(k_refill, 1, 2, refill_b) + mfma_2n(_acc_idx(1, 2, 0), a02, as02, b10, b11, bs10, bs11, 0) + + stage_a_subtile_pass(k_refill, 1, 2, refill_a) + mfma_2n(_acc_idx(1, 2, 2), a02, as02, b12, b13, bs12, bs13, 2) + + stage_b_subtile_pass(k_refill, 1, 3, refill_b) + mfma_2n(_acc_idx(1, 3, 0), a03, as03, b10, b11, bs10, bs11, 0) + + stage_a_subtile_pass(k_refill, 1, 3, refill_a) + mfma_2n(_acc_idx(1, 3, 2), a03, as03, b12, b13, bs12, bs13, 2) + hot_loop_scheduler_q_refill_2n() + + # Leave exactly the K+2 refill and scale loads outstanding. The following + # LDS reads consume the already-ready next page, not the page being refilled. + rocdl.sched_barrier(0) + barrier( + vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE + LOAD_PASSES_SCALES, + lgkmcnt=0, + ) + rocdl.sched_barrier(0) - a12_x0 = load_a_subtile_mi_half(cur_a, 1, 2, 0) - stage_a_subtile_pass(k_refill, 0, 2, refill_a) - mfma_2n(_acc_idx(0, 2, 0), a02, as02, b00, b01, bs00, bs01, 0) + next_a00, next_as00 = load_a_subtile_mi_regs(next_a, next_scales_ready, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, as10, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - a12_x1 = load_a_subtile_mi_half(cur_a, 1, 2, 1) - stage_b_subtile_pass(k_refill, 0, 2, refill_b) - mfma_2n(_acc_idx(0, 2, 2), a02, as02, b02, b03, bs02, bs03, 2) + next_a01, next_as01 = load_a_subtile_mi_regs(next_a, next_scales_ready, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, as11, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - a13_x0 = load_a_subtile_mi_half(cur_a, 1, 3, 0) - stage_a_subtile_pass(k_refill, 0, 3, refill_a) - mfma_2n(_acc_idx(0, 3, 0), a03, as03, b00, b01, bs00, bs01, 0) + next_a02, next_as02 = load_a_subtile_mi_regs(next_a, next_scales_ready, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, as12, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - a13_x1 = load_a_subtile_mi_half(cur_a, 1, 3, 1) - stage_b_subtile_pass(k_refill, 0, 3, refill_b) - mfma_2n(_acc_idx(0, 3, 2), a03, as03, b02, b03, bs02, bs03, 2) + next_a03, next_as03 = load_a_subtile_mi_regs(next_a, next_scales_ready, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, as13, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - hot_loop_scheduler_q0_refill_a1_2n() + next_b00, next_bs00 = load_b_subtile_ni_regs(next_b, next_scales_ready, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, as10, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - # Retire the eight distributed A-bottom LDS reads before K+2 refills - # overwrite the current page's A-bottom half-page. Keep this wait as - # late as possible to maximize read/compute overlap. - rocdl.sched_barrier(0) - barrier(lgkmcnt=0) - rocdl.sched_barrier(0) + next_b01, next_bs01 = load_b_subtile_ni_regs(next_b, next_scales_ready, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, as11, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - a10 = pack_frag_halves(a10_x0, a10_x1) - a11 = pack_frag_halves(a11_x0, a11_x1) - a12 = pack_frag_halves(a12_x0, a12_x1) - a13 = pack_frag_halves(a13_x0, a13_x1) - as10 = cur_scales[1][0] - as11 = cur_scales[1][1] - as12 = cur_scales[1][2] - as13 = cur_scales[1][3] + next_b02, next_bs02 = load_b_subtile_ni_regs(next_b, next_scales_ready, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, as12, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - rocdl.sched_barrier(0) - stage_b_subtile_pass(k_refill, 1, 0, refill_b) - mfma_2n(_acc_idx(1, 0, 0), a00, as00, b10, b11, bs10, bs11, 0) + next_b03, next_bs03 = load_b_subtile_ni_regs(next_b, next_scales_ready, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, as13, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - stage_a_subtile_pass(k_refill, 1, 0, refill_a) - mfma_2n(_acc_idx(1, 0, 2), a00, as00, b12, b13, bs12, bs13, 2) + hot_loop_scheduler_q_prefetch_4n() - stage_b_subtile_pass(k_refill, 1, 1, refill_b) - mfma_2n(_acc_idx(1, 1, 0), a01, as01, b10, b11, bs10, bs11, 0) + next_a0_regs = ( + next_a00, + next_a01, + next_a02, + next_a03, + next_as00, + next_as01, + next_as02, + next_as03, + ) + next_b0_regs = ( + next_b00, + next_b01, + next_b02, + next_b03, + next_bs00, + next_bs01, + next_bs02, + next_bs03, + ) - stage_a_subtile_pass(k_refill, 1, 1, refill_a) - mfma_2n(_acc_idx(1, 1, 2), a01, as01, b12, b13, bs12, bs13, 2) + return next_a0_regs, next_b0_regs, next_scales_ready, refill_scales - stage_b_subtile_pass(k_refill, 1, 2, refill_b) - mfma_2n(_acc_idx(1, 2, 0), a02, as02, b10, b11, bs10, bs11, 0) + def hk_one_k_tail_with_next( + cur_a, cur_b, next_a, next_b, a0_regs, b0_regs, cur_scales, next_scales + ): + barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) - stage_a_subtile_pass(k_refill, 1, 2, refill_a) - mfma_2n(_acc_idx(1, 2, 2), a02, as02, b12, b13, bs12, bs13, 2) + a00, a01, a02, a03, as00, as01, as02, as03 = a0_regs + b00, b01, b02, b03, bs00, bs01, bs02, bs03 = b0_regs - stage_b_subtile_pass(k_refill, 1, 3, refill_b) - mfma_2n(_acc_idx(1, 3, 0), a03, as03, b10, b11, bs10, bs11, 0) + b10, bs10 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 0) + b11, bs11 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 1) + b12, bs12 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 2) + b13, bs13 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 3) - stage_a_subtile_pass(k_refill, 1, 3, refill_a) - mfma_2n(_acc_idx(1, 3, 2), a03, as03, b12, b13, bs12, bs13, 2) - hot_loop_scheduler_q_refill_2n() + mfma_4n(_acc_idx(0, 0, 0), a00, as00, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + mfma_4n(_acc_idx(0, 1, 0), a01, as01, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + mfma_4n(_acc_idx(0, 2, 0), a02, as02, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + mfma_4n(_acc_idx(0, 3, 0), a03, as03, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - # Leave exactly the K+2 refill and scale loads outstanding. The following - # LDS reads consume the already-ready next page, not the page being refilled. - rocdl.sched_barrier(0) - barrier( - vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE + LOAD_PASSES_SCALES, - lgkmcnt=0, - ) - rocdl.sched_barrier(0) + rocdl.sched_barrier(0) + barrier(lgkmcnt=0) + rocdl.sched_barrier(0) - next_a00, next_as00 = load_a_subtile_mi_regs(next_a, next_scales_ready, 0, 0) - mfma_4n(_acc_idx(2, 0, 0), a10, as10, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + a10, as10 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 0) + a11, as11 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 1) + a12, as12 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 2) + a13, as13 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 3) - next_a01, next_as01 = load_a_subtile_mi_regs(next_a, next_scales_ready, 0, 1) - mfma_4n(_acc_idx(2, 1, 0), a11, as11, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + mfma_4n(_acc_idx(1, 0, 0), a00, as00, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + mfma_4n(_acc_idx(1, 1, 0), a01, as01, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + mfma_4n(_acc_idx(1, 2, 0), a02, as02, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + mfma_4n(_acc_idx(1, 3, 0), a03, as03, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - next_a02, next_as02 = load_a_subtile_mi_regs(next_a, next_scales_ready, 0, 2) - mfma_4n(_acc_idx(2, 2, 0), a12, as12, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + rocdl.sched_barrier(0) + barrier(vmcnt=LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) - next_a03, next_as03 = load_a_subtile_mi_regs(next_a, next_scales_ready, 0, 3) - mfma_4n(_acc_idx(2, 3, 0), a13, as13, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + next_a00, next_as00 = load_a_subtile_mi_regs(next_a, next_scales, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, as10, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - next_b00, next_bs00 = load_b_subtile_ni_regs(next_b, next_scales_ready, 0, 0) - mfma_4n(_acc_idx(3, 0, 0), a10, as10, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + next_a01, next_as01 = load_a_subtile_mi_regs(next_a, next_scales, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, as11, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - next_b01, next_bs01 = load_b_subtile_ni_regs(next_b, next_scales_ready, 0, 1) - mfma_4n(_acc_idx(3, 1, 0), a11, as11, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + next_a02, next_as02 = load_a_subtile_mi_regs(next_a, next_scales, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, as12, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - next_b02, next_bs02 = load_b_subtile_ni_regs(next_b, next_scales_ready, 0, 2) - mfma_4n(_acc_idx(3, 2, 0), a12, as12, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + next_a03, next_as03 = load_a_subtile_mi_regs(next_a, next_scales, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, as13, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - next_b03, next_bs03 = load_b_subtile_ni_regs(next_b, next_scales_ready, 0, 3) - mfma_4n(_acc_idx(3, 3, 0), a13, as13, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + next_b00, next_bs00 = load_b_subtile_ni_regs(next_b, next_scales, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, as10, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - hot_loop_scheduler_q_prefetch_4n() + next_b01, next_bs01 = load_b_subtile_ni_regs(next_b, next_scales, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, as11, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - next_a0_regs = ( - next_a00, - next_a01, - next_a02, - next_a03, - next_as00, - next_as01, - next_as02, - next_as03, - ) - next_b0_regs = ( - next_b00, - next_b01, - next_b02, - next_b03, - next_bs00, - next_bs01, - next_bs02, - next_bs03, - ) + next_b02, next_bs02 = load_b_subtile_ni_regs(next_b, next_scales, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, as12, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - return next_a0_regs, next_b0_regs, next_scales_ready, refill_scales + next_b03, next_bs03 = load_b_subtile_ni_regs(next_b, next_scales, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, as13, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - def hk_one_k_tail_with_next( - cur_a, cur_b, next_a, next_b, a0_regs, b0_regs, cur_scales, next_scales - ): - barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + hot_loop_scheduler_q_prefetch_4n() - a00, a01, a02, a03, as00, as01, as02, as03 = a0_regs - b00, b01, b02, b03, bs00, bs01, bs02, bs03 = b0_regs + next_a0_regs = ( + next_a00, + next_a01, + next_a02, + next_a03, + next_as00, + next_as01, + next_as02, + next_as03, + ) + next_b0_regs = ( + next_b00, + next_b01, + next_b02, + next_b03, + next_bs00, + next_bs01, + next_bs02, + next_bs03, + ) - b10, bs10 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 0) - b11, bs11 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 1) - b12, bs12 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 2) - b13, bs13 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 3) + return next_a0_regs, next_b0_regs + + def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs, cur_scales): + barrier(vmcnt=0, lgkmcnt=0) + + a00, a01, a02, a03, as00, as01, as02, as03 = a0_regs + b00, b01, b02, b03, bs00, bs01, bs02, bs03 = b0_regs + + # Materialize the remaining final-page A/B fragments once. The + # subsequent schedule is entirely register/AGPR traffic. + b10, bs10 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 0) + b11, bs11 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 1) + b12, bs12 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 2) + b13, bs13 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 3) + + rocdl.sched_barrier(0) + barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10, as10 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 0) + a11, as11 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 1) + a12, as12 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 2) + a13, as13 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 3) + + rocdl.sched_barrier(0) + barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) + a_scales = (as00, as01, as02, as03, as10, as11, as12, as13) + b_frags = (b00, b01, b02, b03, b10, b11, b12, b13) + b_scales = (bs00, bs01, bs02, bs03, bs10, bs11, bs12, bs13) + + # Rolling final-page epilogue. + # + # Finalize accumulators in their own physical AGPR slots, but delay + # each AGPR read/store until several independent final MFMAs have + # been issued. + # + # MFMA 0, MFMA 1, MFMA 2, MFMA 3, drain 0, + # MFMA 4, drain 1, MFMA 5, drain 2, ... + # + # The buffer stores are only issued here; they may remain in flight + # while later MFMAs and accumulator drains continue. + FINAL_EPILOGUE_DEPTH = 4 + pending = [] + + for old_acc_idx in range_constexpr(ACCS_PER_WAVE): + subtile_id = old_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = old_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE + + a_frag_idx = sm * MFMA_M_PER_SUBTILE + mi + b_frag_idx = sn * MFMA_N_PER_SUBTILE + ni + + # Final MFMA remains in-place. The logical accumulator's own + # AGPR slot is unique and cannot conflict with another pending + # result, so no ad-hoc physical-slot permutation is needed. + pinned_final_mfma( + old_acc_idx, + old_acc_idx, + a_frags[a_frag_idx], + b_frags[b_frag_idx], + a_scales[a_frag_idx], + b_scales[b_frag_idx], + mi, + ni, + ) + pending.append(old_acc_idx) + + # Drain the oldest completed result only after enough newer + # independent MFMAs have supplied the MFMA->AGPR-read spacing. + if len(pending) == FINAL_EPILOGUE_DEPTH: + drain_acc_idx = pending.pop(0) + acc = read_physical_accumulator_slot(drain_acc_idx) + store_acc_vector_for_logical_idx(drain_acc_idx, acc) + + # Flush the final results after all final-page MFMAs have issued. + for drain_acc_idx in pending: + acc = read_physical_accumulator_slot(drain_acc_idx) + store_acc_vector_for_logical_idx(drain_acc_idx, acc) - mfma_4n(_acc_idx(0, 0, 0), a00, as00, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - mfma_4n(_acc_idx(0, 1, 0), a01, as01, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - mfma_4n(_acc_idx(0, 2, 0), a02, as02, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - mfma_4n(_acc_idx(0, 3, 0), a03, as03, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + # Prologue: stage K0/K1 data into ping-pong LDS pages. Scales are not staged in + # LDS: As/Bs are already MFMA-ready preshuffled packed uint32 [K128, row], + # and load_scale_tile returns the current wave's scale operands in VGPRs. - rocdl.sched_barrier(0) - barrier(lgkmcnt=0) - rocdl.sched_barrier(0) + # Load scales first, so that they become the oldest VMEM ops. + scales0 = load_scale_tile(fx.Index(0)) + scales1 = load_scale_tile(fx.Index(1)) - a10, as10 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 0) - a11, as11 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 1) - a12, as12 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 2) - a13, as13 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 3) + stage_a_subtile(fx.Index(0), 0, lds_a0) + stage_b_subtile(fx.Index(0), 0, lds_b0) + stage_b_subtile(fx.Index(0), 1, lds_b0) + stage_a_subtile(fx.Index(0), 1, lds_a0) - mfma_4n(_acc_idx(1, 0, 0), a00, as00, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - mfma_4n(_acc_idx(1, 1, 0), a01, as01, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - mfma_4n(_acc_idx(1, 2, 0), a02, as02, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - mfma_4n(_acc_idx(1, 3, 0), a03, as03, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + stage_a_subtile(fx.Index(BLOCK_K), 0, lds_a1) + stage_b_subtile(fx.Index(BLOCK_K), 0, lds_b1) + stage_b_subtile(fx.Index(BLOCK_K), 1, lds_b1) + stage_a_subtile(fx.Index(BLOCK_K), 1, lds_a1) rocdl.sched_barrier(0) - barrier(vmcnt=LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 4 * LOAD_PASSES_B_SUBTILE) rocdl.sched_barrier(0) - next_a00, next_as00 = load_a_subtile_mi_regs(next_a, next_scales, 0, 0) - mfma_4n(_acc_idx(2, 0, 0), a10, as10, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - - next_a01, next_as01 = load_a_subtile_mi_regs(next_a, next_scales, 0, 1) - mfma_4n(_acc_idx(2, 1, 0), a11, as11, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - - next_a02, next_as02 = load_a_subtile_mi_regs(next_a, next_scales, 0, 2) - mfma_4n(_acc_idx(2, 2, 0), a12, as12, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - - next_a03, next_as03 = load_a_subtile_mi_regs(next_a, next_scales, 0, 3) - mfma_4n(_acc_idx(2, 3, 0), a13, as13, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - - next_b00, next_bs00 = load_b_subtile_ni_regs(next_b, next_scales, 0, 0) - mfma_4n(_acc_idx(3, 0, 0), a10, as10, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - - next_b01, next_bs01 = load_b_subtile_ni_regs(next_b, next_scales, 0, 1) - mfma_4n(_acc_idx(3, 1, 0), a11, as11, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - - next_b02, next_bs02 = load_b_subtile_ni_regs(next_b, next_scales, 0, 2) - mfma_4n(_acc_idx(3, 2, 0), a12, as12, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - - next_b03, next_bs03 = load_b_subtile_ni_regs(next_b, next_scales, 0, 3) - mfma_4n(_acc_idx(3, 3, 0), a13, as13, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - - hot_loop_scheduler_q_prefetch_4n() - - next_a0_regs = ( - next_a00, - next_a01, - next_a02, - next_a03, - next_as00, - next_as01, - next_as02, - next_as03, - ) - next_b0_regs = ( - next_b00, - next_b01, - next_b02, - next_b03, - next_bs00, - next_bs01, - next_bs02, - next_bs03, - ) - - return next_a0_regs, next_b0_regs - - def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs, cur_scales): - barrier(vmcnt=0, lgkmcnt=0) + # scales0 is already MFMA-ready; no byte extraction or broadcast is needed. + # Keep the hot loop consistent for k=0 and k>0: + # K0 is consumed directly. K1 MFMA-ready scales are carried as + # prev_refill_scales and become next_scales_ready at loop entry. - a00, a01, a02, a03, as00, as01, as02, as03 = a0_regs - b00, b01, b02, b03, bs00, bs01, bs02, bs03 = b0_regs - - # Materialize the remaining final-page A/B fragments once. The - # subsequent schedule is entirely register/AGPR traffic. - b10, bs10 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 0) - b11, bs11 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 1) - b12, bs12 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 2) - b13, bs13 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 3) - - rocdl.sched_barrier(0) - barrier(lgkmcnt=0) - rocdl.sched_barrier(0) - - a10, as10 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 0) - a11, as11 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 1) - a12, as12 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 2) - a13, as13 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 3) + # Seed the carried-register pipeline with K0 A-top. In later steady-state + # iterations, Q2/Q3 of the preceding iteration prefetch the next tile's + # A-top and B-left register tiles before their LDS half-pages are reused. + a0_regs = load_a_subtile_regs(lds_a0, scales0, 0) rocdl.sched_barrier(0) - barrier(lgkmcnt=0) + barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) rocdl.sched_barrier(0) - a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) - a_scales = (as00, as01, as02, as03, as10, as11, as12, as13) - b_frags = (b00, b01, b02, b03, b10, b11, b12, b13) - b_scales = (bs00, bs01, bs02, bs03, bs10, bs11, bs12, bs13) - - # Rolling final-page epilogue. - # - # Finalize accumulators in their own physical AGPR slots, but delay - # each AGPR read/store until several independent final MFMAs have - # been issued. - # - # MFMA 0, MFMA 1, MFMA 2, MFMA 3, drain 0, - # MFMA 4, drain 1, MFMA 5, drain 2, ... - # - # The buffer stores are only issued here; they may remain in flight - # while later MFMAs and accumulator drains continue. - FINAL_EPILOGUE_DEPTH = 4 - pending = [] - - for old_acc_idx in range_constexpr(ACCS_PER_WAVE): - subtile_id = old_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) - local_idx = old_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) - sm = subtile_id // 2 - sn = subtile_id % 2 - mi = local_idx // MFMA_N_PER_SUBTILE - ni = local_idx % MFMA_N_PER_SUBTILE - - a_frag_idx = sm * MFMA_M_PER_SUBTILE + mi - b_frag_idx = sn * MFMA_N_PER_SUBTILE + ni - - # Final MFMA remains in-place. The logical accumulator's own - # AGPR slot is unique and cannot conflict with another pending - # result, so no ad-hoc physical-slot permutation is needed. - pinned_final_mfma( - old_acc_idx, - old_acc_idx, - a_frags[a_frag_idx], - b_frags[b_frag_idx], - a_scales[a_frag_idx], - b_scales[b_frag_idx], - mi, - ni, - ) - pending.append(old_acc_idx) - - # Drain the oldest completed result only after enough newer - # independent MFMAs have supplied the MFMA->AGPR-read spacing. - if len(pending) == FINAL_EPILOGUE_DEPTH: - drain_acc_idx = pending.pop(0) - acc = read_physical_accumulator_slot(drain_acc_idx) - store_acc_vector_for_logical_idx(drain_acc_idx, acc) - - # Flush the final results after all final-page MFMAs have issued. - for drain_acc_idx in pending: - acc = read_physical_accumulator_slot(drain_acc_idx) - store_acc_vector_for_logical_idx(drain_acc_idx, acc) - - # Prologue: stage K0/K1 data into ping-pong LDS pages. Scales are not staged in - # LDS: As/Bs are already MFMA-ready preshuffled packed uint32 [K128, row], - # and load_scale_tile returns the current wave's scale operands in VGPRs. - - # Load scales first, so that they become the oldest VMEM ops. - scales0 = load_scale_tile(fx.Index(0)) - scales1 = load_scale_tile(fx.Index(1)) - - stage_a_subtile(fx.Index(0), 0, lds_a0) - stage_b_subtile(fx.Index(0), 0, lds_b0) - stage_b_subtile(fx.Index(0), 1, lds_b0) - stage_a_subtile(fx.Index(0), 1, lds_a0) - - stage_a_subtile(fx.Index(BLOCK_K), 0, lds_a1) - stage_b_subtile(fx.Index(BLOCK_K), 0, lds_b1) - stage_b_subtile(fx.Index(BLOCK_K), 1, lds_b1) - stage_a_subtile(fx.Index(BLOCK_K), 1, lds_a1) - - rocdl.sched_barrier(0) - barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 4 * LOAD_PASSES_B_SUBTILE) - rocdl.sched_barrier(0) - - # scales0 is already MFMA-ready; no byte extraction or broadcast is needed. - # Keep the hot loop consistent for k=0 and k>0: - # K0 is consumed directly. K1 MFMA-ready scales are carried as - # prev_refill_scales and become next_scales_ready at loop entry. - - # Seed the carried-register pipeline with K0 A-top. In later steady-state - # iterations, Q2/Q3 of the preceding iteration prefetch the next tile's - # A-top and B-left register tiles before their LDS half-pages are reused. - a0_regs = load_a_subtile_regs(lds_a0, scales0, 0) - - rocdl.sched_barrier(0) - barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) - rocdl.sched_barrier(0) - - # Complete the K0 carried-register seed with B-left. - b0_regs = load_b_subtile_regs(lds_b0, scales0, 0) - - # Main HK loop: exactly one logical K128 per iteration. - # Even k consumes and refills LDS0; odd k does the same for LDS1. - # Scale tiles follow the same K128 progression but remain in VGPRs. - refill_scales = scales1 # K1 scales become the next ready scale tile at loop entry - for k128 in range_constexpr(NUM_K_TILES - 2): - if (k128 % 2) == 0: - a0_regs, b0_regs, scales1, refill_scales = hk_one_k_with_refill( - k128, + # Complete the K0 carried-register seed with B-left. + b0_regs = load_b_subtile_regs(lds_b0, scales0, 0) + + # Main HK loop: exactly one logical K128 per iteration. + # Even k consumes and refills LDS0; odd k does the same for LDS1. + # Scale tiles follow the same K128 progression but remain in VGPRs. + refill_scales = scales1 # K1 scales become the next ready scale tile at loop entry + for k128 in range_constexpr(NUM_K_TILES - 2): + if (k128 % 2) == 0: + a0_regs, b0_regs, scales1, refill_scales = hk_one_k_with_refill( + k128, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + scales0, + refill_scales, + ) + else: + a0_regs, b0_regs, scales0, refill_scales = hk_one_k_with_refill( + k128, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + scales1, + refill_scales, + ) + + # Common two-page tail. The penultimate tile still uses the Q2/Q3 + # carry-prefetch scheduler to prepare A-top/B-left for the final tile, + # but it performs no K+2 data or scale refill. The final tile performs + # compute only. After the steady loop, a0_regs/b0_regs belong to the + # next tile to consume, while refill_scales belongs to the page most + # recently refilled; therefore tail page order depends on parity: + # even NUM_K_TILES: consume LDS0 then final LDS1 + # odd NUM_K_TILES: consume LDS1 then final LDS0 + if (NUM_K_TILES % 2) == 0: + scales1 = refill_scales + a0_regs, b0_regs = hk_one_k_tail_with_next( lds_a0, lds_b0, lds_a1, lds_b1, - lds_a0, - lds_b0, a0_regs, b0_regs, scales0, - refill_scales, + scales1, ) + hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs, scales1) else: - a0_regs, b0_regs, scales0, refill_scales = hk_one_k_with_refill( - k128, + scales0 = refill_scales + a0_regs, b0_regs = hk_one_k_tail_with_next( lds_a1, lds_b1, lds_a0, lds_b0, - lds_a1, - lds_b1, a0_regs, b0_regs, scales1, - refill_scales, + scales0, ) - - # Common two-page tail. The penultimate tile still uses the Q2/Q3 - # carry-prefetch scheduler to prepare A-top/B-left for the final tile, - # but it performs no K+2 data or scale refill. The final tile performs - # compute only. After the steady loop, a0_regs/b0_regs belong to the - # next tile to consume, while refill_scales belongs to the page most - # recently refilled; therefore tail page order depends on parity: - # even NUM_K_TILES: consume LDS0 then final LDS1 - # odd NUM_K_TILES: consume LDS1 then final LDS0 - if (NUM_K_TILES % 2) == 0: - scales1 = refill_scales - a0_regs, b0_regs = hk_one_k_tail_with_next( - lds_a0, - lds_b0, - lds_a1, - lds_b1, - a0_regs, - b0_regs, - scales0, - scales1, - ) - hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs, scales1) + hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs, scales0) + + # Dispatch. Dense runs one XCD-swizzled tile. Grouped runs a persistent + # loop over precomputed (pid_m, pid_n, expert) descriptors, selecting the + # schedule at compile time (``sched``): + # "worksteal": each iteration claims the next global slot via one atomic + # fetch-add (dynamic load balancing across CUs). + # "static": each workgroup owns a fixed grid-stride slice of slots + # (slot = blockIdx.x, blockIdx.x + gridDim.x, ...). No atomic + # or LDS broadcast, so it is deadlock-proof; used as a + # heuristic fast path and a correctness reference. + def _run_grouped_tile(slot): + pid_m = _desc_field(slot, fx.Int32(0)) + pid_n = _desc_field(slot, fx.Int32(1)) + expert = _desc_field(slot, fx.Int32(2)) + process_tile(pid_m, pid_n, expert) + + if const_expr(group): + if const_expr(sched == "static"): + stride = fx.Int32(gpu.grid_dim.x) + slot = fx.Int32(gpu.block_id("x")) + while slot < num_tiles: + _run_grouped_tile(slot) + slot = slot + stride + else: + # Single persistent loop, mirroring mega_moe_stage1: claim the next + # slot at the TOP of the loop body (atomic + barrier + LDS broadcast, + # so every thread sees the same slot), run the tile only when the + # slot is in range, and derive the loop-continue flag from that same + # broadcast slot. Keeping the barrier at one site with a uniform + # continue condition guarantees convergent barriers. + active = fx.Int32(1) == fx.Int32(1) + while active: + slot = _claim_next_slot() + in_range = slot < num_tiles + if in_range: + _run_grouped_tile(slot) + active = in_range else: - scales0 = refill_scales - a0_regs, b0_regs = hk_one_k_tail_with_next( - lds_a1, - lds_b1, - lds_a0, - lds_b0, - a0_regs, - b0_regs, - scales1, - scales0, - ) - hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs, scales0) + pid_m, pid_n = xcd_swizzle(num_blocks_m, num_blocks_n) + process_tile(pid_m, pid_n, fx.Int32(0)) @flyc.jit def launch_gemm( @@ -1500,10 +1631,18 @@ def launch_gemm( Aux: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32, + num_groups: fx.Int32 = 0, + work_queue: fx.Tensor = None, + work_counter: fx.Tensor = None, stream: fx.Stream = fx.Stream(None), ): # The integration only dispatches aligned shapes; no partial-tile masking exists. + # Dense launch: one workgroup per output tile. Grouped launch: a fixed + # num-CU grid of persistent workgroups that steal tiles from the atomic + # work queue (num_tiles spans the packed [M_total, N] output). grid_x = (c_m // BLOCK_M) * (c_n // BLOCK_N) + if const_expr(group): + grid_x = get_num_cus() kernel_gemm( A, As, @@ -1514,6 +1653,9 @@ def launch_gemm( Aux, c_m, c_n, + num_groups, + work_queue, + work_counter, value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, ).launch(grid=(grid_x, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) @@ -1532,6 +1674,11 @@ def do_gemm( epilogue: str = "DEFAULT", bias: torch.Tensor = None, aux: torch.Tensor = None, + group: bool = False, + num_groups: int = 0, + work_queue: fx.Tensor = None, + work_counter: fx.Tensor = None, + sched: str = "worksteal", ): """Launch one cached compile-time MXFP8 layout specialization. @@ -1551,6 +1698,12 @@ def do_gemm( raise ValueError(f"Unsupported MXFP8 kernel layout: {layout}") assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" + + # Grouped TN: B is packed [G*N, K], so B's leading dim is num_groups * N_out. + # The kernel launches over the per-expert output width (C's N) and re-derives + # the packed B / B-scale stride from num_groups internally. + if group: + N_runtime = C.shape[-1] supported_fp8_dtypes = (torch.float8_e4m3fn, torch.float8_e5m2) assert A.dtype in supported_fp8_dtypes, f"unsupported A FP8 dtype: {A.dtype}" assert B.dtype in supported_fp8_dtypes, f"unsupported B FP8 dtype: {B.dtype}" @@ -1574,7 +1727,10 @@ def do_gemm( ) expected_as = (K_runtime // _BLOCK_K, M_runtime) - expected_bs = (K_runtime // _BLOCK_K, N_runtime) + # Grouped B / B-scale are packed over all experts (num_groups * N_out); the + # kernel indexes them by expert. N_runtime here is the per-expert output N. + n_packed = N_runtime * num_groups if group else N_runtime + expected_bs = (K_runtime // _BLOCK_K, n_packed) assert As.dtype == torch.int32, f"As dtype {As.dtype} != torch.int32 packed scales" assert Bs.dtype == torch.int32, f"Bs dtype {Bs.dtype} != torch.int32 packed scales" assert tuple(As.shape) == expected_as, f"As shape {tuple(As.shape)} != {expected_as}" @@ -1648,6 +1804,15 @@ def do_gemm( else: Aux_arg = torch.zeros(1, dtype=C.dtype, device=A.device) + # Grouped runtime tensors are required by the kernel signature; the dense + # binary never dereferences them but still needs placeholder buffers. + if group: + work_queue_arg = work_queue.contiguous().view(-1) + work_counter_arg = work_counter.contiguous().view(-1) + else: + work_queue_arg = torch.zeros(1, dtype=torch.int32, device=A.device) + work_counter_arg = torch.zeros(1, dtype=torch.int32, device=A.device) + _cached_launch( K_runtime, A.dtype, @@ -1655,6 +1820,8 @@ def do_gemm( C.dtype, layout, epilogue, + group, + sched, )( A_arg, As_arg, @@ -1665,6 +1832,9 @@ def do_gemm( Aux_arg, M_runtime, N_runtime, + num_groups, + work_queue_arg, + work_counter_arg, stream=stream, ) @@ -1677,6 +1847,8 @@ def _cached_launch( output_dtype: torch.dtype, layout: str, epilogue: str = "DEFAULT", + group: bool = False, + sched: str = "worksteal", ): """Cache independent TN/NN/NT binaries with no runtime layout argument.""" return _compile_kernel( @@ -1686,6 +1858,8 @@ def _cached_launch( output_dtype, layout, epilogue, + group, + sched, ) @@ -1727,12 +1901,18 @@ def mxfp8_matmul( epilogue: str = "DEFAULT", bias: torch.Tensor = None, aux: torch.Tensor = None, + group: bool = False, + num_groups: int = 0, + work_queue: fx.Tensor = None, + work_counter: fx.Tensor = None, + sched: str = "worksteal", ): """Normalize scale orientation and launch a compile-time layout binary. Wrapper-visible contracts: TN: a [M,K], b [N,K], scales [M,K/32] and [N,K/32] + (group): a [M_total,K], b [G*N,K] scales [M_total,K/32] and [G*N,K/32] NN: a [M,K], b [K,N], scales [M,K/32] and [K/32,N] NT: a [K,M], b [K,N], scales [K/32,M] and [K/32,N] @@ -1740,6 +1920,10 @@ def mxfp8_matmul( K-major payloads and use ``ds_read_b64_tr_b8`` inside their compile-time-specialized kernels. """ + + if group and layout != "TN": + raise ValueError("FlyDSL MXFP8 Group GEMM only supports the TN layout") + if layout not in ("TN", "NN", "NT"): raise ValueError(f"Unsupported MXFP8 kernel layout: {layout}") @@ -1759,8 +1943,22 @@ def mxfp8_matmul( raise ValueError( f"Incompatible MXFP8 {layout} operands: A{tuple(a.shape)} and B{tuple(b.shape)}" ) - if tuple(D.shape) != (m, n): - raise ValueError(f"D shape {tuple(D.shape)} != expected {(m, n)}") + + # Grouped TN packs the per-expert weights along N, so ``n`` (and the B scale) + # span ``num_groups * N_out``. The output and the kernel's ``c_n`` use the + # per-expert ``n_out``; the kernel re-derives the packed B-scale stride + # (num_groups * c_n) internally. + if group: + if n % num_groups != 0: + raise ValueError( + f"Grouped MXFP8 packed N={n} is not divisible by num_groups={num_groups}" + ) + n_out = n // num_groups + else: + n_out = n + + if tuple(D.shape) != (m, n_out): + raise ValueError(f"D shape {tuple(D.shape)} != expected {(m, n_out)}") if k % SCALE_GROUP_SIZE != 0: raise FlyDSLUnsupportedError( f"K={k} must be divisible by MXFP8 scale group size {SCALE_GROUP_SIZE}" @@ -1843,12 +2041,17 @@ def mxfp8_matmul( a_scale_hk, b_kernel, b_scale_hk, - D.view(m, n), + D.view(m, n_out), layout=layout, stream=stream, epilogue=epilogue, bias=bias, - aux=aux.view(m, n) if aux is not None else None, + aux=aux.view(m, n_out) if aux is not None else None, + group=group, + num_groups=num_groups, + work_queue=work_queue, + work_counter=work_counter, + sched=sched, ) return D diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 059c3a4ce6..fe72c63a97 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -577,6 +577,106 @@ def _is_blockwise_fp8_triton_grouped_gemm_supported( ) ) + @staticmethod + def _is_flydsl_mxfp8_grouped_gemm_supported( + *, + fp8, + recipe, + use_bias, + backward_override, + cpu_offloading, + save_original_input, + debug, + unpad_output, + actual_m_splits, + in_features, + out_features, + ) -> bool: + """FlyDSL persistent work-stealing MXFP8 grouped GEMM: HIP, gfx950, env, MXFP8 recipe. + + Opt-in via ``NVTE_USE_GROUPED_GEMM_FLYDSL=1``. The kernel is a TN-only packed + launch (see :func:`_run_mxfp8_grouped`); N (out_features) and K (in_features) + must be 128-aligned for the tile/scale geometry, and the same feature-set + exclusions as the blockwise-triton path apply. + """ + return ( + IS_HIP_EXTENSION + and os.getenv("NVTE_USE_GROUPED_GEMM_FLYDSL", "0") == "1" + and get_device_compute_capability() == (9, 5) + and fp8 + and recipe is not None + and recipe.mxfp8() + and in_features % 128 == 0 + and out_features % 128 == 0 + and not ( + use_bias + or backward_override + or cpu_offloading + or save_original_input + or debug + or unpad_output + or actual_m_splits is not None + ) + ) + + @staticmethod + def _grouped_gemm_flydsl_mxfp8( + *, + inputmats, + weights_fp8, + out, + m_splits, + m_splits_tensor, + ): + """Dispatch the FlyDSL MXFP8 grouped GEMM over already-quantized operands. + + Called inline at the forward GEMM dispatch (like ``general_grouped_gemm_triton``), + riding off the quantization TE already performed -- ``inputmats`` are the + per-group MXFP8 activations from ``tex.split_quantize`` and ``weights_fp8`` the + per-expert MXFP8 weights. This first iteration does not quantize: it packs the + rowwise MXFP8 payloads into the ``[M_total, K]`` / ``[G*N, K]`` layout, derives + ``group_offs`` from ``m_splits``, and issues one packed TN launch via + :func:`_run_mxfp8_grouped`, writing into ``out``. + """ + from ..flydsl_kernels.gemm import _run_mxfp8_grouped + + num_gemms = len(inputmats) + # TN packed operands: activations concatenated along the token (M) axis, + # weights along the expert*N axis. Each MXFP8 operand carries its rowwise + # fp8 payload and E8M0 scale; concatenate both. + a_data = torch.cat([im._rowwise_data for im in inputmats], dim=0) + a_scale = torch.cat([im._rowwise_scale_inv for im in inputmats], dim=0) + b_data = torch.cat([w._rowwise_data for w in weights_fp8], dim=0) + b_scale = torch.cat([w._rowwise_scale_inv for w in weights_fp8], dim=0) + + # Preserve the native MXFP8 dtype so the wrapper reinterprets the uint8 + # payload to the correct fp8 view at the kernel boundary. + a_data._fp8_dtype = inputmats[0]._fp8_dtype + b_data._fp8_dtype = weights_fp8[0]._fp8_dtype + + n = weights_fp8[0]._rowwise_data.size(0) + k = weights_fp8[0]._rowwise_data.size(-1) + m_total = a_data.size(0) + + # group_offs: [num_gemms+1] exclusive scan over per-group token counts. + group_offs = torch.zeros(num_gemms + 1, dtype=torch.int64, device=out.device) + group_offs[1:] = torch.cumsum( + m_splits_tensor.to(device=out.device, dtype=torch.int64), 0 + ) + + _run_mxfp8_grouped( + a_data, + a_scale, + b_data, + b_scale, + out, + group_offs=group_offs, + m_total=m_total, + n=n, + k=k, + output_dtype=out.dtype, + ) + @staticmethod def _forward_blockwise_fp8_triton( ctx, @@ -981,6 +1081,23 @@ def forward( pow2_grad=blockwise_recipe.fp8_quant_bwd_grad.power_2_scale, ) + # FlyDSL MXFP8 grouped GEMM (ROCm gfx950) opt-in. Unlike the blockwise-triton + # path this does NOT quantize: it rides off TE's existing quantization and is + # dispatched inline at the forward GEMM (see below), like the triton grouped path. + use_grouped_gemm_flydsl = _GroupedLinear._is_flydsl_mxfp8_grouped_gemm_supported( + fp8=fp8, + recipe=blockwise_recipe, + use_bias=use_bias, + backward_override=backward_override, + cpu_offloading=cpu_offloading, + save_original_input=save_original_input, + debug=debug, + unpad_output=unpad_output, + actual_m_splits=actual_m_splits, + in_features=weights[0].size(-1), + out_features=weights[0].size(0), + ) + # Configure quantizers if save_original_input and isinstance(input_quantizers[0], Float8Quantizer): if FP8GlobalStateManager.get_fp8_recipe().custom(): @@ -1150,25 +1267,34 @@ def forward( use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator # Perform GEMM - if use_grouped_gemm_triton: - general_grouped_gemm_func = general_grouped_gemm_triton - kwargs = {"m_splits_tensor": m_splits_tensor} + if use_grouped_gemm_flydsl: + _GroupedLinear._grouped_gemm_flydsl_mxfp8( + inputmats=inputmats, + weights_fp8=weights_fp8, + out=out, + m_splits=m_splits, + m_splits_tensor=m_splits_tensor, + ) else: - general_grouped_gemm_func = general_grouped_gemm - kwargs = {} - general_grouped_gemm_func( - weights_fp8, - inputmats, - [out], - output_quantizers, - activation_dtype, - single_output=True, - m_splits=m_splits, - bias=biases, - use_bias=use_bias, - use_split_accumulator=use_split_accumulator, - **kwargs, - ) + if use_grouped_gemm_triton: + general_grouped_gemm_func = general_grouped_gemm_triton + kwargs = {"m_splits_tensor": m_splits_tensor} + else: + general_grouped_gemm_func = general_grouped_gemm + kwargs = {} + general_grouped_gemm_func( + weights_fp8, + inputmats, + [out], + output_quantizers, + activation_dtype, + single_output=True, + m_splits=m_splits, + bias=biases, + use_bias=use_bias, + use_split_accumulator=use_split_accumulator, + **kwargs, + ) output_unpadded = False if ( From e4fc6bd62b4118c79e7d4f3f6a00c5c388073b6b Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 30 Sep 2026 14:26:40 +0000 Subject: [PATCH 02/22] Add FlyDSL grouped MXFP8 dgrad (NN) backward path Extend the persistent grouped MXFP8 GEMM to the NN layout so GroupedLinear's dgrad (dX = dY @ W) routes through FlyDSL, mirroring the TN forward. The core MFMA pipeline is untouched: NN only adds a const_expr(group) expert offset on the transpose-read B path. Unlike TN forward (experts pack on B's output-N axis), NN's contraction axis is N_out, so the packed weight [G*N_out, K_in] stacks experts on the CONTRACTION axis. The expert offset therefore folds into the k_base term (expert*N_out) and the B-scale k128 term (expert*N_out/128), not the output-N base. Shape-contract handling in mxfp8_matmul/do_gemm is made layout-aware: NN's B contraction dim is G*N_out (not equal to A's per-expert N_out) and the output width K_in is not packed. _run_mxfp8_grouped gains a layout param; its descriptor-queue and work-stealing logic are layout-independent and reused. Backward plumbing: forward records ctx.use_flydsl_mxfp8_grouped; the shared backward branches on it to pack dY (rowwise) + W (columnwise, already ensured) and call the grouped NN launch. wgrad still uses the default path. dgrad is bitwise-equal to the default TE grouped MXFP8 backward on both the worksteal and static schedules; dense (non-grouped) NN is unaffected. Co-Authored-By: Claude --- .../flydsl_kernels/gemm/gemm_wrappers.py | 36 +++-- .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 148 ++++++++++++++---- .../pytorch/module/grouped_linear.py | 108 ++++++++++--- 3 files changed, 230 insertions(+), 62 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 8f581c6f9d..849f2399d6 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -1043,29 +1043,37 @@ def _run_mxfp8_grouped( n, k, output_dtype: torch.dtype, + layout: str = "TN", bias=None, ): - """Dispatch a packed MXFP8 grouped GEMM through the persistent kernel (TN only). + """Dispatch a packed MXFP8 grouped GEMM through the persistent kernel. - Operand contract (TN, single packed launch spanning all groups): + Two layouts, both single packed launches spanning all groups: + TN (forward, C = A @ B^T): a = packed activations [M_total, K] scales [M_total, K/32] b = packed weights [G*N, K] scales [G*N, K/32] - group_offs [G+1] exclusive scan of per-group row counts - - Each group ``g`` owns rows ``[group_offs[g], group_offs[g+1])`` of ``a`` and the - weight block ``b[g*N:(g+1)*N]``. ``group_offs`` is consumed on the host to build - the persistent kernel's per-tile ``(pid_m, pid_n, expert)`` descriptor queue, so - the kernel itself does no group lookup. This iteration assumes the operands arrive - already MXFP8-quantized and packed -- it takes the pre-selected rowwise payloads - directly and performs no quantization, packing, or BLAS operand-ownership swap - (unlike :func:`_run_mxfp8`). + output width n = N_out (per expert); experts pack on B's output-N axis. + + NN (dgrad, dX = dY @ W): + a = packed grad_output [M_total, N_out] scales [M_total, N_out/32] + b = packed weights [G*N_out, K_in] scales [G*N_out/32, K_in] + output width n = K_in; experts pack on B's contraction (N_out) axis. + + In both, ``k`` is A's per-expert contraction (TN: K_in; NN: N_out) and ``n`` is + the output width. Each group ``g`` owns activation/grad rows + ``[group_offs[g], group_offs[g+1])``; ``group_offs`` is consumed on the host to + build the persistent kernel's per-tile ``(pid_m, pid_n, expert)`` descriptor + queue, so the kernel does no group lookup. Operands arrive already + MXFP8-quantized and packed -- no quantization, no BLAS ownership swap. """ + if layout not in ("TN", "NN"): + raise FlyDSLUnsupportedError(f"FlyDSL MXFP8 grouped GEMM supports TN/NN, got {layout}") if k % 32 != 0: raise FlyDSLUnsupportedError(f"K={k} must be divisible by MXFP8 scale group size 32") if n % 128 != 0: raise FlyDSLUnsupportedError( - f"FlyDSL MXFP8 grouped GEMM requires N divisible by 128, got N={n}" + f"FlyDSL MXFP8 grouped GEMM requires output width N divisible by 128, got N={n}" ) a_fp8_dtype = getattr(a_data, "_fp8_dtype", None) @@ -1083,7 +1091,7 @@ def _run_mxfp8_grouped( shape=torch.Size((m_total, n)), dtype=output_dtype, device=a_data.device, - backend_name="MXFP8 grouped TN", + backend_name=f"MXFP8 grouped {layout}", ) # Persistent-kernel work distribution: @@ -1130,7 +1138,7 @@ def _run_mxfp8_grouped( b_data, b_scale, D.view(m_total, n), - layout="TN", + layout=layout, epilogue=epilogue, bias=bias_arg, group=True, diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index b06fc1443a..66f2d76c56 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -616,16 +616,47 @@ def process_tile(pid_m, pid_n, expert): bx_m_idx = fx.Index(bx_m) by_n_idx = fx.Index(by_n) - # B-operand column base and B-scale row stride. In the dense path these - # collapse to the C column base / c_n; in the grouped path they carry the - # expert offset (from the descriptor) into the packed [G*N, K] weight and - # [K128, G*N] scale. + # Per-expert offsets into the packed B operand / B-scale. Where the + # expert index lands depends on the layout, because the packed weight + # stacks experts on B's *first* logical axis, which is a different + # kernel axis for TN vs the transpose-read (NN/NT) case: + # + # TN (forward, C = A @ B^T): B is packed [G*N, K], experts on the + # OUTPUT-N axis. Expert offset folds into the B column base + # (by_n_b_idx = expert*N + n) and the B-scale row stride grows to + # G*N (c_n_bs). The B-scale k128 term is untouched. + # + # NN (dgrad, C = A @ B where A=dY[M,N_out], B=W[N_out,K_in]): the + # kernel contraction axis IS N_out, and B is packed [G*N_out, K_in] + # with experts on that CONTRACTION axis. Expert offset instead folds + # into the contraction (k_base) term as expert*N_out elements + # (b_expert_k_off), and into the B-scale's k128 (contraction-128) + # term as expert*(N_out/128) (bs_expert_k128_off). The B column base + # and B-scale stride stay per-expert (c_n = K_in, unchanged). + # + # Dense (group is False) zeroes every expert term. if const_expr(group): - by_n_b_idx = fx.Index(expert * fx.Int32(c_n) + by_n) - c_n_bs = num_groups * c_n + if const_expr(b_transpose_read): + # NN/NT: expert on the contraction axis. The kernel's + # compile-time ``K`` IS the contraction dim (= N_out for dgrad), + # so the packed weight [G*N_out, K_in] shifts by expert*K + # elements along k_base, and its scale [G*N_out/128, K_in] by + # expert*NUM_K_TILES along the k128 term. + by_n_b_idx = by_n_idx + c_n_bs = c_n + b_expert_k_off = fx.Index(expert * fx.Int32(K)) + bs_expert_k128_off = expert * fx.Int32(NUM_K_TILES) + else: + # TN: expert on the output-N axis. + by_n_b_idx = fx.Index(expert * fx.Int32(c_n) + by_n) + c_n_bs = num_groups * c_n + b_expert_k_off = fx.Index(0) + bs_expert_k128_off = fx.Int32(0) else: by_n_b_idx = by_n_idx c_n_bs = c_n + b_expert_k_off = fx.Index(0) + bs_expert_k128_off = fx.Int32(0) wave_id = tx_i32 // fx.Int32(WARP_SIZE) lane = tx_i32 % fx.Int32(WARP_SIZE) @@ -838,7 +869,9 @@ def load_a_scale_row(k128, row): return fx.memref_load_vec(reg)[0] def load_b_scale_row(k128, row): - off = k128 * c_n_bs_idx + by_n_b_idx + row + # bs_expert_k128_off shifts the contraction-128 base to this tile's + # expert scale block for grouped NN/NT; 0 for dense and grouped TN. + off = (k128 + fx.Index(bs_expert_k128_off)) * c_n_bs_idx + by_n_b_idx + row reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Int32) fx.copy(scale_ld_atom, fx.slice(bs_div, (None, fx.Int32(off))), reg) return fx.memref_load_vec(reg)[0] @@ -873,9 +906,16 @@ def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): ) def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): + # b_expert_k_off shifts the contraction base to this tile's expert + # weight block for the grouped transpose-read (NN/NT) path; it is 0 + # for dense and for grouped TN (which offsets via by_n_b_idx). b_g2s.load_one( lds_b[subtile], - fx.Int32(_b_global_base(k_base, subtile, c_n, by_n_b_idx)), + fx.Int32( + _b_global_base( + k_base + b_expert_k_off, subtile, c_n, by_n_b_idx + ) + ), pass_in_subtile, ) @@ -1697,11 +1737,19 @@ def do_gemm( else: raise ValueError(f"Unsupported MXFP8 kernel layout: {layout}") - assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" + # Grouped NN packs B along the contraction axis, so B's contraction dim is + # G x A's (Kb_runtime = G*N_out vs K_runtime = N_out). Every other case has + # them equal. + if group and layout == "NN": + assert Kb_runtime == num_groups * K_runtime, ( + f"grouped NN B.K={Kb_runtime} != num_groups*A.K={num_groups * K_runtime}" + ) + else: + assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" - # Grouped TN: B is packed [G*N, K], so B's leading dim is num_groups * N_out. - # The kernel launches over the per-expert output width (C's N) and re-derives - # the packed B / B-scale stride from num_groups internally. + # Grouped: the kernel launches over the per-expert output width, which is + # always C's last dim (TN: N_out; NN: K_in). The kernel re-derives the packed + # B / B-scale stride from num_groups internally. if group: N_runtime = C.shape[-1] supported_fp8_dtypes = (torch.float8_e4m3fn, torch.float8_e5m2) @@ -1727,10 +1775,15 @@ def do_gemm( ) expected_as = (K_runtime // _BLOCK_K, M_runtime) - # Grouped B / B-scale are packed over all experts (num_groups * N_out); the - # kernel indexes them by expert. N_runtime here is the per-expert output N. - n_packed = N_runtime * num_groups if group else N_runtime - expected_bs = (K_runtime // _BLOCK_K, n_packed) + # Packed B-scale shape. TN packs B along N, so the scale's row count is the + # per-K128 count and its column count is the packed G*N_out. NN packs B along + # the contraction axis, so the packed count is on the K128 (contraction) axis + # instead: [G*N_out/128, K_in] = (Kb_runtime // BLOCK_K, N_runtime). + if group and layout == "NN": + expected_bs = (Kb_runtime // _BLOCK_K, N_runtime) + else: + n_packed = N_runtime * num_groups if group else N_runtime + expected_bs = (K_runtime // _BLOCK_K, n_packed) assert As.dtype == torch.int32, f"As dtype {As.dtype} != torch.int32 packed scales" assert Bs.dtype == torch.int32, f"Bs dtype {Bs.dtype} != torch.int32 packed scales" assert tuple(As.shape) == expected_as, f"As shape {tuple(As.shape)} != {expected_as}" @@ -1914,6 +1967,9 @@ def mxfp8_matmul( TN: a [M,K], b [N,K], scales [M,K/32] and [N,K/32] (group): a [M_total,K], b [G*N,K] scales [M_total,K/32] and [G*N,K/32] NN: a [M,K], b [K,N], scales [M,K/32] and [K/32,N] + (group): a [M_total,Kc], b [G*Kc,N] scales [M_total,Kc/32] and + [G*Kc/32,N] -- grouped NN is the dgrad dX=dY@W: the kernel + contraction Kc is N_out and experts pack on that axis. NT: a [K,M], b [K,N], scales [K/32,M] and [K/32,N] TN consumes the selected rowwise payloads directly. NN and NT preserve @@ -1921,8 +1977,8 @@ def mxfp8_matmul( compile-time-specialized kernels. """ - if group and layout != "TN": - raise ValueError("FlyDSL MXFP8 Group GEMM only supports the TN layout") + if group and layout not in ("TN", "NN"): + raise ValueError("FlyDSL MXFP8 Group GEMM only supports the TN and NN layouts") if layout not in ("TN", "NN", "NT"): raise ValueError(f"Unsupported MXFP8 kernel layout: {layout}") @@ -1939,24 +1995,53 @@ def mxfp8_matmul( k, m = a.shape kb, n = b.shape - if kb != k: - raise ValueError( - f"Incompatible MXFP8 {layout} operands: A{tuple(a.shape)} and B{tuple(b.shape)}" - ) - - # Grouped TN packs the per-expert weights along N, so ``n`` (and the B scale) - # span ``num_groups * N_out``. The output and the kernel's ``c_n`` use the - # per-expert ``n_out``; the kernel re-derives the packed B-scale stride - # (num_groups * c_n) internally. - if group: + # Grouped GEMM is G INDEPENDENT matmuls stacked into one launch, not a dense + # A @ B_packed. Where the packed B axis lands, and hence which operand shape + # unpacks, depends on the layout: + # + # TN forward (C = A @ B^T): B is packed [G*N_out, K] on the OUTPUT-N axis. + # The kernel-K contraction (K_in) matches A's K, so kb == k. The output + # width unpacks: n = G*N_out -> n_out = N_out. D = [M_total, N_out]. + # (Per expert g: A_g[m_g,K] @ W_g^T[K,N_out] = D_g[m_g,N_out]; stacking the + # rows gives [M_total, N_out] -- the packed G*N_out rows are storage only, + # never the output width; a dense A@B_packed^T would give [M_total,G*N_out], + # G x too much compute, which grouped GEMM avoids.) + # + # NN dgrad (dX = dY @ W): the kernel CONTRACTION is N_out, and B=W is packed + # [G*N_out, K_in] on that CONTRACTION axis. So B's first dim kb = G*N_out + # is G x A's per-expert contraction k = N_out (kb != k by design), and the + # output width n = K_in is NOT packed. D = [M_total, K_in] = (m, n). + if group and layout == "TN": if n % num_groups != 0: raise ValueError( - f"Grouped MXFP8 packed N={n} is not divisible by num_groups={num_groups}" + f"Grouped MXFP8 TN packed N={n} is not divisible by num_groups={num_groups}" ) n_out = n // num_groups + expected_kb = k # contraction unpacked; B rows already per-K + elif group and layout == "NN": + if kb % num_groups != 0: + raise ValueError( + f"Grouped MXFP8 NN packed contraction={kb} not divisible by " + f"num_groups={num_groups}" + ) + if kb // num_groups != k: + raise ValueError( + f"Grouped MXFP8 NN per-expert contraction {kb // num_groups} " + f"!= A contraction {k}" + ) + n_out = n # output width (K_in) is not packed + expected_kb = kb # B rows are G*N_out, checked above else: n_out = n + expected_kb = k + + if kb != expected_kb: + raise ValueError( + f"Incompatible MXFP8 {layout} operands: A{tuple(a.shape)} and B{tuple(b.shape)}" + ) + # D is the full stacked output [M_total, n_out]; expert g writes its + # [m_g, n_out] slice into rows [group_offs[g], group_offs[g+1]). if tuple(D.shape) != (m, n_out): raise ValueError(f"D shape {tuple(D.shape)} != expected {(m, n_out)}") if k % SCALE_GROUP_SIZE != 0: @@ -1970,9 +2055,12 @@ def mxfp8_matmul( expected_a_scale = (m, k // SCALE_GROUP_SIZE) if layout == "TN": + # Grouped TN packs B along N, so the B scale spans n = G*N_out rows. expected_b_scale = (n, k // SCALE_GROUP_SIZE) else: - expected_b_scale = (k // SCALE_GROUP_SIZE, n) + # NN/NT B scale is [contraction/32, N]. Grouped NN packs B along the + # contraction axis, so use kb (= G*N_out) not the per-expert k. + expected_b_scale = (kb // SCALE_GROUP_SIZE, n) if tuple(a_scale.shape) != expected_a_scale: raise ValueError( diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index fe72c63a97..19add49d3b 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -677,6 +677,63 @@ def _grouped_gemm_flydsl_mxfp8( output_dtype=out.dtype, ) + @staticmethod + def _backward_flydsl_mxfp8_grouped( + *, + grad_output, + weights_for_dgrad, + dgrad, + m_splits, + m_splits_tensor, + ): + """FlyDSL grouped dgrad: dX = dY @ W through the packed NN kernel. + + Rides off backward's existing quantization -- ``grad_output`` are the + per-group MXFP8 activations (rowwise) and ``weights_for_dgrad`` the + per-expert MXFP8 weights with columnwise payloads ensured by the caller. + Packs dY as the NN kernel-A ``[M_total, N_out]`` and W (columnwise) as the + contraction-packed kernel-B ``[G*N_out, K_in]``, then issues one grouped NN + launch writing into ``dgrad`` ``[M_total, K_in]``. No quantization here. + """ + from ..flydsl_kernels.gemm import _run_mxfp8_grouped + + num_gemms = len(grad_output) + # NN kernel-A = dY, rowwise (contraction over N_out): concat on the token + # (M) axis. NN kernel-B = W, columnwise (transpose-read, contraction on the + # N_out axis): concat on the packed N_out axis. Each MXFP8 operand carries + # its selected payload + E8M0 scale. + a_data = torch.cat([g._rowwise_data for g in grad_output], dim=0) + a_scale = torch.cat([g._rowwise_scale_inv for g in grad_output], dim=0) + b_data = torch.cat([w._columnwise_data for w in weights_for_dgrad], dim=0) + b_scale = torch.cat([w._columnwise_scale_inv for w in weights_for_dgrad], dim=0) + + a_data._fp8_dtype = grad_output[0]._fp8_dtype + b_data._fp8_dtype = weights_for_dgrad[0]._fp8_dtype + + n_out = grad_output[0]._rowwise_data.size(-1) # per-expert contraction (N_out) + k_in = dgrad.size(-1) # output width (K_in) + m_total = a_data.size(0) + + # group_offs: [num_gemms+1] exclusive scan over per-group token counts. + group_offs = torch.zeros(num_gemms + 1, dtype=torch.int64, device=dgrad.device) + group_offs[1:] = torch.cumsum( + m_splits_tensor.to(device=dgrad.device, dtype=torch.int64), 0 + ) + + _run_mxfp8_grouped( + a_data, + a_scale, + b_data, + b_scale, + dgrad, + group_offs=group_offs, + m_total=m_total, + n=k_in, # NN output width is K_in + k=n_out, # NN per-expert contraction is N_out + output_dtype=dgrad.dtype, + layout="NN", + ) + @staticmethod def _forward_blockwise_fp8_triton( ctx, @@ -1322,6 +1379,9 @@ def forward( if is_grad_enabled: ctx.use_grouped_tensor_path = False + # FlyDSL grouped forward routes dgrad through the grouped NN path; the + # shared backward branches on this flag (wgrad still uses the default). + ctx.use_flydsl_mxfp8_grouped = use_grouped_gemm_flydsl ctx.weight_quantizers = weight_quantizers ctx.weights_shape_1 = weights[0].shape[1] @@ -1770,25 +1830,37 @@ def backward( rowwise_usage=quantizer.rowwise_usage, columnwise_usage=quantizer.columnwise_usage, ) - if ctx.use_grouped_gemm_triton: - general_grouped_gemm_func = general_grouped_gemm_triton - kwargs = {"m_splits_tensor": ctx.m_splits_tensor} + if getattr(ctx, "use_flydsl_mxfp8_grouped", False): + # dgrad dX = dY @ W through the FlyDSL grouped NN kernel, riding + # off the already-quantized grad_output (rowwise) and weights + # (columnwise, ensured above). wgrad still uses the default path. + _GroupedLinear._backward_flydsl_mxfp8_grouped( + grad_output=grad_output, + weights_for_dgrad=weights_for_dgrad, + dgrad=dgrad, + m_splits=ctx.m_splits, + m_splits_tensor=ctx.m_splits_tensor, + ) else: - general_grouped_gemm_func = general_grouped_gemm - kwargs = {} - general_grouped_gemm_func( - weights_for_dgrad, - grad_output, - [dgrad], - ctx.grad_input_quantizers, - ctx.activation_dtype, - single_output=True, - layout="NN", - m_splits=ctx.m_splits, - grad=True, - use_split_accumulator=dgrad_gemm_use_split_accumulator, - **kwargs, - ) + if ctx.use_grouped_gemm_triton: + general_grouped_gemm_func = general_grouped_gemm_triton + kwargs = {"m_splits_tensor": ctx.m_splits_tensor} + else: + general_grouped_gemm_func = general_grouped_gemm + kwargs = {} + general_grouped_gemm_func( + weights_for_dgrad, + grad_output, + [dgrad], + ctx.grad_input_quantizers, + ctx.activation_dtype, + single_output=True, + layout="NN", + m_splits=ctx.m_splits, + grad=True, + use_split_accumulator=dgrad_gemm_use_split_accumulator, + **kwargs, + ) if ( ctx.actual_m_splits is not None From 10f48f81d9c53d1f0034c9cb3d11f02d7e850db0 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 30 Sep 2026 15:32:02 +0000 Subject: [PATCH 03/22] Add FlyDSL grouped MXFP8 wgrad (NT) via uniform padding Complete the grouped MXFP8 backward: wgrad (dW = dY^T @ X per expert) now routes through a dedicated FlyDSL NT kernel, so GroupedLinear forward + dgrad + wgrad all run on FlyDSL. wgrad contracts over the token axis m_e (variable per expert). Uniform padding: every expert's tokens are padded up to a common M_pad and one workgroup contracts the full padded token axis for each output tile (a single-WG contraction, NOT split-K -- no cross-WG reduction). Padded rows are zero and contribute nothing to the contraction, so the NT MFMA pipeline runs byte-identical with no atomics or partial-tile masking. dY and X pack on the contraction axis to [G*M_pad, *]; dW packs on the output-M axis to [G*N_out, K_in]. The kernel is a compile-time `wgrad` specialization of the same generator (distinct binary, no pipeline duplication): both operand contraction bases and their scale k128 terms shift by expert*M_pad, and the store row shifts by expert*N_out (store_m_expert_off) since the operands use local per-expert coords while the output is packed. Reuses the persistent descriptor-queue / work-stealing driver via `grouped = group or wgrad`. Plumbing: _run_mxfp8_grouped_wgrad (pad/pack + descriptor build), _mxfp8_matmul_grouped_wgrad (NT scale pack + do_gemm(wgrad=True)), _wgrad_flydsl_mxfp8_grouped module method, and a backward dispatch gated on ctx.use_flydsl_mxfp8_grouped. fuse_wgrad_accumulation, bias grad, and delayed wgrad still fall through to the default path. wgrad is bitwise-equal to the default TE grouped backward on both worksteal and static schedules; dense NT/NN/TN are unaffected. Co-Authored-By: Claude --- .../pytorch/flydsl_kernels/gemm/__init__.py | 7 +- .../flydsl_kernels/gemm/gemm_wrappers.py | 101 +++++++++ .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 213 +++++++++++++++--- .../pytorch/module/grouped_linear.py | 140 +++++++++--- 4 files changed, 395 insertions(+), 66 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py b/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py index 60c6592c8f..a8f0ca1342 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py @@ -41,10 +41,15 @@ def _check_flydsl_version() -> None: _check_flydsl_version() from .exceptions import FlyDSLUnsupportedError -from .gemm_wrappers import te_generic_gemm_flydsl, _run_mxfp8_grouped +from .gemm_wrappers import ( + te_generic_gemm_flydsl, + _run_mxfp8_grouped, + _run_mxfp8_grouped_wgrad, +) __all__ = [ "FlyDSLUnsupportedError", "te_generic_gemm_flydsl", "_run_mxfp8_grouped", + "_run_mxfp8_grouped_wgrad", ] diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 849f2399d6..b1aa0b3e48 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -1150,6 +1150,107 @@ def _run_mxfp8_grouped( return D +def _run_mxfp8_grouped_wgrad( + dy_cols, # list[G] of per-expert dY columnwise fp8 data, each [m_e, N_out] + dy_scales, # list[G] of per-expert dY columnwise scale, each [m_e/32, N_out] + x_cols, # list[G] of per-expert X columnwise fp8 data, each [m_e, K_in] + x_scales, # list[G] of per-expert X columnwise scale, each [m_e/32, K_in] + dW, # output [G*N_out, K_in] + *, + m_splits, # list[int] per-expert token counts + n_out, + k_in, + dy_fp8_dtype, + x_fp8_dtype, + output_dtype, +): + """Grouped MXFP8 wgrad (NT): dW_e = dY_e^T @ X_e, per expert, one launch. + + Fixed uniform padding (phase 1): every expert's token count m_e is padded up to + a common M_pad (a multiple of BLOCK_M and 32); the padded rows are zero, so they + contribute nothing to the token-axis contraction -- no atomics or masking needed, + the dense NT pipeline runs unchanged. dY and X are packed on the contraction + (token) axis to [G*M_pad, *]; dW is packed on the output-M axis to [G*N_out, K_in] + and each expert's output tile is written via the kernel's store_m_expert_off. + """ + from .mxfp8_gemm import BLOCK_M, BLOCK_N, SCALE_GROUP_SIZE + + num_groups = len(m_splits) + device = dW.device + if n_out % BLOCK_M != 0: + raise FlyDSLUnsupportedError( + f"wgrad requires N_out divisible by BLOCK_M={BLOCK_M}, got {n_out}" + ) + if k_in % BLOCK_N != 0: + raise FlyDSLUnsupportedError( + f"wgrad requires K_in divisible by BLOCK_N={BLOCK_N}, got {k_in}" + ) + + # Uniform pad target: the largest per-expert token count, rounded up to BLOCK_M + # (also a multiple of 32). Enforce the two-page pipeline minimum (>= 4 K128 + # tiles => M_pad >= 512). + max_m = max(int(m) for m in m_splits) + m_pad = ((max_m + BLOCK_M - 1) // BLOCK_M) * BLOCK_M + m_pad = max(m_pad, 512) + + def _pad_rows(t, rows): + if t.size(0) == rows: + return t.contiguous() + out = t.new_zeros((rows, *t.shape[1:])) + out[: t.size(0)].copy_(t) + return out + + scale_rows = m_pad // SCALE_GROUP_SIZE + dy_data = torch.cat([_pad_rows(t, m_pad) for t in dy_cols], dim=0) + dy_scale = torch.cat([_pad_rows(t, scale_rows) for t in dy_scales], dim=0) + x_data = torch.cat([_pad_rows(t, m_pad) for t in x_cols], dim=0) + x_scale = torch.cat([_pad_rows(t, scale_rows) for t in x_scales], dim=0) + + if dy_data.dtype == torch.uint8: + dy_data = reinterpret_as_fp8_tensor(dy_data, dy_fp8_dtype) + if x_data.dtype == torch.uint8: + x_data = reinterpret_as_fp8_tensor(x_data, x_fp8_dtype) + + dW = _validate_or_allocate_output( + dW, + shape=torch.Size((num_groups * n_out, k_in)), + dtype=output_dtype, + device=device, + backend_name="MXFP8 grouped wgrad NT", + ) + + # Descriptor queue: one (pid_m, pid_n, expert) per output tile of every expert's + # [N_out, K_in] weight-gradient block. pid_m/pid_n are LOCAL (per-expert); the + # kernel re-packs the output row via expert*N_out. + num_pid_m = n_out // BLOCK_M + num_pid_n = k_in // BLOCK_N + descriptors = [ + (pid_m, pid_n, g) + for g in range(num_groups) + for pid_m in range(num_pid_m) + for pid_n in range(num_pid_n) + ] + num_tiles = len(descriptors) + work_queue = torch.tensor(descriptors, dtype=torch.int32, device=device).reshape(-1) + work_counter = torch.zeros(1, dtype=torch.int32, device=device) + sched = _select_grouped_sched(num_tiles, device) + + mxfp8_matmul( + dy_data, + dy_scale, + x_data, + x_scale, + dW, + group=False, + num_groups=num_groups, + work_queue=work_queue, + work_counter=work_counter, + sched=sched, + wgrad=True, + ) + return dW + + def _select_fp8_storage_for_layout(A, transa, B, transb): """Select existing TE FP8 backings and normalize to one core contract. diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index 66f2d76c56..df612094a7 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -260,6 +260,7 @@ def _compile_kernel( epilogue: str = "DEFAULT", group: bool = False, sched: str = "worksteal", + wgrad: bool = False, ): """Build one compile-time-specialized TN, NN, or NT kernel. @@ -518,10 +519,14 @@ def kernel_gemm( work_queue: fx.Tensor = None, work_counter: fx.Tensor = None, ): - # Total output tiles over the packed [M_total, N] result. Dense uses this - # as the grid extent; grouped uses it as the drain bound for the persistent - # work-stealing loop (see the dispatch at the end of the kernel). + # Total output tiles. Dense uses this as the grid extent; grouped uses it as + # the drain bound for the persistent loop. For fwd/dgrad, c_m = M_total spans + # all experts, so this is already the total. For wgrad, c_m = N_out is the + # PER-EXPERT output-row count (the packed output [G*N_out, K_in] is addressed + # via store_m_expert_off), so multiply by num_groups for the full count. num_tiles = (c_m // BLOCK_M) * (c_n // BLOCK_N) + if const_expr(wgrad): + num_tiles = num_tiles * num_groups lds = fx.SharedAllocator().allocate(SharedStorage).peek() lds_a0 = (lds.a0_0, lds.a0_1) @@ -546,6 +551,10 @@ def kernel_gemm( num_blocks_m = c_m // BLOCK_M num_blocks_n = c_n // BLOCK_N + # Both the grouped fwd/dgrad path and the wgrad specialization use the same + # persistent descriptor-queue driver; ``grouped`` gates that shared setup. + grouped = group or wgrad + # Persistent work-stealing driver. # # Dense path (group is False): the XCD-swizzled block id owns exactly one @@ -562,7 +571,7 @@ def kernel_gemm( # without any kernel change. B is packed [G*N, K] and B-scale [K128, G*N], # so B reads use the expert-shifted column ``expert*N + n`` with stride # ``G*N`` while A / A-scale / C / bias keep absolute-row / local-N coords. - if const_expr(group): + if const_expr(grouped): wq_rsrc = fx.rocdl.make_buffer_tensor(work_queue, max_size=True) wq_div = fx.logical_divide(wq_rsrc, fx.make_layout(1, 1)) @@ -635,28 +644,55 @@ def process_tile(pid_m, pid_n, expert): # and B-scale stride stay per-expert (c_n = K_in, unchanged). # # Dense (group is False) zeroes every expert term. - if const_expr(group): + # + # wgrad (NT, dW = dY^T @ X): a separate specialization. Both operands + # transpose-read and BOTH contract over the token axis M_tok (the + # kernel's compile-time K = the per-expert padded token count M_pad). + # Experts are padded to a uniform M_pad and packed on the contraction + # axis, so both a (dY [G*M_pad, N_out]) and b (X [G*M_pad, K_in]) + # shift by expert*K along k_base (a_expert_k_off / b_expert_k_off) and + # expert*NUM_K_TILES along their scale k128 terms. The output is packed + # [G*N_out, K_in] on the OUTPUT-M axis, so the store row gets an extra + # expert*N_out (= expert*c_m) offset (store_m_expert_off); the a/b + # operand row/col coords stay per-expert (local pid_m over N_out). + if const_expr(wgrad): + by_n_b_idx = by_n_idx + c_n_bs = c_n + a_expert_k_off = fx.Index(expert * fx.Int32(K)) + b_expert_k_off = fx.Index(expert * fx.Int32(K)) + as_expert_k128_off = expert * fx.Int32(NUM_K_TILES) + bs_expert_k128_off = expert * fx.Int32(NUM_K_TILES) + store_m_expert_off = fx.Index(expert * fx.Int32(c_m)) + elif const_expr(group): if const_expr(b_transpose_read): - # NN/NT: expert on the contraction axis. The kernel's - # compile-time ``K`` IS the contraction dim (= N_out for dgrad), - # so the packed weight [G*N_out, K_in] shifts by expert*K - # elements along k_base, and its scale [G*N_out/128, K_in] by - # expert*NUM_K_TILES along the k128 term. + # NN: expert on the contraction axis. The kernel's compile-time + # ``K`` IS the contraction dim (= N_out for dgrad), so the packed + # weight [G*N_out, K_in] shifts by expert*K elements along k_base, + # and its scale [G*N_out/128, K_in] by expert*NUM_K_TILES. by_n_b_idx = by_n_idx c_n_bs = c_n + a_expert_k_off = fx.Index(0) b_expert_k_off = fx.Index(expert * fx.Int32(K)) + as_expert_k128_off = fx.Int32(0) bs_expert_k128_off = expert * fx.Int32(NUM_K_TILES) + store_m_expert_off = fx.Index(0) else: # TN: expert on the output-N axis. by_n_b_idx = fx.Index(expert * fx.Int32(c_n) + by_n) c_n_bs = num_groups * c_n + a_expert_k_off = fx.Index(0) b_expert_k_off = fx.Index(0) + as_expert_k128_off = fx.Int32(0) bs_expert_k128_off = fx.Int32(0) + store_m_expert_off = fx.Index(0) else: by_n_b_idx = by_n_idx c_n_bs = c_n + a_expert_k_off = fx.Index(0) b_expert_k_off = fx.Index(0) + as_expert_k128_off = fx.Int32(0) bs_expert_k128_off = fx.Int32(0) + store_m_expert_off = fx.Index(0) wave_id = tx_i32 // fx.Int32(WARP_SIZE) lane = tx_i32 % fx.Int32(WARP_SIZE) @@ -710,7 +746,12 @@ def process_tile(pid_m, pid_n, expert): # Per-CTA tile base in C elements, folded into each store's linear # coordinate below (matching the scale-load addressing on this build; # add_offset on a dynamic Index is unsupported here). - c_tile_base_elems = bx_m_idx * fx.Index(c_n) + by_n_idx + # store_m_expert_off (grouped wgrad only) shifts the output row into this + # tile's expert block of the packed [G*N_out, K_in] weight-gradient; the + # a/b operands used local (per-expert) pid_m, so the store re-packs here. + c_tile_base_elems = ( + (bx_m_idx + store_m_expert_off) * fx.Index(c_n) + by_n_idx + ) gC = fx.rocdl.make_buffer_tensor(C, max_size=True) c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) c_store_atom = fx.make_copy_atom( @@ -863,7 +904,10 @@ def hot_loop_scheduler_q_prefetch_4n(): rocdl.sched_barrier(0) def load_a_scale_row(k128, row): - off = k128 * c_m_idx + bx_m_idx + row + # as_expert_k128_off shifts the contraction-128 base to this tile's + # expert A-scale block for grouped wgrad (both operands packed on the + # token/contraction axis); 0 for every other path. + off = (k128 + fx.Index(as_expert_k128_off)) * c_m_idx + bx_m_idx + row reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Int32) fx.copy(scale_ld_atom, fx.slice(as_div, (None, fx.Int32(off))), reg) return fx.memref_load_vec(reg)[0] @@ -899,9 +943,14 @@ def load_scale_tile(k128): ) def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): + # a_expert_k_off shifts the contraction base to this tile's expert + # A block for grouped wgrad (dY packed on the token axis); 0 for + # every other path. a_g2s.load_one( lds_a[subtile], - fx.Int32(_a_global_base(k_base, subtile, c_m, bx_m_idx)), + fx.Int32( + _a_global_base(k_base + a_expert_k_off, subtile, c_m, bx_m_idx) + ), pass_in_subtile, ) @@ -1635,7 +1684,7 @@ def _run_grouped_tile(slot): expert = _desc_field(slot, fx.Int32(2)) process_tile(pid_m, pid_n, expert) - if const_expr(group): + if const_expr(grouped): if const_expr(sched == "static"): stride = fx.Int32(gpu.grid_dim.x) slot = fx.Int32(gpu.block_id("x")) @@ -1681,7 +1730,7 @@ def launch_gemm( # num-CU grid of persistent workgroups that steal tiles from the atomic # work queue (num_tiles spans the packed [M_total, N] output). grid_x = (c_m // BLOCK_M) * (c_n // BLOCK_N) - if const_expr(group): + if const_expr(group or wgrad): grid_x = get_num_cus() kernel_gemm( A, @@ -1719,6 +1768,7 @@ def do_gemm( work_queue: fx.Tensor = None, work_counter: fx.Tensor = None, sched: str = "worksteal", + wgrad: bool = False, ): """Launch one cached compile-time MXFP8 layout specialization. @@ -1747,9 +1797,20 @@ def do_gemm( else: assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" - # Grouped: the kernel launches over the per-expert output width, which is - # always C's last dim (TN: N_out; NN: K_in). The kernel re-derives the packed - # B / B-scale stride from num_groups internally. + # Grouped wgrad (NT): both operands (dY, X) are packed on the token/contraction + # axis to G*M_pad, but the kernel's compile-time contraction is the PER-EXPERT + # padded token count M_pad. M_runtime (=N_out) and N_runtime (=K_in) are already + # per-expert; only K_runtime must be de-packed. + if wgrad: + assert K_runtime % num_groups == 0, ( + f"grouped wgrad contraction {K_runtime} not divisible by num_groups={num_groups}" + ) + K_runtime = K_runtime // num_groups + + # Grouped fwd/dgrad launch over the per-expert output width = C's last dim (TN: + # N_out; NN: K_in); the kernel re-derives the packed stride from num_groups. + # wgrad keeps M_runtime/N_runtime per-expert (C is packed on the output-M axis, + # handled in-kernel via store_m_expert_off). if group: N_runtime = C.shape[-1] supported_fp8_dtypes = (torch.float8_e4m3fn, torch.float8_e5m2) @@ -1774,24 +1835,30 @@ def do_gemm( ("C", C), ) - expected_as = (K_runtime // _BLOCK_K, M_runtime) - # Packed B-scale shape. TN packs B along N, so the scale's row count is the - # per-K128 count and its column count is the packed G*N_out. NN packs B along - # the contraction axis, so the packed count is on the K128 (contraction) axis - # instead: [G*N_out/128, K_in] = (Kb_runtime // BLOCK_K, N_runtime). - if group and layout == "NN": - expected_bs = (Kb_runtime // _BLOCK_K, N_runtime) + # Packed scale shapes. + if wgrad: + # wgrad (NT): both A (dY) and B (X) scales are packed on the contraction + # axis to G*M_pad/128. A-scale [G*M_pad/128, N_out], B-scale [.., K_in]. + packed_k128 = (K_runtime * num_groups) // _BLOCK_K + expected_as = (packed_k128, M_runtime) + expected_bs = (packed_k128, N_runtime) else: - n_packed = N_runtime * num_groups if group else N_runtime - expected_bs = (K_runtime // _BLOCK_K, n_packed) + expected_as = (K_runtime // _BLOCK_K, M_runtime) + # TN packs B along N (column count = G*N_out); NN packs B along the + # contraction axis (row count on K128 = G*N_out/128). + if group and layout == "NN": + expected_bs = (Kb_runtime // _BLOCK_K, N_runtime) + else: + n_packed = N_runtime * num_groups if group else N_runtime + expected_bs = (K_runtime // _BLOCK_K, n_packed) assert As.dtype == torch.int32, f"As dtype {As.dtype} != torch.int32 packed scales" assert Bs.dtype == torch.int32, f"Bs dtype {Bs.dtype} != torch.int32 packed scales" assert tuple(As.shape) == expected_as, f"As shape {tuple(As.shape)} != {expected_as}" assert tuple(Bs.shape) == expected_bs, f"Bs shape {tuple(Bs.shape)} != {expected_bs}" - assert tuple(C.shape) == ( - M_runtime, - N_runtime, - ), f"C shape {tuple(C.shape)} != {(M_runtime, N_runtime)}" + # wgrad's C is packed [G*N_out, K_in] on the output-M axis; fwd/dgrad C is the + # unpacked (M_runtime, N_runtime). + expected_c = (num_groups * M_runtime, N_runtime) if wgrad else (M_runtime, N_runtime) + assert tuple(C.shape) == expected_c, f"C shape {tuple(C.shape)} != {expected_c}" if C.dtype not in (torch.float16, torch.bfloat16, torch.float32): raise TypeError( f"C dtype must be torch.float16, torch.bfloat16, or torch.float32, got {C.dtype}" @@ -1858,8 +1925,9 @@ def do_gemm( Aux_arg = torch.zeros(1, dtype=C.dtype, device=A.device) # Grouped runtime tensors are required by the kernel signature; the dense - # binary never dereferences them but still needs placeholder buffers. - if group: + # binary never dereferences them but still needs placeholder buffers. wgrad + # uses the same persistent descriptor queue as the grouped fwd/dgrad path. + if group or wgrad: work_queue_arg = work_queue.contiguous().view(-1) work_counter_arg = work_counter.contiguous().view(-1) else: @@ -1875,6 +1943,7 @@ def do_gemm( epilogue, group, sched, + wgrad, )( A_arg, As_arg, @@ -1902,6 +1971,7 @@ def _cached_launch( epilogue: str = "DEFAULT", group: bool = False, sched: str = "worksteal", + wgrad: bool = False, ): """Cache independent TN/NN/NT binaries with no runtime layout argument.""" return _compile_kernel( @@ -1913,6 +1983,7 @@ def _cached_launch( epilogue, group, sched, + wgrad, ) @@ -1942,6 +2013,67 @@ def _validate_common_payloads( ) +def _mxfp8_matmul_grouped_wgrad( + a, # dY packed [G*M_pad, N_out], columnwise fp8 + a_scale, # [G*M_pad/32, N_out] raw E8M0 uint8 + b, # X packed [G*M_pad, K_in], columnwise fp8 + b_scale, # [G*M_pad/32, K_in] raw E8M0 uint8 + D, # dW packed [G*N_out, K_in] + *, + stream, + num_groups, + work_queue, + work_counter, + sched, +): + """Grouped wgrad (NT) launch: dW = dY^T @ X, contracting over padded tokens. + + Both operands transpose-read (columnwise) and contract over the token axis; the + kernel's compile-time contraction is the per-expert padded token count + M_pad = (G*M_pad)/G. Experts are packed on the contraction axis for a/b and on + the output-M axis for D. Operands arrive already packed + MXFP8-quantized. + """ + gm, n_out = a.shape # gm = G*M_pad, n_out = per-expert N_out + gm_b, k_in = b.shape + if gm != gm_b: + raise ValueError(f"wgrad a/b token dims differ: {gm} vs {gm_b}") + if gm % num_groups != 0: + raise ValueError(f"wgrad packed tokens {gm} not divisible by num_groups={num_groups}") + m_pad = gm // num_groups + if m_pad % SCALE_GROUP_SIZE != 0: + raise FlyDSLUnsupportedError( + f"wgrad padded per-expert tokens M_pad={m_pad} must be divisible by " + f"{SCALE_GROUP_SIZE}" + ) + if tuple(D.shape) != (num_groups * n_out, k_in): + raise ValueError( + f"wgrad D shape {tuple(D.shape)} != expected {(num_groups * n_out, k_in)}" + ) + if a_scale.dtype != torch.uint8 or b_scale.dtype != torch.uint8: + raise TypeError("FlyDSL MXFP8 wgrad expects raw E8M0 scales as torch.uint8") + + # NT transpose-read: both scales are packed from their columnwise (K-major) + # source into the kernel's [K/128, dim] int32 layout. + a_scale_hk = pack_mx32_scales_for_hk(a_scale, source_colwise=True, stream=stream) + b_scale_hk = pack_mx32_scales_for_hk(b_scale, source_colwise=True, stream=stream) + + do_gemm( + a, + a_scale_hk, + b, + b_scale_hk, + D, + layout="NT", + stream=stream, + num_groups=num_groups, + work_queue=work_queue, + work_counter=work_counter, + sched=sched, + wgrad=True, + ) + return D + + def mxfp8_matmul( a: torch.Tensor, a_scale: torch.Tensor, @@ -1959,6 +2091,7 @@ def mxfp8_matmul( work_queue: fx.Tensor = None, work_counter: fx.Tensor = None, sched: str = "worksteal", + wgrad: bool = False, ): """Normalize scale orientation and launch a compile-time layout binary. @@ -1977,6 +2110,20 @@ def mxfp8_matmul( compile-time-specialized kernels. """ + if wgrad: + return _mxfp8_matmul_grouped_wgrad( + a, + a_scale, + b, + b_scale, + D, + stream=stream, + num_groups=num_groups, + work_queue=work_queue, + work_counter=work_counter, + sched=sched, + ) + if group and layout not in ("TN", "NN"): raise ValueError("FlyDSL MXFP8 Group GEMM only supports the TN and NN layouts") diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 19add49d3b..b48dc39bc6 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -734,6 +734,57 @@ def _backward_flydsl_mxfp8_grouped( layout="NN", ) + @staticmethod + def _wgrad_flydsl_mxfp8_grouped( + *, + grad_output, + inputmats, + wgrad_list, + m_splits, + ): + """FlyDSL grouped wgrad: dW = dY^T @ X per expert, through the packed NT kernel. + + Both operands transpose-read and contract over the token axis, so both use + their COLUMNWISE MXFP8 payloads (dY from ``grad_output``, X from + ``inputmats``). The wrapper pads each expert's tokens to a uniform M_pad and + writes into the packed ``[G*N_out, K_in]`` gradient; per-expert results land + in ``wgrad_list[g]`` (a view of that packed buffer). No quantization here. + """ + from ..flydsl_kernels.gemm import _run_mxfp8_grouped_wgrad + + num_gemms = len(inputmats) + dy_cols = [g._columnwise_data for g in grad_output] + dy_scales = [g._columnwise_scale_inv for g in grad_output] + x_cols = [x._columnwise_data for x in inputmats] + x_scales = [x._columnwise_scale_inv for x in inputmats] + + n_out = grad_output[0]._columnwise_data.size(-1) # per-expert output feature N + k_in = inputmats[0]._columnwise_data.size(-1) # per-expert input feature K + + # Pack the per-expert [N_out, K_in] outputs into one [G*N_out, K_in] buffer; + # wgrad_list[g] are the per-expert slices the caller expects written. + dW = torch.empty( + num_gemms * n_out, k_in, dtype=wgrad_list[0].dtype, device=wgrad_list[0].device + ) + + _run_mxfp8_grouped_wgrad( + dy_cols, + dy_scales, + x_cols, + x_scales, + dW, + m_splits=m_splits, + n_out=n_out, + k_in=k_in, + dy_fp8_dtype=grad_output[0]._fp8_dtype, + x_fp8_dtype=inputmats[0]._fp8_dtype, + output_dtype=wgrad_list[0].dtype, + ) + + # Scatter the packed result back into the caller's per-expert wgrad tensors. + for g in range(num_gemms): + wgrad_list[g].copy_(dW[g * n_out : (g + 1) * n_out]) + @staticmethod def _forward_blockwise_fp8_triton( ctx, @@ -1956,42 +2007,67 @@ def backward( inputmats_dequant.append(cast_if_needed(inputmat, ctx.activation_dtype)) inputmats = inputmats_dequant - if ctx.use_grouped_gemm_triton: - general_grouped_gemm_func = general_grouped_gemm_triton - kwargs = {"m_splits_tensor": ctx.m_splits_tensor} - else: - general_grouped_gemm_func = general_grouped_gemm - kwargs = {} - grouped_gemm_wgrad = functools.partial( - general_grouped_gemm_func, - quantization_params=ctx.grad_weight_quantizers, - out_dtype=ctx.activation_dtype, - layout="NT", - grad=True, - m_splits=ctx.m_splits, - use_bias=ctx.use_bias if grad_biases[0] is None else None, - bias=biases, - use_split_accumulator=wgrad_gemm_use_split_accumulator, - accumulate=( - accumulate_wgrad_into_param_main_grad - if not getattr(ctx, "origin_weights_overwrite_main_grad", False) - else False - ), - **kwargs, + # FlyDSL grouped wgrad (NT) opt-in. Phase 1 handles the plain path + # only; fused wgrad accumulation, bias grad, and delayed wgrad still + # fall through to the default grouped GEMM below. + use_flydsl_wgrad = ( + getattr(ctx, "use_flydsl_mxfp8_grouped", False) + and not ctx.fuse_wgrad_accumulation + and not ctx.use_bias # bias-grad-in-wgrad-GEMM not supported yet + and not ( + ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute() + ) ) - # WGRAD - if ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute(): - ctx.wgrad_store.put([inputmats, grad_output, wgrad_list], grouped_gemm_wgrad) + if use_flydsl_wgrad: + _GroupedLinear._wgrad_flydsl_mxfp8_grouped( + grad_output=grad_output, + inputmats=inputmats, + wgrad_list=wgrad_list, + m_splits=ctx.m_splits, + ) + clear_tensor_data(*inputmats) + grouped_gemm_wgrad = None else: - _, grad_biases_, _ = grouped_gemm_wgrad(inputmats, grad_output, wgrad_list) + if ctx.use_grouped_gemm_triton: + general_grouped_gemm_func = general_grouped_gemm_triton + kwargs = {"m_splits_tensor": ctx.m_splits_tensor} + else: + general_grouped_gemm_func = general_grouped_gemm + kwargs = {} + grouped_gemm_wgrad = functools.partial( + general_grouped_gemm_func, + quantization_params=ctx.grad_weight_quantizers, + out_dtype=ctx.activation_dtype, + layout="NT", + grad=True, + m_splits=ctx.m_splits, + use_bias=ctx.use_bias if grad_biases[0] is None else None, + bias=biases, + use_split_accumulator=wgrad_gemm_use_split_accumulator, + accumulate=( + accumulate_wgrad_into_param_main_grad + if not getattr(ctx, "origin_weights_overwrite_main_grad", False) + else False + ), + **kwargs, + ) + # WGRAD + if ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute(): + ctx.wgrad_store.put( + [inputmats, grad_output, wgrad_list], grouped_gemm_wgrad + ) + else: + _, grad_biases_, _ = grouped_gemm_wgrad( + inputmats, grad_output, wgrad_list + ) - for i in range(ctx.num_gemms): - if grad_biases[i] is None: - grad_biases[i] = grad_biases_[i] - del grad_biases_ + for i in range(ctx.num_gemms): + if grad_biases[i] is None: + grad_biases[i] = grad_biases_[i] + del grad_biases_ - # Deallocate input tensor - clear_tensor_data(*inputmats) + # Deallocate input tensor + clear_tensor_data(*inputmats) def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): if ctx.weights_requires_grad: From 6ca17d56e63252ba27362983de4372a769930a4d Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 14 Oct 2026 14:33:39 +0000 Subject: [PATCH 04/22] Generalize FlyDSL grouped MXFP8 GEMM to BF16 and tensor-wise FP8 --- .../pytorch/flydsl_kernels/gemm/__init__.py | 8 + .../pytorch/flydsl_kernels/gemm/fp8_gemm.py | 1427 ++++++++------- .../flydsl_kernels/gemm/gemm_common_utils.py | 6 + .../flydsl_kernels/gemm/gemm_wrappers.py | 362 +++- .../flydsl_kernels/gemm/grouped_gemm_utils.py | 173 ++ .../flydsl_kernels/gemm/half_prec_gemm.py | 1545 +++++++++-------- .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 39 +- .../pytorch/module/grouped_linear.py | 411 ++++- 8 files changed, 2536 insertions(+), 1435 deletions(-) create mode 100644 transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py b/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py index a8f0ca1342..23a3947d4c 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py @@ -45,6 +45,10 @@ def _check_flydsl_version() -> None: te_generic_gemm_flydsl, _run_mxfp8_grouped, _run_mxfp8_grouped_wgrad, + _run_bf16_grouped, + _run_bf16_grouped_wgrad, + _run_fp8_grouped, + _run_fp8_grouped_wgrad, ) __all__ = [ @@ -52,4 +56,8 @@ def _check_flydsl_version() -> None: "te_generic_gemm_flydsl", "_run_mxfp8_grouped", "_run_mxfp8_grouped_wgrad", + "_run_bf16_grouped", + "_run_bf16_grouped_wgrad", + "_run_fp8_grouped", + "_run_fp8_grouped_wgrad", ] diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py index 7a9427f95a..c8e40ad446 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -25,8 +25,9 @@ from flydsl.expr.typing import T from flydsl.expr.typing import Vector as Vec +from . import grouped_gemm_utils as gd from .exceptions import FlyDSLUnsupportedError -from .gemm_common_utils import require_block_tiling, require_launch_size +from .gemm_common_utils import get_num_cus, require_block_tiling, require_launch_size # Transformer Engine-local FlyDSL utilities. from .fp8_gemm_utils import ( @@ -95,6 +96,9 @@ def _compile_kernel( output_dtype: torch.dtype, use_xcd_remap: bool = True, epilogue: str = "DEFAULT", + group: bool = False, + sched: str = "worksteal", + wgrad: bool = False, ): """Build the specialized kernel for compile-time K, A/B FP8 types, and output dtype. @@ -194,6 +198,9 @@ class SharedStorage: b0_1: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] b1_0: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] b1_1: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + # Persistent grouped path: one-int32 scratch to broadcast the atomically + # claimed queue slot from thread 0 to the whole workgroup. + claim_slot: fx.Array[fx.Int32, 1, 16] @flyc.kernel(known_block_size=[NUM_THREADS, 1, 1]) def kernel_gemm( @@ -206,6 +213,9 @@ def kernel_gemm( Aux: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32, + num_groups: fx.Int32 = 0, + work_queue: fx.Tensor = None, + work_counter: fx.Tensor = None, ): lds = fx.SharedAllocator().allocate(SharedStorage).peek() lds_a0 = (lds.a0_0, lds.a0_1) @@ -225,735 +235,844 @@ def kernel_gemm( b_scale_div = fx.logical_divide(b_scale_rsrc, fx.make_layout(1, 1)) scale_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) - def _load_scale(scale_div): + def _load_scale(scale_div, index): + # Dense passes a 1-element scale buffer (index 0). Grouped passes a + # length-num_groups vector indexed by the tile's expert, so each tile + # recomputes its own per-expert output_scale inside process_tile. reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) - fx.copy(scale_ld_atom, fx.slice(scale_div, (None, fx.Int32(0))), reg) + fx.copy(scale_ld_atom, fx.slice(scale_div, (None, index)), reg) return fx.memref_load_vec(reg)[0] - output_scale = _load_scale(a_scale_div) * _load_scale(b_scale_div) tx = gpu.thread_id("x") + tx_i32 = fx.Int32(tx) num_blocks_m = c_m // BLOCK_M num_blocks_n = c_n // BLOCK_N - if const_expr(use_xcd_remap): - pid_m, pid_n = xcd_swizzle(num_blocks_m, num_blocks_n) - else: - pid_m, pid_n = divmod(fx.block_idx.x, num_blocks_n) - - bx_m = pid_m * BLOCK_M - by_n = pid_n * BLOCK_N - - # The flattened/XCD-swizzled block coordinates are i32, while global - # address arithmetic below is expressed in MLIR index type. Convert - # once here and use these index-typed tile bases for every address. - bx_m_idx = fx.Index(bx_m) - by_n_idx = fx.Index(by_n) - - # Keep wave/lane arithmetic in i32. compute_global_swizzle() combines - # these values with i32 constants, so Index-typed coordinates would make - # arith.addi receive mixed operand types. - tx_i32 = fx.Int32(tx) - wave_id = tx_i32 // fx.Int32(WARP_SIZE) - lane = tx_i32 % fx.Int32(WARP_SIZE) - - # The utility mapping is identical to the previous manual staging: - # each step contributes one contiguous 16-byte vector per thread, while - # the global K coordinate is XOR-unswizzled for the physical LDS slot. - gl_off_a = compute_global_swizzle(lane, wave_id, K, LOAD_PASSES_HALF, preshuffled=False) - gl_off_b = compute_global_swizzle(lane, wave_id, K, LOAD_PASSES_HALF, preshuffled=False) - a_g2s = G2SLoader(a_div, gl_off_a, LOAD_PASSES_HALF, a_f8_ir_t, wave_id) - b_g2s = G2SLoader(b_div, gl_off_b, LOAD_PASSES_HALF, b_f8_ir_t, wave_id) - s2r = S2RLoader(fx.Int32(0), 1) - - layout_lane16 = fx.make_layout((4, 16), (16, 1)) - coord_lane16 = fx.idx2crd(fx.Int32(lane), layout_lane16) - lane_div_16 = fx.get(coord_lane16, 0) - lane_mod_16 = fx.get(coord_lane16, 1) - - # Per-CTA tile base in C elements, folded into each store's linear - # coordinate below (matching the scale-load addressing on this build; - # add_offset on a dynamic Index is unsupported here). - c_tile_base_elems = bx_m_idx * fx.Index(c_n) + by_n_idx - gC = fx.rocdl.make_buffer_tensor(C, max_size=True) - c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) - c_store_atom = fx.make_copy_atom( - fx.rocdl.BufferCopy32b() if output_element_bytes == 4 else fx.rocdl.BufferCopy16b(), - output_fx_dtype, - ) - - # Bias is a length-N fp32 vector indexed by the global output-feature - # (N) coordinate and broadcast across the M/token rows. const_expr folds - # this compile-time flag at trace time so the setup (and load_bias) are - # inlined into the kernel scope with no runtime dispatch branch. - if const_expr(has_bias): - gBias = fx.rocdl.make_buffer_tensor(Bias, max_size=True) - bias_div = fx.logical_divide(gBias, fx.make_layout(1, 1)) - bias_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) - - def load_bias(col): - reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) - fx.copy(bias_ld_atom, fx.slice(bias_div, (None, fx.Int32(col))), reg) - return fx.memref_load_vec(reg)[0] - - # GELU_AUX saves the pre-activation value (A@B[+bias]) to a second M x N - # output so the backward pass can recompute the GELU gradient. Same tile - # base / store atom shape as C; DEFAULT gets a dummy 1-element tensor. - if const_expr(has_gelu): - gAux = fx.rocdl.make_buffer_tensor(Aux, max_size=True) - aux_div = fx.logical_divide(gAux, fx.make_layout(1, 1)) - aux_store_atom = fx.make_copy_atom( + # Total output tiles. Dense uses this as the grid extent; grouped uses it as + # the persistent-loop drain bound. ``grouped`` gates the shared descriptor + # queue / work-steal driver; const_expr folds it out of the dense binary. + grouped = group or wgrad + num_tiles = gd.num_tiles_bound(c_m, c_n, BLOCK_M, BLOCK_N, num_groups, const_expr(wgrad)) + + if const_expr(grouped): + wq_rsrc = fx.rocdl.make_buffer_tensor(work_queue, max_size=True) + wq_div = fx.logical_divide(wq_rsrc, fx.make_layout(1, 1)) + desc_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32) + claim_iter, claim_view = gd.make_claim_broadcast(lds.claim_slot) + + def _claim_next_slot(): + if tx_i32 == fx.Int32(0): + old = llvm.AtomicRMWOp( + llvm.AtomicBinOp.add, + gd.global_i32_ptr(work_counter, fx.Int32(0)), + fx.as_ir_value(fx.Int32(1)), + llvm.AtomicOrdering.monotonic, + syncscope="agent", + alignment=4, + ).result + fx.ptr_store(fx.Vector.from_elements([fx.Int32(old)], fx.Int32), claim_iter) + fx.barrier() + return fx.Vector(claim_view.load())[0] + + def _desc_field(slot, field): + return gd.read_desc_field(wq_div, desc_ld_atom, slot, field) + + def process_tile(pid_m, pid_n, expert): + # Per-expert scalar scale: dense reads index 0 (length-1 buffer); + # grouped reads scale[expert] from the length-num_groups vector. + if const_expr(grouped): + scale_idx = fx.Int32(expert) + else: + scale_idx = fx.Int32(0) + output_scale = _load_scale(a_scale_div, scale_idx) * _load_scale(b_scale_div, scale_idx) + + bx_m = pid_m * BLOCK_M + by_n = pid_n * BLOCK_N + + # The flattened/XCD-swizzled block coordinates are i32, while global + # address arithmetic below is expressed in MLIR index type. Convert + # once here and use these index-typed tile bases for every address. + bx_m_idx = fx.Index(bx_m) + by_n_idx = fx.Index(by_n) + + # Per-expert offsets fold the expert index into the packed operands / + # output (const_expr-zeroed for dense). FP8 is TN-NORMAL-read, so the + # contraction axis (K) is the trailing dim and is NEVER packed; experts + # always stack on an OUTPUT/feature axis: + # + # fwd/dgrad: A is per-token (absolute M rows, no offset); B (weights, + # or W^T for dgrad) is packed on its N-row by expert + # (by_n_b_idx = expert*c_n + n). Output per-token (no store offset). + # + # wgrad (dW_e = dY_e^T @ X_e): host pre-transposes both operands to + # [feature, m_pad] and token-pads the trailing (contraction) axis, so + # A = dY^T packed [G*N_out, m_pad] on the M-row (a_expert_m_off = + # expert*c_m), B = X^T packed [G*K_in, m_pad] on the N-row + # (by_n_b_idx = expert*c_n + n). Output dW packed [G*N_out, K_in] on + # the M-row (store_m_expert_off = expert*c_m). + if const_expr(wgrad): + by_n_b_idx = fx.Index(expert * fx.Int32(c_n) + by_n) + a_expert_m_off = fx.Index(expert * fx.Int32(c_m)) + store_m_expert_off = fx.Index(expert * fx.Int32(c_m)) + elif const_expr(group): + by_n_b_idx = fx.Index(expert * fx.Int32(c_n) + by_n) + a_expert_m_off = fx.Index(0) + store_m_expert_off = fx.Index(0) + else: + by_n_b_idx = by_n_idx + a_expert_m_off = fx.Index(0) + store_m_expert_off = fx.Index(0) + + # Keep wave/lane arithmetic in i32. compute_global_swizzle() combines + # these values with i32 constants, so Index-typed coordinates would make + # arith.addi receive mixed operand types. + wave_id = tx_i32 // fx.Int32(WARP_SIZE) + lane = tx_i32 % fx.Int32(WARP_SIZE) + + # The utility mapping is identical to the previous manual staging: + # each step contributes one contiguous 16-byte vector per thread, while + # the global K coordinate is XOR-unswizzled for the physical LDS slot. + gl_off_a = compute_global_swizzle(lane, wave_id, K, LOAD_PASSES_HALF, preshuffled=False) + gl_off_b = compute_global_swizzle(lane, wave_id, K, LOAD_PASSES_HALF, preshuffled=False) + a_g2s = G2SLoader(a_div, gl_off_a, LOAD_PASSES_HALF, a_f8_ir_t, wave_id) + b_g2s = G2SLoader(b_div, gl_off_b, LOAD_PASSES_HALF, b_f8_ir_t, wave_id) + s2r = S2RLoader(fx.Int32(0), 1) + + layout_lane16 = fx.make_layout((4, 16), (16, 1)) + coord_lane16 = fx.idx2crd(fx.Int32(lane), layout_lane16) + lane_div_16 = fx.get(coord_lane16, 0) + lane_mod_16 = fx.get(coord_lane16, 1) + + # Per-CTA tile base in C elements, folded into each store's linear + # coordinate below (matching the scale-load addressing on this build; + # add_offset on a dynamic Index is unsupported here). + # store_m_expert_off (grouped wgrad only) repacks the output row to + # this expert's block; 0 for dense and grouped fwd/dgrad. + c_tile_base_elems = (bx_m_idx + store_m_expert_off) * fx.Index(c_n) + by_n_idx + gC = fx.rocdl.make_buffer_tensor(C, max_size=True) + c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) + c_store_atom = fx.make_copy_atom( fx.rocdl.BufferCopy32b() if output_element_bytes == 4 else fx.rocdl.BufferCopy16b(), output_fx_dtype, ) - def gelu_tanh(x): - # tanh-approx GELU (matches PyTorch approximate='tanh' and the - # FlyDSL preshuffle reference), expressed through a non-positive - # exponent so exp() cannot overflow: - # 0.5*x*(1 + tanh(y)), y = sqrt(2/pi)*(x + 0.044715*x^3) - half_f32 = fx.Float32(0.5) - one_f32 = fx.Float32(1.0) - zero_f32 = fx.Float32(0.0) - two_f32 = fx.Float32(2.0) - x3 = x * x * x - y = fx.Float32(0.7978845608) * (x + fx.Float32(0.044715) * x3) - abs_y = fx.Float32(y).maximumf(zero_f32 - y) - e_neg2abs = math.exp(fx.Float32(-2.0) * abs_y) - denom = one_f32 + e_neg2abs - numerator = (y > zero_f32).select(two_f32, two_f32 * e_neg2abs) - return half_f32 * x * (numerator * (one_f32 / denom)) - - PIN_ACC_BASE = 0 - - def _reg_list(prefix, start, end): - return ",".join(f"~{{{prefix}{r}}}" for r in range(start, end + 1)) - - def reserve_pinned_accumulators(): - # Reserve a fixed physical AGPR bank for all accumulators. In the - # SSA-lowered path, the compiler generated heavy AGPR <-> VGPR traffic, - # including v_accvgpr_mov/read sequences, s_nop stalls, and accumulator - # spills. Pinning each f32x4 accumulator to a stable AGPR range keeps the - # scaled MFMA accumulation in place and avoids those transfers and spills. - # - # ACCS_PER_WAVE = 64 accumulator objects and each object is f32x4, - # so the physical bank is exactly 64 * 4 = 256 AGPRs: a[0:255]. - clobbers = _reg_list("a", PIN_ACC_BASE, PIN_ACC_BASE + ACCS_PER_WAVE * 4 - 1) - llvm.InlineAsmOp( - None, - [], - "", - clobbers, - has_side_effects=True, - ) + # Bias is a length-N fp32 vector indexed by the global output-feature + # (N) coordinate and broadcast across the M/token rows. const_expr folds + # this compile-time flag at trace time so the setup (and load_bias) are + # inlined into the kernel scope with no runtime dispatch branch. + if const_expr(has_bias): + gBias = fx.rocdl.make_buffer_tensor(Bias, max_size=True) + bias_div = fx.logical_divide(gBias, fx.make_layout(1, 1)) + bias_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) + + def load_bias(col): + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) + fx.copy(bias_ld_atom, fx.slice(bias_div, (None, fx.Int32(col))), reg) + return fx.memref_load_vec(reg)[0] + + # GELU_AUX saves the pre-activation value (A@B[+bias]) to a second M x N + # output so the backward pass can recompute the GELU gradient. Same tile + # base / store atom shape as C; DEFAULT gets a dummy 1-element tensor. + if const_expr(has_gelu): + gAux = fx.rocdl.make_buffer_tensor(Aux, max_size=True) + aux_div = fx.logical_divide(gAux, fx.make_layout(1, 1)) + aux_store_atom = fx.make_copy_atom( + fx.rocdl.BufferCopy32b() if output_element_bytes == 4 else fx.rocdl.BufferCopy16b(), + output_fx_dtype, + ) - def zero_pinned_accumulators(): - for ai in range_constexpr(ACCS_PER_WAVE * 4): + def gelu_tanh(x): + # tanh-approx GELU (matches PyTorch approximate='tanh' and the + # FlyDSL preshuffle reference), expressed through a non-positive + # exponent so exp() cannot overflow: + # 0.5*x*(1 + tanh(y)), y = sqrt(2/pi)*(x + 0.044715*x^3) + half_f32 = fx.Float32(0.5) + one_f32 = fx.Float32(1.0) + zero_f32 = fx.Float32(0.0) + two_f32 = fx.Float32(2.0) + x3 = x * x * x + y = fx.Float32(0.7978845608) * (x + fx.Float32(0.044715) * x3) + abs_y = fx.Float32(y).maximumf(zero_f32 - y) + e_neg2abs = math.exp(fx.Float32(-2.0) * abs_y) + denom = one_f32 + e_neg2abs + numerator = (y > zero_f32).select(two_f32, two_f32 * e_neg2abs) + return half_f32 * x * (numerator * (one_f32 / denom)) + + PIN_ACC_BASE = 0 + + def _reg_list(prefix, start, end): + return ",".join(f"~{{{prefix}{r}}}" for r in range(start, end + 1)) + + def reserve_pinned_accumulators(): + # Reserve a fixed physical AGPR bank for all accumulators. In the + # SSA-lowered path, the compiler generated heavy AGPR <-> VGPR traffic, + # including v_accvgpr_mov/read sequences, s_nop stalls, and accumulator + # spills. Pinning each f32x4 accumulator to a stable AGPR range keeps the + # scaled MFMA accumulation in place and avoids those transfers and spills. + # + # ACCS_PER_WAVE = 64 accumulator objects and each object is f32x4, + # so the physical bank is exactly 64 * 4 = 256 AGPRs: a[0:255]. + clobbers = _reg_list("a", PIN_ACC_BASE, PIN_ACC_BASE + ACCS_PER_WAVE * 4 - 1) llvm.InlineAsmOp( None, [], - f"v_accvgpr_write_b32 a[{PIN_ACC_BASE + ai}], 0", - f"~{{a{PIN_ACC_BASE + ai}}}", + "", + clobbers, has_side_effects=True, ) - def _inline_asm_i32(asm_string, constraints, operands=None): - op = llvm.InlineAsmOp( - T.i32, - operands or [], - asm_string, - constraints, - has_side_effects=True, - ) - return _one_i32_result(op) - - def _one_i32_result(op): - # Accept the result attribute names exposed by the supported MLIR Python bindings. - return getattr(op, "result", getattr(op, "res", op.results[0])) - - def read_pinned_accumulator(acc_idx): - acc_pin = PIN_ACC_BASE + acc_idx * 4 - r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") - r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") - r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") - r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") - return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) - - def read_physical_accumulator_slot(slot_idx): - acc_pin = PIN_ACC_BASE + slot_idx * 4 - r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") - r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") - r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") - r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") - return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) - - def hot_loop_scheduler_q_refill_2n(): - for _ in range_constexpr(8): - rocdl.sched_vmem(1) - rocdl.sched_mfma(2) - rocdl.sched_barrier(0) - - def hot_loop_scheduler_q0_refill_a1_2n(): - for _ in range_constexpr(8): - rocdl.sched_vmem(1) - rocdl.sched_dsrd(1) - rocdl.sched_mfma(2) - rocdl.sched_barrier(0) + def zero_pinned_accumulators(): + for ai in range_constexpr(ACCS_PER_WAVE * 4): + llvm.InlineAsmOp( + None, + [], + f"v_accvgpr_write_b32 a[{PIN_ACC_BASE + ai}], 0", + f"~{{a{PIN_ACC_BASE + ai}}}", + has_side_effects=True, + ) + + def _inline_asm_i32(asm_string, constraints, operands=None): + op = llvm.InlineAsmOp( + T.i32, + operands or [], + asm_string, + constraints, + has_side_effects=True, + ) + return _one_i32_result(op) + + def _one_i32_result(op): + # Accept the result attribute names exposed by the supported MLIR Python bindings. + return getattr(op, "result", getattr(op, "res", op.results[0])) + + def read_pinned_accumulator(acc_idx): + acc_pin = PIN_ACC_BASE + acc_idx * 4 + r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") + r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") + r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") + r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") + return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) + + def read_physical_accumulator_slot(slot_idx): + acc_pin = PIN_ACC_BASE + slot_idx * 4 + r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") + r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") + r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") + r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") + return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) + + def hot_loop_scheduler_q_refill_2n(): + for _ in range_constexpr(8): + rocdl.sched_vmem(1) + rocdl.sched_mfma(2) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q0_refill_a1_2n(): + for _ in range_constexpr(8): + rocdl.sched_vmem(1) + rocdl.sched_dsrd(1) + rocdl.sched_mfma(2) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q_prefetch_4n(): + for _ in range_constexpr(8): + rocdl.sched_dsrd(2) + rocdl.sched_mfma(4) + rocdl.sched_barrier(0) + + def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): + # One pass writes 256 threads * 16 B = 4 KiB. Four passes fill one + # 128x128 half-page (16 KiB). Each half has its own LDS base. + # a_expert_m_off shifts the A row to this tile's expert block (grouped + # wgrad, where A = dY^T is packed [G*N_out, m_pad] on the M-row); 0 + # for dense and grouped fwd/dgrad (A is per-token, absolute rows). + global_base = ( + (bx_m_idx + a_expert_m_off + fx.Index(subtile * (BLOCK_M // 2))) * fx.Index(K) + + k_base + ) + a_g2s.load_one(lds_a[subtile], fx.Int32(global_base), pass_in_subtile) + + def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): + # by_n_b_idx shifts the B row to this tile's expert block (grouped TN + # fwd/dgrad weights, and grouped wgrad X^T packed [G*K_in, m_pad]); + # reduces to the dense by_n for the dense path. + global_base = ( + (by_n_b_idx + fx.Index(subtile * (BLOCK_N // 2))) * fx.Index(K) + + k_base + ) + b_g2s.load_one(lds_b[subtile], fx.Int32(global_base), pass_in_subtile) + + def stage_a_subtile(k_base, subtile, lds_a): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a) + + def stage_b_subtile(k_base, subtile, lds_b): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b) + + def load_frag_half_at_byte_base(lds_page, row_byte_base, half): + # Issue exactly one 16-byte LDS read for one K64 half of an MFMA operand. + # Keeping the halves separate allows steady-state Q0 to schedule one + # A-bottom ds_read_b128 in each refill/MFMA chunk. + k_col = reg_lds_k_col0 if half == 0 else reg_lds_k_col1 + return s2r.load_one(lds_page, fx.Int32(row_byte_base + k_col)) + + def pack_frag_halves(x0, x1): + return pack_i32x4_i32x8(x0, x1) + + def load_frag_at_byte_base(lds_page, row_byte_base): + # Default complete-fragment path used outside the dedicated Q0 schedule. + x0 = load_frag_half_at_byte_base(lds_page, row_byte_base, 0) + x1 = load_frag_half_at_byte_base(lds_page, row_byte_base, 1) + return pack_frag_halves(x0, x1) + + def load_b_frag(lds_b, local_row, half): + # B is [N, K]. Each 128-row half-page has a local row origin of 0. + half_row = local_row - fx.Index(half * (BLOCK_N // 2)) + return load_frag_at_byte_base(lds_b[half], half_row * fx.Index(BLOCK_K)) + + def _acc_idx(subtile_id, mi, ni): + return ( + subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + mi * MFMA_N_PER_SUBTILE + ni + ) - def hot_loop_scheduler_q_prefetch_4n(): - for _ in range_constexpr(8): - rocdl.sched_dsrd(2) - rocdl.sched_mfma(4) - rocdl.sched_barrier(0) + def pinned_mfma(acc_idx, a_frag, b_frag): + """Issue ordinary FP8 MFMA into the fixed physical accumulator bank.""" + acc_pin = PIN_ACC_BASE + acc_idx * 4 + llvm.InlineAsmOp( + None, + [ + arith._to_raw(a_frag), + arith._to_raw(b_frag), + ], + "v_mfma_f32_16x16x128_f8f6f4 " + f"a[{acc_pin}:{acc_pin + 3}], " + "$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}] " + f"cbsz:{a_matrix_format} blgp:{b_matrix_format}", + f"v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}},~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}", + has_side_effects=True, + ) - def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): - # One pass writes 256 threads * 16 B = 4 KiB. Four passes fill one - # 128x128 half-page (16 KiB). Each half has its own LDS base. - global_base = (bx_m_idx + fx.Index(subtile * (BLOCK_M // 2))) * fx.Index(K) + k_base - a_g2s.load_one(lds_a[subtile], fx.Int32(global_base), pass_in_subtile) - - def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): - global_base = (by_n_idx + fx.Index(subtile * (BLOCK_N // 2))) * fx.Index(K) + k_base - b_g2s.load_one(lds_b[subtile], fx.Int32(global_base), pass_in_subtile) - - def stage_a_subtile(k_base, subtile, lds_a): - for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): - stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a) - - def stage_b_subtile(k_base, subtile, lds_b): - for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): - stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b) - - def load_frag_half_at_byte_base(lds_page, row_byte_base, half): - # Issue exactly one 16-byte LDS read for one K64 half of an MFMA operand. - # Keeping the halves separate allows steady-state Q0 to schedule one - # A-bottom ds_read_b128 in each refill/MFMA chunk. - k_col = reg_lds_k_col0 if half == 0 else reg_lds_k_col1 - return s2r.load_one(lds_page, fx.Int32(row_byte_base + k_col)) - - def pack_frag_halves(x0, x1): - return pack_i32x4_i32x8(x0, x1) - - def load_frag_at_byte_base(lds_page, row_byte_base): - # Default complete-fragment path used outside the dedicated Q0 schedule. - x0 = load_frag_half_at_byte_base(lds_page, row_byte_base, 0) - x1 = load_frag_half_at_byte_base(lds_page, row_byte_base, 1) - return pack_frag_halves(x0, x1) - - def load_b_frag(lds_b, local_row, half): - # B is [N, K]. Each 128-row half-page has a local row origin of 0. - half_row = local_row - fx.Index(half * (BLOCK_N // 2)) - return load_frag_at_byte_base(lds_b[half], half_row * fx.Index(BLOCK_K)) - - def _acc_idx(subtile_id, mi, ni): - return ( - subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + mi * MFMA_N_PER_SUBTILE + ni - ) + def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag): + """Final-page ordinary FP8 MFMA with independently named AGPR source/destination.""" + dst_pin = PIN_ACC_BASE + dst_slot * 4 + old_pin = PIN_ACC_BASE + old_acc_idx * 4 + llvm.InlineAsmOp( + None, + [ + arith._to_raw(a_frag), + arith._to_raw(b_frag), + ], + "v_mfma_f32_16x16x128_f8f6f4 " + f"a[{dst_pin}:{dst_pin + 3}], " + "$0, $1, " + f"a[{old_pin}:{old_pin + 3}] " + f"cbsz:{a_matrix_format} blgp:{b_matrix_format}", + f"v,v,~{{a{dst_pin}}},~{{a{dst_pin + 1}}},~{{a{dst_pin + 2}}},~{{a{dst_pin + 3}}}", + has_side_effects=True, + ) - def pinned_mfma(acc_idx, a_frag, b_frag): - """Issue ordinary FP8 MFMA into the fixed physical accumulator bank.""" - acc_pin = PIN_ACC_BASE + acc_idx * 4 - llvm.InlineAsmOp( - None, - [ - arith._to_raw(a_frag), - arith._to_raw(b_frag), - ], - "v_mfma_f32_16x16x128_f8f6f4 " - f"a[{acc_pin}:{acc_pin + 3}], " - "$0, $1, " - f"a[{acc_pin}:{acc_pin + 3}] " - f"cbsz:{a_matrix_format} blgp:{b_matrix_format}", - f"v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}},~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}", - has_side_effects=True, - ) + def mfma_4n(acc_base, a_frag, b0, b1, b2, b3): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + pinned_mfma(acc_base + 2, a_frag, b2) + pinned_mfma(acc_base + 3, a_frag, b3) - def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag): - """Final-page ordinary FP8 MFMA with independently named AGPR source/destination.""" - dst_pin = PIN_ACC_BASE + dst_slot * 4 - old_pin = PIN_ACC_BASE + old_acc_idx * 4 - llvm.InlineAsmOp( - None, - [ - arith._to_raw(a_frag), - arith._to_raw(b_frag), - ], - "v_mfma_f32_16x16x128_f8f6f4 " - f"a[{dst_pin}:{dst_pin + 3}], " - "$0, $1, " - f"a[{old_pin}:{old_pin + 3}] " - f"cbsz:{a_matrix_format} blgp:{b_matrix_format}", - f"v,v,~{{a{dst_pin}}},~{{a{dst_pin + 1}}},~{{a{dst_pin + 2}}},~{{a{dst_pin + 3}}}", - has_side_effects=True, - ) + def mfma_2n(acc_base, a_frag, b0, b1): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) - def mfma_4n(acc_base, a_frag, b0, b1, b2, b3): - pinned_mfma(acc_base + 0, a_frag, b0) - pinned_mfma(acc_base + 1, a_frag, b1) - pinned_mfma(acc_base + 2, a_frag, b2) - pinned_mfma(acc_base + 3, a_frag, b3) - - def mfma_2n(acc_base, a_frag, b0, b1): - pinned_mfma(acc_base + 0, a_frag, b0) - pinned_mfma(acc_base + 1, a_frag, b1) - - def store_acc_vector_for_logical_idx(logical_acc_idx, acc): - subtile_id = logical_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) - local_idx = logical_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) - sm = subtile_id // 2 - sn = subtile_id % 2 - mi = local_idx // MFMA_N_PER_SUBTILE - ni = local_idx % MFMA_N_PER_SUBTILE - - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - row_base = subtile_m_idx * SUBTILE_M + fx.Index(mi * MFMA_M) + lane_div_16 * 4 - col = subtile_n_idx * SUBTILE_N + fx.Index(ni * MFMA_N) + lane_mod_16 - - # Bias depends only on the output-feature (N) coordinate, so read it - # once per column (global index by_n_idx + col) and reuse across the - # four M rows below. - if const_expr(has_bias): - bias_value = load_bias(by_n_idx + col) + def store_acc_vector_for_logical_idx(logical_acc_idx, acc): + subtile_id = logical_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = logical_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE - for ii in range_constexpr(4): - row = row_base + fx.Index(ii) - c_idx = c_tile_base_elems + row * fx.Index(c_n) + col + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + row_base = subtile_m_idx * SUBTILE_M + fx.Index(mi * MFMA_M) + lane_div_16 * 4 + col = subtile_n_idx * SUBTILE_N + fx.Index(ni * MFMA_N) + lane_mod_16 - # Epilogue stages run on the fp32 accumulator, in order, before - # the output-dtype narrowing. Bias is added after the tensor-wise - # output scale, matching the reference epilogue ordering; GELU_AUX - # then saves the pre-activation and applies GELU to the C value. - value = Vec(acc)[ii] * output_scale + # Bias depends only on the output-feature (N) coordinate, so read it + # once per column (global index by_n_idx + col) and reuse across the + # four M rows below. if const_expr(has_bias): - value = value + bias_value + bias_value = load_bias(by_n_idx + col) + + for ii in range_constexpr(4): + row = row_base + fx.Index(ii) + c_idx = c_tile_base_elems + row * fx.Index(c_n) + col + + # Epilogue stages run on the fp32 accumulator, in order, before + # the output-dtype narrowing. Bias is added after the tensor-wise + # output scale, matching the reference epilogue ordering; GELU_AUX + # then saves the pre-activation and applies GELU to the C value. + value = Vec(acc)[ii] * output_scale + if const_expr(has_bias): + value = value + bias_value + + if const_expr(has_gelu): + aux_val = value + if output_dtype != torch.float32: + aux_val = aux_val.to(output_fx_dtype) + aux_reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) + fx.memref_store_vec(Vec.filled(1, aux_val, output_fx_dtype), aux_reg) + fx.copy(aux_store_atom, aux_reg, fx.slice(aux_div, (None, fx.Int32(c_idx)))) + value = gelu_tanh(value) - if const_expr(has_gelu): - aux_val = value if output_dtype != torch.float32: - aux_val = aux_val.to(output_fx_dtype) - aux_reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) - fx.memref_store_vec(Vec.filled(1, aux_val, output_fx_dtype), aux_reg) - fx.copy(aux_store_atom, aux_reg, fx.slice(aux_div, (None, fx.Int32(c_idx)))) - value = gelu_tanh(value) - - if output_dtype != torch.float32: - value = value.to(output_fx_dtype) - reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) - fx.memref_store_vec(Vec.filled(1, value, output_fx_dtype), reg) - fx.copy(c_store_atom, reg, fx.slice(c_div, (None, fx.Int32(c_idx)))) - - # Explicit register coordinates for HK-style four-quadrant mapping. - # BLOCK_M/BLOCK_N are 256x256. Four waves map to warp positions - # inside each 128x128 quadrant: - # cA: (warp_m, warp_n) - # cB: (warp_m, warp_n + 2) - # cC: (warp_m + 2, warp_n) - # cD: (warp_m + 2, warp_n + 2) - reg_k_col0 = lane_div_16 * 16 - reg_k_col1 = 64 + lane_div_16 * 16 - - # Every fragment row differs only by multiples of 16, so row % 16 is - # always lane_mod_16. Hoist the logical->physical XOR mapping once. - _, reg_lds_k_col0 = swizzle_128(lane_mod_16, reg_k_col0) - _, reg_lds_k_col1 = swizzle_128(lane_mod_16, reg_k_col1) - - reg_subtile_m_idx0 = wave_id // 2 - reg_subtile_n_idx0 = wave_id % 2 - - reserve_pinned_accumulators() - zero_pinned_accumulators() - - def load_b_subtile_ni_regs(lds_b, sn, ni): - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - b_row_addr = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 - return load_b_frag(lds_b, b_row_addr, sn) - - def load_b_subtile_regs(lds_b, sn): - return ( - load_b_subtile_ni_regs(lds_b, sn, 0), - load_b_subtile_ni_regs(lds_b, sn, 1), - load_b_subtile_ni_regs(lds_b, sn, 2), - load_b_subtile_ni_regs(lds_b, sn, 3), - ) - - def load_a_subtile_mi_half(lds_a, sm, mi, half): - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - a_row_addr = subtile_m_idx * fx.Index(SUBTILE_M) + fx.Index(mi * MFMA_M) + lane_mod_16 - half_row = a_row_addr - fx.Index(sm * (BLOCK_M // 2)) - row_byte_base = half_row * fx.Index(BLOCK_K) - return load_frag_half_at_byte_base(lds_a[sm], row_byte_base, half) - - def load_a_subtile_mi_regs(lds_a, sm, mi): - x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) - x1 = load_a_subtile_mi_half(lds_a, sm, mi, 1) - return pack_frag_halves(x0, x1) - - def load_a_subtile_regs(lds_a, sm): - return ( - load_a_subtile_mi_regs(lds_a, sm, 0), - load_a_subtile_mi_regs(lds_a, sm, 1), - load_a_subtile_mi_regs(lds_a, sm, 2), - load_a_subtile_mi_regs(lds_a, sm, 3), - ) - - def hk_one_k_with_refill( - k128, - cur_a, - cur_b, - next_a, - next_b, - refill_a, - refill_b, - a0_regs, - b0_regs, - ): - - # Wait only far enough for the current page; the next-page refill may remain in flight. - barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) - rocdl.sched_barrier(0) - - # A-top and B-left are both carried as complete 64-row register tiles, - # so their LDS half-pages can be refilled immediately. - a00, a01, a02, a03 = a0_regs - b00, b01, b02, b03 = b0_regs - - b10 = load_b_subtile_ni_regs(cur_b, 1, 0) - b11 = load_b_subtile_ni_regs(cur_b, 1, 1) - b12 = load_b_subtile_ni_regs(cur_b, 1, 2) - b13 = load_b_subtile_ni_regs(cur_b, 1, 3) - - # Refill the current ping-pong page with K+2, alternating A and B passes. - k_refill = fx.Index((k128 + 2) * BLOCK_K) - - # Q0: interleave the current tile's A-bottom LDS reads with K+2 - # refills and Q0 compute. Each complete A-bottom fragment is assembled - # from two independently scheduled K64 halves. - rocdl.sched_barrier(0) - a10_x0 = load_a_subtile_mi_half(cur_a, 1, 0, 0) - stage_a_subtile_pass(k_refill, 0, 0, refill_a) - mfma_2n(_acc_idx(0, 0, 0), a00, b00, b01) + value = value.to(output_fx_dtype) + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) + fx.memref_store_vec(Vec.filled(1, value, output_fx_dtype), reg) + fx.copy(c_store_atom, reg, fx.slice(c_div, (None, fx.Int32(c_idx)))) + + # Explicit register coordinates for HK-style four-quadrant mapping. + # BLOCK_M/BLOCK_N are 256x256. Four waves map to warp positions + # inside each 128x128 quadrant: + # cA: (warp_m, warp_n) + # cB: (warp_m, warp_n + 2) + # cC: (warp_m + 2, warp_n) + # cD: (warp_m + 2, warp_n + 2) + reg_k_col0 = lane_div_16 * 16 + reg_k_col1 = 64 + lane_div_16 * 16 + + # Every fragment row differs only by multiples of 16, so row % 16 is + # always lane_mod_16. Hoist the logical->physical XOR mapping once. + _, reg_lds_k_col0 = swizzle_128(lane_mod_16, reg_k_col0) + _, reg_lds_k_col1 = swizzle_128(lane_mod_16, reg_k_col1) + + reg_subtile_m_idx0 = wave_id // 2 + reg_subtile_n_idx0 = wave_id % 2 + + reserve_pinned_accumulators() + zero_pinned_accumulators() + + def load_b_subtile_ni_regs(lds_b, sn, ni): + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + b_row_addr = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 + return load_b_frag(lds_b, b_row_addr, sn) + + def load_b_subtile_regs(lds_b, sn): + return ( + load_b_subtile_ni_regs(lds_b, sn, 0), + load_b_subtile_ni_regs(lds_b, sn, 1), + load_b_subtile_ni_regs(lds_b, sn, 2), + load_b_subtile_ni_regs(lds_b, sn, 3), + ) - a10_x1 = load_a_subtile_mi_half(cur_a, 1, 0, 1) - stage_b_subtile_pass(k_refill, 0, 0, refill_b) - mfma_2n(_acc_idx(0, 0, 2), a00, b02, b03) + def load_a_subtile_mi_half(lds_a, sm, mi, half): + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + a_row_addr = subtile_m_idx * fx.Index(SUBTILE_M) + fx.Index(mi * MFMA_M) + lane_mod_16 + half_row = a_row_addr - fx.Index(sm * (BLOCK_M // 2)) + row_byte_base = half_row * fx.Index(BLOCK_K) + return load_frag_half_at_byte_base(lds_a[sm], row_byte_base, half) + + def load_a_subtile_mi_regs(lds_a, sm, mi): + x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) + x1 = load_a_subtile_mi_half(lds_a, sm, mi, 1) + return pack_frag_halves(x0, x1) + + def load_a_subtile_regs(lds_a, sm): + return ( + load_a_subtile_mi_regs(lds_a, sm, 0), + load_a_subtile_mi_regs(lds_a, sm, 1), + load_a_subtile_mi_regs(lds_a, sm, 2), + load_a_subtile_mi_regs(lds_a, sm, 3), + ) - a11_x0 = load_a_subtile_mi_half(cur_a, 1, 1, 0) - stage_a_subtile_pass(k_refill, 0, 1, refill_a) - mfma_2n(_acc_idx(0, 1, 0), a01, b00, b01) + def hk_one_k_with_refill( + k128, + cur_a, + cur_b, + next_a, + next_b, + refill_a, + refill_b, + a0_regs, + b0_regs, + ): - a11_x1 = load_a_subtile_mi_half(cur_a, 1, 1, 1) - stage_b_subtile_pass(k_refill, 0, 1, refill_b) - mfma_2n(_acc_idx(0, 1, 2), a01, b02, b03) + # Wait only far enough for the current page; the next-page refill may remain in flight. + barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) - a12_x0 = load_a_subtile_mi_half(cur_a, 1, 2, 0) - stage_a_subtile_pass(k_refill, 0, 2, refill_a) - mfma_2n(_acc_idx(0, 2, 0), a02, b00, b01) + # A-top and B-left are both carried as complete 64-row register tiles, + # so their LDS half-pages can be refilled immediately. + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs - a12_x1 = load_a_subtile_mi_half(cur_a, 1, 2, 1) - stage_b_subtile_pass(k_refill, 0, 2, refill_b) - mfma_2n(_acc_idx(0, 2, 2), a02, b02, b03) + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) - a13_x0 = load_a_subtile_mi_half(cur_a, 1, 3, 0) - stage_a_subtile_pass(k_refill, 0, 3, refill_a) - mfma_2n(_acc_idx(0, 3, 0), a03, b00, b01) + # Refill the current ping-pong page with K+2, alternating A and B passes. + k_refill = fx.Index((k128 + 2) * BLOCK_K) - a13_x1 = load_a_subtile_mi_half(cur_a, 1, 3, 1) - stage_b_subtile_pass(k_refill, 0, 3, refill_b) - mfma_2n(_acc_idx(0, 3, 2), a03, b02, b03) + # Q0: interleave the current tile's A-bottom LDS reads with K+2 + # refills and Q0 compute. Each complete A-bottom fragment is assembled + # from two independently scheduled K64 halves. + rocdl.sched_barrier(0) + a10_x0 = load_a_subtile_mi_half(cur_a, 1, 0, 0) + stage_a_subtile_pass(k_refill, 0, 0, refill_a) + mfma_2n(_acc_idx(0, 0, 0), a00, b00, b01) - hot_loop_scheduler_q0_refill_a1_2n() + a10_x1 = load_a_subtile_mi_half(cur_a, 1, 0, 1) + stage_b_subtile_pass(k_refill, 0, 0, refill_b) + mfma_2n(_acc_idx(0, 0, 2), a00, b02, b03) - # Retire the eight distributed A-bottom LDS reads before K+2 refills - # overwrite the current page's A-bottom half-page. Keep this wait as - # late as possible to maximize read/compute overlap. - rocdl.sched_barrier(0) - barrier(lgkmcnt=0) - rocdl.sched_barrier(0) + a11_x0 = load_a_subtile_mi_half(cur_a, 1, 1, 0) + stage_a_subtile_pass(k_refill, 0, 1, refill_a) + mfma_2n(_acc_idx(0, 1, 0), a01, b00, b01) - a10 = pack_frag_halves(a10_x0, a10_x1) - a11 = pack_frag_halves(a11_x0, a11_x1) - a12 = pack_frag_halves(a12_x0, a12_x1) - a13 = pack_frag_halves(a13_x0, a13_x1) + a11_x1 = load_a_subtile_mi_half(cur_a, 1, 1, 1) + stage_b_subtile_pass(k_refill, 0, 1, refill_b) + mfma_2n(_acc_idx(0, 1, 2), a01, b02, b03) - rocdl.sched_barrier(0) - stage_b_subtile_pass(k_refill, 1, 0, refill_b) - mfma_2n(_acc_idx(1, 0, 0), a00, b10, b11) + a12_x0 = load_a_subtile_mi_half(cur_a, 1, 2, 0) + stage_a_subtile_pass(k_refill, 0, 2, refill_a) + mfma_2n(_acc_idx(0, 2, 0), a02, b00, b01) - stage_a_subtile_pass(k_refill, 1, 0, refill_a) - mfma_2n(_acc_idx(1, 0, 2), a00, b12, b13) + a12_x1 = load_a_subtile_mi_half(cur_a, 1, 2, 1) + stage_b_subtile_pass(k_refill, 0, 2, refill_b) + mfma_2n(_acc_idx(0, 2, 2), a02, b02, b03) - stage_b_subtile_pass(k_refill, 1, 1, refill_b) - mfma_2n(_acc_idx(1, 1, 0), a01, b10, b11) + a13_x0 = load_a_subtile_mi_half(cur_a, 1, 3, 0) + stage_a_subtile_pass(k_refill, 0, 3, refill_a) + mfma_2n(_acc_idx(0, 3, 0), a03, b00, b01) - stage_a_subtile_pass(k_refill, 1, 1, refill_a) - mfma_2n(_acc_idx(1, 1, 2), a01, b12, b13) + a13_x1 = load_a_subtile_mi_half(cur_a, 1, 3, 1) + stage_b_subtile_pass(k_refill, 0, 3, refill_b) + mfma_2n(_acc_idx(0, 3, 2), a03, b02, b03) - stage_b_subtile_pass(k_refill, 1, 2, refill_b) - mfma_2n(_acc_idx(1, 2, 0), a02, b10, b11) + hot_loop_scheduler_q0_refill_a1_2n() - stage_a_subtile_pass(k_refill, 1, 2, refill_a) - mfma_2n(_acc_idx(1, 2, 2), a02, b12, b13) + # Retire the eight distributed A-bottom LDS reads before K+2 refills + # overwrite the current page's A-bottom half-page. Keep this wait as + # late as possible to maximize read/compute overlap. + rocdl.sched_barrier(0) + barrier(lgkmcnt=0) + rocdl.sched_barrier(0) - stage_b_subtile_pass(k_refill, 1, 3, refill_b) - mfma_2n(_acc_idx(1, 3, 0), a03, b10, b11) + a10 = pack_frag_halves(a10_x0, a10_x1) + a11 = pack_frag_halves(a11_x0, a11_x1) + a12 = pack_frag_halves(a12_x0, a12_x1) + a13 = pack_frag_halves(a13_x0, a13_x1) - stage_a_subtile_pass(k_refill, 1, 3, refill_a) - mfma_2n(_acc_idx(1, 3, 2), a03, b12, b13) - hot_loop_scheduler_q_refill_2n() + rocdl.sched_barrier(0) + stage_b_subtile_pass(k_refill, 1, 0, refill_b) + mfma_2n(_acc_idx(1, 0, 0), a00, b10, b11) - # Leave exactly the K+2 refill and scale loads outstanding. The following - # LDS reads consume the already-ready next page, not the page being refilled. - rocdl.sched_barrier(0) - barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) - rocdl.sched_barrier(0) + stage_a_subtile_pass(k_refill, 1, 0, refill_a) + mfma_2n(_acc_idx(1, 0, 2), a00, b12, b13) - next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) - mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) + stage_b_subtile_pass(k_refill, 1, 1, refill_b) + mfma_2n(_acc_idx(1, 1, 0), a01, b10, b11) - next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) - mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) + stage_a_subtile_pass(k_refill, 1, 1, refill_a) + mfma_2n(_acc_idx(1, 1, 2), a01, b12, b13) - next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) - mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) + stage_b_subtile_pass(k_refill, 1, 2, refill_b) + mfma_2n(_acc_idx(1, 2, 0), a02, b10, b11) - next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) - mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) + stage_a_subtile_pass(k_refill, 1, 2, refill_a) + mfma_2n(_acc_idx(1, 2, 2), a02, b12, b13) - next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) - mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) + stage_b_subtile_pass(k_refill, 1, 3, refill_b) + mfma_2n(_acc_idx(1, 3, 0), a03, b10, b11) - next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) - mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) + stage_a_subtile_pass(k_refill, 1, 3, refill_a) + mfma_2n(_acc_idx(1, 3, 2), a03, b12, b13) + hot_loop_scheduler_q_refill_2n() - next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) - mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) + # Leave exactly the K+2 refill and scale loads outstanding. The following + # LDS reads consume the already-ready next page, not the page being refilled. + rocdl.sched_barrier(0) + barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) - next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) - mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) - hot_loop_scheduler_q_prefetch_4n() + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) - next_a0_regs = (next_a00, next_a01, next_a02, next_a03) - next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) - return next_a0_regs, next_b0_regs + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) - def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): - barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) - a00, a01, a02, a03 = a0_regs - b00, b01, b02, b03 = b0_regs + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) - b10 = load_b_subtile_ni_regs(cur_b, 1, 0) - b11 = load_b_subtile_ni_regs(cur_b, 1, 1) - b12 = load_b_subtile_ni_regs(cur_b, 1, 2) - b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) - mfma_4n(_acc_idx(0, 0, 0), a00, b00, b01, b02, b03) - mfma_4n(_acc_idx(0, 1, 0), a01, b00, b01, b02, b03) - mfma_4n(_acc_idx(0, 2, 0), a02, b00, b01, b02, b03) - mfma_4n(_acc_idx(0, 3, 0), a03, b00, b01, b02, b03) + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) - rocdl.sched_barrier(0) - barrier(lgkmcnt=0) - rocdl.sched_barrier(0) + hot_loop_scheduler_q_prefetch_4n() - a10 = load_a_subtile_mi_regs(cur_a, 1, 0) - a11 = load_a_subtile_mi_regs(cur_a, 1, 1) - a12 = load_a_subtile_mi_regs(cur_a, 1, 2) - a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) - mfma_4n(_acc_idx(1, 0, 0), a00, b10, b11, b12, b13) - mfma_4n(_acc_idx(1, 1, 0), a01, b10, b11, b12, b13) - mfma_4n(_acc_idx(1, 2, 0), a02, b10, b11, b12, b13) - mfma_4n(_acc_idx(1, 3, 0), a03, b10, b11, b12, b13) + return next_a0_regs, next_b0_regs - rocdl.sched_barrier(0) - barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) - rocdl.sched_barrier(0) + def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): + barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) - next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) - mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs - next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) - mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) - next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) - mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 0, 0), a00, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 1, 0), a01, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 2, 0), a02, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 3, 0), a03, b00, b01, b02, b03) - next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) - mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) + rocdl.sched_barrier(0) + barrier(lgkmcnt=0) + rocdl.sched_barrier(0) - next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) - mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) - next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) - mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 0, 0), a00, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 1, 0), a01, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 2, 0), a02, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 3, 0), a03, b10, b11, b12, b13) - next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) - mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) + rocdl.sched_barrier(0) + barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) - next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) - mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) - hot_loop_scheduler_q_prefetch_4n() + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) - next_a0_regs = (next_a00, next_a01, next_a02, next_a03) - next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) - return next_a0_regs, next_b0_regs + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) - def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): - barrier(vmcnt=0, lgkmcnt=0) + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) - a00, a01, a02, a03 = a0_regs - b00, b01, b02, b03 = b0_regs + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) - # Materialize the remaining final-page A/B fragments once. The - # subsequent schedule is entirely register/AGPR traffic. - b10 = load_b_subtile_ni_regs(cur_b, 1, 0) - b11 = load_b_subtile_ni_regs(cur_b, 1, 1) - b12 = load_b_subtile_ni_regs(cur_b, 1, 2) - b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) - rocdl.sched_barrier(0) - barrier(lgkmcnt=0) - rocdl.sched_barrier(0) + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) - a10 = load_a_subtile_mi_regs(cur_a, 1, 0) - a11 = load_a_subtile_mi_regs(cur_a, 1, 1) - a12 = load_a_subtile_mi_regs(cur_a, 1, 2) - a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) - rocdl.sched_barrier(0) - barrier(lgkmcnt=0) - rocdl.sched_barrier(0) + return next_a0_regs, next_b0_regs - a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) - b_frags = (b00, b01, b02, b03, b10, b11, b12, b13) + def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): + barrier(vmcnt=0, lgkmcnt=0) - # Rolling final-page epilogue. - # - # Finalize accumulators in their own physical AGPR slots, but delay - # each AGPR read/store until several independent final MFMAs have - # been issued. - # - # MFMA 0, MFMA 1, MFMA 2, MFMA 3, drain 0, - # MFMA 4, drain 1, MFMA 5, drain 2, ... - # - # The buffer stores are only issued here; they may remain in flight - # while later MFMAs and accumulator drains continue. - FINAL_EPILOGUE_DEPTH = 4 - pending = [] - - for old_acc_idx in range_constexpr(ACCS_PER_WAVE): - subtile_id = old_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) - local_idx = old_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) - sm = subtile_id // 2 - sn = subtile_id % 2 - mi = local_idx // MFMA_N_PER_SUBTILE - ni = local_idx % MFMA_N_PER_SUBTILE + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs - a_frag_idx = sm * MFMA_M_PER_SUBTILE + mi - b_frag_idx = sn * MFMA_N_PER_SUBTILE + ni - - # Final MFMA remains in-place. The logical accumulator's own - # AGPR slot is unique and cannot conflict with another pending - # result, so no ad-hoc physical-slot permutation is needed. - pinned_final_mfma( - old_acc_idx, - old_acc_idx, - a_frags[a_frag_idx], - b_frags[b_frag_idx], - ) - pending.append(old_acc_idx) + # Materialize the remaining final-page A/B fragments once. The + # subsequent schedule is entirely register/AGPR traffic. + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) - # Drain the oldest completed result only after enough newer - # independent MFMAs have supplied the MFMA->AGPR-read spacing. - if len(pending) == FINAL_EPILOGUE_DEPTH: - drain_acc_idx = pending.pop(0) + rocdl.sched_barrier(0) + barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + + rocdl.sched_barrier(0) + barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) + b_frags = (b00, b01, b02, b03, b10, b11, b12, b13) + + # Rolling final-page epilogue. + # + # Finalize accumulators in their own physical AGPR slots, but delay + # each AGPR read/store until several independent final MFMAs have + # been issued. + # + # MFMA 0, MFMA 1, MFMA 2, MFMA 3, drain 0, + # MFMA 4, drain 1, MFMA 5, drain 2, ... + # + # The buffer stores are only issued here; they may remain in flight + # while later MFMAs and accumulator drains continue. + FINAL_EPILOGUE_DEPTH = 4 + pending = [] + + for old_acc_idx in range_constexpr(ACCS_PER_WAVE): + subtile_id = old_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = old_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE + + a_frag_idx = sm * MFMA_M_PER_SUBTILE + mi + b_frag_idx = sn * MFMA_N_PER_SUBTILE + ni + + # Final MFMA remains in-place. The logical accumulator's own + # AGPR slot is unique and cannot conflict with another pending + # result, so no ad-hoc physical-slot permutation is needed. + pinned_final_mfma( + old_acc_idx, + old_acc_idx, + a_frags[a_frag_idx], + b_frags[b_frag_idx], + ) + pending.append(old_acc_idx) + + # Drain the oldest completed result only after enough newer + # independent MFMAs have supplied the MFMA->AGPR-read spacing. + if len(pending) == FINAL_EPILOGUE_DEPTH: + drain_acc_idx = pending.pop(0) + acc = read_physical_accumulator_slot(drain_acc_idx) + store_acc_vector_for_logical_idx(drain_acc_idx, acc) + + # Flush the final results after all final-page MFMAs have issued. + for drain_acc_idx in pending: acc = read_physical_accumulator_slot(drain_acc_idx) store_acc_vector_for_logical_idx(drain_acc_idx, acc) - # Flush the final results after all final-page MFMAs have issued. - for drain_acc_idx in pending: - acc = read_physical_accumulator_slot(drain_acc_idx) - store_acc_vector_for_logical_idx(drain_acc_idx, acc) - - # Prologue: stage K0/K1 data into ping-pong LDS pages. - stage_a_subtile(fx.Index(0), 0, lds_a0) - stage_b_subtile(fx.Index(0), 0, lds_b0) - stage_b_subtile(fx.Index(0), 1, lds_b0) - stage_a_subtile(fx.Index(0), 1, lds_a0) + # Prologue: stage K0/K1 data into ping-pong LDS pages. + stage_a_subtile(fx.Index(0), 0, lds_a0) + stage_b_subtile(fx.Index(0), 0, lds_b0) + stage_b_subtile(fx.Index(0), 1, lds_b0) + stage_a_subtile(fx.Index(0), 1, lds_a0) - stage_a_subtile(fx.Index(BLOCK_K), 0, lds_a1) - stage_b_subtile(fx.Index(BLOCK_K), 0, lds_b1) - stage_b_subtile(fx.Index(BLOCK_K), 1, lds_b1) - stage_a_subtile(fx.Index(BLOCK_K), 1, lds_a1) + stage_a_subtile(fx.Index(BLOCK_K), 0, lds_a1) + stage_b_subtile(fx.Index(BLOCK_K), 0, lds_b1) + stage_b_subtile(fx.Index(BLOCK_K), 1, lds_b1) + stage_a_subtile(fx.Index(BLOCK_K), 1, lds_a1) - rocdl.sched_barrier(0) - barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 4 * LOAD_PASSES_B_SUBTILE) - rocdl.sched_barrier(0) - - a0_regs = load_a_subtile_regs(lds_a0, 0) + rocdl.sched_barrier(0) + barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 4 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) - rocdl.sched_barrier(0) - barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) - rocdl.sched_barrier(0) + a0_regs = load_a_subtile_regs(lds_a0, 0) - b0_regs = load_b_subtile_regs(lds_b0, 0) + rocdl.sched_barrier(0) + barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) - # Main HK loop: exactly one logical K128 per iteration. - # Even k consumes and refills LDS0; odd k does the same for LDS1. - for k128 in range_constexpr(NUM_K_TILES - 2): - if (k128 % 2) == 0: - a0_regs, b0_regs = hk_one_k_with_refill( - k128, + b0_regs = load_b_subtile_regs(lds_b0, 0) + + # Main HK loop: exactly one logical K128 per iteration. + # Even k consumes and refills LDS0; odd k does the same for LDS1. + for k128 in range_constexpr(NUM_K_TILES - 2): + if (k128 % 2) == 0: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + else: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) + + # Common two-page tail. The penultimate tile uses Q2/Q3 carry-prefetch + # to prepare A-top/B-left for the final tile, but performs no K+2 refill. + if (NUM_K_TILES % 2) == 0: + a0_regs, b0_regs = hk_one_k_tail_with_next( lds_a0, lds_b0, lds_a1, lds_b1, - lds_a0, - lds_b0, a0_regs, b0_regs, ) + hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs) else: - a0_regs, b0_regs = hk_one_k_with_refill( - k128, + a0_regs, b0_regs = hk_one_k_tail_with_next( lds_a1, lds_b1, lds_a0, lds_b0, - lds_a1, - lds_b1, a0_regs, b0_regs, ) - - # Common two-page tail. The penultimate tile uses Q2/Q3 carry-prefetch - # to prepare A-top/B-left for the final tile, but performs no K+2 refill. - if (NUM_K_TILES % 2) == 0: - a0_regs, b0_regs = hk_one_k_tail_with_next( - lds_a0, - lds_b0, - lds_a1, - lds_b1, - a0_regs, - b0_regs, - ) - hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs) + hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs) + + # Dispatch. Dense runs one XCD-swizzled tile. Grouped runs a persistent + # loop over precomputed (pid_m, pid_n, expert) descriptors, selecting the + # schedule at compile time (``sched``): "worksteal" atomically claims the + # next global slot each iteration; "static" owns a fixed grid-stride slice. + def _run_grouped_tile(slot): + pid_m = _desc_field(slot, fx.Int32(0)) + pid_n = _desc_field(slot, fx.Int32(1)) + expert = _desc_field(slot, fx.Int32(2)) + process_tile(pid_m, pid_n, expert) + + if const_expr(grouped): + if const_expr(sched == "static"): + stride = fx.Int32(gpu.grid_dim.x) + slot = fx.Int32(gpu.block_id("x")) + while slot < num_tiles: + _run_grouped_tile(slot) + slot = slot + stride + else: + active = fx.Int32(1) == fx.Int32(1) + while active: + slot = _claim_next_slot() + in_range = slot < num_tiles + if in_range: + _run_grouped_tile(slot) + active = in_range else: - a0_regs, b0_regs = hk_one_k_tail_with_next( - lds_a1, - lds_b1, - lds_a0, - lds_b0, - a0_regs, - b0_regs, - ) - hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs) + if const_expr(use_xcd_remap): + pid_m, pid_n = xcd_swizzle(num_blocks_m, num_blocks_n) + else: + pid_m, pid_n = divmod(fx.block_idx.x, num_blocks_n) + process_tile(pid_m, pid_n, fx.Int32(0)) @flyc.jit def launch_gemm( @@ -966,10 +1085,17 @@ def launch_gemm( Aux: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32, + num_groups: fx.Int32 = 0, + work_queue: fx.Tensor = None, + work_counter: fx.Tensor = None, stream: fx.Stream = fx.Stream(None), ): - # The integration only dispatches aligned shapes; no partial-tile masking exists. + # Dense: one workgroup per output tile. Grouped/wgrad: a fixed num-CU + # persistent grid that drains the descriptor queue. Placeholder work + # tensors for the dense path are allocated host-side in ``doGemm``. grid_x = (c_m // BLOCK_M) * (c_n // BLOCK_N) + if const_expr(group or wgrad): + grid_x = get_num_cus() kernel_gemm( A, B, @@ -980,6 +1106,9 @@ def launch_gemm( Aux, c_m, c_n, + num_groups, + work_queue, + work_counter, value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, ).launch(grid=(grid_x, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) @@ -994,6 +1123,9 @@ def _cached_launch( output_dtype: torch.dtype, use_xcd_remap: bool = True, epilogue: str = "DEFAULT", + group: bool = False, + sched: str = "worksteal", + wgrad: bool = False, ): return _compile_kernel( K, @@ -1002,6 +1134,9 @@ def _cached_launch( output_dtype, use_xcd_remap=use_xcd_remap, epilogue=epilogue, + group=group, + sched=sched, + wgrad=wgrad, ) @@ -1101,10 +1236,28 @@ def doGemm( epilogue: str = "DEFAULT", bias: torch.Tensor = None, aux: torch.Tensor = None, + group: bool = False, + num_groups: int = 0, + work_queue: torch.Tensor = None, + work_counter: torch.Tensor = None, + sched: str = "worksteal", + wgrad: bool = False, + m: int = None, + n: int = None, + k: int = None, ): - """Launch tensor-wise FP8 GEMM with TE-style inverse input scales.""" - M_runtime, K_runtime = A.shape - N_runtime, Kb_runtime = B.shape + """Launch tensor-wise FP8 GEMM with TE-style inverse input scales. + + Dense derives M/N/K from operand shapes. Grouped/wgrad pass explicit per-expert + ``m``/``n``/``k`` (the packed operands do not reveal the per-expert geometry), + a length-``num_groups`` scale vector for each operand, and the descriptor queue. + """ + if group or wgrad: + M_runtime, N_runtime, K_runtime = int(m), int(n), int(k) + else: + M_runtime, K_runtime = A.shape + N_runtime, Kb_runtime = B.shape + assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" supported_fp8_dtypes = ( torch.float8_e4m3fn, torch.float8_e5m2, @@ -1116,7 +1269,6 @@ def doGemm( torch.bfloat16, torch.float32, ), f"C dtype must be torch.float16, torch.bfloat16, or torch.float32, got {C.dtype}" - assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" require_block_tiling( M_runtime, N_runtime, @@ -1127,12 +1279,14 @@ def doGemm( label="FP8 GEMM", ) require_launch_size("FP8 GEMM", ("A", A), ("B", B), ("C", C)) - assert A_scale_inv.dtype == torch.float32 and A_scale_inv.numel() == 1 - assert B_scale_inv.dtype == torch.float32 and B_scale_inv.numel() == 1 - assert C.shape == ( - M_runtime, - N_runtime, - ), f"C shape {tuple(C.shape)} != ({M_runtime}, {N_runtime})" + # Dense: one scalar inverse scale per operand. Grouped: a per-expert vector. + expected_scale_numel = num_groups if (group or wgrad) else 1 + assert A_scale_inv.dtype == torch.float32 and A_scale_inv.numel() == expected_scale_numel + assert B_scale_inv.dtype == torch.float32 and B_scale_inv.numel() == expected_scale_numel + # wgrad's C is packed [G*N_out, K_in] on the output-M axis; fwd/dgrad/dense C is + # the unpacked (M_runtime, N_runtime). + expected_c = (num_groups * M_runtime, N_runtime) if wgrad else (M_runtime, N_runtime) + assert tuple(C.shape) == expected_c, f"C shape {tuple(C.shape)} != {expected_c}" if epilogue not in ("DEFAULT", "BIAS", "GELU_AUX", "GELU_AUX_BIAS"): raise ValueError(f"Unsupported FP8 epilogue: {epilogue}") @@ -1183,6 +1337,15 @@ def doGemm( else: Aux_arg = torch.zeros(1, dtype=C.dtype, device=A.device) + # Grouped runtime tensors are required by the kernel signature; the dense + # binary never dereferences them but still needs placeholder buffers. + if group or wgrad: + work_queue_arg = work_queue.contiguous().view(-1) + work_counter_arg = work_counter.contiguous().view(-1) + else: + work_queue_arg = torch.zeros(1, dtype=torch.int32, device=A.device) + work_counter_arg = torch.zeros(1, dtype=torch.int32, device=A.device) + launch = _cached_launch( int(K_runtime), A.dtype, @@ -1190,6 +1353,9 @@ def doGemm( C.dtype, bool(use_xcd_remap), epilogue, + group, + sched, + wgrad, ) launch( A_arg, @@ -1201,5 +1367,8 @@ def doGemm( Aux_arg, M_runtime, N_runtime, + num_groups, + work_queue_arg, + work_counter_arg, stream=stream, ) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py index f91aa6e043..36136f08a9 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py @@ -12,12 +12,18 @@ G2S/S2R loaders, buffer-tensor makers) stays in the per-dtype modules. """ +import torch + import flydsl.expr as fx from flydsl.expr import arith, const_expr, range_constexpr, rocdl from .exceptions import FlyDSLUnsupportedError +def get_num_cus() -> int: + return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + + def require_block_tiling(m, n, k, *, block_m, block_n, block_k, label, min_k_tiles=4): """Validate the host-side M/N/K tiling contract shared by every GEMM core. diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index b1aa0b3e48..8fee753fa7 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -9,6 +9,7 @@ import torch import transformer_engine_torch as tex +from . import grouped_gemm_utils as gd from .exceptions import FlyDSLUnsupportedError from .half_prec_gemm import bf16_matmul, fp16_matmul @@ -1008,27 +1009,8 @@ def _run_mxfp8( return D, aux -# Above this many persistent-loop rounds, dynamic atomic work-stealing pays for -# itself; below it a single grid-stride sweep already balances the tiles, so the -# atomic-free static schedule is preferred. Matches the HK reference's SCHED_ROUNDS. -_GROUPED_SCHED_ROUNDS = 4 - - -def _select_grouped_sched(num_tiles: int, device) -> str: - """Pick the grouped-GEMM schedule: 'worksteal' (atomic) or 'static' (grid-stride). - - Honors the NVTE_FLYDSL_GROUPED_SCHED override ('auto', 'static', 'worksteal') - for testing; otherwise applies the round-count heuristic against the CU grid. - """ - override = os.getenv("NVTE_FLYDSL_GROUPED_SCHED", "auto").lower() - if override in ("static", "worksteal"): - return override - if override not in ("auto", ""): - raise ValueError( - f"NVTE_FLYDSL_GROUPED_SCHED must be auto|static|worksteal, got {override!r}" - ) - num_cus = torch.cuda.get_device_properties(device).multi_processor_count - return "static" if num_tiles <= _GROUPED_SCHED_ROUNDS * num_cus else "worksteal" +# Schedule selection now lives in the shared grouped driver (dtype-independent). +_select_grouped_sched = gd.select_grouped_sched def _run_mxfp8_grouped( @@ -1111,17 +1093,8 @@ def _run_mxfp8_grouped( # Per-expert M-block ranges over the packed [M_total, N] output. Each expert's # token count must be BLOCK_M-aligned (enforced upstream by the tiling check), # so tiles never straddle a group boundary. - descriptors = [] - for g in range(num_groups): - m_start_blk = int(group_offs_cpu[g]) // BLOCK_M - m_end_blk = int(group_offs_cpu[g + 1]) // BLOCK_M - for pid_m in range(m_start_blk, m_end_blk): - for pid_n in range(num_pid_n): - descriptors.append((pid_m, pid_n, g)) - - num_tiles = len(descriptors) - work_queue = torch.tensor(descriptors, dtype=torch.int32, device=a_data.device).reshape(-1) - work_counter = torch.zeros(1, dtype=torch.int32, device=a_data.device) + descriptors = gd.build_fwd_dgrad_descriptors(group_offs_cpu, num_groups, num_pid_n, BLOCK_M) + work_queue, work_counter, num_tiles = gd.make_work_queue(descriptors, a_data.device) # Schedule selection (compile-time specialization on the kernel side): # "worksteal": persistent atomic tile-claiming, best for imbalanced experts. @@ -1193,18 +1166,11 @@ def _run_mxfp8_grouped_wgrad( m_pad = ((max_m + BLOCK_M - 1) // BLOCK_M) * BLOCK_M m_pad = max(m_pad, 512) - def _pad_rows(t, rows): - if t.size(0) == rows: - return t.contiguous() - out = t.new_zeros((rows, *t.shape[1:])) - out[: t.size(0)].copy_(t) - return out - scale_rows = m_pad // SCALE_GROUP_SIZE - dy_data = torch.cat([_pad_rows(t, m_pad) for t in dy_cols], dim=0) - dy_scale = torch.cat([_pad_rows(t, scale_rows) for t in dy_scales], dim=0) - x_data = torch.cat([_pad_rows(t, m_pad) for t in x_cols], dim=0) - x_scale = torch.cat([_pad_rows(t, scale_rows) for t in x_scales], dim=0) + dy_data = gd.uniform_pad_concat(dy_cols, m_pad) + dy_scale = gd.uniform_pad_concat(dy_scales, scale_rows) + x_data = gd.uniform_pad_concat(x_cols, m_pad) + x_scale = gd.uniform_pad_concat(x_scales, scale_rows) if dy_data.dtype == torch.uint8: dy_data = reinterpret_as_fp8_tensor(dy_data, dy_fp8_dtype) @@ -1224,15 +1190,8 @@ def _pad_rows(t, rows): # kernel re-packs the output row via expert*N_out. num_pid_m = n_out // BLOCK_M num_pid_n = k_in // BLOCK_N - descriptors = [ - (pid_m, pid_n, g) - for g in range(num_groups) - for pid_m in range(num_pid_m) - for pid_n in range(num_pid_n) - ] - num_tiles = len(descriptors) - work_queue = torch.tensor(descriptors, dtype=torch.int32, device=device).reshape(-1) - work_counter = torch.zeros(1, dtype=torch.int32, device=device) + descriptors = gd.build_wgrad_descriptors(num_groups, num_pid_m, num_pid_n) + work_queue, work_counter, num_tiles = gd.make_work_queue(descriptors, device) sched = _select_grouped_sched(num_tiles, device) mxfp8_matmul( @@ -1251,6 +1210,305 @@ def _pad_rows(t, rows): return dW +def _run_bf16_grouped( + a_data, + b_data, + D, + *, + group_offs, + m_total, + n, + k, + output_dtype: torch.dtype, + layout: str = "TN", + bias=None, +): + """Dispatch a packed BF16 grouped GEMM through the persistent kernel. + + Scaleless analog of :func:`_run_mxfp8_grouped`; the packed operand contracts + match the MXFP8 TN/NN layouts exactly (experts packed on B's output-N axis for + TN, on B's contraction axis for NN), only without E8M0 scales: + + TN (forward, C = A @ B^T): + a = packed activations [M_total, K] b = packed weights [G*N, K] + NN (dgrad, dX = dY @ W): + a = packed grad_output [M_total, N_out] b = packed weights [G*N_out, K_in] + """ + from .half_prec_gemm import doGemm as _half_prec_doGemm + + if layout not in ("TN", "NN"): + raise FlyDSLUnsupportedError(f"FlyDSL BF16 grouped GEMM supports TN/NN, got {layout}") + + D = _validate_or_allocate_output( + D, + shape=torch.Size((m_total, n)), + dtype=output_dtype, + device=a_data.device, + backend_name=f"BF16 grouped {layout}", + ) + + from .half_prec_gemm import _BLOCK_M, _BLOCK_N + + num_groups = int(group_offs.numel()) - 1 + num_pid_n = n // _BLOCK_N + group_offs_cpu = group_offs.to(device="cpu", dtype=torch.int64) + descriptors = gd.build_fwd_dgrad_descriptors(group_offs_cpu, num_groups, num_pid_n, _BLOCK_M) + work_queue, work_counter, num_tiles = gd.make_work_queue(descriptors, a_data.device) + sched = _select_grouped_sched(num_tiles, a_data.device) + + epilogue, bias_arg = _resolve_epilogue(bias, n) + _half_prec_doGemm( + a_data, + b_data, + D.view(m_total, n), + layout=layout, + m=m_total, + n=n, + k=k, + input_dtype=torch.bfloat16, + label="BF16 grouped", + epilogue=epilogue, + bias=bias_arg, + group=True, + num_groups=num_groups, + work_queue=work_queue, + work_counter=work_counter, + sched=sched, + ) + return D + + +def _run_bf16_grouped_wgrad( + dy_cols, # list[G] of per-expert dY, each [m_e, N_out] + x_cols, # list[G] of per-expert X, each [m_e, K_in] + dW, # output [G*N_out, K_in] + *, + m_splits, + n_out, + k_in, + output_dtype, +): + """Grouped BF16 wgrad (NT): dW_e = dY_e^T @ X_e, per expert, one launch. + + Scaleless analog of :func:`_run_mxfp8_grouped_wgrad`: pad each expert's tokens + to a uniform M_pad on the contraction axis, pack dY/X to [G*M_pad, *], and emit + one NT launch writing the packed [G*N_out, K_in] gradient (each expert's output + tile shifted via store_m_expert_off). + """ + from .half_prec_gemm import doGemm as _half_prec_doGemm, _BLOCK_M, _BLOCK_N + + num_groups = len(m_splits) + device = dW.device + if n_out % _BLOCK_M != 0: + raise FlyDSLUnsupportedError(f"wgrad requires N_out divisible by {_BLOCK_M}, got {n_out}") + if k_in % _BLOCK_N != 0: + raise FlyDSLUnsupportedError(f"wgrad requires K_in divisible by {_BLOCK_N}, got {k_in}") + + max_m = max(int(m) for m in m_splits) + m_pad = ((max_m + _BLOCK_M - 1) // _BLOCK_M) * _BLOCK_M + m_pad = max(m_pad, 512) + + dy_data = gd.uniform_pad_concat(dy_cols, m_pad) + x_data = gd.uniform_pad_concat(x_cols, m_pad) + + dW = _validate_or_allocate_output( + dW, + shape=torch.Size((num_groups * n_out, k_in)), + dtype=output_dtype, + device=device, + backend_name="BF16 grouped wgrad NT", + ) + + num_pid_m = n_out // _BLOCK_M + num_pid_n = k_in // _BLOCK_N + descriptors = gd.build_wgrad_descriptors(num_groups, num_pid_m, num_pid_n) + work_queue, work_counter, num_tiles = gd.make_work_queue(descriptors, device) + sched = _select_grouped_sched(num_tiles, device) + + # The kernel's compile-time contraction K is the PER-EXPERT padded token count; + # both operands are packed [G*M_pad, *] so doGemm de-packs K via num_groups. + _half_prec_doGemm( + dy_data, + x_data, + dW, + layout="NT", + m=n_out, + n=k_in, + k=m_pad, + input_dtype=torch.bfloat16, + label="BF16 grouped wgrad", + group=False, + num_groups=num_groups, + work_queue=work_queue, + work_counter=work_counter, + sched=sched, + wgrad=True, + ) + return dW + + +def _run_fp8_grouped( + a_data, + a_scale_vec, # [num_groups] fp32 per-expert inverse scales + b_data, + b_scale_vec, # [num_groups] fp32 per-expert inverse scales + D, + *, + group_offs, + m_total, + n, + k, + output_dtype: torch.dtype, + bias=None, +): + """Dispatch a packed tensor-wise FP8 grouped GEMM (TN) through the persistent kernel. + + Tensor-wise FP8 is TN-only; grouped dgrad/wgrad reach this same TN kernel via + host-materialized transpose backings (see the backward helpers). Experts pack on + B's output-N axis ([G*N, K]); each expert carries one scalar inverse scale, so the + scales arrive as length-``num_groups`` vectors the kernel indexes by ``expert``. + """ + from .fp8_gemm import doGemm as _fp8_doGemm + + num_groups = int(group_offs.numel()) - 1 + if a_scale_vec.numel() != num_groups or b_scale_vec.numel() != num_groups: + raise ValueError( + f"FP8 grouped expects length-{num_groups} scale vectors, got " + f"A={a_scale_vec.numel()}, B={b_scale_vec.numel()}" + ) + + a_fp8_dtype = getattr(a_data, "_fp8_dtype", None) + b_fp8_dtype = getattr(b_data, "_fp8_dtype", None) + if a_data.dtype == torch.uint8: + a_data = reinterpret_as_fp8_tensor(a_data, a_fp8_dtype) + if b_data.dtype == torch.uint8: + b_data = reinterpret_as_fp8_tensor(b_data, b_fp8_dtype) + + D = _validate_or_allocate_output( + D, + shape=torch.Size((m_total, n)), + dtype=output_dtype, + device=a_data.device, + backend_name="FP8 grouped TN", + ) + + from .fp8_gemm import _BLOCK_M, _BLOCK_N + + num_pid_n = n // _BLOCK_N + group_offs_cpu = group_offs.to(device="cpu", dtype=torch.int64) + descriptors = gd.build_fwd_dgrad_descriptors(group_offs_cpu, num_groups, num_pid_n, _BLOCK_M) + work_queue, work_counter, num_tiles = gd.make_work_queue(descriptors, a_data.device) + sched = _select_grouped_sched(num_tiles, a_data.device) + + epilogue, bias_arg = _resolve_epilogue(bias, n) + _fp8_doGemm( + a_data, + b_data, + D.view(m_total, n), + a_scale_vec.to(torch.float32).contiguous(), + b_scale_vec.to(torch.float32).contiguous(), + epilogue=epilogue, + bias=bias_arg, + group=True, + num_groups=num_groups, + work_queue=work_queue, + work_counter=work_counter, + sched=sched, + m=m_total, + n=n, + k=k, + ) + return D + + +def _run_fp8_grouped_wgrad( + dyt_cols, # list[G] of per-expert dY^T (host-transposed) [N_out, m_e] + dy_scale_vec, # [num_groups] fp32 per-expert inverse scales for dY + xt_cols, # list[G] of per-expert X^T (host-transposed) [K_in, m_e] + x_scale_vec, # [num_groups] fp32 per-expert inverse scales for X + dW, # output [G*N_out, K_in] + *, + m_splits, + n_out, + k_in, + dy_fp8_dtype, + x_fp8_dtype, + output_dtype, +): + """Grouped tensor-wise FP8 wgrad: dW_e = dY_e^T @ X_e, one TN launch. + + Tensor-wise FP8 is TN-NORMAL-read: the contraction axis (tokens m) is the + trailing dim and is never packed, so -- unlike the MXFP8/BF16 NT path -- the + operands arrive already transposed (the FP8 ``_transpose`` backing is a + materialized row-major transpose). Per expert: A = dY^T [N_out, m_e], + B = X^T [K_in, m_e]. Each expert's token count m_e is zero-padded up to a + common ``m_pad`` on the trailing (contraction) axis; experts stack on the + feature (output) axis to A [G*N_out, m_pad] / B [G*K_in, m_pad]. Output dW is + packed [G*N_out, K_in] on the M-row. Per-expert scalar scales are length-G + vectors indexed by expert in the kernel. + """ + from .fp8_gemm import doGemm as _fp8_doGemm, _BLOCK_M, _BLOCK_N + + num_groups = len(m_splits) + device = dW.device + if n_out % _BLOCK_M != 0: + raise FlyDSLUnsupportedError(f"wgrad requires N_out divisible by {_BLOCK_M}, got {n_out}") + if k_in % _BLOCK_N != 0: + raise FlyDSLUnsupportedError(f"wgrad requires K_in divisible by {_BLOCK_N}, got {k_in}") + + max_m = max(int(m) for m in m_splits) + m_pad = ((max_m + _BLOCK_M - 1) // _BLOCK_M) * _BLOCK_M + m_pad = max(m_pad, 512) + + # Zero-pad the trailing (token/contraction) axis of each [feature, m_e] + # transpose up to m_pad, then concat on the feature axis. + def _pad_cols(t): + if t.size(1) == m_pad: + return t.contiguous() + out = t.new_zeros((t.size(0), m_pad)) + out[:, : t.size(1)].copy_(t) + return out + + dy_data = torch.cat([_pad_cols(t) for t in dyt_cols], dim=0) + x_data = torch.cat([_pad_cols(t) for t in xt_cols], dim=0) + if dy_data.dtype == torch.uint8: + dy_data = reinterpret_as_fp8_tensor(dy_data, dy_fp8_dtype) + if x_data.dtype == torch.uint8: + x_data = reinterpret_as_fp8_tensor(x_data, x_fp8_dtype) + + dW = _validate_or_allocate_output( + dW, + shape=torch.Size((num_groups * n_out, k_in)), + dtype=output_dtype, + device=device, + backend_name="FP8 grouped wgrad TN", + ) + + num_pid_m = n_out // _BLOCK_M + num_pid_n = k_in // _BLOCK_N + descriptors = gd.build_wgrad_descriptors(num_groups, num_pid_m, num_pid_n) + work_queue, work_counter, num_tiles = gd.make_work_queue(descriptors, device) + sched = _select_grouped_sched(num_tiles, device) + + _fp8_doGemm( + dy_data, + x_data, + dW, + dy_scale_vec.to(torch.float32).contiguous(), + x_scale_vec.to(torch.float32).contiguous(), + group=False, + num_groups=num_groups, + work_queue=work_queue, + work_counter=work_counter, + sched=sched, + wgrad=True, + m=n_out, + n=k_in, + k=m_pad, + ) + return dW + + def _select_fp8_storage_for_layout(A, transa, B, transb): """Select existing TE FP8 backings and normalize to one core contract. diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py new file mode 100644 index 0000000000..f83f84d7f3 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py @@ -0,0 +1,173 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# License for AMD contributions = MIT. See LICENSE for more information + +"""Shared utilities for the FlyDSL persistent work-stealing grouped GEMM. + +The grouped (MoE) GEMM path is dtype-independent: one packed kernel launch spans +all local experts, a fixed num-CU persistent grid claims tiles from a precomputed +descriptor queue, and the per-expert offset is folded into operand/scale/output +addressing. MXFP8, BF16/FP16, and tensor-wise FP8 all share this machinery. + +This module holds the reusable pieces that can live outside the ``@flyc.kernel`` +body. The persistent loop itself -- the ``if tx==0`` atomic claim and the ``while`` +drain loops -- is NOT here: the FlyDSL compiler AST-rewrites only the decorated +kernel's own source, so that runtime-SSA control flow MUST stay inline in each +``kernel_gemm``. What's shared instead: + +* **Device-side IR emitters** (``global_i32_ptr``, ``make_claim_broadcast``, + ``read_desc_field``, ``num_tiles_bound``) -- control-flow-free functions that + append IR at the current insertion point, so they emit identically whether + defined inline or imported. +* **Host-side builders** (descriptor-queue construction, schedule selection, + work-counter allocation, uniform token-axis padding) -- plain Python/torch, + identical across dtypes; they differ only in whether scales are also padded. +""" + +import os + +import torch + +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import fly as _fly +from flydsl._mlir.dialects import llvm +from flydsl.expr.typing import T + + +# --------------------------------------------------------------------------- +# Device-side IR emitters (called while tracing a kernel; pure, no control flow) +# --------------------------------------------------------------------------- + + +def global_i32_ptr(tensor, index): + """Build an address-space-1 ``!llvm.ptr`` to ``tensor[index]`` (int32). + + ``llvm.atomicrmw`` requires an ``!llvm.ptr`` (not an ``!fly.ptr``); this + recovers the raw base pointer and offsets it by ``index`` int32 elements. + """ + ptr_ty = ir.Type.parse("!llvm.ptr<1>") + base = _fly.extract_aligned_pointer_as_index(ptr_ty, tensor) + base = llvm.PtrToIntOp(T.i64, base).result + byte_off = fx.as_ir_value(fx.Int64(fx.Index(index) * fx.Index(4))) + addr = llvm.AddOp(base, byte_off, llvm.IntegerOverflowFlags(0)).result + return llvm.IntToPtrOp(ptr_ty, addr).result + + +def make_claim_broadcast(claim_slot_array): + """Recast a one-int32 LDS scratch array into a claim iterator + 1-elem view. + + Thread 0 writes the atomically-claimed slot via ``claim_iter``; after a + workgroup barrier every thread reads the same value through ``claim_view``. + """ + claim_iter = fx.recast_iter(fx.Int32, claim_slot_array.ptr) + claim_view = fx.make_view(claim_iter, fx.make_layout(1, 1)) + return claim_iter, claim_view + + +def read_desc_field(wq_div, desc_ld_atom, slot, field): + """Read one int32 of the packed ``(pid_m, pid_n, expert)`` descriptor. + + ``wq_div`` is the ``logical_divide`` view of the work-queue buffer tensor and + ``desc_ld_atom`` an int32 ``BufferCopy32b`` copy atom. + """ + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Int32) + fx.copy(desc_ld_atom, fx.slice(wq_div, (None, slot * fx.Int32(3) + field)), reg) + return fx.memref_load_vec(reg)[0] + + +def num_tiles_bound(c_m, c_n, block_m, block_n, num_groups, wgrad): + """Drain bound for the persistent loop. + + For dense/grouped fwd/dgrad, ``c_m`` already spans all experts, so the tile + count is complete. For wgrad, ``c_m = N_out`` is the PER-EXPERT output-row + count (the packed ``[G*N_out, K_in]`` output is addressed via + ``store_m_expert_off``), so multiply by ``num_groups``. + """ + num_tiles = (c_m // block_m) * (c_n // block_n) + if wgrad: + num_tiles = num_tiles * num_groups + return num_tiles + + +# --------------------------------------------------------------------------- +# Host-side builders (plain torch; dtype-independent) +# --------------------------------------------------------------------------- + +# Above this many persistent-loop rounds, dynamic atomic work-stealing pays for +# itself; below it a single grid-stride sweep already balances the tiles, so the +# atomic-free static schedule is preferred. Matches the HK reference's SCHED_ROUNDS. +_GROUPED_SCHED_ROUNDS = 4 + + +def select_grouped_sched(num_tiles: int, device) -> str: + """Pick the grouped-GEMM schedule: 'worksteal' (atomic) or 'static' (grid-stride). + + Honors the ``NVTE_FLYDSL_GROUPED_SCHED`` override ('auto', 'static', + 'worksteal') for testing; otherwise applies the round-count heuristic against + the CU grid. + """ + override = os.getenv("NVTE_FLYDSL_GROUPED_SCHED", "auto").lower() + if override in ("static", "worksteal"): + return override + if override not in ("auto", ""): + raise ValueError( + f"NVTE_FLYDSL_GROUPED_SCHED must be auto|static|worksteal, got {override!r}" + ) + num_cus = torch.cuda.get_device_properties(device).multi_processor_count + return "static" if num_tiles <= _GROUPED_SCHED_ROUNDS * num_cus else "worksteal" + + +def build_fwd_dgrad_descriptors(group_offs_cpu, num_groups, num_pid_n, block_m): + """Build the fwd/dgrad ``(pid_m, pid_n, expert)`` descriptor list. + + Each expert ``g`` owns the packed output M-block range + ``[group_offs[g], group_offs[g+1])`` (BLOCK_M-aligned so tiles never straddle + a group boundary), tiled over ``num_pid_n`` columns. + """ + descriptors = [] + for g in range(num_groups): + m_start_blk = int(group_offs_cpu[g]) // block_m + m_end_blk = int(group_offs_cpu[g + 1]) // block_m + for pid_m in range(m_start_blk, m_end_blk): + for pid_n in range(num_pid_n): + descriptors.append((pid_m, pid_n, g)) + return descriptors + + +def build_wgrad_descriptors(num_groups, num_pid_m, num_pid_n): + """Build the wgrad descriptor list: one tile per expert's local [N_out, K_in].""" + return [ + (pid_m, pid_n, g) + for g in range(num_groups) + for pid_m in range(num_pid_m) + for pid_n in range(num_pid_n) + ] + + +def make_work_queue(descriptors, device): + """Flatten descriptors into the kernel's queue buffer + zeroed atomic counter. + + Returns ``(work_queue, work_counter, num_tiles)``. ``work_queue`` is a flat + int32 ``[num_tiles*3]`` tensor; ``work_counter`` a zeroed one-int32 dispenser. + """ + num_tiles = len(descriptors) + work_queue = torch.tensor(descriptors, dtype=torch.int32, device=device).reshape(-1) + work_counter = torch.zeros(1, dtype=torch.int32, device=device) + return work_queue, work_counter, num_tiles + + +def uniform_pad_concat(tensors, rows): + """Zero-pad each tensor's leading dim up to ``rows`` and concat on dim 0. + + Used by the wgrad uniform-pad path to stack per-expert operands (or their + scales) on the contraction/token axis. Padded rows are zero, so they + contribute nothing to the token-axis contraction. + """ + def _pad_rows(t): + if t.size(0) == rows: + return t.contiguous() + out = t.new_zeros((rows, *t.shape[1:])) + out[: t.size(0)].copy_(t) + return out + + return torch.cat([_pad_rows(t) for t in tensors], dim=0) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py index 36ed68d5b1..8ecf89aa1f 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py @@ -37,8 +37,9 @@ from flydsl.expr.typing import Vector as Vec # Transformer Engine-local FlyDSL utilities. +from . import grouped_gemm_utils as gd from .exceptions import FlyDSLUnsupportedError -from .gemm_common_utils import require_block_tiling, require_launch_size +from .gemm_common_utils import get_num_cus, require_block_tiling, require_launch_size from .fp16_gemm_utils import ( G2SLoader, S2RLoader, @@ -110,6 +111,9 @@ def _compile_kernel( mfma_suffix: str, use_xcd_remap: bool = True, epilogue: str = "DEFAULT", + group: bool = False, + sched: str = "worksteal", + wgrad: bool = False, ): """Build one compile-time-specialized TN, NN, or NT half-precision kernel. @@ -372,6 +376,9 @@ class SharedStorage: b0_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] b1_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] b1_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + # Persistent grouped path: one-int32 scratch to broadcast the atomically + # claimed queue slot from thread 0 to the whole workgroup. + claim_slot: fx.Array[fx.Int32, 1, 16] @flyc.kernel(known_block_size=[NUM_THREADS, 1, 1]) def kernel_gemm( @@ -382,6 +389,9 @@ def kernel_gemm( Aux: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32, + num_groups: fx.Int32 = 0, + work_queue: fx.Tensor = None, + work_counter: fx.Tensor = None, ): lds = fx.SharedAllocator().allocate(SharedStorage).peek() lds_a0 = (lds.a0_0, lds.a0_1) @@ -397,821 +407,922 @@ def kernel_gemm( a_div = fx.logical_divide(gA, fx.make_layout(1, 1)) b_div = fx.logical_divide(gB, fx.make_layout(1, 1)) tx = gpu.thread_id("x") + tx_i32 = fx.Int32(tx) num_blocks_m = c_m // BLOCK_M num_blocks_n = c_n // BLOCK_N - if const_expr(use_xcd_remap): - pid_m, pid_n = xcd_swizzle(num_blocks_m, num_blocks_n) - else: - pid_m, pid_n = divmod(fx.block_idx.x, num_blocks_n) - - bx_m = pid_m * BLOCK_M - by_n = pid_n * BLOCK_N - - # The flattened/XCD-swizzled block coordinates are i32, while global - # address arithmetic below is expressed in MLIR index type. Convert - # once here and use these index-typed tile bases for every address. - bx_m_idx = fx.Index(bx_m) - by_n_idx = fx.Index(by_n) - - # Keep wave/lane arithmetic in i32. compute_global_swizzle() combines - # these values with i32 constants, so Index-typed coordinates would make - # arith.addi receive mixed operand types. - tx_i32 = fx.Int32(tx) - wave_id = tx_i32 // fx.Int32(WARP_SIZE) - lane = tx_i32 % fx.Int32(WARP_SIZE) - - # Offsets are always bytes. TN uses the original 128-byte XOR - # swizzle. NN/NT stage K-major 16-bit data as two [K64, X64] slices for - # ds_read_b64_tr_b16; the layout choice was resolved before capture. - gl_off_a = _a_global_offsets(lane, wave_id, c_m) - gl_off_b = _b_global_offsets(lane, wave_id, c_n) - - a_g2s = G2SLoader( - a_div, - gl_off_a, - LOAD_PASSES_HALF, - fx.Uint8.ir_type, - wave_id, - ) - b_g2s = G2SLoader( - b_div, - gl_off_b, - LOAD_PASSES_HALF, - fx.Uint8.ir_type, - wave_id, - ) - s2r = S2RLoader(fx.Int32(0), 1) - - layout_lane16 = fx.make_layout((4, 16), (16, 1)) - coord_lane16 = fx.idx2crd(fx.Int32(lane), layout_lane16) - lane_div_16 = fx.get(coord_lane16, 0) - lane_mod_16 = fx.get(coord_lane16, 1) - - # Per-CTA tile base in C elements, folded into each store's linear - # coordinate below (matching the scale-load addressing on this build; - # add_offset on a dynamic Index is unsupported here). - c_tile_base_elems = bx_m_idx * fx.Index(c_n) + by_n_idx - gC = fx.rocdl.make_buffer_tensor(C, max_size=True) - c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) - c_store_atom = fx.make_copy_atom( - fx.rocdl.BufferCopy32b() if output_element_bytes == 4 else fx.rocdl.BufferCopy16b(), - output_fx_dtype, - ) - - # Bias is a length-N fp32 vector indexed by the global output-feature - # (N) coordinate and broadcast across the M/token rows. const_expr folds - # this compile-time flag at trace time so the setup (and load_bias) are - # inlined into the kernel scope with no runtime dispatch branch. - if const_expr(has_bias): - gBias = fx.rocdl.make_buffer_tensor(Bias, max_size=True) - bias_div = fx.logical_divide(gBias, fx.make_layout(1, 1)) - bias_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) - - def load_bias(col): - reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) - fx.copy(bias_ld_atom, fx.slice(bias_div, (None, fx.Int32(col))), reg) - return fx.memref_load_vec(reg)[0] - - # GELU_AUX saves the pre-activation value (A@B[+bias]) to a second M x N - # output so the backward pass can recompute the GELU gradient. Same tile - # base / store atom shape as C; DEFAULT gets a dummy 1-element tensor. - if const_expr(has_gelu): - gAux = fx.rocdl.make_buffer_tensor(Aux, max_size=True) - aux_div = fx.logical_divide(gAux, fx.make_layout(1, 1)) - aux_store_atom = fx.make_copy_atom( + # Total output tiles. Dense uses this as the grid extent; grouped uses it as + # the persistent-loop drain bound. ``grouped`` gates the shared descriptor + # queue / work-steal driver (group fwd/dgrad or the wgrad specialization); + # const_expr folds it all out of the dense binary. + grouped = group or wgrad + num_tiles = gd.num_tiles_bound(c_m, c_n, BLOCK_M, BLOCK_N, num_groups, const_expr(wgrad)) + + if const_expr(grouped): + wq_rsrc = fx.rocdl.make_buffer_tensor(work_queue, max_size=True) + wq_div = fx.logical_divide(wq_rsrc, fx.make_layout(1, 1)) + desc_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32) + + # LDS broadcast scratch (one int32): thread 0 writes the atomically + # claimed slot, a barrier broadcasts it to the whole workgroup. The + # ``if tx==0`` claim control flow must stay inline (AST-rewritten here). + claim_iter, claim_view = gd.make_claim_broadcast(lds.claim_slot) + + def _claim_next_slot(): + if tx_i32 == fx.Int32(0): + old = llvm.AtomicRMWOp( + llvm.AtomicBinOp.add, + gd.global_i32_ptr(work_counter, fx.Int32(0)), + fx.as_ir_value(fx.Int32(1)), + llvm.AtomicOrdering.monotonic, + syncscope="agent", + alignment=4, + ).result + fx.ptr_store(fx.Vector.from_elements([fx.Int32(old)], fx.Int32), claim_iter) + fx.barrier() + return fx.Vector(claim_view.load())[0] + + def _desc_field(slot, field): + return gd.read_desc_field(wq_div, desc_ld_atom, slot, field) + + def process_tile(pid_m, pid_n, expert): + bx_m = pid_m * BLOCK_M + by_n = pid_n * BLOCK_N + + # The flattened/XCD-swizzled block coordinates are i32, while global + # address arithmetic below is expressed in MLIR index type. Convert + # once here and use these index-typed tile bases for every address. + bx_m_idx = fx.Index(bx_m) + by_n_idx = fx.Index(by_n) + + # Per-expert offsets fold the expert index into the packed operands / + # output (const_expr-zeroed for dense, so the dense binary is unchanged): + # TN fwd -- experts on B's output-N axis (by_n_b_idx = expert*c_n + n) + # NN dgrad-- experts on B's contraction axis (b_expert_k_off = expert*K) + # NT wgrad-- both operands on the contraction (token) axis + # (a/b_expert_k_off = expert*K); output packed on M (store_m_expert_off) + if const_expr(wgrad): + by_n_b_idx = by_n_idx + a_expert_k_off = fx.Index(expert * fx.Int32(K)) + b_expert_k_off = fx.Index(expert * fx.Int32(K)) + store_m_expert_off = fx.Index(expert * fx.Int32(c_m)) + elif const_expr(group): + if const_expr(b_transpose_read): + by_n_b_idx = by_n_idx + a_expert_k_off = fx.Index(0) + b_expert_k_off = fx.Index(expert * fx.Int32(K)) + store_m_expert_off = fx.Index(0) + else: + by_n_b_idx = fx.Index(expert * fx.Int32(c_n) + by_n) + a_expert_k_off = fx.Index(0) + b_expert_k_off = fx.Index(0) + store_m_expert_off = fx.Index(0) + else: + by_n_b_idx = by_n_idx + a_expert_k_off = fx.Index(0) + b_expert_k_off = fx.Index(0) + store_m_expert_off = fx.Index(0) + + # Keep wave/lane arithmetic in i32. compute_global_swizzle() combines + # these values with i32 constants, so Index-typed coordinates would make + # arith.addi receive mixed operand types. + wave_id = tx_i32 // fx.Int32(WARP_SIZE) + lane = tx_i32 % fx.Int32(WARP_SIZE) + + # Offsets are always bytes. TN uses the original 128-byte XOR + # swizzle. NN/NT stage K-major 16-bit data as two [K64, X64] slices for + # ds_read_b64_tr_b16; the layout choice was resolved before capture. + gl_off_a = _a_global_offsets(lane, wave_id, c_m) + gl_off_b = _b_global_offsets(lane, wave_id, c_n) + + a_g2s = G2SLoader( + a_div, + gl_off_a, + LOAD_PASSES_HALF, + fx.Uint8.ir_type, + wave_id, + ) + b_g2s = G2SLoader( + b_div, + gl_off_b, + LOAD_PASSES_HALF, + fx.Uint8.ir_type, + wave_id, + ) + s2r = S2RLoader(fx.Int32(0), 1) + + layout_lane16 = fx.make_layout((4, 16), (16, 1)) + coord_lane16 = fx.idx2crd(fx.Int32(lane), layout_lane16) + lane_div_16 = fx.get(coord_lane16, 0) + lane_mod_16 = fx.get(coord_lane16, 1) + + # Per-CTA tile base in C elements, folded into each store's linear + # coordinate below (matching the scale-load addressing on this build; + # add_offset on a dynamic Index is unsupported here). store_m_expert_off + # (grouped wgrad only) repacks the output row to this expert's block; + # it is 0 for dense and for grouped fwd/dgrad. + c_tile_base_elems = (bx_m_idx + store_m_expert_off) * fx.Index(c_n) + by_n_idx + gC = fx.rocdl.make_buffer_tensor(C, max_size=True) + c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) + c_store_atom = fx.make_copy_atom( fx.rocdl.BufferCopy32b() if output_element_bytes == 4 else fx.rocdl.BufferCopy16b(), output_fx_dtype, ) - def gelu_tanh(x): - # tanh-approx GELU (matches PyTorch approximate='tanh' and the - # FlyDSL preshuffle reference), expressed through a non-positive - # exponent so exp() cannot overflow: - # 0.5*x*(1 + tanh(y)), y = sqrt(2/pi)*(x + 0.044715*x^3) - half_f32 = fx.Float32(0.5) - one_f32 = fx.Float32(1.0) - zero_f32 = fx.Float32(0.0) - two_f32 = fx.Float32(2.0) - x3 = x * x * x - y = fx.Float32(0.7978845608) * (x + fx.Float32(0.044715) * x3) - abs_y = fx.Float32(y).maximumf(zero_f32 - y) - e_neg2abs = math.exp(fx.Float32(-2.0) * abs_y) - denom = one_f32 + e_neg2abs - numerator = (y > zero_f32).select(two_f32, two_f32 * e_neg2abs) - return half_f32 * x * (numerator * (one_f32 / denom)) - - PIN_ACC_BASE = 0 - - def _reg_list(prefix, start, end): - return ",".join(f"~{{{prefix}{r}}}" for r in range(start, end + 1)) - - def reserve_pinned_accumulators(): - # Reserve a fixed physical AGPR bank for all accumulators. In the - # SSA-lowered path, the compiler generated heavy AGPR <-> VGPR traffic, - # including v_accvgpr_mov/read sequences, s_nop stalls, and accumulator - # spills. Pinning each f32x4 accumulator to a stable AGPR range keeps the - # scaled MFMA accumulation in place and avoids those transfers and spills. - # - # ACCS_PER_WAVE = 64 accumulator objects and each object is f32x4, - # so the physical bank is exactly 64 * 4 = 256 AGPRs: a[0:255]. - clobbers = _reg_list("a", PIN_ACC_BASE, PIN_ACC_BASE + ACCS_PER_WAVE * 4 - 1) - llvm.InlineAsmOp( - None, - [], - "", - clobbers, - has_side_effects=True, - ) + # Bias is a length-N fp32 vector indexed by the global output-feature + # (N) coordinate and broadcast across the M/token rows. const_expr folds + # this compile-time flag at trace time so the setup (and load_bias) are + # inlined into the kernel scope with no runtime dispatch branch. + if const_expr(has_bias): + gBias = fx.rocdl.make_buffer_tensor(Bias, max_size=True) + bias_div = fx.logical_divide(gBias, fx.make_layout(1, 1)) + bias_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) + + def load_bias(col): + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) + fx.copy(bias_ld_atom, fx.slice(bias_div, (None, fx.Int32(col))), reg) + return fx.memref_load_vec(reg)[0] + + # GELU_AUX saves the pre-activation value (A@B[+bias]) to a second M x N + # output so the backward pass can recompute the GELU gradient. Same tile + # base / store atom shape as C; DEFAULT gets a dummy 1-element tensor. + if const_expr(has_gelu): + gAux = fx.rocdl.make_buffer_tensor(Aux, max_size=True) + aux_div = fx.logical_divide(gAux, fx.make_layout(1, 1)) + aux_store_atom = fx.make_copy_atom( + fx.rocdl.BufferCopy32b() if output_element_bytes == 4 else fx.rocdl.BufferCopy16b(), + output_fx_dtype, + ) - def zero_pinned_accumulators(): - for ai in range_constexpr(ACCS_PER_WAVE * 4): + def gelu_tanh(x): + # tanh-approx GELU (matches PyTorch approximate='tanh' and the + # FlyDSL preshuffle reference), expressed through a non-positive + # exponent so exp() cannot overflow: + # 0.5*x*(1 + tanh(y)), y = sqrt(2/pi)*(x + 0.044715*x^3) + half_f32 = fx.Float32(0.5) + one_f32 = fx.Float32(1.0) + zero_f32 = fx.Float32(0.0) + two_f32 = fx.Float32(2.0) + x3 = x * x * x + y = fx.Float32(0.7978845608) * (x + fx.Float32(0.044715) * x3) + abs_y = fx.Float32(y).maximumf(zero_f32 - y) + e_neg2abs = math.exp(fx.Float32(-2.0) * abs_y) + denom = one_f32 + e_neg2abs + numerator = (y > zero_f32).select(two_f32, two_f32 * e_neg2abs) + return half_f32 * x * (numerator * (one_f32 / denom)) + + PIN_ACC_BASE = 0 + + def _reg_list(prefix, start, end): + return ",".join(f"~{{{prefix}{r}}}" for r in range(start, end + 1)) + + def reserve_pinned_accumulators(): + # Reserve a fixed physical AGPR bank for all accumulators. In the + # SSA-lowered path, the compiler generated heavy AGPR <-> VGPR traffic, + # including v_accvgpr_mov/read sequences, s_nop stalls, and accumulator + # spills. Pinning each f32x4 accumulator to a stable AGPR range keeps the + # scaled MFMA accumulation in place and avoids those transfers and spills. + # + # ACCS_PER_WAVE = 64 accumulator objects and each object is f32x4, + # so the physical bank is exactly 64 * 4 = 256 AGPRs: a[0:255]. + clobbers = _reg_list("a", PIN_ACC_BASE, PIN_ACC_BASE + ACCS_PER_WAVE * 4 - 1) llvm.InlineAsmOp( None, [], - f"v_accvgpr_write_b32 a[{PIN_ACC_BASE + ai}], 0", - f"~{{a{PIN_ACC_BASE + ai}}}", + "", + clobbers, has_side_effects=True, ) - def _inline_asm_i32(asm_string, constraints, operands=None): - op = llvm.InlineAsmOp( - T.i32, - operands or [], - asm_string, - constraints, - has_side_effects=True, - ) - return _one_i32_result(op) - - def _one_i32_result(op): - # Accept the result attribute names exposed by the supported MLIR Python bindings. - return getattr(op, "result", getattr(op, "res", op.results[0])) - - def read_pinned_accumulator(acc_idx): - acc_pin = PIN_ACC_BASE + acc_idx * 4 - r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") - r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") - r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") - r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") - return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) - - def read_physical_accumulator_slot(slot_idx): - acc_pin = PIN_ACC_BASE + slot_idx * 4 - r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") - r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") - r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") - r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") - return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) - - def hot_loop_scheduler_q_refill_2n(): - # Eight refill VMEM operations overlap four independent 8-MFMA - # groups (two K32 slices x two N-halves). - for _ in range_constexpr(4): - rocdl.sched_vmem(2) - rocdl.sched_mfma(8) - rocdl.sched_barrier(0) - - def hot_loop_scheduler_q0_refill_a1_2n(): - # Eight refill VMEM operations and eight distributed A-bottom LDS - # reads overlap four independent 8-MFMA K32 groups. - for _ in range_constexpr(4): - rocdl.sched_vmem(2) - rocdl.sched_dsrd(Q0_SCHED_DSRD) - rocdl.sched_mfma(8) - rocdl.sched_barrier(0) - - def hot_loop_scheduler_q_prefetch_4n(): - # Eight two-read prefetch groups overlap four complete-quadrant - # 16-MFMA groups (two K32 slices for each of Q2 and Q3). - for _ in range_constexpr(4): - rocdl.sched_dsrd(PREFETCH_SCHED_DSRD) - rocdl.sched_mfma(16) - rocdl.sched_barrier(0) - - def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): - # One pass writes 256 threads * 16 B = 4 KiB. Four passes fill one - # 128x64 half-page (16 KiB). Each half has its own LDS base. - global_base = _a_global_base_bytes(k_base, subtile, c_m, bx_m_idx) - a_g2s.load_one(lds_a[subtile], fx.Int32(global_base), pass_in_subtile) - - def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): - global_base = _b_global_base_bytes(k_base, subtile, c_n, by_n_idx) - b_g2s.load_one(lds_b[subtile], fx.Int32(global_base), pass_in_subtile) - - def stage_a_subtile(k_base, subtile, lds_a): - for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): - stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a) - - def stage_b_subtile(k_base, subtile, lds_b): - for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): - stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b) - - def load_frag_half_at_byte_base(lds_page, row_byte_base, half): - # Issue exactly one 16-byte LDS read for one 16-byte half of the wave operand tile. - # Keeping the halves separate allows steady-state Q0 to schedule one - # A-bottom ds_read_b128 in each refill/MFMA chunk. - k_col = reg_lds_k_col0 if half == 0 else reg_lds_k_col1 - return s2r.load_one(lds_page, fx.Int32(row_byte_base + k_col)) - - def pack_frag_halves(x0, x1): - return pack_i32x4_i32x8(x0, x1) - - def load_frag_at_byte_base(lds_page, row_byte_base): - # Default complete-fragment path used outside the dedicated Q0 schedule. - x0 = load_frag_half_at_byte_base(lds_page, row_byte_base, 0) - x1 = load_frag_half_at_byte_base(lds_page, row_byte_base, 1) - return pack_frag_halves(x0, x1) - - def load_b_frag(lds_b, local_row, half): - # B is [N, K]. Each 128-row half-page has a local row origin of 0. - half_row = local_row - fx.Index(half * (BLOCK_N // 2)) - return load_frag_at_byte_base( - lds_b[half], - half_row * fx.Index(BLOCK_K * ELEM_BYTES), - ) - - def load_transposed_frag_half(lds_page, local_x_tile, half): - # Half-precision uses v_mfma_f32_16x16x32_{f16,bf16}, not the MXFP8 K128 - # instruction. A 128-X half-page is therefore two independent - # swizzled [K64, X64] half-precision slices. One ds_read_b64_tr_b16 returns - # four half-precision values/lane; two reads form one K32 MFMA fragment. - local_x_i32 = fx.Int32(local_x_tile) - slice_idx = local_x_i32 // fx.Int32(64) - x_in_slice = local_x_i32 % fx.Int32(64) - lane_div16_i32 = fx.Int32(lane_div_16) - lane_in16_i32 = fx.Int32(lane_mod_16) - - source_k = lane_div16_i32 * fx.Int32(8) + lane_in16_i32 // fx.Int32(4) - source_x_byte = x_in_slice * fx.Int32(ELEM_BYTES) + ( - lane_in16_i32 % fx.Int32(4) - ) * fx.Int32(8) - - physical_k, physical_x = swizzle_128(source_k, source_x_byte) - slice_base = slice_idx * fx.Int32(64 * 128) - base = slice_base + physical_k * fx.Int32(128) + physical_x - other = base ^ fx.Int32(0x220) - immediate_offset = 0 if half == 0 else 0x1000 - return s2r.load_one_transpose( - lds_page, - base, - other, - immediate_offset=immediate_offset, - ) + def zero_pinned_accumulators(): + for ai in range_constexpr(ACCS_PER_WAVE * 4): + llvm.InlineAsmOp( + None, + [], + f"v_accvgpr_write_b32 a[{PIN_ACC_BASE + ai}], 0", + f"~{{a{PIN_ACC_BASE + ai}}}", + has_side_effects=True, + ) - def load_transposed_frag(lds_page, local_x_tile): - x0 = load_transposed_frag_half(lds_page, local_x_tile, 0) - x1 = load_transposed_frag_half(lds_page, local_x_tile, 1) - return pack_frag_halves(x0, x1) + def _inline_asm_i32(asm_string, constraints, operands=None): + op = llvm.InlineAsmOp( + T.i32, + operands or [], + asm_string, + constraints, + has_side_effects=True, + ) + return _one_i32_result(op) + + def _one_i32_result(op): + # Accept the result attribute names exposed by the supported MLIR Python bindings. + return getattr(op, "result", getattr(op, "res", op.results[0])) + + def read_pinned_accumulator(acc_idx): + acc_pin = PIN_ACC_BASE + acc_idx * 4 + r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") + r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") + r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") + r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") + return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) + + def read_physical_accumulator_slot(slot_idx): + acc_pin = PIN_ACC_BASE + slot_idx * 4 + r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") + r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") + r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") + r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") + return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) + + def hot_loop_scheduler_q_refill_2n(): + # Eight refill VMEM operations overlap four independent 8-MFMA + # groups (two K32 slices x two N-halves). + for _ in range_constexpr(4): + rocdl.sched_vmem(2) + rocdl.sched_mfma(8) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q0_refill_a1_2n(): + # Eight refill VMEM operations and eight distributed A-bottom LDS + # reads overlap four independent 8-MFMA K32 groups. + for _ in range_constexpr(4): + rocdl.sched_vmem(2) + rocdl.sched_dsrd(Q0_SCHED_DSRD) + rocdl.sched_mfma(8) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q_prefetch_4n(): + # Eight two-read prefetch groups overlap four complete-quadrant + # 16-MFMA groups (two K32 slices for each of Q2 and Q3). + for _ in range_constexpr(4): + rocdl.sched_dsrd(PREFETCH_SCHED_DSRD) + rocdl.sched_mfma(16) + rocdl.sched_barrier(0) + + def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): + # One pass writes 256 threads * 16 B = 4 KiB. Four passes fill one + # 128x64 half-page (16 KiB). Each half has its own LDS base. + # a_expert_k_off shifts the contraction base to this tile's expert A + # block for grouped wgrad (dY packed on the token axis); 0 otherwise. + global_base = _a_global_base_bytes(k_base + a_expert_k_off, subtile, c_m, bx_m_idx) + a_g2s.load_one(lds_a[subtile], fx.Int32(global_base), pass_in_subtile) + + def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): + # b_expert_k_off shifts the contraction base (grouped NN/NT transpose + # read); by_n_b_idx shifts the B column (grouped TN). Both 0 for dense. + global_base = _b_global_base_bytes(k_base + b_expert_k_off, subtile, c_n, by_n_b_idx) + b_g2s.load_one(lds_b[subtile], fx.Int32(global_base), pass_in_subtile) + + def stage_a_subtile(k_base, subtile, lds_a): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a) + + def stage_b_subtile(k_base, subtile, lds_b): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b) + + def load_frag_half_at_byte_base(lds_page, row_byte_base, half): + # Issue exactly one 16-byte LDS read for one 16-byte half of the wave operand tile. + # Keeping the halves separate allows steady-state Q0 to schedule one + # A-bottom ds_read_b128 in each refill/MFMA chunk. + k_col = reg_lds_k_col0 if half == 0 else reg_lds_k_col1 + return s2r.load_one(lds_page, fx.Int32(row_byte_base + k_col)) + + def pack_frag_halves(x0, x1): + return pack_i32x4_i32x8(x0, x1) + + def load_frag_at_byte_base(lds_page, row_byte_base): + # Default complete-fragment path used outside the dedicated Q0 schedule. + x0 = load_frag_half_at_byte_base(lds_page, row_byte_base, 0) + x1 = load_frag_half_at_byte_base(lds_page, row_byte_base, 1) + return pack_frag_halves(x0, x1) + + def load_b_frag(lds_b, local_row, half): + # B is [N, K]. Each 128-row half-page has a local row origin of 0. + half_row = local_row - fx.Index(half * (BLOCK_N // 2)) + return load_frag_at_byte_base( + lds_b[half], + half_row * fx.Index(BLOCK_K * ELEM_BYTES), + ) - def _acc_idx(subtile_id, mi, ni): - return ( - subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + mi * MFMA_N_PER_SUBTILE + ni - ) + def load_transposed_frag_half(lds_page, local_x_tile, half): + # Half-precision uses v_mfma_f32_16x16x32_{f16,bf16}, not the MXFP8 K128 + # instruction. A 128-X half-page is therefore two independent + # swizzled [K64, X64] half-precision slices. One ds_read_b64_tr_b16 returns + # four half-precision values/lane; two reads form one K32 MFMA fragment. + local_x_i32 = fx.Int32(local_x_tile) + slice_idx = local_x_i32 // fx.Int32(64) + x_in_slice = local_x_i32 % fx.Int32(64) + lane_div16_i32 = fx.Int32(lane_div_16) + lane_in16_i32 = fx.Int32(lane_mod_16) + + source_k = lane_div16_i32 * fx.Int32(8) + lane_in16_i32 // fx.Int32(4) + source_x_byte = x_in_slice * fx.Int32(ELEM_BYTES) + ( + lane_in16_i32 % fx.Int32(4) + ) * fx.Int32(8) + + physical_k, physical_x = swizzle_128(source_k, source_x_byte) + slice_base = slice_idx * fx.Int32(64 * 128) + base = slice_base + physical_k * fx.Int32(128) + physical_x + other = base ^ fx.Int32(0x220) + immediate_offset = 0 if half == 0 else 0x1000 + return s2r.load_one_transpose( + lds_page, + base, + other, + immediate_offset=immediate_offset, + ) - def _k32_frag(full_frag, k32): - # A/B 16x64 half-precision wave fragments are i32x8. Each K32 MFMA - # consumes one contiguous i32x4 slice (eight half-precision values/lane). - lo = k32 * 4 - v = Vec(full_frag) - return Vec.from_elements( - [v[lo], v[lo + 1], v[lo + 2], v[lo + 3]], - fx.Int32, - ) + def load_transposed_frag(lds_page, local_x_tile): + x0 = load_transposed_frag_half(lds_page, local_x_tile, 0) + x1 = load_transposed_frag_half(lds_page, local_x_tile, 1) + return pack_frag_halves(x0, x1) - def _pinned_mfma_once(acc_idx, a_k32, b_k32): - acc_pin = PIN_ACC_BASE + acc_idx * 4 - llvm.InlineAsmOp( - None, - [arith._to_raw(a_k32), arith._to_raw(b_k32)], - f"v_mfma_f32_16x16x32_{mfma_suffix} " - f"a[{acc_pin}:{acc_pin + 3}], " - "$0, $1, " - f"a[{acc_pin}:{acc_pin + 3}]", - f"v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}},~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}", - has_side_effects=True, - ) + def _acc_idx(subtile_id, mi, ni): + return ( + subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + mi * MFMA_N_PER_SUBTILE + ni + ) - def pinned_mfma(acc_idx, a_frag, b_frag): - """Accumulate one logical 16x16x64 half-precision product into pinned AGPRs.""" - for k32 in range_constexpr(2): - _pinned_mfma_once( - acc_idx, - _k32_frag(a_frag, k32), - _k32_frag(b_frag, k32), + def _k32_frag(full_frag, k32): + # A/B 16x64 half-precision wave fragments are i32x8. Each K32 MFMA + # consumes one contiguous i32x4 slice (eight half-precision values/lane). + lo = k32 * 4 + v = Vec(full_frag) + return Vec.from_elements( + [v[lo], v[lo + 1], v[lo + 2], v[lo + 3]], + fx.Int32, ) - def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag): - # The final logical K64 update is two in-place K32 MFMAs. - assert dst_slot == old_acc_idx - pinned_mfma(old_acc_idx, a_frag, b_frag) - - def mfma_4n(acc_base, a_frag, b0, b1, b2, b3): - pinned_mfma(acc_base + 0, a_frag, b0) - pinned_mfma(acc_base + 1, a_frag, b1) - pinned_mfma(acc_base + 2, a_frag, b2) - pinned_mfma(acc_base + 3, a_frag, b3) - - def mfma_2n(acc_base, a_frag, b0, b1): - pinned_mfma(acc_base + 0, a_frag, b0) - pinned_mfma(acc_base + 1, a_frag, b1) - - def mfma_2n_4mi_k32(subtile_id, n_base, k32, a0, a1, a2, a3, b0, b1): - """Issue one K32 slice for a 4x2 accumulator slab.""" - a_frags = (a0, a1, a2, a3) - b_frags = (b0, b1) - for mi in range_constexpr(4): - a_k32 = _k32_frag(a_frags[mi], k32) - for nj in range_constexpr(2): - _pinned_mfma_once( - _acc_idx(subtile_id, mi, n_base + nj), - a_k32, - _k32_frag(b_frags[nj], k32), - ) + def _pinned_mfma_once(acc_idx, a_k32, b_k32): + acc_pin = PIN_ACC_BASE + acc_idx * 4 + llvm.InlineAsmOp( + None, + [arith._to_raw(a_k32), arith._to_raw(b_k32)], + f"v_mfma_f32_16x16x32_{mfma_suffix} " + f"a[{acc_pin}:{acc_pin + 3}], " + "$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}]", + f"v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}},~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}", + has_side_effects=True, + ) - def mfma_4n_4mi_k32(subtile_id, k32, a0, a1, a2, a3, b0, b1, b2, b3): - """Issue one K32 slice for a complete 4x4 quadrant.""" - a_frags = (a0, a1, a2, a3) - b_frags = (b0, b1, b2, b3) - for mi in range_constexpr(4): - a_k32 = _k32_frag(a_frags[mi], k32) - for ni in range_constexpr(4): + def pinned_mfma(acc_idx, a_frag, b_frag): + """Accumulate one logical 16x16x64 half-precision product into pinned AGPRs.""" + for k32 in range_constexpr(2): _pinned_mfma_once( - _acc_idx(subtile_id, mi, ni), - a_k32, - _k32_frag(b_frags[ni], k32), + acc_idx, + _k32_frag(a_frag, k32), + _k32_frag(b_frag, k32), ) - def store_acc_vector_for_logical_idx(logical_acc_idx, acc): - subtile_id = logical_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) - local_idx = logical_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) - sm = subtile_id // 2 - sn = subtile_id % 2 - mi = local_idx // MFMA_N_PER_SUBTILE - ni = local_idx % MFMA_N_PER_SUBTILE - - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - row_base = subtile_m_idx * SUBTILE_M + fx.Index(mi * MFMA_M) + lane_div_16 * 4 - col = subtile_n_idx * SUBTILE_N + fx.Index(ni * MFMA_N) + lane_mod_16 - - # Bias depends only on the output-feature (N) coordinate, so read it - # once per column (global index by_n_idx + col) and reuse across the - # four M rows below. - if const_expr(has_bias): - bias_value = load_bias(by_n_idx + col) + def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag): + # The final logical K64 update is two in-place K32 MFMAs. + assert dst_slot == old_acc_idx + pinned_mfma(old_acc_idx, a_frag, b_frag) + + def mfma_4n(acc_base, a_frag, b0, b1, b2, b3): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + pinned_mfma(acc_base + 2, a_frag, b2) + pinned_mfma(acc_base + 3, a_frag, b3) + + def mfma_2n(acc_base, a_frag, b0, b1): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + + def mfma_2n_4mi_k32(subtile_id, n_base, k32, a0, a1, a2, a3, b0, b1): + """Issue one K32 slice for a 4x2 accumulator slab.""" + a_frags = (a0, a1, a2, a3) + b_frags = (b0, b1) + for mi in range_constexpr(4): + a_k32 = _k32_frag(a_frags[mi], k32) + for nj in range_constexpr(2): + _pinned_mfma_once( + _acc_idx(subtile_id, mi, n_base + nj), + a_k32, + _k32_frag(b_frags[nj], k32), + ) + + def mfma_4n_4mi_k32(subtile_id, k32, a0, a1, a2, a3, b0, b1, b2, b3): + """Issue one K32 slice for a complete 4x4 quadrant.""" + a_frags = (a0, a1, a2, a3) + b_frags = (b0, b1, b2, b3) + for mi in range_constexpr(4): + a_k32 = _k32_frag(a_frags[mi], k32) + for ni in range_constexpr(4): + _pinned_mfma_once( + _acc_idx(subtile_id, mi, ni), + a_k32, + _k32_frag(b_frags[ni], k32), + ) + + def store_acc_vector_for_logical_idx(logical_acc_idx, acc): + subtile_id = logical_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = logical_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE - for ii in range_constexpr(4): - row = row_base + fx.Index(ii) - c_idx = c_tile_base_elems + row * fx.Index(c_n) + col + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + row_base = subtile_m_idx * SUBTILE_M + fx.Index(mi * MFMA_M) + lane_div_16 * 4 + col = subtile_n_idx * SUBTILE_N + fx.Index(ni * MFMA_N) + lane_mod_16 - # Epilogue stages run on the fp32 accumulator, in order, before - # the output-dtype narrowing: - # value = acc [+ bias] (pre-activation) - # GELU_AUX: save pre-activation to Aux, then value = gelu(value) - value = Vec(acc)[ii] + # Bias depends only on the output-feature (N) coordinate, so read it + # once per column (global index by_n_idx + col) and reuse across the + # four M rows below. if const_expr(has_bias): - value = value + bias_value + bias_value = load_bias(by_n_idx + col) + + for ii in range_constexpr(4): + row = row_base + fx.Index(ii) + c_idx = c_tile_base_elems + row * fx.Index(c_n) + col + + # Epilogue stages run on the fp32 accumulator, in order, before + # the output-dtype narrowing: + # value = acc [+ bias] (pre-activation) + # GELU_AUX: save pre-activation to Aux, then value = gelu(value) + value = Vec(acc)[ii] + if const_expr(has_bias): + value = value + bias_value + + if const_expr(has_gelu): + aux_val = value + if const_expr(output_dtype != torch.float32): + aux_val = aux_val.to(output_fx_dtype) + aux_reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) + fx.memref_store_vec(Vec.filled(1, aux_val, output_fx_dtype), aux_reg) + fx.copy(aux_store_atom, aux_reg, fx.slice(aux_div, (None, fx.Int32(c_idx)))) + value = gelu_tanh(value) - if const_expr(has_gelu): - aux_val = value if const_expr(output_dtype != torch.float32): - aux_val = aux_val.to(output_fx_dtype) - aux_reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) - fx.memref_store_vec(Vec.filled(1, aux_val, output_fx_dtype), aux_reg) - fx.copy(aux_store_atom, aux_reg, fx.slice(aux_div, (None, fx.Int32(c_idx)))) - value = gelu_tanh(value) - - if const_expr(output_dtype != torch.float32): - value = value.to(output_fx_dtype) - reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) - fx.memref_store_vec(Vec.filled(1, value, output_fx_dtype), reg) - fx.copy(c_store_atom, reg, fx.slice(c_div, (None, fx.Int32(c_idx)))) - - # Explicit register coordinates for HK-style four-quadrant mapping. - # BLOCK_M/BLOCK_N are 256x256. Four waves map to warp positions - # inside each 128x128 quadrant: - # cA: (warp_m, warp_n) - # cB: (warp_m, warp_n + 2) - # cC: (warp_m + 2, warp_n) - # cD: (warp_m + 2, warp_n + 2) - reg_k_col0 = lane_div_16 * 16 - reg_k_col1 = 64 + lane_div_16 * 16 - - # Every fragment row differs only by multiples of 16, so row % 16 is - # always lane_mod_16. Hoist the logical->physical XOR mapping once. - _, reg_lds_k_col0 = swizzle_128(lane_mod_16, reg_k_col0) - _, reg_lds_k_col1 = swizzle_128(lane_mod_16, reg_k_col1) - - reg_subtile_m_idx0 = wave_id // 2 - reg_subtile_n_idx0 = wave_id % 2 - - reserve_pinned_accumulators() - zero_pinned_accumulators() - - def load_b_subtile_ni_regs(lds_b, sn, ni): - return _load_b_ni( - load_transposed_frag, - load_b_frag, - lds_b, - sn, - ni, - reg_subtile_n_idx0, - lane_mod_16, - ) - - def load_b_subtile_regs(lds_b, sn): - return ( - load_b_subtile_ni_regs(lds_b, sn, 0), - load_b_subtile_ni_regs(lds_b, sn, 1), - load_b_subtile_ni_regs(lds_b, sn, 2), - load_b_subtile_ni_regs(lds_b, sn, 3), - ) + value = value.to(output_fx_dtype) + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) + fx.memref_store_vec(Vec.filled(1, value, output_fx_dtype), reg) + fx.copy(c_store_atom, reg, fx.slice(c_div, (None, fx.Int32(c_idx)))) + + # Explicit register coordinates for HK-style four-quadrant mapping. + # BLOCK_M/BLOCK_N are 256x256. Four waves map to warp positions + # inside each 128x128 quadrant: + # cA: (warp_m, warp_n) + # cB: (warp_m, warp_n + 2) + # cC: (warp_m + 2, warp_n) + # cD: (warp_m + 2, warp_n + 2) + reg_k_col0 = lane_div_16 * 16 + reg_k_col1 = 64 + lane_div_16 * 16 + + # Every fragment row differs only by multiples of 16, so row % 16 is + # always lane_mod_16. Hoist the logical->physical XOR mapping once. + _, reg_lds_k_col0 = swizzle_128(lane_mod_16, reg_k_col0) + _, reg_lds_k_col1 = swizzle_128(lane_mod_16, reg_k_col1) + + reg_subtile_m_idx0 = wave_id // 2 + reg_subtile_n_idx0 = wave_id % 2 + + reserve_pinned_accumulators() + zero_pinned_accumulators() + + def load_b_subtile_ni_regs(lds_b, sn, ni): + return _load_b_ni( + load_transposed_frag, + load_b_frag, + lds_b, + sn, + ni, + reg_subtile_n_idx0, + lane_mod_16, + ) - def load_a_subtile_mi_half(lds_a, sm, mi, half): - return _load_a_half( - load_transposed_frag_half, - load_frag_half_at_byte_base, - lds_a, - sm, - mi, - half, - reg_subtile_m_idx0, - lane_mod_16, - ) + def load_b_subtile_regs(lds_b, sn): + return ( + load_b_subtile_ni_regs(lds_b, sn, 0), + load_b_subtile_ni_regs(lds_b, sn, 1), + load_b_subtile_ni_regs(lds_b, sn, 2), + load_b_subtile_ni_regs(lds_b, sn, 3), + ) - def load_a_subtile_mi_regs(lds_a, sm, mi): - x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) - x1 = load_a_subtile_mi_half(lds_a, sm, mi, 1) - return pack_frag_halves(x0, x1) - - def load_a_subtile_regs(lds_a, sm): - return ( - load_a_subtile_mi_regs(lds_a, sm, 0), - load_a_subtile_mi_regs(lds_a, sm, 1), - load_a_subtile_mi_regs(lds_a, sm, 2), - load_a_subtile_mi_regs(lds_a, sm, 3), - ) + def load_a_subtile_mi_half(lds_a, sm, mi, half): + return _load_a_half( + load_transposed_frag_half, + load_frag_half_at_byte_base, + lds_a, + sm, + mi, + half, + reg_subtile_m_idx0, + lane_mod_16, + ) - def hk_one_k_with_refill( - k128, - cur_a, - cur_b, - next_a, - next_b, - refill_a, - refill_b, - a0_regs, - b0_regs, - ): + def load_a_subtile_mi_regs(lds_a, sm, mi): + x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) + x1 = load_a_subtile_mi_half(lds_a, sm, mi, 1) + return pack_frag_halves(x0, x1) + + def load_a_subtile_regs(lds_a, sm): + return ( + load_a_subtile_mi_regs(lds_a, sm, 0), + load_a_subtile_mi_regs(lds_a, sm, 1), + load_a_subtile_mi_regs(lds_a, sm, 2), + load_a_subtile_mi_regs(lds_a, sm, 3), + ) - # Wait only far enough for the current page; the next-page refill may remain in flight. - barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) - rocdl.sched_barrier(0) + def hk_one_k_with_refill( + k128, + cur_a, + cur_b, + next_a, + next_b, + refill_a, + refill_b, + a0_regs, + b0_regs, + ): - # A-top and B-left are both carried as complete 64-row register tiles, - # so their LDS half-pages can be refilled immediately. - a00, a01, a02, a03 = a0_regs - b00, b01, b02, b03 = b0_regs + # Wait only far enough for the current page; the next-page refill may remain in flight. + barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) - b10 = load_b_subtile_ni_regs(cur_b, 1, 0) - b11 = load_b_subtile_ni_regs(cur_b, 1, 1) - b12 = load_b_subtile_ni_regs(cur_b, 1, 2) - b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + # A-top and B-left are both carried as complete 64-row register tiles, + # so their LDS half-pages can be refilled immediately. + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs - # Refill the current ping-pong page with K+2, alternating A and B passes. - k_refill = fx.Index((k128 + 2) * BLOCK_K) + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) - # Q0: interleave the current tile's A-bottom LDS reads with K+2 - # refills and Q0 compute. Compute is K32-major across all 16 - # independent accumulators, eliminating the two-deep same-AGPR - # dependency chains produced by pinned_mfma(). - rocdl.sched_barrier(0) - a10_x0 = load_a_subtile_mi_half(cur_a, 1, 0, 0) - stage_a_subtile_pass(k_refill, 0, 0, refill_a) - mfma_2n_4mi_k32(0, 0, 0, a00, a01, a02, a03, b00, b01) + # Refill the current ping-pong page with K+2, alternating A and B passes. + k_refill = fx.Index((k128 + 2) * BLOCK_K) - a10_x1 = load_a_subtile_mi_half(cur_a, 1, 0, 1) - stage_b_subtile_pass(k_refill, 0, 0, refill_b) - mfma_2n_4mi_k32(0, 2, 0, a00, a01, a02, a03, b02, b03) + # Q0: interleave the current tile's A-bottom LDS reads with K+2 + # refills and Q0 compute. Compute is K32-major across all 16 + # independent accumulators, eliminating the two-deep same-AGPR + # dependency chains produced by pinned_mfma(). + rocdl.sched_barrier(0) + a10_x0 = load_a_subtile_mi_half(cur_a, 1, 0, 0) + stage_a_subtile_pass(k_refill, 0, 0, refill_a) + mfma_2n_4mi_k32(0, 0, 0, a00, a01, a02, a03, b00, b01) - a11_x0 = load_a_subtile_mi_half(cur_a, 1, 1, 0) - stage_a_subtile_pass(k_refill, 0, 1, refill_a) - # K32 slice 0 already covers K[0:32]. + a10_x1 = load_a_subtile_mi_half(cur_a, 1, 0, 1) + stage_b_subtile_pass(k_refill, 0, 0, refill_b) + mfma_2n_4mi_k32(0, 2, 0, a00, a01, a02, a03, b02, b03) - a11_x1 = load_a_subtile_mi_half(cur_a, 1, 1, 1) - stage_b_subtile_pass(k_refill, 0, 1, refill_b) - # Keep this refill/LDS-read slot compute-free. + a11_x0 = load_a_subtile_mi_half(cur_a, 1, 1, 0) + stage_a_subtile_pass(k_refill, 0, 1, refill_a) + # K32 slice 0 already covers K[0:32]. - a12_x0 = load_a_subtile_mi_half(cur_a, 1, 2, 0) - stage_a_subtile_pass(k_refill, 0, 2, refill_a) - mfma_2n_4mi_k32(0, 0, 1, a00, a01, a02, a03, b00, b01) + a11_x1 = load_a_subtile_mi_half(cur_a, 1, 1, 1) + stage_b_subtile_pass(k_refill, 0, 1, refill_b) + # Keep this refill/LDS-read slot compute-free. - a12_x1 = load_a_subtile_mi_half(cur_a, 1, 2, 1) - stage_b_subtile_pass(k_refill, 0, 2, refill_b) - mfma_2n_4mi_k32(0, 2, 1, a00, a01, a02, a03, b02, b03) + a12_x0 = load_a_subtile_mi_half(cur_a, 1, 2, 0) + stage_a_subtile_pass(k_refill, 0, 2, refill_a) + mfma_2n_4mi_k32(0, 0, 1, a00, a01, a02, a03, b00, b01) - a13_x0 = load_a_subtile_mi_half(cur_a, 1, 3, 0) - stage_a_subtile_pass(k_refill, 0, 3, refill_a) - # K32 slice 1 already covers K[32:64]. + a12_x1 = load_a_subtile_mi_half(cur_a, 1, 2, 1) + stage_b_subtile_pass(k_refill, 0, 2, refill_b) + mfma_2n_4mi_k32(0, 2, 1, a00, a01, a02, a03, b02, b03) - a13_x1 = load_a_subtile_mi_half(cur_a, 1, 3, 1) - stage_b_subtile_pass(k_refill, 0, 3, refill_b) - # Keep this refill/LDS-read slot compute-free. + a13_x0 = load_a_subtile_mi_half(cur_a, 1, 3, 0) + stage_a_subtile_pass(k_refill, 0, 3, refill_a) + # K32 slice 1 already covers K[32:64]. - hot_loop_scheduler_q0_refill_a1_2n() + a13_x1 = load_a_subtile_mi_half(cur_a, 1, 3, 1) + stage_b_subtile_pass(k_refill, 0, 3, refill_b) + # Keep this refill/LDS-read slot compute-free. - # Retire the eight distributed A-bottom LDS reads before K+2 refills - # overwrite the current page's A-bottom half-page. Keep this wait as - # late as possible to maximize read/compute overlap. - rocdl.sched_barrier(0) - barrier(lgkmcnt=0) - rocdl.sched_barrier(0) + hot_loop_scheduler_q0_refill_a1_2n() - a10 = pack_frag_halves(a10_x0, a10_x1) - a11 = pack_frag_halves(a11_x0, a11_x1) - a12 = pack_frag_halves(a12_x0, a12_x1) - a13 = pack_frag_halves(a13_x0, a13_x1) + # Retire the eight distributed A-bottom LDS reads before K+2 refills + # overwrite the current page's A-bottom half-page. Keep this wait as + # late as possible to maximize read/compute overlap. + rocdl.sched_barrier(0) + barrier(lgkmcnt=0) + rocdl.sched_barrier(0) - rocdl.sched_barrier(0) - stage_b_subtile_pass(k_refill, 1, 0, refill_b) - mfma_2n_4mi_k32(1, 0, 0, a00, a01, a02, a03, b10, b11) + a10 = pack_frag_halves(a10_x0, a10_x1) + a11 = pack_frag_halves(a11_x0, a11_x1) + a12 = pack_frag_halves(a12_x0, a12_x1) + a13 = pack_frag_halves(a13_x0, a13_x1) - stage_a_subtile_pass(k_refill, 1, 0, refill_a) - mfma_2n_4mi_k32(1, 2, 0, a00, a01, a02, a03, b12, b13) + rocdl.sched_barrier(0) + stage_b_subtile_pass(k_refill, 1, 0, refill_b) + mfma_2n_4mi_k32(1, 0, 0, a00, a01, a02, a03, b10, b11) - stage_b_subtile_pass(k_refill, 1, 1, refill_b) - # K32 slice 0 already covers K[0:32]. + stage_a_subtile_pass(k_refill, 1, 0, refill_a) + mfma_2n_4mi_k32(1, 2, 0, a00, a01, a02, a03, b12, b13) - stage_a_subtile_pass(k_refill, 1, 1, refill_a) - # Keep this refill slot compute-free. + stage_b_subtile_pass(k_refill, 1, 1, refill_b) + # K32 slice 0 already covers K[0:32]. - stage_b_subtile_pass(k_refill, 1, 2, refill_b) - mfma_2n_4mi_k32(1, 0, 1, a00, a01, a02, a03, b10, b11) + stage_a_subtile_pass(k_refill, 1, 1, refill_a) + # Keep this refill slot compute-free. - stage_a_subtile_pass(k_refill, 1, 2, refill_a) - mfma_2n_4mi_k32(1, 2, 1, a00, a01, a02, a03, b12, b13) + stage_b_subtile_pass(k_refill, 1, 2, refill_b) + mfma_2n_4mi_k32(1, 0, 1, a00, a01, a02, a03, b10, b11) - stage_b_subtile_pass(k_refill, 1, 3, refill_b) - # K32 slice 1 already covers K[32:64]. + stage_a_subtile_pass(k_refill, 1, 2, refill_a) + mfma_2n_4mi_k32(1, 2, 1, a00, a01, a02, a03, b12, b13) - stage_a_subtile_pass(k_refill, 1, 3, refill_a) - # Keep this refill slot compute-free. - hot_loop_scheduler_q_refill_2n() + stage_b_subtile_pass(k_refill, 1, 3, refill_b) + # K32 slice 1 already covers K[32:64]. - # Leave exactly the K+2 refill and scale loads outstanding. The following - # LDS reads consume the already-ready next page, not the page being refilled. - rocdl.sched_barrier(0) - barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) - rocdl.sched_barrier(0) + stage_a_subtile_pass(k_refill, 1, 3, refill_a) + # Keep this refill slot compute-free. + hot_loop_scheduler_q_refill_2n() - next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) - mfma_4n_4mi_k32(2, 0, a10, a11, a12, a13, b00, b01, b02, b03) + # Leave exactly the K+2 refill and scale loads outstanding. The following + # LDS reads consume the already-ready next page, not the page being refilled. + rocdl.sched_barrier(0) + barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) - next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) - # K32 slice 0 already covers K[0:32]. + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n_4mi_k32(2, 0, a10, a11, a12, a13, b00, b01, b02, b03) - next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) - mfma_4n_4mi_k32(2, 1, a10, a11, a12, a13, b00, b01, b02, b03) + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + # K32 slice 0 already covers K[0:32]. - next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) - # K32 slice 1 already covers K[32:64]. + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n_4mi_k32(2, 1, a10, a11, a12, a13, b00, b01, b02, b03) - next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) - mfma_4n_4mi_k32(3, 0, a10, a11, a12, a13, b10, b11, b12, b13) + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + # K32 slice 1 already covers K[32:64]. - next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) - # K32 slice 0 already covers K[0:32]. + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n_4mi_k32(3, 0, a10, a11, a12, a13, b10, b11, b12, b13) - next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) - mfma_4n_4mi_k32(3, 1, a10, a11, a12, a13, b10, b11, b12, b13) + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + # K32 slice 0 already covers K[0:32]. - next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) - # K32 slice 1 already covers K[32:64]. + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n_4mi_k32(3, 1, a10, a11, a12, a13, b10, b11, b12, b13) - hot_loop_scheduler_q_prefetch_4n() + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + # K32 slice 1 already covers K[32:64]. - next_a0_regs = (next_a00, next_a01, next_a02, next_a03) - next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + hot_loop_scheduler_q_prefetch_4n() - return next_a0_regs, next_b0_regs + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) - def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): - barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + return next_a0_regs, next_b0_regs - a00, a01, a02, a03 = a0_regs - b00, b01, b02, b03 = b0_regs + def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): + barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) - b10 = load_b_subtile_ni_regs(cur_b, 1, 0) - b11 = load_b_subtile_ni_regs(cur_b, 1, 1) - b12 = load_b_subtile_ni_regs(cur_b, 1, 2) - b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs - mfma_4n(_acc_idx(0, 0, 0), a00, b00, b01, b02, b03) - mfma_4n(_acc_idx(0, 1, 0), a01, b00, b01, b02, b03) - mfma_4n(_acc_idx(0, 2, 0), a02, b00, b01, b02, b03) - mfma_4n(_acc_idx(0, 3, 0), a03, b00, b01, b02, b03) + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) - rocdl.sched_barrier(0) - barrier(lgkmcnt=0) - rocdl.sched_barrier(0) + mfma_4n(_acc_idx(0, 0, 0), a00, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 1, 0), a01, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 2, 0), a02, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 3, 0), a03, b00, b01, b02, b03) - a10 = load_a_subtile_mi_regs(cur_a, 1, 0) - a11 = load_a_subtile_mi_regs(cur_a, 1, 1) - a12 = load_a_subtile_mi_regs(cur_a, 1, 2) - a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + rocdl.sched_barrier(0) + barrier(lgkmcnt=0) + rocdl.sched_barrier(0) - mfma_4n(_acc_idx(1, 0, 0), a00, b10, b11, b12, b13) - mfma_4n(_acc_idx(1, 1, 0), a01, b10, b11, b12, b13) - mfma_4n(_acc_idx(1, 2, 0), a02, b10, b11, b12, b13) - mfma_4n(_acc_idx(1, 3, 0), a03, b10, b11, b12, b13) + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) - rocdl.sched_barrier(0) - barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) - rocdl.sched_barrier(0) + mfma_4n(_acc_idx(1, 0, 0), a00, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 1, 0), a01, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 2, 0), a02, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 3, 0), a03, b10, b11, b12, b13) - next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) - mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) + rocdl.sched_barrier(0) + barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) - next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) - mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) - next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) - mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) - next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) - mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) - next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) - mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) - next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) - mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) - next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) - mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) - next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) - mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) - hot_loop_scheduler_q_prefetch_4n() + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) - next_a0_regs = (next_a00, next_a01, next_a02, next_a03) - next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + hot_loop_scheduler_q_prefetch_4n() - return next_a0_regs, next_b0_regs + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) - def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): - barrier(vmcnt=0, lgkmcnt=0) + return next_a0_regs, next_b0_regs - a00, a01, a02, a03 = a0_regs - b00, b01, b02, b03 = b0_regs + def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): + barrier(vmcnt=0, lgkmcnt=0) - # Materialize the remaining final-page A/B fragments once. The - # subsequent schedule is entirely register/AGPR traffic. - b10 = load_b_subtile_ni_regs(cur_b, 1, 0) - b11 = load_b_subtile_ni_regs(cur_b, 1, 1) - b12 = load_b_subtile_ni_regs(cur_b, 1, 2) - b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs - rocdl.sched_barrier(0) - barrier(lgkmcnt=0) - rocdl.sched_barrier(0) + # Materialize the remaining final-page A/B fragments once. The + # subsequent schedule is entirely register/AGPR traffic. + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) - a10 = load_a_subtile_mi_regs(cur_a, 1, 0) - a11 = load_a_subtile_mi_regs(cur_a, 1, 1) - a12 = load_a_subtile_mi_regs(cur_a, 1, 2) - a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + rocdl.sched_barrier(0) + barrier(lgkmcnt=0) + rocdl.sched_barrier(0) - rocdl.sched_barrier(0) - barrier(lgkmcnt=0) - rocdl.sched_barrier(0) + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) - a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) - b_frags = (b00, b01, b02, b03, b10, b11, b12, b13) - - # Rolling final-page epilogue. - # - # Finalize accumulators in their own physical AGPR slots, but delay - # each AGPR read/store until several independent final MFMAs have - # been issued. - # - # MFMA 0, MFMA 1, MFMA 2, MFMA 3, drain 0, - # MFMA 4, drain 1, MFMA 5, drain 2, ... - # - # The buffer stores are only issued here; they may remain in flight - # while later MFMAs and accumulator drains continue. - FINAL_EPILOGUE_DEPTH = 4 - pending = [] - - for old_acc_idx in range_constexpr(ACCS_PER_WAVE): - subtile_id = old_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) - local_idx = old_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) - sm = subtile_id // 2 - sn = subtile_id % 2 - mi = local_idx // MFMA_N_PER_SUBTILE - ni = local_idx % MFMA_N_PER_SUBTILE + rocdl.sched_barrier(0) + barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) + b_frags = (b00, b01, b02, b03, b10, b11, b12, b13) + + # Rolling final-page epilogue. + # + # Finalize accumulators in their own physical AGPR slots, but delay + # each AGPR read/store until several independent final MFMAs have + # been issued. + # + # MFMA 0, MFMA 1, MFMA 2, MFMA 3, drain 0, + # MFMA 4, drain 1, MFMA 5, drain 2, ... + # + # The buffer stores are only issued here; they may remain in flight + # while later MFMAs and accumulator drains continue. + FINAL_EPILOGUE_DEPTH = 4 + pending = [] + + for old_acc_idx in range_constexpr(ACCS_PER_WAVE): + subtile_id = old_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = old_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE + + a_frag_idx = sm * MFMA_M_PER_SUBTILE + mi + b_frag_idx = sn * MFMA_N_PER_SUBTILE + ni + + # Final MFMA remains in-place. The logical accumulator's own + # AGPR slot is unique and cannot conflict with another pending + # result, so no ad-hoc physical-slot permutation is needed. + pinned_final_mfma( + old_acc_idx, + old_acc_idx, + a_frags[a_frag_idx], + b_frags[b_frag_idx], + ) + pending.append(old_acc_idx) - a_frag_idx = sm * MFMA_M_PER_SUBTILE + mi - b_frag_idx = sn * MFMA_N_PER_SUBTILE + ni - - # Final MFMA remains in-place. The logical accumulator's own - # AGPR slot is unique and cannot conflict with another pending - # result, so no ad-hoc physical-slot permutation is needed. - pinned_final_mfma( - old_acc_idx, - old_acc_idx, - a_frags[a_frag_idx], - b_frags[b_frag_idx], - ) - pending.append(old_acc_idx) + # Drain the oldest completed result only after enough newer + # independent MFMAs have supplied the MFMA->AGPR-read spacing. + if len(pending) == FINAL_EPILOGUE_DEPTH: + drain_acc_idx = pending.pop(0) + acc = read_physical_accumulator_slot(drain_acc_idx) + store_acc_vector_for_logical_idx(drain_acc_idx, acc) - # Drain the oldest completed result only after enough newer - # independent MFMAs have supplied the MFMA->AGPR-read spacing. - if len(pending) == FINAL_EPILOGUE_DEPTH: - drain_acc_idx = pending.pop(0) + # Flush the final results after all final-page MFMAs have issued. + for drain_acc_idx in pending: acc = read_physical_accumulator_slot(drain_acc_idx) store_acc_vector_for_logical_idx(drain_acc_idx, acc) - # Flush the final results after all final-page MFMAs have issued. - for drain_acc_idx in pending: - acc = read_physical_accumulator_slot(drain_acc_idx) - store_acc_vector_for_logical_idx(drain_acc_idx, acc) - - # Prologue: stage K0/K1 data into ping-pong LDS pages. - stage_a_subtile(fx.Index(0), 0, lds_a0) - stage_b_subtile(fx.Index(0), 0, lds_b0) - stage_b_subtile(fx.Index(0), 1, lds_b0) - stage_a_subtile(fx.Index(0), 1, lds_a0) + # Prologue: stage K0/K1 data into ping-pong LDS pages. + stage_a_subtile(fx.Index(0), 0, lds_a0) + stage_b_subtile(fx.Index(0), 0, lds_b0) + stage_b_subtile(fx.Index(0), 1, lds_b0) + stage_a_subtile(fx.Index(0), 1, lds_a0) - stage_a_subtile(fx.Index(BLOCK_K), 0, lds_a1) - stage_b_subtile(fx.Index(BLOCK_K), 0, lds_b1) - stage_b_subtile(fx.Index(BLOCK_K), 1, lds_b1) - stage_a_subtile(fx.Index(BLOCK_K), 1, lds_a1) + stage_a_subtile(fx.Index(BLOCK_K), 0, lds_a1) + stage_b_subtile(fx.Index(BLOCK_K), 0, lds_b1) + stage_b_subtile(fx.Index(BLOCK_K), 1, lds_b1) + stage_a_subtile(fx.Index(BLOCK_K), 1, lds_a1) - rocdl.sched_barrier(0) - barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 4 * LOAD_PASSES_B_SUBTILE) - rocdl.sched_barrier(0) + rocdl.sched_barrier(0) + barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 4 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) - a0_regs = load_a_subtile_regs(lds_a0, 0) + a0_regs = load_a_subtile_regs(lds_a0, 0) - rocdl.sched_barrier(0) - barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) - rocdl.sched_barrier(0) + rocdl.sched_barrier(0) + barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) - b0_regs = load_b_subtile_regs(lds_b0, 0) + b0_regs = load_b_subtile_regs(lds_b0, 0) + + # Main HK loop: exactly one logical K64 per iteration. + # Even k consumes and refills LDS0; odd k does the same for LDS1. + for k128 in range_constexpr(NUM_K_TILES - 2): + if (k128 % 2) == 0: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + else: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) - # Main HK loop: exactly one logical K64 per iteration. - # Even k consumes and refills LDS0; odd k does the same for LDS1. - for k128 in range_constexpr(NUM_K_TILES - 2): - if (k128 % 2) == 0: - a0_regs, b0_regs = hk_one_k_with_refill( - k128, + # Common two-page tail. The penultimate tile uses Q2/Q3 carry-prefetch + # to prepare A-top/B-left for the final tile, but performs no K+2 refill. + if (NUM_K_TILES % 2) == 0: + a0_regs, b0_regs = hk_one_k_tail_with_next( lds_a0, lds_b0, lds_a1, lds_b1, - lds_a0, - lds_b0, a0_regs, b0_regs, ) + hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs) else: - a0_regs, b0_regs = hk_one_k_with_refill( - k128, + a0_regs, b0_regs = hk_one_k_tail_with_next( lds_a1, lds_b1, lds_a0, lds_b0, - lds_a1, - lds_b1, a0_regs, b0_regs, ) - - # Common two-page tail. The penultimate tile uses Q2/Q3 carry-prefetch - # to prepare A-top/B-left for the final tile, but performs no K+2 refill. - if (NUM_K_TILES % 2) == 0: - a0_regs, b0_regs = hk_one_k_tail_with_next( - lds_a0, - lds_b0, - lds_a1, - lds_b1, - a0_regs, - b0_regs, - ) - hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs) + hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs) + + # Dispatch. Dense runs one XCD-swizzled tile. Grouped runs a persistent + # loop over precomputed (pid_m, pid_n, expert) descriptors, selecting the + # schedule at compile time (``sched``): "worksteal" atomically claims the + # next global slot each iteration; "static" owns a fixed grid-stride slice + # (no atomic or LDS broadcast, deadlock-proof). + def _run_grouped_tile(slot): + pid_m = _desc_field(slot, fx.Int32(0)) + pid_n = _desc_field(slot, fx.Int32(1)) + expert = _desc_field(slot, fx.Int32(2)) + process_tile(pid_m, pid_n, expert) + + if const_expr(grouped): + if const_expr(sched == "static"): + stride = fx.Int32(gpu.grid_dim.x) + slot = fx.Int32(gpu.block_id("x")) + while slot < num_tiles: + _run_grouped_tile(slot) + slot = slot + stride + else: + # Single persistent loop (mega_moe_stage1 idiom): claim the next + # slot at the TOP (atomic + barrier + LDS broadcast), run the tile + # only when in range, derive the continue flag from that same + # broadcast slot -- one barrier site, workgroup-uniform condition. + active = fx.Int32(1) == fx.Int32(1) + while active: + slot = _claim_next_slot() + in_range = slot < num_tiles + if in_range: + _run_grouped_tile(slot) + active = in_range else: - a0_regs, b0_regs = hk_one_k_tail_with_next( - lds_a1, - lds_b1, - lds_a0, - lds_b0, - a0_regs, - b0_regs, - ) - hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs) + if const_expr(use_xcd_remap): + pid_m, pid_n = xcd_swizzle(num_blocks_m, num_blocks_n) + else: + pid_m, pid_n = divmod(fx.block_idx.x, num_blocks_n) + process_tile(pid_m, pid_n, fx.Int32(0)) @flyc.jit def launch_gemm( @@ -1222,10 +1333,17 @@ def launch_gemm( Aux: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32, + num_groups: fx.Int32 = 0, + work_queue: fx.Tensor = None, + work_counter: fx.Tensor = None, stream: fx.Stream = fx.Stream(None), ): - # The integration only dispatches aligned shapes; no partial-tile masking exists. + # Dense: one workgroup per output tile. Grouped/wgrad: a fixed num-CU + # persistent grid that drains the descriptor queue. Placeholder work + # tensors for the dense path are allocated host-side in ``doGemm``. grid_x = (c_m // BLOCK_M) * (c_n // BLOCK_N) + if const_expr(group or wgrad): + grid_x = get_num_cus() kernel_gemm( A, B, @@ -1234,6 +1352,9 @@ def launch_gemm( Aux, c_m, c_n, + num_groups, + work_queue, + work_counter, value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, ).launch(grid=(grid_x, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) @@ -1248,6 +1369,9 @@ def _cached_launch( mfma_suffix: str, use_xcd_remap: bool = True, epilogue: str = "DEFAULT", + group: bool = False, + sched: str = "worksteal", + wgrad: bool = False, ): return _compile_kernel( K, @@ -1256,6 +1380,9 @@ def _cached_launch( mfma_suffix, use_xcd_remap=use_xcd_remap, epilogue=epilogue, + group=group, + sched=sched, + wgrad=wgrad, ) @@ -1424,6 +1551,12 @@ def doGemm( epilogue: str = "DEFAULT", bias: torch.Tensor = None, aux: torch.Tensor = None, + group: bool = False, + num_groups: int = 0, + work_queue: torch.Tensor = None, + work_counter: torch.Tensor = None, + sched: str = "worksteal", + wgrad: bool = False, ): """Launch one cached K/output/layout-specialized half-precision core. @@ -1462,8 +1595,11 @@ def doGemm( ) require_launch_size(f"{label} GEMM", ("A", A), ("B", B), ("C", C)) - if tuple(C.shape) != (M_runtime, N_runtime): - raise ValueError(f"C shape {tuple(C.shape)} != expected {(M_runtime, N_runtime)}") + # wgrad's C is packed [G*N_out, K_in] on the output-M axis; fwd/dgrad/dense C is + # the unpacked (M_runtime, N_runtime). + expected_c = (num_groups * M_runtime, N_runtime) if wgrad else (M_runtime, N_runtime) + if tuple(C.shape) != expected_c: + raise ValueError(f"C shape {tuple(C.shape)} != expected {expected_c}") if epilogue not in ("DEFAULT", "BIAS", "GELU_AUX", "GELU_AUX_BIAS"): raise ValueError(f"Unsupported {label} epilogue: {epilogue}") @@ -1506,6 +1642,9 @@ def doGemm( _MFMA_SUFFIX[input_dtype], bool(use_xcd_remap), epilogue, + group, + sched, + wgrad, ) # Preserve the original validated byte-addressed G2L path. These are # metadata-only dtype/flatten views of the already-contiguous row-major @@ -1523,6 +1662,15 @@ def doGemm( else: Aux_arg = torch.zeros(1, dtype=C.dtype, device=A.device) + # Grouped runtime tensors are required by the kernel signature; the dense + # binary never dereferences them but still needs placeholder buffers. + if group or wgrad: + work_queue_arg = work_queue.contiguous().view(-1) + work_counter_arg = work_counter.contiguous().view(-1) + else: + work_queue_arg = torch.zeros(1, dtype=torch.int32, device=A.device) + work_counter_arg = torch.zeros(1, dtype=torch.int32, device=A.device) + launch( A_arg, B_arg, @@ -1531,5 +1679,8 @@ def doGemm( Aux_arg, M_runtime, N_runtime, + num_groups, + work_queue_arg, + work_counter_arg, stream=stream, ) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index df612094a7..14d13e7d7a 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -28,16 +28,15 @@ import flydsl.compiler as flyc import flydsl.expr as fx -from flydsl._mlir import ir -from flydsl._mlir.dialects import fly as _fly from flydsl._mlir.dialects import llvm from flydsl.expr import arith, const_expr, gpu, math, range_constexpr, rocdl from flydsl.expr.typing import T from flydsl.expr.typing import Vector as Vec # Transformer Engine-local FlyDSL utilities. +from . import grouped_gemm_utils as gd from .exceptions import FlyDSLUnsupportedError -from .gemm_common_utils import require_block_tiling, require_launch_size +from .gemm_common_utils import get_num_cus, require_block_tiling, require_launch_size from .fp8_gemm_utils import ( G2SLoader, S2RLoader, @@ -61,10 +60,6 @@ SCALE_GROUP_SIZE = 32 -def get_num_cus() -> int: - return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count - - def _debug_enabled() -> bool: value = os.getenv("DEBUG_FLYDSL_MXFP8_GEMM", "") return value.lower() not in ("", "0", "false", "no", "off") @@ -524,9 +519,7 @@ def kernel_gemm( # all experts, so this is already the total. For wgrad, c_m = N_out is the # PER-EXPERT output-row count (the packed output [G*N_out, K_in] is addressed # via store_m_expert_off), so multiply by num_groups for the full count. - num_tiles = (c_m // BLOCK_M) * (c_n // BLOCK_N) - if const_expr(wgrad): - num_tiles = num_tiles * num_groups + num_tiles = gd.num_tiles_bound(c_m, c_n, BLOCK_M, BLOCK_N, num_groups, const_expr(wgrad)) lds = fx.SharedAllocator().allocate(SharedStorage).peek() lds_a0 = (lds.a0_0, lds.a0_1) @@ -578,28 +571,19 @@ def kernel_gemm( # LDS broadcast scratch (one int32), following the persistent-kernel # idiom in mega_moe_stage1: recast the LDS slot to an int32 iterator and # a 1-element view. Thread 0 writes via ptr_store; every thread reads via - # the view after a workgroup barrier. - claim_iter = fx.recast_iter(fx.Int32, lds.claim_slot.ptr) - claim_view = fx.make_view(claim_iter, fx.make_layout(1, 1)) - - def _global_i32_ptr(tensor, index): - # Build an address-space-1 !llvm.ptr to tensor[index] (int32). - # llvm.atomicrmw requires an !llvm.ptr (not !fly.ptr). - ptr_ty = ir.Type.parse("!llvm.ptr<1>") - base = _fly.extract_aligned_pointer_as_index(ptr_ty, tensor) - base = llvm.PtrToIntOp(T.i64, base).result - byte_off = fx.as_ir_value(fx.Int64(fx.Index(index) * fx.Index(4))) - addr = llvm.AddOp(base, byte_off, llvm.IntegerOverflowFlags(0)).result - return llvm.IntToPtrOp(ptr_ty, addr).result + # the view after a workgroup barrier. (Descriptor reads reuse the int32 + # ``scale_ld_atom`` as their copy atom.) + claim_iter, claim_view = gd.make_claim_broadcast(lds.claim_slot) def _claim_next_slot(): # One atomic per workgroup: thread 0 fetch-adds the shared counter # and stores the pre-update value into the LDS scratch; a barrier - # then broadcasts it so every thread returns the same slot. + # then broadcasts it so every thread returns the same slot. The + # ``if tx==0`` control flow must stay inline (AST-rewritten here). if tx_i32 == fx.Int32(0): old = llvm.AtomicRMWOp( llvm.AtomicBinOp.add, - _global_i32_ptr(work_counter, fx.Int32(0)), + gd.global_i32_ptr(work_counter, fx.Int32(0)), fx.as_ir_value(fx.Int32(1)), llvm.AtomicOrdering.monotonic, syncscope="agent", @@ -610,10 +594,7 @@ def _claim_next_slot(): return fx.Vector(claim_view.load())[0] def _desc_field(slot, field): - # Read one int32 of the packed (pid_m, pid_n, expert) descriptor. - reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Int32) - fx.copy(scale_ld_atom, fx.slice(wq_div, (None, slot * fx.Int32(3) + field)), reg) - return fx.memref_load_vec(reg)[0] + return gd.read_desc_field(wq_div, scale_ld_atom, slot, field) def process_tile(pid_m, pid_n, expert): bx_m = pid_m * BLOCK_M diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index b48dc39bc6..af158e7b45 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -578,10 +578,8 @@ def _is_blockwise_fp8_triton_grouped_gemm_supported( ) @staticmethod - def _is_flydsl_mxfp8_grouped_gemm_supported( + def _flydsl_grouped_common_ok( *, - fp8, - recipe, use_bias, backward_override, cpu_offloading, @@ -592,20 +590,17 @@ def _is_flydsl_mxfp8_grouped_gemm_supported( in_features, out_features, ) -> bool: - """FlyDSL persistent work-stealing MXFP8 grouped GEMM: HIP, gfx950, env, MXFP8 recipe. + """Dtype-independent gate shared by every FlyDSL grouped GEMM path. - Opt-in via ``NVTE_USE_GROUPED_GEMM_FLYDSL=1``. The kernel is a TN-only packed - launch (see :func:`_run_mxfp8_grouped`); N (out_features) and K (in_features) - must be 128-aligned for the tile/scale geometry, and the same feature-set - exclusions as the blockwise-triton path apply. + HIP extension, gfx950, the ``NVTE_USE_GROUPED_GEMM_FLYDSL=1`` opt-in, + 128-aligned N/K for the tile geometry, and the same feature-set exclusions + as the blockwise-triton path. The per-dtype gates add only the quantization + condition (BF16: none; tensor-wise FP8: current/delayed scaling; MXFP8). """ return ( IS_HIP_EXTENSION and os.getenv("NVTE_USE_GROUPED_GEMM_FLYDSL", "0") == "1" and get_device_compute_capability() == (9, 5) - and fp8 - and recipe is not None - and recipe.mxfp8() and in_features % 128 == 0 and out_features % 128 == 0 and not ( @@ -619,6 +614,84 @@ def _is_flydsl_mxfp8_grouped_gemm_supported( ) ) + @staticmethod + def _flydsl_grouped_kind( + *, + fp8, + recipe, + use_bias, + backward_override, + cpu_offloading, + save_original_input, + debug, + unpad_output, + actual_m_splits, + in_features, + out_features, + ): + """Return which FlyDSL grouped path applies: 'mxfp8' | 'fp8' | 'bf16' | None. + + One dispatch signal replacing the former MXFP8-only boolean. All kinds share + the persistent work-stealing driver; they differ only in operand packing and + (fp8/mxfp8) the scale payload. BF16 is the unquantized path (fp8 is False). + """ + if not _GroupedLinear._flydsl_grouped_common_ok( + use_bias=use_bias, + backward_override=backward_override, + cpu_offloading=cpu_offloading, + save_original_input=save_original_input, + debug=debug, + unpad_output=unpad_output, + actual_m_splits=actual_m_splits, + in_features=in_features, + out_features=out_features, + ): + return None + if not fp8: + # Unquantized BF16 grouped GEMM (no recipe). FP16 is not routed here; + # the activation dtype is checked at the dispatch site. + return "bf16" + if recipe is None: + return None + if recipe.mxfp8(): + return "mxfp8" + if recipe.float8_current_scaling() or recipe.delayed(): + return "fp8" + return None + + @staticmethod + def _is_flydsl_mxfp8_grouped_gemm_supported( + *, + fp8, + recipe, + use_bias, + backward_override, + cpu_offloading, + save_original_input, + debug, + unpad_output, + actual_m_splits, + in_features, + out_features, + ) -> bool: + """Back-compat shim: True only for the MXFP8 FlyDSL grouped path.""" + return ( + _GroupedLinear._flydsl_grouped_kind( + fp8=fp8, + recipe=recipe, + use_bias=use_bias, + backward_override=backward_override, + cpu_offloading=cpu_offloading, + save_original_input=save_original_input, + debug=debug, + unpad_output=unpad_output, + actual_m_splits=actual_m_splits, + in_features=in_features, + out_features=out_features, + ) + == "mxfp8" + ) + @staticmethod def _grouped_gemm_flydsl_mxfp8( *, @@ -785,6 +858,232 @@ def _wgrad_flydsl_mxfp8_grouped( for g in range(num_gemms): wgrad_list[g].copy_(dW[g * n_out : (g + 1) * n_out]) + @staticmethod + def _flydsl_group_offs(num_gemms, m_splits, m_splits_tensor, device): + """Build the [num_gemms+1] exclusive-scan token offsets on ``device``. + + Prefers the GPU ``m_splits_tensor`` (avoids a blocking H2D of pageable + ``m_splits``) but falls back to the Python ``m_splits`` list when the tensor + is absent (the BF16 path does not always carry it). + """ + group_offs = torch.zeros(num_gemms + 1, dtype=torch.int64, device=device) + if m_splits_tensor is not None: + group_offs[1:] = torch.cumsum( + m_splits_tensor.to(device=device, dtype=torch.int64), 0 + ) + else: + counts = torch.tensor(list(m_splits), dtype=torch.int64, device=device) + group_offs[1:] = torch.cumsum(counts, 0) + return group_offs + + @staticmethod + def _grouped_gemm_flydsl_bf16(*, inputmats, weights, out, m_splits, m_splits_tensor): + """FlyDSL BF16 grouped forward (TN): Y_e = X_e @ W_e^T, one packed launch. + + Unquantized analog of :meth:`_grouped_gemm_flydsl_mxfp8`: activations are the + per-group BF16 splits (concat on the token axis), weights the per-expert BF16 + tensors (concat on the output-N axis). No scales. + """ + from ..flydsl_kernels.gemm import _run_bf16_grouped + + num_gemms = len(inputmats) + a_data = torch.cat([im.reshape(-1, im.shape[-1]) for im in inputmats], dim=0) + b_data = torch.cat([w for w in weights], dim=0) + n = weights[0].size(0) + k = weights[0].size(-1) + m_total = a_data.size(0) + + group_offs = _GroupedLinear._flydsl_group_offs( + num_gemms, m_splits, m_splits_tensor, out.device + ) + + _run_bf16_grouped( + a_data, + b_data, + out, + group_offs=group_offs, + m_total=m_total, + n=n, + k=k, + output_dtype=out.dtype, + ) + + @staticmethod + def _backward_flydsl_bf16_grouped(*, grad_output, weights, dgrad, m_splits, m_splits_tensor): + """FlyDSL BF16 grouped dgrad (NN): dX_e = dY_e @ W_e, one packed launch.""" + from ..flydsl_kernels.gemm import _run_bf16_grouped + + num_gemms = len(grad_output) + a_data = torch.cat([g.reshape(-1, g.shape[-1]) for g in grad_output], dim=0) + b_data = torch.cat([w for w in weights], dim=0) + n_out = grad_output[0].shape[-1] # per-expert contraction (N_out) + k_in = dgrad.size(-1) # output width (K_in) + m_total = a_data.size(0) + + group_offs = _GroupedLinear._flydsl_group_offs( + num_gemms, m_splits, m_splits_tensor, dgrad.device + ) + + _run_bf16_grouped( + a_data, + b_data, + dgrad, + group_offs=group_offs, + m_total=m_total, + n=k_in, + k=n_out, + output_dtype=dgrad.dtype, + layout="NN", + ) + + @staticmethod + def _wgrad_flydsl_bf16_grouped(*, grad_output, inputmats, wgrad_list, m_splits): + """FlyDSL BF16 grouped wgrad (NT): dW_e = dY_e^T @ X_e, one packed launch.""" + from ..flydsl_kernels.gemm import _run_bf16_grouped_wgrad + + num_gemms = len(inputmats) + dy_cols = [g.reshape(-1, g.shape[-1]) for g in grad_output] + x_cols = [x.reshape(-1, x.shape[-1]) for x in inputmats] + n_out = dy_cols[0].shape[-1] + k_in = x_cols[0].shape[-1] + + dW = torch.empty( + num_gemms * n_out, k_in, dtype=wgrad_list[0].dtype, device=wgrad_list[0].device + ) + _run_bf16_grouped_wgrad( + dy_cols, + x_cols, + dW, + m_splits=m_splits, + n_out=n_out, + k_in=k_in, + output_dtype=wgrad_list[0].dtype, + ) + for g in range(num_gemms): + wgrad_list[g].copy_(dW[g * n_out : (g + 1) * n_out]) + + @staticmethod + def _grouped_gemm_flydsl_fp8(*, inputmats, weights_fp8, out, m_splits, m_splits_tensor): + """FlyDSL tensor-wise FP8 grouped forward (TN): Y_e = X_e @ W_e^T. + + Experts pack weights on the output-N axis ([G*N, K]); each operand carries + one scalar inverse scale, gathered into length-G per-expert vectors. + """ + from ..flydsl_kernels.gemm import _run_fp8_grouped + + num_gemms = len(inputmats) + a_data = torch.cat([im._data for im in inputmats], dim=0) + b_data = torch.cat([w._data for w in weights_fp8], dim=0) + a_data._fp8_dtype = inputmats[0]._fp8_dtype + b_data._fp8_dtype = weights_fp8[0]._fp8_dtype + a_scale = torch.cat([im._scale_inv.reshape(1) for im in inputmats], dim=0) + b_scale = torch.cat([w._scale_inv.reshape(1) for w in weights_fp8], dim=0) + + n = weights_fp8[0]._data.size(0) + k = weights_fp8[0]._data.size(-1) + m_total = a_data.size(0) + + group_offs = _GroupedLinear._flydsl_group_offs( + num_gemms, m_splits, m_splits_tensor, out.device + ) + + _run_fp8_grouped( + a_data, + a_scale, + b_data, + b_scale, + out, + group_offs=group_offs, + m_total=m_total, + n=n, + k=k, + output_dtype=out.dtype, + ) + + @staticmethod + def _backward_flydsl_fp8_grouped( + *, grad_output, weights_for_dgrad, dgrad, m_splits, m_splits_tensor + ): + """FlyDSL tensor-wise FP8 grouped dgrad: dX_e = dY_e @ W_e, via the TN kernel. + + FP8 is TN-only, so dgrad reaches the same TN kernel using materialized + transpose backings: A = dY rowwise ``_data`` [M_total, N_out], B = W^T + columnwise ``_transpose`` [G*N_out... ]. The TN kernel computes A @ B^T with + B packed on the output axis = K_in, i.e. per expert dY @ (W^T)^T = dY @ W. + """ + from ..flydsl_kernels.gemm import _run_fp8_grouped + + num_gemms = len(grad_output) + # A = dY (rowwise fp8 data); B = W columnwise transpose, which is the + # materialized [K_in, N_out] -> packed on K_in (output axis of the TN GEMM). + a_data = torch.cat([g._data for g in grad_output], dim=0) + b_data = torch.cat([w._transpose for w in weights_for_dgrad], dim=0) + a_data._fp8_dtype = grad_output[0]._fp8_dtype + b_data._fp8_dtype = weights_for_dgrad[0]._fp8_dtype + a_scale = torch.cat([g._scale_inv.reshape(1) for g in grad_output], dim=0) + b_scale = torch.cat([w._scale_inv.reshape(1) for w in weights_for_dgrad], dim=0) + + n_out = grad_output[0]._data.size(-1) # per-expert contraction (N_out) + k_in = dgrad.size(-1) # output width (K_in) + m_total = a_data.size(0) + + group_offs = _GroupedLinear._flydsl_group_offs( + num_gemms, m_splits, m_splits_tensor, dgrad.device + ) + + _run_fp8_grouped( + a_data, + a_scale, + b_data, + b_scale, + dgrad, + group_offs=group_offs, + m_total=m_total, + n=k_in, + k=n_out, + output_dtype=dgrad.dtype, + ) + + @staticmethod + def _wgrad_flydsl_fp8_grouped(*, grad_output, inputmats, wgrad_list, m_splits): + """FlyDSL tensor-wise FP8 grouped wgrad: dW_e = dY_e^T @ X_e, via the TN kernel. + + Both operands reach the TN kernel already transposed: A = dY^T (the fp8 + ``_transpose`` backing, [N_out, m_e]), B = X^T ([K_in, m_e]). The kernel + computes A @ B^T = dY^T @ X = dW. + """ + from ..flydsl_kernels.gemm import _run_fp8_grouped_wgrad + + num_gemms = len(inputmats) + dyt_cols = [g._transpose for g in grad_output] + xt_cols = [x._transpose for x in inputmats] + dy_scale = torch.cat([g._scale_inv.reshape(1) for g in grad_output], dim=0) + x_scale = torch.cat([x._scale_inv.reshape(1) for x in inputmats], dim=0) + + # Operands may carry only the columnwise (_transpose) backing. The + # transpose is [feature, m_e]: dY^T is [N_out, m_e], X^T is [K_in, m_e]. + n_out = dyt_cols[0].size(0) # per-expert output feature N + k_in = xt_cols[0].size(0) # per-expert input feature K + + dW = torch.empty( + num_gemms * n_out, k_in, dtype=wgrad_list[0].dtype, device=wgrad_list[0].device + ) + _run_fp8_grouped_wgrad( + dyt_cols, + dy_scale, + xt_cols, + x_scale, + dW, + m_splits=m_splits, + n_out=n_out, + k_in=k_in, + dy_fp8_dtype=grad_output[0]._fp8_dtype, + x_fp8_dtype=inputmats[0]._fp8_dtype, + output_dtype=wgrad_list[0].dtype, + ) + for g in range(num_gemms): + wgrad_list[g].copy_(dW[g * n_out : (g + 1) * n_out]) + @staticmethod def _forward_blockwise_fp8_triton( ctx, @@ -1192,7 +1491,7 @@ def forward( # FlyDSL MXFP8 grouped GEMM (ROCm gfx950) opt-in. Unlike the blockwise-triton # path this does NOT quantize: it rides off TE's existing quantization and is # dispatched inline at the forward GEMM (see below), like the triton grouped path. - use_grouped_gemm_flydsl = _GroupedLinear._is_flydsl_mxfp8_grouped_gemm_supported( + flydsl_grouped_kind = _GroupedLinear._flydsl_grouped_kind( fp8=fp8, recipe=blockwise_recipe, use_bias=use_bias, @@ -1205,6 +1504,11 @@ def forward( in_features=weights[0].size(-1), out_features=weights[0].size(0), ) + # The BF16 kernel is bfloat16-only (FP16 grouped is not wired through + # FlyDSL); drop the BF16 kind for any other activation dtype. + if flydsl_grouped_kind == "bf16" and activation_dtype != torch.bfloat16: + flydsl_grouped_kind = None + use_grouped_gemm_flydsl = flydsl_grouped_kind is not None # Configure quantizers if save_original_input and isinstance(input_quantizers[0], Float8Quantizer): @@ -1376,13 +1680,30 @@ def forward( # Perform GEMM if use_grouped_gemm_flydsl: - _GroupedLinear._grouped_gemm_flydsl_mxfp8( - inputmats=inputmats, - weights_fp8=weights_fp8, - out=out, - m_splits=m_splits, - m_splits_tensor=m_splits_tensor, - ) + if flydsl_grouped_kind == "mxfp8": + _GroupedLinear._grouped_gemm_flydsl_mxfp8( + inputmats=inputmats, + weights_fp8=weights_fp8, + out=out, + m_splits=m_splits, + m_splits_tensor=m_splits_tensor, + ) + elif flydsl_grouped_kind == "fp8": + _GroupedLinear._grouped_gemm_flydsl_fp8( + inputmats=inputmats, + weights_fp8=weights_fp8, + out=out, + m_splits=m_splits, + m_splits_tensor=m_splits_tensor, + ) + else: # bf16 + _GroupedLinear._grouped_gemm_flydsl_bf16( + inputmats=inputmats, + weights=weights_fp8, + out=out, + m_splits=m_splits, + m_splits_tensor=m_splits_tensor, + ) else: if use_grouped_gemm_triton: general_grouped_gemm_func = general_grouped_gemm_triton @@ -1433,6 +1754,7 @@ def forward( # FlyDSL grouped forward routes dgrad through the grouped NN path; the # shared backward branches on this flag (wgrad still uses the default). ctx.use_flydsl_mxfp8_grouped = use_grouped_gemm_flydsl + ctx.flydsl_grouped_kind = flydsl_grouped_kind ctx.weight_quantizers = weight_quantizers ctx.weights_shape_1 = weights[0].shape[1] @@ -1881,10 +2203,11 @@ def backward( rowwise_usage=quantizer.rowwise_usage, columnwise_usage=quantizer.columnwise_usage, ) - if getattr(ctx, "use_flydsl_mxfp8_grouped", False): + _flydsl_kind = getattr(ctx, "flydsl_grouped_kind", None) + if _flydsl_kind == "mxfp8": # dgrad dX = dY @ W through the FlyDSL grouped NN kernel, riding # off the already-quantized grad_output (rowwise) and weights - # (columnwise, ensured above). wgrad still uses the default path. + # (columnwise, ensured above). _GroupedLinear._backward_flydsl_mxfp8_grouped( grad_output=grad_output, weights_for_dgrad=weights_for_dgrad, @@ -1892,6 +2215,22 @@ def backward( m_splits=ctx.m_splits, m_splits_tensor=ctx.m_splits_tensor, ) + elif _flydsl_kind == "fp8": + _GroupedLinear._backward_flydsl_fp8_grouped( + grad_output=grad_output, + weights_for_dgrad=weights_for_dgrad, + dgrad=dgrad, + m_splits=ctx.m_splits, + m_splits_tensor=ctx.m_splits_tensor, + ) + elif _flydsl_kind == "bf16": + _GroupedLinear._backward_flydsl_bf16_grouped( + grad_output=grad_output, + weights=weights_for_dgrad, + dgrad=dgrad, + m_splits=ctx.m_splits, + m_splits_tensor=ctx.m_splits_tensor, + ) else: if ctx.use_grouped_gemm_triton: general_grouped_gemm_func = general_grouped_gemm_triton @@ -2010,8 +2349,9 @@ def backward( # FlyDSL grouped wgrad (NT) opt-in. Phase 1 handles the plain path # only; fused wgrad accumulation, bias grad, and delayed wgrad still # fall through to the default grouped GEMM below. + _flydsl_kind = getattr(ctx, "flydsl_grouped_kind", None) use_flydsl_wgrad = ( - getattr(ctx, "use_flydsl_mxfp8_grouped", False) + _flydsl_kind is not None and not ctx.fuse_wgrad_accumulation and not ctx.use_bias # bias-grad-in-wgrad-GEMM not supported yet and not ( @@ -2019,12 +2359,27 @@ def backward( ) ) if use_flydsl_wgrad: - _GroupedLinear._wgrad_flydsl_mxfp8_grouped( - grad_output=grad_output, - inputmats=inputmats, - wgrad_list=wgrad_list, - m_splits=ctx.m_splits, - ) + if _flydsl_kind == "mxfp8": + _GroupedLinear._wgrad_flydsl_mxfp8_grouped( + grad_output=grad_output, + inputmats=inputmats, + wgrad_list=wgrad_list, + m_splits=ctx.m_splits, + ) + elif _flydsl_kind == "fp8": + _GroupedLinear._wgrad_flydsl_fp8_grouped( + grad_output=grad_output, + inputmats=inputmats, + wgrad_list=wgrad_list, + m_splits=ctx.m_splits, + ) + else: # bf16 + _GroupedLinear._wgrad_flydsl_bf16_grouped( + grad_output=grad_output, + inputmats=inputmats, + wgrad_list=wgrad_list, + m_splits=ctx.m_splits, + ) clear_tensor_data(*inputmats) grouped_gemm_wgrad = None else: From ddee2cf939509c9697e1a8f2aed040156f29cd15 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 14 Oct 2026 17:35:54 +0000 Subject: [PATCH 05/22] Cache grouped-GEMM descriptor queue, drop per-call D2H sync --- .../flydsl_kernels/gemm/gemm_wrappers.py | 65 ++++++++++------- .../flydsl_kernels/gemm/grouped_gemm_utils.py | 65 ++++++++++++----- .../pytorch/module/grouped_linear.py | 70 +++---------------- 3 files changed, 94 insertions(+), 106 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 8fee753fa7..4457234d82 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -1020,7 +1020,7 @@ def _run_mxfp8_grouped( b_scale, D, *, - group_offs, + m_splits, m_total, n, k, @@ -1043,10 +1043,11 @@ def _run_mxfp8_grouped( output width n = K_in; experts pack on B's contraction (N_out) axis. In both, ``k`` is A's per-expert contraction (TN: K_in; NN: N_out) and ``n`` is - the output width. Each group ``g`` owns activation/grad rows - ``[group_offs[g], group_offs[g+1])``; ``group_offs`` is consumed on the host to - build the persistent kernel's per-tile ``(pid_m, pid_n, expert)`` descriptor - queue, so the kernel does no group lookup. Operands arrive already + the output width. ``m_splits`` is the Python list of per-expert token counts + (already host-resident); the per-tile ``(pid_m, pid_n, expert)`` descriptor + queue is built + shape-cached from it directly -- no device->host copy of a + ``group_offs`` tensor (that copy forced a blocking ``hipDeviceSynchronize`` + that stranded the GPU idle before the kernel). Operands arrive already MXFP8-quantized and packed -- no quantization, no BLAS ownership swap. """ if layout not in ("TN", "NN"): @@ -1086,15 +1087,15 @@ def _run_mxfp8_grouped( # and repeats until slot >= num_tiles. from .mxfp8_gemm import BLOCK_M, BLOCK_N - num_groups = int(group_offs.numel()) - 1 + num_groups = len(m_splits) num_pid_n = n // BLOCK_N - group_offs_cpu = group_offs.to(device="cpu", dtype=torch.int64) - # Per-expert M-block ranges over the packed [M_total, N] output. Each expert's # token count must be BLOCK_M-aligned (enforced upstream by the tiling check), - # so tiles never straddle a group boundary. - descriptors = gd.build_fwd_dgrad_descriptors(group_offs_cpu, num_groups, num_pid_n, BLOCK_M) - work_queue, work_counter, num_tiles = gd.make_work_queue(descriptors, a_data.device) + # so tiles never straddle a group boundary. Built + cached from the host-side + # m_splits list (no GPU->CPU sync). + work_queue, work_counter, num_tiles = gd.build_fwd_dgrad_work_queue( + m_splits, num_pid_n, BLOCK_M, a_data.device + ) # Schedule selection (compile-time specialization on the kernel side): # "worksteal": persistent atomic tile-claiming, best for imbalanced experts. @@ -1190,8 +1191,9 @@ def _run_mxfp8_grouped_wgrad( # kernel re-packs the output row via expert*N_out. num_pid_m = n_out // BLOCK_M num_pid_n = k_in // BLOCK_N - descriptors = gd.build_wgrad_descriptors(num_groups, num_pid_m, num_pid_n) - work_queue, work_counter, num_tiles = gd.make_work_queue(descriptors, device) + work_queue, work_counter, num_tiles = gd.build_wgrad_work_queue( + num_groups, num_pid_m, num_pid_n, device + ) sched = _select_grouped_sched(num_tiles, device) mxfp8_matmul( @@ -1215,7 +1217,7 @@ def _run_bf16_grouped( b_data, D, *, - group_offs, + m_splits, m_total, n, k, @@ -1233,6 +1235,11 @@ def _run_bf16_grouped( a = packed activations [M_total, K] b = packed weights [G*N, K] NN (dgrad, dX = dY @ W): a = packed grad_output [M_total, N_out] b = packed weights [G*N_out, K_in] + + ``m_splits`` is the Python list of per-expert token counts (already on the + host). The descriptor queue is built + cached from it directly -- no + device->host copy of a ``group_offs`` tensor (that copy forced a blocking + ``hipDeviceSynchronize`` that stranded the GPU idle before the kernel). """ from .half_prec_gemm import doGemm as _half_prec_doGemm @@ -1249,11 +1256,11 @@ def _run_bf16_grouped( from .half_prec_gemm import _BLOCK_M, _BLOCK_N - num_groups = int(group_offs.numel()) - 1 + num_groups = len(m_splits) num_pid_n = n // _BLOCK_N - group_offs_cpu = group_offs.to(device="cpu", dtype=torch.int64) - descriptors = gd.build_fwd_dgrad_descriptors(group_offs_cpu, num_groups, num_pid_n, _BLOCK_M) - work_queue, work_counter, num_tiles = gd.make_work_queue(descriptors, a_data.device) + work_queue, work_counter, num_tiles = gd.build_fwd_dgrad_work_queue( + m_splits, num_pid_n, _BLOCK_M, a_data.device + ) sched = _select_grouped_sched(num_tiles, a_data.device) epilogue, bias_arg = _resolve_epilogue(bias, n) @@ -1321,8 +1328,9 @@ def _run_bf16_grouped_wgrad( num_pid_m = n_out // _BLOCK_M num_pid_n = k_in // _BLOCK_N - descriptors = gd.build_wgrad_descriptors(num_groups, num_pid_m, num_pid_n) - work_queue, work_counter, num_tiles = gd.make_work_queue(descriptors, device) + work_queue, work_counter, num_tiles = gd.build_wgrad_work_queue( + num_groups, num_pid_m, num_pid_n, device + ) sched = _select_grouped_sched(num_tiles, device) # The kernel's compile-time contraction K is the PER-EXPERT padded token count; @@ -1354,7 +1362,7 @@ def _run_fp8_grouped( b_scale_vec, # [num_groups] fp32 per-expert inverse scales D, *, - group_offs, + m_splits, m_total, n, k, @@ -1367,10 +1375,12 @@ def _run_fp8_grouped( host-materialized transpose backings (see the backward helpers). Experts pack on B's output-N axis ([G*N, K]); each expert carries one scalar inverse scale, so the scales arrive as length-``num_groups`` vectors the kernel indexes by ``expert``. + ``m_splits`` is the host-side per-expert token list; the descriptor queue is built + + shape-cached from it (no GPU->CPU sync). """ from .fp8_gemm import doGemm as _fp8_doGemm - num_groups = int(group_offs.numel()) - 1 + num_groups = len(m_splits) if a_scale_vec.numel() != num_groups or b_scale_vec.numel() != num_groups: raise ValueError( f"FP8 grouped expects length-{num_groups} scale vectors, got " @@ -1395,9 +1405,9 @@ def _run_fp8_grouped( from .fp8_gemm import _BLOCK_M, _BLOCK_N num_pid_n = n // _BLOCK_N - group_offs_cpu = group_offs.to(device="cpu", dtype=torch.int64) - descriptors = gd.build_fwd_dgrad_descriptors(group_offs_cpu, num_groups, num_pid_n, _BLOCK_M) - work_queue, work_counter, num_tiles = gd.make_work_queue(descriptors, a_data.device) + work_queue, work_counter, num_tiles = gd.build_fwd_dgrad_work_queue( + m_splits, num_pid_n, _BLOCK_M, a_data.device + ) sched = _select_grouped_sched(num_tiles, a_data.device) epilogue, bias_arg = _resolve_epilogue(bias, n) @@ -1486,8 +1496,9 @@ def _pad_cols(t): num_pid_m = n_out // _BLOCK_M num_pid_n = k_in // _BLOCK_N - descriptors = gd.build_wgrad_descriptors(num_groups, num_pid_m, num_pid_n) - work_queue, work_counter, num_tiles = gd.make_work_queue(descriptors, device) + work_queue, work_counter, num_tiles = gd.build_wgrad_work_queue( + num_groups, num_pid_m, num_pid_n, device + ) sched = _select_grouped_sched(num_tiles, device) _fp8_doGemm( diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py index f83f84d7f3..046ef77550 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py @@ -23,6 +23,7 @@ identical across dtypes; they differ only in whether scales are also padded. """ +import functools import os import torch @@ -117,41 +118,69 @@ def select_grouped_sched(num_tiles: int, device) -> str: return "static" if num_tiles <= _GROUPED_SCHED_ROUNDS * num_cus else "worksteal" -def build_fwd_dgrad_descriptors(group_offs_cpu, num_groups, num_pid_n, block_m): - """Build the fwd/dgrad ``(pid_m, pid_n, expert)`` descriptor list. +@functools.lru_cache(maxsize=512) +def _fwd_dgrad_queue_cached(m_splits, num_pid_n, block_m, device): + """Cached flat ``(pid_m, pid_n, expert)`` work queue for a fwd/dgrad shape. - Each expert ``g`` owns the packed output M-block range - ``[group_offs[g], group_offs[g+1])`` (BLOCK_M-aligned so tiles never straddle - a group boundary), tiled over ``num_pid_n`` columns. + Keyed by the per-expert token counts + tiling, so repeated calls with the same + ``m_splits`` (every microbatch in training, every bench iteration) reuse one + already-uploaded queue instead of rebuilding + re-H2D-copying it each forward. + Built from the Python ``m_splits`` tuple directly -- NO device->host copy of + ``group_offs`` (that copy forced a ``hipDeviceSynchronize`` that stranded the + GPU idle for hundreds of microseconds before the kernel). """ descriptors = [] - for g in range(num_groups): - m_start_blk = int(group_offs_cpu[g]) // block_m - m_end_blk = int(group_offs_cpu[g + 1]) // block_m + off = 0 + for g, m in enumerate(m_splits): + m_start_blk = off // block_m + m_end_blk = (off + m) // block_m for pid_m in range(m_start_blk, m_end_blk): for pid_n in range(num_pid_n): descriptors.append((pid_m, pid_n, g)) - return descriptors + off += m + num_tiles = len(descriptors) + work_queue = torch.tensor(descriptors, dtype=torch.int32, device=device).reshape(-1) + return work_queue, num_tiles + + +def build_fwd_dgrad_work_queue(m_splits, num_pid_n, block_m, device): + """Return ``(work_queue, work_counter, num_tiles)`` for a fwd/dgrad launch. + + The descriptor queue is cached (shape-only, data-independent) so only the first + call with a given ``m_splits`` pays the build+upload; the zeroed atomic counter + is cheap and allocated fresh each call. ``m_splits`` is a Python list/tuple of + per-expert token counts -- no GPU->CPU sync. + """ + work_queue, num_tiles = _fwd_dgrad_queue_cached( + tuple(int(m) for m in m_splits), int(num_pid_n), int(block_m), device + ) + work_counter = torch.zeros(1, dtype=torch.int32, device=device) + return work_queue, work_counter, num_tiles -def build_wgrad_descriptors(num_groups, num_pid_m, num_pid_n): - """Build the wgrad descriptor list: one tile per expert's local [N_out, K_in].""" - return [ +@functools.lru_cache(maxsize=512) +def _wgrad_queue_cached(num_groups, num_pid_m, num_pid_n, device): + descriptors = [ (pid_m, pid_n, g) for g in range(num_groups) for pid_m in range(num_pid_m) for pid_n in range(num_pid_n) ] + num_tiles = len(descriptors) + work_queue = torch.tensor(descriptors, dtype=torch.int32, device=device).reshape(-1) + return work_queue, num_tiles -def make_work_queue(descriptors, device): - """Flatten descriptors into the kernel's queue buffer + zeroed atomic counter. +def build_wgrad_work_queue(num_groups, num_pid_m, num_pid_n, device): + """Return cached ``(work_queue, work_counter, num_tiles)`` for a wgrad launch. - Returns ``(work_queue, work_counter, num_tiles)``. ``work_queue`` is a flat - int32 ``[num_tiles*3]`` tensor; ``work_counter`` a zeroed one-int32 dispenser. + The wgrad descriptor queue depends only on ``(num_groups, num_pid_m, num_pid_n)`` + (per-expert tile geometry), so it is shape-cached and reused across calls -- no + per-call Python rebuild or H2D re-upload. Fresh zeroed atomic counter each call. """ - num_tiles = len(descriptors) - work_queue = torch.tensor(descriptors, dtype=torch.int32, device=device).reshape(-1) + work_queue, num_tiles = _wgrad_queue_cached( + int(num_groups), int(num_pid_m), int(num_pid_n), device + ) work_counter = torch.zeros(1, dtype=torch.int32, device=device) return work_queue, work_counter, num_tiles diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 8e4ad94630..2e57aad1fe 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -1237,13 +1237,12 @@ def _grouped_gemm_flydsl_mxfp8( riding off the quantization TE already performed -- ``inputmats`` are the per-group MXFP8 activations from ``tex.split_quantize`` and ``weights_fp8`` the per-expert MXFP8 weights. This first iteration does not quantize: it packs the - rowwise MXFP8 payloads into the ``[M_total, K]`` / ``[G*N, K]`` layout, derives - ``group_offs`` from ``m_splits``, and issues one packed TN launch via - :func:`_run_mxfp8_grouped`, writing into ``out``. + rowwise MXFP8 payloads into the ``[M_total, K]`` / ``[G*N, K]`` layout and + issues one packed TN launch via :func:`_run_mxfp8_grouped`, writing into + ``out``. The descriptor queue is built host-side from ``m_splits``. """ from ..flydsl_kernels.gemm import _run_mxfp8_grouped - num_gemms = len(inputmats) # TN packed operands: activations concatenated along the token (M) axis, # weights along the expert*N axis. Each MXFP8 operand carries its rowwise # fp8 payload and E8M0 scale; concatenate both. @@ -1261,19 +1260,13 @@ def _grouped_gemm_flydsl_mxfp8( k = weights_fp8[0]._rowwise_data.size(-1) m_total = a_data.size(0) - # group_offs: [num_gemms+1] exclusive scan over per-group token counts. - group_offs = torch.zeros(num_gemms + 1, dtype=torch.int64, device=out.device) - group_offs[1:] = torch.cumsum( - m_splits_tensor.to(device=out.device, dtype=torch.int64), 0 - ) - _run_mxfp8_grouped( a_data, a_scale, b_data, b_scale, out, - group_offs=group_offs, + m_splits=m_splits, m_total=m_total, n=n, k=k, @@ -1300,7 +1293,6 @@ def _backward_flydsl_mxfp8_grouped( """ from ..flydsl_kernels.gemm import _run_mxfp8_grouped - num_gemms = len(grad_output) # NN kernel-A = dY, rowwise (contraction over N_out): concat on the token # (M) axis. NN kernel-B = W, columnwise (transpose-read, contraction on the # N_out axis): concat on the packed N_out axis. Each MXFP8 operand carries @@ -1317,19 +1309,13 @@ def _backward_flydsl_mxfp8_grouped( k_in = dgrad.size(-1) # output width (K_in) m_total = a_data.size(0) - # group_offs: [num_gemms+1] exclusive scan over per-group token counts. - group_offs = torch.zeros(num_gemms + 1, dtype=torch.int64, device=dgrad.device) - group_offs[1:] = torch.cumsum( - m_splits_tensor.to(device=dgrad.device, dtype=torch.int64), 0 - ) - _run_mxfp8_grouped( a_data, a_scale, b_data, b_scale, dgrad, - group_offs=group_offs, + m_splits=m_splits, m_total=m_total, n=k_in, # NN output width is K_in k=n_out, # NN per-expert contraction is N_out @@ -1388,24 +1374,6 @@ def _wgrad_flydsl_mxfp8_grouped( for g in range(num_gemms): wgrad_list[g].copy_(dW[g * n_out : (g + 1) * n_out]) - @staticmethod - def _flydsl_group_offs(num_gemms, m_splits, m_splits_tensor, device): - """Build the [num_gemms+1] exclusive-scan token offsets on ``device``. - - Prefers the GPU ``m_splits_tensor`` (avoids a blocking H2D of pageable - ``m_splits``) but falls back to the Python ``m_splits`` list when the tensor - is absent (the BF16 path does not always carry it). - """ - group_offs = torch.zeros(num_gemms + 1, dtype=torch.int64, device=device) - if m_splits_tensor is not None: - group_offs[1:] = torch.cumsum( - m_splits_tensor.to(device=device, dtype=torch.int64), 0 - ) - else: - counts = torch.tensor(list(m_splits), dtype=torch.int64, device=device) - group_offs[1:] = torch.cumsum(counts, 0) - return group_offs - @staticmethod def _grouped_gemm_flydsl_bf16(*, inputmats, weights, out, m_splits, m_splits_tensor): """FlyDSL BF16 grouped forward (TN): Y_e = X_e @ W_e^T, one packed launch. @@ -1416,22 +1384,17 @@ def _grouped_gemm_flydsl_bf16(*, inputmats, weights, out, m_splits, m_splits_ten """ from ..flydsl_kernels.gemm import _run_bf16_grouped - num_gemms = len(inputmats) a_data = torch.cat([im.reshape(-1, im.shape[-1]) for im in inputmats], dim=0) b_data = torch.cat([w for w in weights], dim=0) n = weights[0].size(0) k = weights[0].size(-1) m_total = a_data.size(0) - group_offs = _GroupedLinear._flydsl_group_offs( - num_gemms, m_splits, m_splits_tensor, out.device - ) - _run_bf16_grouped( a_data, b_data, out, - group_offs=group_offs, + m_splits=m_splits, m_total=m_total, n=n, k=k, @@ -1443,22 +1406,17 @@ def _backward_flydsl_bf16_grouped(*, grad_output, weights, dgrad, m_splits, m_sp """FlyDSL BF16 grouped dgrad (NN): dX_e = dY_e @ W_e, one packed launch.""" from ..flydsl_kernels.gemm import _run_bf16_grouped - num_gemms = len(grad_output) a_data = torch.cat([g.reshape(-1, g.shape[-1]) for g in grad_output], dim=0) b_data = torch.cat([w for w in weights], dim=0) n_out = grad_output[0].shape[-1] # per-expert contraction (N_out) k_in = dgrad.size(-1) # output width (K_in) m_total = a_data.size(0) - group_offs = _GroupedLinear._flydsl_group_offs( - num_gemms, m_splits, m_splits_tensor, dgrad.device - ) - _run_bf16_grouped( a_data, b_data, dgrad, - group_offs=group_offs, + m_splits=m_splits, m_total=m_total, n=k_in, k=n_out, @@ -1501,7 +1459,6 @@ def _grouped_gemm_flydsl_fp8(*, inputmats, weights_fp8, out, m_splits, m_splits_ """ from ..flydsl_kernels.gemm import _run_fp8_grouped - num_gemms = len(inputmats) a_data = torch.cat([im._data for im in inputmats], dim=0) b_data = torch.cat([w._data for w in weights_fp8], dim=0) a_data._fp8_dtype = inputmats[0]._fp8_dtype @@ -1513,17 +1470,13 @@ def _grouped_gemm_flydsl_fp8(*, inputmats, weights_fp8, out, m_splits, m_splits_ k = weights_fp8[0]._data.size(-1) m_total = a_data.size(0) - group_offs = _GroupedLinear._flydsl_group_offs( - num_gemms, m_splits, m_splits_tensor, out.device - ) - _run_fp8_grouped( a_data, a_scale, b_data, b_scale, out, - group_offs=group_offs, + m_splits=m_splits, m_total=m_total, n=n, k=k, @@ -1543,7 +1496,6 @@ def _backward_flydsl_fp8_grouped( """ from ..flydsl_kernels.gemm import _run_fp8_grouped - num_gemms = len(grad_output) # A = dY (rowwise fp8 data); B = W columnwise transpose, which is the # materialized [K_in, N_out] -> packed on K_in (output axis of the TN GEMM). a_data = torch.cat([g._data for g in grad_output], dim=0) @@ -1557,17 +1509,13 @@ def _backward_flydsl_fp8_grouped( k_in = dgrad.size(-1) # output width (K_in) m_total = a_data.size(0) - group_offs = _GroupedLinear._flydsl_group_offs( - num_gemms, m_splits, m_splits_tensor, dgrad.device - ) - _run_fp8_grouped( a_data, a_scale, b_data, b_scale, dgrad, - group_offs=group_offs, + m_splits=m_splits, m_total=m_total, n=k_in, k=n_out, From d5e00fd15b367a58a2bab107a5101f8a8c7aaa78 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 14 Oct 2026 20:25:21 +0000 Subject: [PATCH 06/22] add K-bucketing for bf16 wgrad grouped gemm --- .../flydsl_kernels/gemm/gemm_wrappers.py | 106 ++++++++++++------ .../flydsl_kernels/gemm/grouped_gemm_utils.py | 46 ++++++++ 2 files changed, 116 insertions(+), 36 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 4457234d82..d8feb31f49 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -1285,6 +1285,14 @@ def _run_bf16_grouped( return D +def _wgrad_max_buckets(): + """Launch-count cap for the K-bucketed wgrad (env override, default uncapped).""" + v = os.getenv("NVTE_FLYDSL_WGRAD_MAX_BUCKETS", "") + if v.strip() in ("", "0"): + return None + return int(v) + + def _run_bf16_grouped_wgrad( dy_cols, # list[G] of per-expert dY, each [m_e, N_out] x_cols, # list[G] of per-expert X, each [m_e, K_in] @@ -1295,12 +1303,23 @@ def _run_bf16_grouped_wgrad( k_in, output_dtype, ): - """Grouped BF16 wgrad (NT): dW_e = dY_e^T @ X_e, per expert, one launch. - - Scaleless analog of :func:`_run_mxfp8_grouped_wgrad`: pad each expert's tokens - to a uniform M_pad on the contraction axis, pack dY/X to [G*M_pad, *], and emit - one NT launch writing the packed [G*N_out, K_in] gradient (each expert's output - tile shifted via store_m_expert_off). + """Grouped BF16 wgrad (NT): dW_e = dY_e^T @ X_e, via K-bucketed launches. + + Scaleless. The kernel's contraction K (the per-expert token count) is baked in + at compile time, so experts are grouped into **K-buckets** by their own padded + token count (``bucketize_wgrad_tokens``): each expert rounds up to its own + BLOCK_M multiple (floor 512), and experts sharing a rounded K share one launch. + This avoids the old global-max pad -- a single large expert no longer forces + every small expert to contract over its token count (which wasted work and + OOM'd the int32 launch-arg limit under skew). Balanced token counts collapse to + a single bucket, reproducing the prior single-launch path bit-for-bit. + + Per bucket ``b`` with experts ``E_b`` and padded K ``k_pad``: pack those + experts' dY/X to ``[len(E_b)*k_pad, *]`` (bucket-local, so the kernel's + ``expert*K`` operand addressing uses the bucket-local index), launch into a + bucket-local output ``[len(E_b)*N_out, K_in]``, then copy each expert's + ``[N_out, K_in]`` block into its GLOBAL row of ``dW``. The hand-scheduled 4-wave + MFMA pipeline is untouched -- K stays compile-time-constant per launch. """ from .half_prec_gemm import doGemm as _half_prec_doGemm, _BLOCK_M, _BLOCK_N @@ -1311,13 +1330,6 @@ def _run_bf16_grouped_wgrad( if k_in % _BLOCK_N != 0: raise FlyDSLUnsupportedError(f"wgrad requires K_in divisible by {_BLOCK_N}, got {k_in}") - max_m = max(int(m) for m in m_splits) - m_pad = ((max_m + _BLOCK_M - 1) // _BLOCK_M) * _BLOCK_M - m_pad = max(m_pad, 512) - - dy_data = gd.uniform_pad_concat(dy_cols, m_pad) - x_data = gd.uniform_pad_concat(x_cols, m_pad) - dW = _validate_or_allocate_output( dW, shape=torch.Size((num_groups * n_out, k_in)), @@ -1328,30 +1340,52 @@ def _run_bf16_grouped_wgrad( num_pid_m = n_out // _BLOCK_M num_pid_n = k_in // _BLOCK_N - work_queue, work_counter, num_tiles = gd.build_wgrad_work_queue( - num_groups, num_pid_m, num_pid_n, device - ) - sched = _select_grouped_sched(num_tiles, device) + buckets = gd.bucketize_wgrad_tokens(m_splits, _BLOCK_M, max_buckets=_wgrad_max_buckets()) + + for k_pad, ids in buckets: + bsize = len(ids) + dy_data = gd.uniform_pad_concat([dy_cols[e] for e in ids], k_pad) + x_data = gd.uniform_pad_concat([x_cols[e] for e in ids], k_pad) + + # Fast path: when this bucket's experts are exactly rows 0..G-1 in order + # (the common single-bucket / balanced case), the bucket-local expert index + # equals the global index, so the kernel can write straight into dW -- no + # bucket-local buffer, no scatter copy. This reproduces the pre-bucketing + # single-launch path byte-for-byte (and avoids a full-dW-size copy that + # otherwise regresses the balanced case ~15%). Only genuinely multi-bucket + # (skewed) launches pay the scatter, where the K savings dwarf it. + direct = ids == list(range(num_groups)) + dW_b = dW if direct else torch.empty(bsize * n_out, k_in, dtype=output_dtype, + device=device) + + work_queue, work_counter, num_tiles = gd.build_wgrad_work_queue( + bsize, num_pid_m, num_pid_n, device + ) + sched = _select_grouped_sched(num_tiles, device) + + _half_prec_doGemm( + dy_data, + x_data, + dW_b, + layout="NT", + m=n_out, + n=k_in, + k=k_pad, + input_dtype=torch.bfloat16, + label="BF16 grouped wgrad", + group=False, + num_groups=bsize, + work_queue=work_queue, + work_counter=work_counter, + sched=sched, + wgrad=True, + ) + + if not direct: + # Scatter each bucket-local expert block to its global dW row. + for j, e in enumerate(ids): + dW[e * n_out : (e + 1) * n_out].copy_(dW_b[j * n_out : (j + 1) * n_out]) - # The kernel's compile-time contraction K is the PER-EXPERT padded token count; - # both operands are packed [G*M_pad, *] so doGemm de-packs K via num_groups. - _half_prec_doGemm( - dy_data, - x_data, - dW, - layout="NT", - m=n_out, - n=k_in, - k=m_pad, - input_dtype=torch.bfloat16, - label="BF16 grouped wgrad", - group=False, - num_groups=num_groups, - work_queue=work_queue, - work_counter=work_counter, - sched=sched, - wgrad=True, - ) return dW diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py index 046ef77550..85b1910db4 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py @@ -200,3 +200,49 @@ def _pad_rows(t): return out return torch.cat([_pad_rows(t) for t in tensors], dim=0) + + +def bucketize_wgrad_tokens(m_splits, block_m, max_buckets=None): + """Group experts by their per-expert padded contraction length (K-bucketing). + + The grouped wgrad kernel bakes its contraction length ``K`` in at compile time, + so experts that share a padded token count can share one launch. Instead of + padding every expert to the GLOBAL max token count (which wastes work and can + OOM under skew), round each expert's ``m_g`` UP to its own ``block_m`` multiple + (floor 512 = the pipeline's 4-K128-tile minimum), then group experts by that + rounded value. Each distinct value is one bucket -> one fixed-K launch. + + When the number of distinct values exceeds ``max_buckets``, adjacent buckets + are merged UPWARD (experts promoted to the larger K) so the launch count is + capped; promoting only ever adds zero padding, never truncates, so results are + unchanged. Returns a list of ``(k_pad, [global_expert_ids])`` tuples, ordered + by ascending ``k_pad``. + """ + def _round(m): + return max(((int(m) + block_m - 1) // block_m) * block_m, 512) + + k_of = [_round(m) for m in m_splits] + distinct = sorted(set(k_of)) + + if max_buckets is not None and len(distinct) > max_buckets: + # Merge the smallest-K buckets upward onto a coarser ladder until the + # distinct-K count fits. Keep the largest K (one dominant expert) its own + # bucket so it is never dragged smaller (it can't be) nor inflates others. + # Simple policy: collapse the sorted distinct values into `max_buckets` + # contiguous groups, each group promoted to its group-max. + import math + per = math.ceil(len(distinct) / max_buckets) + ladder = {} + for i in range(0, len(distinct), per): + grp = distinct[i : i + per] + top = grp[-1] + for v in grp: + ladder[v] = top + k_of = [ladder[k] for k in k_of] + distinct = sorted(set(k_of)) + + buckets = [] + for k_pad in distinct: + ids = [g for g, k in enumerate(k_of) if k == k_pad] + buckets.append((k_pad, ids)) + return buckets From f394a18ad7eab8a1d16706e9c20e178188c5b1b7 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 14 Oct 2026 21:45:25 +0000 Subject: [PATCH 07/22] Apply K-bucket wgrad to MXFP8 and FP8 grouped GEMM --- .../flydsl_kernels/gemm/gemm_wrappers.py | 208 ++++++++++-------- 1 file changed, 116 insertions(+), 92 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index d8feb31f49..63c85fa73f 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -1138,14 +1138,16 @@ def _run_mxfp8_grouped_wgrad( x_fp8_dtype, output_dtype, ): - """Grouped MXFP8 wgrad (NT): dW_e = dY_e^T @ X_e, per expert, one launch. - - Fixed uniform padding (phase 1): every expert's token count m_e is padded up to - a common M_pad (a multiple of BLOCK_M and 32); the padded rows are zero, so they - contribute nothing to the token-axis contraction -- no atomics or masking needed, - the dense NT pipeline runs unchanged. dY and X are packed on the contraction - (token) axis to [G*M_pad, *]; dW is packed on the output-M axis to [G*N_out, K_in] - and each expert's output tile is written via the kernel's store_m_expert_off. + """Grouped MXFP8 wgrad (NT): dW_e = dY_e^T @ X_e, via K-bucketed launches. + + K-bucketing (see :func:`_run_bf16_grouped_wgrad`): experts are grouped by their + own padded token count (rounded up to BLOCK_M, floor 512) rather than padded to + a single global max, so one large expert no longer forces every small expert to + contract over its token count. Each bucket is one fixed-K launch; MXFP8 also + packs each operand's E8M0 scale (``scale_rows = k_pad / SCALE_GROUP_SIZE``) to + the same bucket K. Per-bucket output goes to a bucket-local buffer then scatters + to global dW rows (fast direct-write when the bucket is all experts in order). + The 4-wave MFMA pipeline is untouched -- K stays compile-time per launch. """ from .mxfp8_gemm import BLOCK_M, BLOCK_N, SCALE_GROUP_SIZE @@ -1160,24 +1162,6 @@ def _run_mxfp8_grouped_wgrad( f"wgrad requires K_in divisible by BLOCK_N={BLOCK_N}, got {k_in}" ) - # Uniform pad target: the largest per-expert token count, rounded up to BLOCK_M - # (also a multiple of 32). Enforce the two-page pipeline minimum (>= 4 K128 - # tiles => M_pad >= 512). - max_m = max(int(m) for m in m_splits) - m_pad = ((max_m + BLOCK_M - 1) // BLOCK_M) * BLOCK_M - m_pad = max(m_pad, 512) - - scale_rows = m_pad // SCALE_GROUP_SIZE - dy_data = gd.uniform_pad_concat(dy_cols, m_pad) - dy_scale = gd.uniform_pad_concat(dy_scales, scale_rows) - x_data = gd.uniform_pad_concat(x_cols, m_pad) - x_scale = gd.uniform_pad_concat(x_scales, scale_rows) - - if dy_data.dtype == torch.uint8: - dy_data = reinterpret_as_fp8_tensor(dy_data, dy_fp8_dtype) - if x_data.dtype == torch.uint8: - x_data = reinterpret_as_fp8_tensor(x_data, x_fp8_dtype) - dW = _validate_or_allocate_output( dW, shape=torch.Size((num_groups * n_out, k_in)), @@ -1191,24 +1175,48 @@ def _run_mxfp8_grouped_wgrad( # kernel re-packs the output row via expert*N_out. num_pid_m = n_out // BLOCK_M num_pid_n = k_in // BLOCK_N - work_queue, work_counter, num_tiles = gd.build_wgrad_work_queue( - num_groups, num_pid_m, num_pid_n, device - ) - sched = _select_grouped_sched(num_tiles, device) + buckets = gd.bucketize_wgrad_tokens(m_splits, BLOCK_M, max_buckets=_wgrad_max_buckets()) + + for k_pad, ids in buckets: + bsize = len(ids) + scale_rows = k_pad // SCALE_GROUP_SIZE + dy_data = gd.uniform_pad_concat([dy_cols[e] for e in ids], k_pad) + dy_scale = gd.uniform_pad_concat([dy_scales[e] for e in ids], scale_rows) + x_data = gd.uniform_pad_concat([x_cols[e] for e in ids], k_pad) + x_scale = gd.uniform_pad_concat([x_scales[e] for e in ids], scale_rows) + + if dy_data.dtype == torch.uint8: + dy_data = reinterpret_as_fp8_tensor(dy_data, dy_fp8_dtype) + if x_data.dtype == torch.uint8: + x_data = reinterpret_as_fp8_tensor(x_data, x_fp8_dtype) + + direct = ids == list(range(num_groups)) + dW_b = dW if direct else torch.empty(bsize * n_out, k_in, dtype=output_dtype, + device=device) + + work_queue, work_counter, num_tiles = gd.build_wgrad_work_queue( + bsize, num_pid_m, num_pid_n, device + ) + sched = _select_grouped_sched(num_tiles, device) + + mxfp8_matmul( + dy_data, + dy_scale, + x_data, + x_scale, + dW_b, + group=False, + num_groups=bsize, + work_queue=work_queue, + work_counter=work_counter, + sched=sched, + wgrad=True, + ) + + if not direct: + for j, e in enumerate(ids): + dW[e * n_out : (e + 1) * n_out].copy_(dW_b[j * n_out : (j + 1) * n_out]) - mxfp8_matmul( - dy_data, - dy_scale, - x_data, - x_scale, - dW, - group=False, - num_groups=num_groups, - work_queue=work_queue, - work_counter=work_counter, - sched=sched, - wgrad=True, - ) return dW @@ -1479,17 +1487,20 @@ def _run_fp8_grouped_wgrad( x_fp8_dtype, output_dtype, ): - """Grouped tensor-wise FP8 wgrad: dW_e = dY_e^T @ X_e, one TN launch. + """Grouped tensor-wise FP8 wgrad: dW_e = dY_e^T @ X_e, via K-bucketed TN launches. Tensor-wise FP8 is TN-NORMAL-read: the contraction axis (tokens m) is the - trailing dim and is never packed, so -- unlike the MXFP8/BF16 NT path -- the - operands arrive already transposed (the FP8 ``_transpose`` backing is a - materialized row-major transpose). Per expert: A = dY^T [N_out, m_e], - B = X^T [K_in, m_e]. Each expert's token count m_e is zero-padded up to a - common ``m_pad`` on the trailing (contraction) axis; experts stack on the - feature (output) axis to A [G*N_out, m_pad] / B [G*K_in, m_pad]. Output dW is - packed [G*N_out, K_in] on the M-row. Per-expert scalar scales are length-G - vectors indexed by expert in the kernel. + trailing dim, so -- unlike the MXFP8/BF16 NT path -- operands arrive already + transposed (the FP8 ``_transpose`` backing is a materialized row-major + transpose). Per expert: A = dY^T [N_out, m_e], B = X^T [K_in, m_e]. + + K-bucketing (see :func:`_run_bf16_grouped_wgrad`): experts are grouped by their + own padded token count instead of a single global max. Each expert's tokens are + zero-padded on the TRAILING (contraction) axis up to the bucket's ``k_pad``, and + the bucket's experts stack on the feature axis to A [bsize*N_out, k_pad] / + B [bsize*K_in, k_pad]. Per-expert scalar scales are sliced to the bucket. Output + goes to a bucket-local buffer then scatters to global dW rows (direct-write when + the bucket is all experts in order). The 4-wave MFMA pipeline is untouched. """ from .fp8_gemm import doGemm as _fp8_doGemm, _BLOCK_M, _BLOCK_N @@ -1500,26 +1511,6 @@ def _run_fp8_grouped_wgrad( if k_in % _BLOCK_N != 0: raise FlyDSLUnsupportedError(f"wgrad requires K_in divisible by {_BLOCK_N}, got {k_in}") - max_m = max(int(m) for m in m_splits) - m_pad = ((max_m + _BLOCK_M - 1) // _BLOCK_M) * _BLOCK_M - m_pad = max(m_pad, 512) - - # Zero-pad the trailing (token/contraction) axis of each [feature, m_e] - # transpose up to m_pad, then concat on the feature axis. - def _pad_cols(t): - if t.size(1) == m_pad: - return t.contiguous() - out = t.new_zeros((t.size(0), m_pad)) - out[:, : t.size(1)].copy_(t) - return out - - dy_data = torch.cat([_pad_cols(t) for t in dyt_cols], dim=0) - x_data = torch.cat([_pad_cols(t) for t in xt_cols], dim=0) - if dy_data.dtype == torch.uint8: - dy_data = reinterpret_as_fp8_tensor(dy_data, dy_fp8_dtype) - if x_data.dtype == torch.uint8: - x_data = reinterpret_as_fp8_tensor(x_data, x_fp8_dtype) - dW = _validate_or_allocate_output( dW, shape=torch.Size((num_groups * n_out, k_in)), @@ -1528,29 +1519,62 @@ def _pad_cols(t): backend_name="FP8 grouped wgrad TN", ) + def _pad_cols(t, cols): + # Zero-pad the trailing (token/contraction) axis of a [feature, m_e] view. + if t.size(1) == cols: + return t.contiguous() + out = t.new_zeros((t.size(0), cols)) + out[:, : t.size(1)].copy_(t) + return out + num_pid_m = n_out // _BLOCK_M num_pid_n = k_in // _BLOCK_N - work_queue, work_counter, num_tiles = gd.build_wgrad_work_queue( - num_groups, num_pid_m, num_pid_n, device - ) - sched = _select_grouped_sched(num_tiles, device) + dy_scale_vec = dy_scale_vec.to(torch.float32).contiguous() + x_scale_vec = x_scale_vec.to(torch.float32).contiguous() + buckets = gd.bucketize_wgrad_tokens(m_splits, _BLOCK_M, max_buckets=_wgrad_max_buckets()) + + for k_pad, ids in buckets: + bsize = len(ids) + dy_data = torch.cat([_pad_cols(dyt_cols[e], k_pad) for e in ids], dim=0) + x_data = torch.cat([_pad_cols(xt_cols[e], k_pad) for e in ids], dim=0) + if dy_data.dtype == torch.uint8: + dy_data = reinterpret_as_fp8_tensor(dy_data, dy_fp8_dtype) + if x_data.dtype == torch.uint8: + x_data = reinterpret_as_fp8_tensor(x_data, x_fp8_dtype) + idx = torch.tensor(ids, device=device) + dy_scale_b = dy_scale_vec[idx] + x_scale_b = x_scale_vec[idx] + + direct = ids == list(range(num_groups)) + dW_b = dW if direct else torch.empty(bsize * n_out, k_in, dtype=output_dtype, + device=device) + + work_queue, work_counter, num_tiles = gd.build_wgrad_work_queue( + bsize, num_pid_m, num_pid_n, device + ) + sched = _select_grouped_sched(num_tiles, device) + + _fp8_doGemm( + dy_data, + x_data, + dW_b, + dy_scale_b, + x_scale_b, + group=False, + num_groups=bsize, + work_queue=work_queue, + work_counter=work_counter, + sched=sched, + wgrad=True, + m=n_out, + n=k_in, + k=k_pad, + ) + + if not direct: + for j, e in enumerate(ids): + dW[e * n_out : (e + 1) * n_out].copy_(dW_b[j * n_out : (j + 1) * n_out]) - _fp8_doGemm( - dy_data, - x_data, - dW, - dy_scale_vec.to(torch.float32).contiguous(), - x_scale_vec.to(torch.float32).contiguous(), - group=False, - num_groups=num_groups, - work_queue=work_queue, - work_counter=work_counter, - sched=sched, - wgrad=True, - m=n_out, - n=k_in, - k=m_pad, - ) return dW From 14d85a47c8b938136d2069aef2a5b515bc0f9a83 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Thu, 15 Oct 2026 15:36:34 +0000 Subject: [PATCH 08/22] Address grouped GEMM weights via per-expert pointers, dropping the host concat --- .../pytorch/flydsl_kernels/gemm/fp8_gemm.py | 47 +++- .../flydsl_kernels/gemm/gemm_wrappers.py | 12 + .../flydsl_kernels/gemm/grouped_gemm_utils.py | 112 +++++++++ .../flydsl_kernels/gemm/half_prec_gemm.py | 44 +++- .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 233 ++++++++++++++---- .../pytorch/module/grouped_linear.py | 113 +++++++-- 6 files changed, 474 insertions(+), 87 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py index c8e40ad446..130cfc250e 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -99,6 +99,7 @@ def _compile_kernel( group: bool = False, sched: str = "worksteal", wgrad: bool = False, + weight_ptrs: bool = False, ): """Build the specialized kernel for compile-time K, A/B FP8 types, and output dtype. @@ -216,6 +217,7 @@ def kernel_gemm( num_groups: fx.Int32 = 0, work_queue: fx.Tensor = None, work_counter: fx.Tensor = None, + b_data_ptrs: fx.Tensor = None, ): lds = fx.SharedAllocator().allocate(SharedStorage).peek() lds_a0 = (lds.a0_0, lds.a0_1) @@ -255,6 +257,13 @@ def _load_scale(scale_div, index): grouped = group or wgrad num_tiles = gd.num_tiles_bound(c_m, c_n, BLOCK_M, BLOCK_N, num_groups, const_expr(wgrad)) + # Grouped-TN forward (weight_ptrs): weights are G independently-allocated + # per-expert slabs, addressed from a per-expert base pointer read from + # ``b_data_ptrs`` (reusing gB's static layout) instead of an expert-shifted + # B row. Forward-only compile flag; dgrad/wgrad keep the packed-B fold. + if const_expr(weight_ptrs): + b_ptr_layout = fx.get_layout(gB) + if const_expr(grouped): wq_rsrc = fx.rocdl.make_buffer_tensor(work_queue, max_size=True) wq_div = fx.logical_divide(wq_rsrc, fx.make_layout(1, 1)) @@ -336,7 +345,17 @@ def process_tile(pid_m, pid_n, expert): gl_off_a = compute_global_swizzle(lane, wave_id, K, LOAD_PASSES_HALF, preshuffled=False) gl_off_b = compute_global_swizzle(lane, wave_id, K, LOAD_PASSES_HALF, preshuffled=False) a_g2s = G2SLoader(a_div, gl_off_a, LOAD_PASSES_HALF, a_f8_ir_t, wave_id) - b_g2s = G2SLoader(b_div, gl_off_b, LOAD_PASSES_HALF, b_f8_ir_t, wave_id) + # Grouped-TN forward B-data: rebase onto this tile's expert weight slab and + # read with a LOCAL column (base selects the expert). Other paths keep the + # shared packed b_div / by_n_b_idx. + if const_expr(weight_ptrs): + b_base_i64 = gd.read_ptr_at(b_data_ptrs, expert) + b_div_tile = gd.rebase_buffer_view(b_base_i64, b_f8_ir_t, b_ptr_layout) + b_data_col = by_n_idx + else: + b_div_tile = b_div + b_data_col = by_n_b_idx + b_g2s = G2SLoader(b_div_tile, gl_off_b, LOAD_PASSES_HALF, b_f8_ir_t, wave_id) s2r = S2RLoader(fx.Int32(0), 1) layout_lane16 = fx.make_layout((4, 16), (16, 1)) @@ -494,11 +513,12 @@ def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): a_g2s.load_one(lds_a[subtile], fx.Int32(global_base), pass_in_subtile) def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): - # by_n_b_idx shifts the B row to this tile's expert block (grouped TN - # fwd/dgrad weights, and grouped wgrad X^T packed [G*K_in, m_pad]); - # reduces to the dense by_n for the dense path. + # Grouped-TN forward selects the expert via the rebased per-expert B + # view (b_g2s) with a LOCAL column (b_data_col). Other paths: by_n_b_idx + # shifts the B row to the expert block (grouped dgrad weights, grouped + # wgrad X^T); reduces to dense by_n for the dense path. global_base = ( - (by_n_b_idx + fx.Index(subtile * (BLOCK_N // 2))) * fx.Index(K) + (b_data_col + fx.Index(subtile * (BLOCK_N // 2))) * fx.Index(K) + k_base ) b_g2s.load_one(lds_b[subtile], fx.Int32(global_base), pass_in_subtile) @@ -1088,6 +1108,7 @@ def launch_gemm( num_groups: fx.Int32 = 0, work_queue: fx.Tensor = None, work_counter: fx.Tensor = None, + b_data_ptrs: fx.Tensor = None, stream: fx.Stream = fx.Stream(None), ): # Dense: one workgroup per output tile. Grouped/wgrad: a fixed num-CU @@ -1109,6 +1130,7 @@ def launch_gemm( num_groups, work_queue, work_counter, + b_data_ptrs, value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, ).launch(grid=(grid_x, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) @@ -1126,6 +1148,7 @@ def _cached_launch( group: bool = False, sched: str = "worksteal", wgrad: bool = False, + weight_ptrs: bool = False, ): return _compile_kernel( K, @@ -1137,6 +1160,7 @@ def _cached_launch( group=group, sched=sched, wgrad=wgrad, + weight_ptrs=weight_ptrs, ) @@ -1240,6 +1264,7 @@ def doGemm( num_groups: int = 0, work_queue: torch.Tensor = None, work_counter: torch.Tensor = None, + b_data_ptrs: torch.Tensor = None, sched: str = "worksteal", wgrad: bool = False, m: int = None, @@ -1346,6 +1371,16 @@ def doGemm( work_queue_arg = torch.zeros(1, dtype=torch.int32, device=A.device) work_counter_arg = torch.zeros(1, dtype=torch.int32, device=A.device) + # Per-expert weight base pointers (grouped TN forward only): compile a kernel that + # reads each expert's weight slab from b_data_ptrs[g] instead of concatenating the + # G weights host-side. weight_ptrs is a compile-time flag keyed into the kernel + # cache; dgrad/wgrad leave it False and pass a placeholder never dereferenced. + weight_ptrs = b_data_ptrs is not None + if weight_ptrs: + b_data_ptrs_arg = b_data_ptrs.contiguous().view(-1) + else: + b_data_ptrs_arg = torch.zeros(1, dtype=torch.int64, device=A.device) + launch = _cached_launch( int(K_runtime), A.dtype, @@ -1356,6 +1391,7 @@ def doGemm( group, sched, wgrad, + weight_ptrs, ) launch( A_arg, @@ -1370,5 +1406,6 @@ def doGemm( num_groups, work_queue_arg, work_counter_arg, + b_data_ptrs_arg, stream=stream, ) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 63c85fa73f..053d9b8025 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -1027,6 +1027,7 @@ def _run_mxfp8_grouped( output_dtype: torch.dtype, layout: str = "TN", bias=None, + b_data_ptrs=None, ): """Dispatch a packed MXFP8 grouped GEMM through the persistent kernel. @@ -1106,6 +1107,12 @@ def _run_mxfp8_grouped( sched = _select_grouped_sched(num_tiles, a_data.device) epilogue, bias_arg = _resolve_epilogue(bias, n) + + # TN forward addresses the weight operand from per-expert base pointers supplied + # by the caller (one int64 base per expert slab); the kernel reads B slab g from + # b_data_ptrs[g] instead of an expert-shifted column into one packed B, so the + # host never concatenates the G per-expert weights. NN dgrad has no pointer array + # and keeps the packed-B contraction fold (phase 2). mxfp8_matmul( a_data, a_scale, @@ -1119,6 +1126,7 @@ def _run_mxfp8_grouped( num_groups=num_groups, work_queue=work_queue, work_counter=work_counter, + b_data_ptrs=b_data_ptrs, sched=sched, ) return D @@ -1232,6 +1240,7 @@ def _run_bf16_grouped( output_dtype: torch.dtype, layout: str = "TN", bias=None, + b_data_ptrs=None, ): """Dispatch a packed BF16 grouped GEMM through the persistent kernel. @@ -1288,6 +1297,7 @@ def _run_bf16_grouped( num_groups=num_groups, work_queue=work_queue, work_counter=work_counter, + b_data_ptrs=b_data_ptrs, sched=sched, ) return D @@ -1410,6 +1420,7 @@ def _run_fp8_grouped( k, output_dtype: torch.dtype, bias=None, + b_data_ptrs=None, ): """Dispatch a packed tensor-wise FP8 grouped GEMM (TN) through the persistent kernel. @@ -1465,6 +1476,7 @@ def _run_fp8_grouped( num_groups=num_groups, work_queue=work_queue, work_counter=work_counter, + b_data_ptrs=b_data_ptrs, sched=sched, m=m_total, n=n, diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py index 85b1910db4..1b660b42f8 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py @@ -54,6 +54,40 @@ def global_i32_ptr(tensor, index): return llvm.IntToPtrOp(ptr_ty, addr).result +def read_ptr_at(ptr_array, index): + """Load the int64 base address stored at ``ptr_array[index]``. + + ``ptr_array`` is a device ``int64[G]`` of per-expert operand base addresses + (each the ``data_ptr()`` of one expert's independently-allocated weight slab). + Recovers the array base, offsets by ``index`` int64 elements (8 bytes), and + loads the stored address as an i64 IR value -- the per-expert operand base. + """ + ptr_ty = ir.Type.parse("!llvm.ptr<1>") + base = _fly.extract_aligned_pointer_as_index(ptr_ty, ptr_array) + base = llvm.PtrToIntOp(T.i64, base).result + byte_off = fx.as_ir_value(fx.Int64(fx.Index(index) * fx.Index(8))) + addr = llvm.AddOp(base, byte_off, llvm.IntegerOverflowFlags(0)).result + elt_ptr = llvm.IntToPtrOp(ptr_ty, addr).result + return llvm.LoadOp(T.i64, elt_ptr).result + + +def rebase_buffer_view(base_i64, elem_ir_type, layout): + """Rebuild a buffer-tensor logical-divide view at a runtime base address. + + Reuses an EXISTING static ``layout`` (compile-time cosize) but swaps the base + pointer to the runtime ``base_i64`` address. This is how a per-expert operand + slice is addressed without constructing a fresh runtime-sized layout + (``fx.make_layout`` requires a static extent): the packed operand's layout is + reused and only the base moves. ``max_size=True`` lets the buffer descriptor + span the per-expert extent; the loader's dynamic ``soffset`` walks within it. + """ + buf_ptr_ty = fx.PointerType.get(elem_ty=elem_ir_type, address_space=1) + base_ptr = fx.inttoptr(buf_ptr_ty, fx.Int64(base_i64)) + view = fx.make_view(base_ptr, layout) + buf = fx.rocdl.make_buffer_tensor(fx.Tensor(view), max_size=True) + return fx.logical_divide(buf, fx.make_layout(1, 1)) + + def make_claim_broadcast(claim_slot_array): """Recast a one-int32 LDS scratch array into a claim iterator + 1-elem view. @@ -185,6 +219,84 @@ def build_wgrad_work_queue(num_groups, num_pid_m, num_pid_n, device): return work_queue, work_counter, num_tiles +def concat_or_view_rows(tensors): + """Concatenate per-expert tensors on dim 0, skipping the copy when possible. + + ``tex.split_quantize`` bulk-allocates its per-expert outputs into one base + buffer (data blocks packed back-to-back, then scale blocks back-to-back), so + the per-expert rowwise payloads are frequently already physically adjacent -- + the ``[M_total, cols]`` concatenation the grouped kernel wants already exists + in memory. When that holds, return a zero-copy spanning view over the shared + region instead of materializing a redundant ``torch.cat`` copy (the activation + and activation-scale cats on the MXFP8 forward/dgrad hot path). + + Adjacency is not guaranteed (independently-allocated operands -- e.g. weights + from separate ``nn.Parameter`` storages -- are scattered), so fall back to + ``torch.cat`` whenever the inputs are not a single contiguous run. The view and + the cat produce bitwise-identical operands; only the copy is elided. + """ + if len(tensors) == 1: + return tensors[0] + + first = tensors[0] + if not first.is_cuda or first.ndim < 1: + return torch.cat(tensors, dim=0) + + # __cuda_array_interface__ is unavailable on grad-tracking tensors. The grouped + # GEMM operand is consumed inside a custom autograd Function (grad handled at the + # module level), so a detached data view is correct; detach once up front so the + # adjacency walk and span wrapper both see plain tensors. + if any(t.requires_grad for t in tensors): + tensors = [t.detach() for t in tensors] + first = tensors[0] + + trailing = first.shape[1:] + elem = first.element_size() + total_rows = 0 + expected_ptr = first.data_ptr() + for t in tensors: + # Adjacency + uniformity: same trailing shape/dtype, row-contiguous, and + # starting exactly where the previous tensor ended. + if ( + t.dtype != first.dtype + or t.shape[1:] != trailing + or not t.is_contiguous() + or t.data_ptr() != expected_ptr + ): + return torch.cat(tensors, dim=0) + total_rows += t.size(0) + expected_ptr += t.numel() * elem + + # One contiguous run: wrap [total_rows*row_bytes] over the shared base pointer as + # raw uint8, then reinterpret back to the operand dtype. Going through bytes keeps + # this dtype-agnostic -- __cuda_array_interface__ has no typestr for bfloat16, so a + # native-dtype CAI would reject it; uint8 is always representable. + # The CAI carries no ownership, so anchor the source tensors on the result to keep + # their bulk buffer alive for the view's lifetime. + row_bytes = 1 + for d in trailing: + row_bytes *= d + row_bytes *= elem + total_bytes = total_rows * row_bytes + + class _SpanCAI: + __slots__ = ("__cuda_array_interface__",) + + def __init__(self): + self.__cuda_array_interface__ = { + "shape": (total_bytes,), + "typestr": "|u1", + "data": (first.data_ptr(), False), + "strides": None, + "version": 3, + } + + byte_view = torch.as_tensor(_SpanCAI(), device=first.device) + view = byte_view.view(first.dtype).view(total_rows, *tuple(trailing)) + view._flydsl_span_anchor = tensors + return view + + def uniform_pad_concat(tensors, rows): """Zero-pad each tensor's leading dim up to ``rows`` and concat on dim 0. diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py index 8ecf89aa1f..8684195551 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py @@ -392,6 +392,7 @@ def kernel_gemm( num_groups: fx.Int32 = 0, work_queue: fx.Tensor = None, work_counter: fx.Tensor = None, + b_data_ptrs: fx.Tensor = None, ): lds = fx.SharedAllocator().allocate(SharedStorage).peek() lds_a0 = (lds.a0_0, lds.a0_1) @@ -424,6 +425,13 @@ def kernel_gemm( wq_div = fx.logical_divide(wq_rsrc, fx.make_layout(1, 1)) desc_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32) + # Grouped-TN forward: weights arrive as G independently-allocated + # per-expert slabs, so B is addressed from a per-expert base pointer read + # from ``b_data_ptrs`` (byte-addressed, reusing gB's static layout) rather + # than an expert-shifted column into one packed B. NN/wgrad keep the + # packed-B folds; dense never enters here. + b_ptr_layout = fx.get_layout(gB) + # LDS broadcast scratch (one int32): thread 0 writes the atomically # claimed slot, a barrier broadcasts it to the whole workgroup. The # ``if tx==0`` claim control flow must stay inline (AST-rewritten here). @@ -503,8 +511,22 @@ def process_tile(pid_m, pid_n, expert): fx.Uint8.ir_type, wave_id, ) + # Grouped B-data: rebase onto this tile's expert weight slab (byte buffer); + # the per-expert base already selects the expert, so in-tile addressing is + # LOCAL. TN reads a local column (by_n_idx); NN (dgrad, transpose-read) + # packs experts on the contraction axis, so the base replaces the expert*K + # contraction shift (b_k_off -> 0). Dense/wgrad keep the shared packed view. + if const_expr(group): + b_base_i64 = gd.read_ptr_at(b_data_ptrs, expert) + b_div_tile = gd.rebase_buffer_view(b_base_i64, fx.Uint8.ir_type, b_ptr_layout) + b_data_col = by_n_idx + b_k_off = fx.Index(0) + else: + b_div_tile = b_div + b_data_col = by_n_b_idx + b_k_off = b_expert_k_off b_g2s = G2SLoader( - b_div, + b_div_tile, gl_off_b, LOAD_PASSES_HALF, fx.Uint8.ir_type, @@ -669,9 +691,11 @@ def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): a_g2s.load_one(lds_a[subtile], fx.Int32(global_base), pass_in_subtile) def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): - # b_expert_k_off shifts the contraction base (grouped NN/NT transpose - # read); by_n_b_idx shifts the B column (grouped TN). Both 0 for dense. - global_base = _b_global_base_bytes(k_base + b_expert_k_off, subtile, c_n, by_n_b_idx) + # Grouped fwd/dgrad select the expert via the rebased per-expert B view + # (b_g2s) with LOCAL coords: TN a local column (b_data_col), NN a zero + # contraction shift (b_k_off). Dense/wgrad use the shared view with the + # packed column by_n_b_idx and expert contraction shift b_k_off. + global_base = _b_global_base_bytes(k_base + b_k_off, subtile, c_n, b_data_col) b_g2s.load_one(lds_b[subtile], fx.Int32(global_base), pass_in_subtile) def stage_a_subtile(k_base, subtile, lds_a): @@ -1336,6 +1360,7 @@ def launch_gemm( num_groups: fx.Int32 = 0, work_queue: fx.Tensor = None, work_counter: fx.Tensor = None, + b_data_ptrs: fx.Tensor = None, stream: fx.Stream = fx.Stream(None), ): # Dense: one workgroup per output tile. Grouped/wgrad: a fixed num-CU @@ -1355,6 +1380,7 @@ def launch_gemm( num_groups, work_queue, work_counter, + b_data_ptrs, value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, ).launch(grid=(grid_x, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) @@ -1555,6 +1581,7 @@ def doGemm( num_groups: int = 0, work_queue: torch.Tensor = None, work_counter: torch.Tensor = None, + b_data_ptrs: torch.Tensor = None, sched: str = "worksteal", wgrad: bool = False, ): @@ -1671,6 +1698,14 @@ def doGemm( work_queue_arg = torch.zeros(1, dtype=torch.int32, device=A.device) work_counter_arg = torch.zeros(1, dtype=torch.int32, device=A.device) + # Per-expert weight base pointers (grouped TN only): an int64[G] device array of + # each expert's weight-slab base address, read in-kernel so the weight operand + # needs no host-side concat. Dense/NN/wgrad pass a placeholder never read there. + if b_data_ptrs is not None: + b_data_ptrs_arg = b_data_ptrs.contiguous().view(-1) + else: + b_data_ptrs_arg = torch.zeros(1, dtype=torch.int64, device=A.device) + launch( A_arg, B_arg, @@ -1682,5 +1717,6 @@ def doGemm( num_groups, work_queue_arg, work_counter_arg, + b_data_ptrs_arg, stream=stream, ) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index 14d13e7d7a..e99dc26ae2 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -60,6 +60,44 @@ SCALE_GROUP_SIZE = 32 +@functools.lru_cache(maxsize=None) +def _dummy_arg(dtype: torch.dtype, device: torch.device): + """One-element placeholder for an unused uniform-signature kernel arg. + + Cached per (dtype, device): the kernel never dereferences these buffers, so a + single shared zero tensor is reused across calls instead of allocating+zeroing + a fresh one each launch (which added a per-call alloc + fill-kernel + launch). + """ + return torch.zeros(1, dtype=dtype, device=device) + + +# Per-key cache of flyc.compile'd launchers. flydsl's JitFunction.__call__ rebuilds +# a full cache key (sig.bind + globals snapshot + arg-signature resolve) on EVERY +# launch -- ~70 us/launch of pure Python, which dominated the host cost on small +# GEMMs. flyc.compile returns a CompiledFunction whose __call__ skips all of that +# (~5 us): constexpr values are baked at compile time, so one CompiledFunction per +# (shape-class, dtype, layout) key serves every pointer/scalar/stream variation. +# NOTE: the fast path intentionally bypasses flydsl's globals-drift safety check; +# these kernel modules never mutate their captured globals between launches, so it +# is safe here. Keyed identically to the underlying @lru_cache launcher builders. +_COMPILED_LAUNCH_CACHE = {} + + +def _compiled_launch(key, jit_fn, args): + """Return (and cache) a flyc.compile'd fast launcher for ``jit_fn``. + + ``key`` uniquely identifies the compiled specialization (same tuple the + JitFunction's lru_cache uses). ``args`` is the full positional arg list for the + first call; it is used once to compile, then the cached CompiledFunction is + reused positionally for all later calls with the same ABI. + """ + cf = _COMPILED_LAUNCH_CACHE.get(key) + if cf is None: + cf = flyc.compile(jit_fn, *args) + _COMPILED_LAUNCH_CACHE[key] = cf + return cf + + def _debug_enabled() -> bool: value = os.getenv("DEBUG_FLYDSL_MXFP8_GEMM", "") return value.lower() not in ("", "0", "false", "no", "off") @@ -232,17 +270,13 @@ def pack_mx32_scales_for_hk( stream = torch.cuda.current_stream(scales_u8.device) stride0, stride1 = (int(x) for x in scales_u8.stride()) - _cached_mx32_scale_pack_launch( - dim, - qk, - bool(source_colwise), - stride0, - stride1, - )( - scales_u8, - packed, - stream=stream, - ) + pack_key = ("mx32_scale_pack", dim, qk, bool(source_colwise), stride0, stride1) + pack_args = (scales_u8, packed, stream) + _compiled_launch( + pack_key, + _cached_mx32_scale_pack_launch(dim, qk, bool(source_colwise), stride0, stride1), + pack_args, + )(*pack_args) return packed @@ -513,6 +547,8 @@ def kernel_gemm( num_groups: fx.Int32 = 0, work_queue: fx.Tensor = None, work_counter: fx.Tensor = None, + b_data_ptrs: fx.Tensor = None, + b_scale_ptrs: fx.Tensor = None, ): # Total output tiles. Dense uses this as the grid extent; grouped uses it as # the drain bound for the persistent loop. For fwd/dgrad, c_m = M_total spans @@ -568,6 +604,14 @@ def kernel_gemm( wq_rsrc = fx.rocdl.make_buffer_tensor(work_queue, max_size=True) wq_div = fx.logical_divide(wq_rsrc, fx.make_layout(1, 1)) + # Grouped-TN forward: the weights arrive as G independently-allocated + # per-expert slabs (not one packed buffer), so B is addressed from a + # per-expert base pointer read from ``b_data_ptrs`` rather than an + # expert-shifted column into a packed B. Reuse ``gB``'s static layout + # (compile-time cosize) and swap only the base per tile. NN/wgrad keep + # the packed-B expert-offset folds; dense never enters here. + b_ptr_layout = fx.get_layout(gB) + # LDS broadcast scratch (one int32), following the persistent-kernel # idiom in mega_moe_stage1: recast the LDS slot to an int32 iterator and # a 1-element view. Thread 0 writes via ptr_store; every thread reads via @@ -658,7 +702,13 @@ def process_tile(pid_m, pid_n, expert): bs_expert_k128_off = expert * fx.Int32(NUM_K_TILES) store_m_expert_off = fx.Index(0) else: - # TN: expert on the output-N axis. + # TN: weights are per-expert slabs addressed via a base-pointer + # swap (b_div_tile below), so B's column coord stays LOCAL + # (by_n_idx) -- the expert shift lives in the base, not the + # column. The B-scale still packs [K128, G*N], so its row stride + # grows to G*N (c_n_bs) and column folds in expert (handled at + # by_n_b_idx for the scale read; B-data uses by_n_idx directly + # via the per-expert base). by_n_b_idx = fx.Index(expert * fx.Int32(c_n) + by_n) c_n_bs = num_groups * c_n a_expert_k_off = fx.Index(0) @@ -710,8 +760,23 @@ def process_tile(pid_m, pid_n, expert): a_f8_ir_t, wave_id, ) + # Grouped B-data source: rebase onto this tile's expert weight slab (the + # per-expert base pointer already selects the expert), so the in-tile + # addressing uses LOCAL coords. TN reads with a local column (by_n_idx); + # NN (dgrad, transpose-read) packs experts on the contraction axis, so the + # per-expert base replaces the expert*K contraction shift (b_k_off -> 0). + # Dense/wgrad keep the shared packed ``b_div`` and its expert folds. + if const_expr(group): + b_base_i64 = gd.read_ptr_at(b_data_ptrs, expert) + b_div_tile = gd.rebase_buffer_view(b_base_i64, b_f8_ir_t, b_ptr_layout) + b_data_col = by_n_idx + b_k_off = fx.Index(0) + else: + b_div_tile = b_div + b_data_col = by_n_b_idx + b_k_off = b_expert_k_off b_g2s = G2SLoader( - b_div, + b_div_tile, gl_off_b, LOAD_PASSES_HALF, b_f8_ir_t, @@ -936,15 +1001,14 @@ def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): ) def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): - # b_expert_k_off shifts the contraction base to this tile's expert - # weight block for the grouped transpose-read (NN/NT) path; it is 0 - # for dense and for grouped TN (which offsets via by_n_b_idx). + # Grouped fwd/dgrad select the expert via the rebased per-expert B view + # (b_g2s) with LOCAL coords: TN a local column (b_data_col), NN a zero + # contraction shift (b_k_off). Dense/wgrad use the shared view with the + # packed column by_n_b_idx and the expert contraction shift b_k_off. b_g2s.load_one( lds_b[subtile], fx.Int32( - _b_global_base( - k_base + b_expert_k_off, subtile, c_n, by_n_b_idx - ) + _b_global_base(k_base + b_k_off, subtile, c_n, b_data_col) ), pass_in_subtile, ) @@ -1704,6 +1768,8 @@ def launch_gemm( num_groups: fx.Int32 = 0, work_queue: fx.Tensor = None, work_counter: fx.Tensor = None, + b_data_ptrs: fx.Tensor = None, + b_scale_ptrs: fx.Tensor = None, stream: fx.Stream = fx.Stream(None), ): # The integration only dispatches aligned shapes; no partial-tile masking exists. @@ -1726,6 +1792,8 @@ def launch_gemm( num_groups, work_queue, work_counter, + b_data_ptrs, + b_scale_ptrs, value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, ).launch(grid=(grid_x, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) @@ -1748,6 +1816,8 @@ def do_gemm( num_groups: int = 0, work_queue: fx.Tensor = None, work_counter: fx.Tensor = None, + b_data_ptrs: torch.Tensor = None, + b_scale_ptrs: torch.Tensor = None, sched: str = "worksteal", wgrad: bool = False, ): @@ -1768,10 +1838,11 @@ def do_gemm( else: raise ValueError(f"Unsupported MXFP8 kernel layout: {layout}") - # Grouped NN packs B along the contraction axis, so B's contraction dim is - # G x A's (Kb_runtime = G*N_out vs K_runtime = N_out). Every other case has - # them equal. - if group and layout == "NN": + # Grouped NN packs B along the contraction axis, so the packed B's contraction + # dim is G x A's (Kb_runtime = G*N_out vs K_runtime = N_out). With per-expert + # weight pointers B is a per-expert [N_out, K_in] representative, so its + # contraction already equals A's. Every other case has them equal. + if group and layout == "NN" and b_data_ptrs is None: assert Kb_runtime == num_groups * K_runtime, ( f"grouped NN B.K={Kb_runtime} != num_groups*A.K={num_groups * K_runtime}" ) @@ -1826,9 +1897,12 @@ def do_gemm( else: expected_as = (K_runtime // _BLOCK_K, M_runtime) # TN packs B along N (column count = G*N_out); NN packs B along the - # contraction axis (row count on K128 = G*N_out/128). + # contraction axis (row count on K128 = G*N_out/128). The B-scale stays packed + # even when the B data is addressed per-expert, so recover the packed + # contraction (Kb_runtime is per-expert N_out when pointers are used). if group and layout == "NN": - expected_bs = (Kb_runtime // _BLOCK_K, N_runtime) + kb_packed = Kb_runtime * num_groups if b_data_ptrs is not None else Kb_runtime + expected_bs = (kb_packed // _BLOCK_K, N_runtime) else: n_packed = N_runtime * num_groups if group else N_runtime expected_bs = (K_runtime // _BLOCK_K, n_packed) @@ -1896,14 +1970,18 @@ def do_gemm( Bs_arg = Bs.contiguous().view(-1) C_arg = C.view(-1) # DEFAULT keeps the kernel signature uniform with dummy 1-element buffers. + # These placeholders are never dereferenced by the kernel, so they are cached + # (per dtype+device) rather than re-allocated+zeroed every call -- each + # torch.zeros is an alloc + a fill kernel + a launch, which on small/fast GEMMs + # dominated the host launch time (4 such allocs per call -> ~4.6x CPU overhead). if needs_bias: Bias_arg = bias.contiguous().view(-1) else: - Bias_arg = torch.zeros(1, dtype=torch.float32, device=A.device) + Bias_arg = _dummy_arg(torch.float32, A.device) if needs_aux: Aux_arg = aux.view(-1) else: - Aux_arg = torch.zeros(1, dtype=C.dtype, device=A.device) + Aux_arg = _dummy_arg(C.dtype, A.device) # Grouped runtime tensors are required by the kernel signature; the dense # binary never dereferences them but still needs placeholder buffers. wgrad @@ -1912,10 +1990,25 @@ def do_gemm( work_queue_arg = work_queue.contiguous().view(-1) work_counter_arg = work_counter.contiguous().view(-1) else: - work_queue_arg = torch.zeros(1, dtype=torch.int32, device=A.device) - work_counter_arg = torch.zeros(1, dtype=torch.int32, device=A.device) + work_queue_arg = _dummy_arg(torch.int32, A.device) + work_counter_arg = _dummy_arg(torch.int32, A.device) + + # Per-expert weight base pointers (grouped TN forward only): an int64[G] device + # array of each expert's B-data / B-scale base address, read in-kernel so the + # weight operand needs no host-side concat. Every other path (dense, NN, wgrad) + # supplies a dtype-correct placeholder to keep the flyc.compile positional ABI + # uniform; the kernel never dereferences it there. + if b_data_ptrs is not None: + b_data_ptrs_arg = b_data_ptrs.contiguous().view(-1) + else: + b_data_ptrs_arg = _dummy_arg(torch.int64, A.device) + if b_scale_ptrs is not None: + b_scale_ptrs_arg = b_scale_ptrs.contiguous().view(-1) + else: + b_scale_ptrs_arg = _dummy_arg(torch.int64, A.device) - _cached_launch( + launch_key = ( + "mxfp8_gemm", K_runtime, A.dtype, B.dtype, @@ -1925,7 +2018,8 @@ def do_gemm( group, sched, wgrad, - )( + ) + launch_args = ( A_arg, As_arg, B_arg, @@ -1938,8 +2032,17 @@ def do_gemm( num_groups, work_queue_arg, work_counter_arg, - stream=stream, + b_data_ptrs_arg, + b_scale_ptrs_arg, + stream, ) + _compiled_launch( + launch_key, + _cached_launch( + K_runtime, A.dtype, B.dtype, C.dtype, layout, epilogue, group, sched, wgrad + ), + launch_args, + )(*launch_args) @functools.lru_cache(maxsize=None) @@ -2071,6 +2174,7 @@ def mxfp8_matmul( num_groups: int = 0, work_queue: fx.Tensor = None, work_counter: fx.Tensor = None, + b_data_ptrs: torch.Tensor = None, sched: str = "worksteal", wgrad: bool = False, ): @@ -2140,25 +2244,41 @@ def mxfp8_matmul( # is G x A's per-expert contraction k = N_out (kb != k by design), and the # output width n = K_in is NOT packed. D = [M_total, K_in] = (m, n). if group and layout == "TN": - if n % num_groups != 0: - raise ValueError( - f"Grouped MXFP8 TN packed N={n} is not divisible by num_groups={num_groups}" - ) - n_out = n // num_groups + # Per-expert weight pointers: B is a per-expert [N_out, K] representative, + # so b.shape[0] already IS N_out. Without pointers B is the packed [G*N_out, + # K] and N_out unpacks. Either way n_packed = G*N_out for the B-scale, which + # stays packed. + if b_data_ptrs is not None: + n_out = n + else: + if n % num_groups != 0: + raise ValueError( + f"Grouped MXFP8 TN packed N={n} is not divisible by num_groups={num_groups}" + ) + n_out = n // num_groups expected_kb = k # contraction unpacked; B rows already per-K elif group and layout == "NN": - if kb % num_groups != 0: - raise ValueError( - f"Grouped MXFP8 NN packed contraction={kb} not divisible by " - f"num_groups={num_groups}" - ) - if kb // num_groups != k: - raise ValueError( - f"Grouped MXFP8 NN per-expert contraction {kb // num_groups} " - f"!= A contraction {k}" - ) + # Per-expert weight pointers: B is a per-expert [N_out, K_in] representative, + # so b.shape[0] IS the per-expert contraction N_out (= k). Without pointers B + # is the packed [G*N_out, K_in] and the contraction unpacks. + if b_data_ptrs is not None: + if kb != k: + raise ValueError( + f"Grouped MXFP8 NN per-expert B contraction {kb} != A contraction {k}" + ) + else: + if kb % num_groups != 0: + raise ValueError( + f"Grouped MXFP8 NN packed contraction={kb} not divisible by " + f"num_groups={num_groups}" + ) + if kb // num_groups != k: + raise ValueError( + f"Grouped MXFP8 NN per-expert contraction {kb // num_groups} " + f"!= A contraction {k}" + ) n_out = n # output width (K_in) is not packed - expected_kb = kb # B rows are G*N_out, checked above + expected_kb = kb # checked above (per-expert or packed) else: n_out = n expected_kb = k @@ -2183,12 +2303,18 @@ def mxfp8_matmul( expected_a_scale = (m, k // SCALE_GROUP_SIZE) if layout == "TN": - # Grouped TN packs B along N, so the B scale spans n = G*N_out rows. - expected_b_scale = (n, k // SCALE_GROUP_SIZE) + # Grouped TN packs the B scale along N, spanning G*N_out rows. With per-expert + # weight pointers B is per-expert so n == N_out; scale rows are n_out*num_groups. + # Without pointers (packed B) n already is G*N_out. + b_scale_rows = n_out * num_groups if (group and b_data_ptrs is not None) else n + expected_b_scale = (b_scale_rows, k // SCALE_GROUP_SIZE) else: - # NN/NT B scale is [contraction/32, N]. Grouped NN packs B along the - # contraction axis, so use kb (= G*N_out) not the per-expert k. - expected_b_scale = (kb // SCALE_GROUP_SIZE, n) + # NN/NT B scale is [contraction/32, N]. Grouped NN packs the scale along the + # contraction axis (G*N_out rows). With per-expert weight pointers kb is the + # per-expert N_out, so the packed scale spans kb*num_groups; without pointers + # kb already is G*N_out. + kb_packed = kb * num_groups if (group and b_data_ptrs is not None) else kb + expected_b_scale = (kb_packed // SCALE_GROUP_SIZE, n) if tuple(a_scale.shape) != expected_a_scale: raise ValueError( @@ -2267,6 +2393,7 @@ def mxfp8_matmul( num_groups=num_groups, work_queue=work_queue, work_counter=work_counter, + b_data_ptrs=b_data_ptrs, sched=sched, ) return D diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 2e57aad1fe..2f2c813201 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -1242,14 +1242,29 @@ def _grouped_gemm_flydsl_mxfp8( ``out``. The descriptor queue is built host-side from ``m_splits``. """ from ..flydsl_kernels.gemm import _run_mxfp8_grouped + from ..flydsl_kernels.gemm.grouped_gemm_utils import concat_or_view_rows # TN packed operands: activations concatenated along the token (M) axis, # weights along the expert*N axis. Each MXFP8 operand carries its rowwise - # fp8 payload and E8M0 scale; concatenate both. - a_data = torch.cat([im._rowwise_data for im in inputmats], dim=0) - a_scale = torch.cat([im._rowwise_scale_inv for im in inputmats], dim=0) - b_data = torch.cat([w._rowwise_data for w in weights_fp8], dim=0) - b_scale = torch.cat([w._rowwise_scale_inv for w in weights_fp8], dim=0) + # fp8 payload and E8M0 scale. Activations are bulk-allocated by + # split_quantize (per-expert payloads physically adjacent), so pack them + # with a zero-copy spanning view when contiguous; weights come from + # separate parameter storages and fall back to a real concat. + a_data = concat_or_view_rows([im._rowwise_data for im in inputmats]) + a_scale = concat_or_view_rows([im._rowwise_scale_inv for im in inputmats]) + # Weights are independently-allocated per-expert slabs. The grouped-TN + # kernel reads each expert's weight from a per-expert base pointer, so the + # host passes an int64[G] address array instead of concatenating the G + # weights into one packed buffer (the dominant host cost). b_data is a + # per-expert representative used only for dtype/shape. The B-scale stays + # packed (cheap) and keeps the in-kernel packed-column fold. + b_data_ptrs = torch.tensor( + [w._rowwise_data.data_ptr() for w in weights_fp8], + dtype=torch.int64, + device=a_data.device, + ) + b_data = weights_fp8[0]._rowwise_data + b_scale = concat_or_view_rows([w._rowwise_scale_inv for w in weights_fp8]) # Preserve the native MXFP8 dtype so the wrapper reinterprets the uint8 # payload to the correct fp8 view at the kernel boundary. @@ -1271,6 +1286,7 @@ def _grouped_gemm_flydsl_mxfp8( n=n, k=k, output_dtype=out.dtype, + b_data_ptrs=b_data_ptrs, ) @staticmethod @@ -1292,15 +1308,21 @@ def _backward_flydsl_mxfp8_grouped( launch writing into ``dgrad`` ``[M_total, K_in]``. No quantization here. """ from ..flydsl_kernels.gemm import _run_mxfp8_grouped + from ..flydsl_kernels.gemm.grouped_gemm_utils import concat_or_view_rows # NN kernel-A = dY, rowwise (contraction over N_out): concat on the token # (M) axis. NN kernel-B = W, columnwise (transpose-read, contraction on the - # N_out axis): concat on the packed N_out axis. Each MXFP8 operand carries - # its selected payload + E8M0 scale. - a_data = torch.cat([g._rowwise_data for g in grad_output], dim=0) - a_scale = torch.cat([g._rowwise_scale_inv for g in grad_output], dim=0) - b_data = torch.cat([w._columnwise_data for w in weights_for_dgrad], dim=0) - b_scale = torch.cat([w._columnwise_scale_inv for w in weights_for_dgrad], dim=0) + # N_out axis). Weights are independent per-expert [N_out, K_in] columnwise + # slabs addressed via per-expert base pointers (no concat); dY still concats. + a_data = concat_or_view_rows([g._rowwise_data for g in grad_output]) + a_scale = concat_or_view_rows([g._rowwise_scale_inv for g in grad_output]) + b_data_ptrs = torch.tensor( + [w._columnwise_data.data_ptr() for w in weights_for_dgrad], + dtype=torch.int64, + device=a_data.device, + ) + b_data = weights_for_dgrad[0]._columnwise_data + b_scale = concat_or_view_rows([w._columnwise_scale_inv for w in weights_for_dgrad]) a_data._fp8_dtype = grad_output[0]._fp8_dtype b_data._fp8_dtype = weights_for_dgrad[0]._fp8_dtype @@ -1321,6 +1343,7 @@ def _backward_flydsl_mxfp8_grouped( k=n_out, # NN per-expert contraction is N_out output_dtype=dgrad.dtype, layout="NN", + b_data_ptrs=b_data_ptrs, ) @staticmethod @@ -1383,9 +1406,19 @@ def _grouped_gemm_flydsl_bf16(*, inputmats, weights, out, m_splits, m_splits_ten tensors (concat on the output-N axis). No scales. """ from ..flydsl_kernels.gemm import _run_bf16_grouped - - a_data = torch.cat([im.reshape(-1, im.shape[-1]) for im in inputmats], dim=0) - b_data = torch.cat([w for w in weights], dim=0) + from ..flydsl_kernels.gemm.grouped_gemm_utils import concat_or_view_rows + + a_data = concat_or_view_rows([im.reshape(-1, im.shape[-1]) for im in inputmats]) + # Weights are independent per-expert [N_out, K] tensors; the grouped-TN kernel + # reads each from a per-expert base pointer, so pass an int64[G] address array + # instead of concatenating the G weights. b_data is a per-expert representative + # (shape only; the kernel addresses weights via the pointers). + b_data_ptrs = torch.tensor( + [w.data_ptr() for w in weights], + dtype=torch.int64, + device=a_data.device, + ) + b_data = weights[0] n = weights[0].size(0) k = weights[0].size(-1) m_total = a_data.size(0) @@ -1399,15 +1432,25 @@ def _grouped_gemm_flydsl_bf16(*, inputmats, weights, out, m_splits, m_splits_ten n=n, k=k, output_dtype=out.dtype, + b_data_ptrs=b_data_ptrs, ) @staticmethod def _backward_flydsl_bf16_grouped(*, grad_output, weights, dgrad, m_splits, m_splits_tensor): """FlyDSL BF16 grouped dgrad (NN): dX_e = dY_e @ W_e, one packed launch.""" from ..flydsl_kernels.gemm import _run_bf16_grouped - - a_data = torch.cat([g.reshape(-1, g.shape[-1]) for g in grad_output], dim=0) - b_data = torch.cat([w for w in weights], dim=0) + from ..flydsl_kernels.gemm.grouped_gemm_utils import concat_or_view_rows + + a_data = concat_or_view_rows([g.reshape(-1, g.shape[-1]) for g in grad_output]) + # Weights are independent per-expert [N_out, K_in] tensors; the grouped-NN + # kernel reads each from a per-expert base pointer (no concat). b_data is a + # per-expert representative. + b_data_ptrs = torch.tensor( + [w.data_ptr() for w in weights], + dtype=torch.int64, + device=a_data.device, + ) + b_data = weights[0] n_out = grad_output[0].shape[-1] # per-expert contraction (N_out) k_in = dgrad.size(-1) # output width (K_in) m_total = a_data.size(0) @@ -1422,6 +1465,7 @@ def _backward_flydsl_bf16_grouped(*, grad_output, weights, dgrad, m_splits, m_sp k=n_out, output_dtype=dgrad.dtype, layout="NN", + b_data_ptrs=b_data_ptrs, ) @staticmethod @@ -1458,9 +1502,18 @@ def _grouped_gemm_flydsl_fp8(*, inputmats, weights_fp8, out, m_splits, m_splits_ one scalar inverse scale, gathered into length-G per-expert vectors. """ from ..flydsl_kernels.gemm import _run_fp8_grouped - - a_data = torch.cat([im._data for im in inputmats], dim=0) - b_data = torch.cat([w._data for w in weights_fp8], dim=0) + from ..flydsl_kernels.gemm.grouped_gemm_utils import concat_or_view_rows + + a_data = concat_or_view_rows([im._data for im in inputmats]) + # Weights are independent per-expert [N, K] FP8 tensors; the grouped-TN kernel + # reads each from a per-expert base pointer, so pass an int64[G] address array + # instead of concatenating the G weights. b_data is a per-expert representative. + b_data_ptrs = torch.tensor( + [w._data.data_ptr() for w in weights_fp8], + dtype=torch.int64, + device=a_data.device, + ) + b_data = weights_fp8[0]._data a_data._fp8_dtype = inputmats[0]._fp8_dtype b_data._fp8_dtype = weights_fp8[0]._fp8_dtype a_scale = torch.cat([im._scale_inv.reshape(1) for im in inputmats], dim=0) @@ -1481,6 +1534,7 @@ def _grouped_gemm_flydsl_fp8(*, inputmats, weights_fp8, out, m_splits, m_splits_ n=n, k=k, output_dtype=out.dtype, + b_data_ptrs=b_data_ptrs, ) @staticmethod @@ -1495,11 +1549,19 @@ def _backward_flydsl_fp8_grouped( B packed on the output axis = K_in, i.e. per expert dY @ (W^T)^T = dY @ W. """ from ..flydsl_kernels.gemm import _run_fp8_grouped - - # A = dY (rowwise fp8 data); B = W columnwise transpose, which is the - # materialized [K_in, N_out] -> packed on K_in (output axis of the TN GEMM). - a_data = torch.cat([g._data for g in grad_output], dim=0) - b_data = torch.cat([w._transpose for w in weights_for_dgrad], dim=0) + from ..flydsl_kernels.gemm.grouped_gemm_utils import concat_or_view_rows + + # A = dY (rowwise fp8 data); B = W columnwise transpose, the materialized + # [K_in, N_out] packed on K_in (the TN GEMM output axis). The transpose + # backings are independent per-expert slabs, addressed via per-expert base + # pointers through the same TN per-expert path (no concat); dY still concats. + a_data = concat_or_view_rows([g._data for g in grad_output]) + b_data_ptrs = torch.tensor( + [w._transpose.data_ptr() for w in weights_for_dgrad], + dtype=torch.int64, + device=a_data.device, + ) + b_data = weights_for_dgrad[0]._transpose a_data._fp8_dtype = grad_output[0]._fp8_dtype b_data._fp8_dtype = weights_for_dgrad[0]._fp8_dtype a_scale = torch.cat([g._scale_inv.reshape(1) for g in grad_output], dim=0) @@ -1520,6 +1582,7 @@ def _backward_flydsl_fp8_grouped( n=k_in, k=n_out, output_dtype=dgrad.dtype, + b_data_ptrs=b_data_ptrs, ) @staticmethod From 4d5d448fcfdceb39aa6acdcb5b18c6328cd93969 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Thu, 15 Oct 2026 20:42:36 +0000 Subject: [PATCH 09/22] Enable FP16 grouped GEMM, add a FlyDSL dispatch table, and wire a test backend --- tests/pytorch/test_grouped_linear.py | 38 ++- .../pytorch/flydsl_kernels/gemm/fp8_gemm.py | 19 +- .../flydsl_kernels/gemm/gemm_wrappers.py | 57 ++-- .../flydsl_kernels/gemm/half_prec_gemm.py | 10 +- .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 50 ++-- .../pytorch/module/grouped_linear.py | 245 +++++++----------- transformer_engine/pytorch/quantization.py | 9 + 7 files changed, 207 insertions(+), 221 deletions(-) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 19e4f4d0b2..085bdaa33d 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -266,6 +266,10 @@ def _test_grouped_linear_accuracy( split_size = 1 if fp8: split_size = get_align_size_for_quantization(recipe) + elif os.environ.get("NVTE_USE_GROUPED_GEMM_FLYDSL", "0") == "1": + # The FlyDSL grouped GEMM tiles every expert with BLOCK_M=256, so even the + # unquantized path needs 256-aligned per-expert token counts. + split_size = 256 m = config.max_seqlen_q // split_size dist = torch.sort(torch.randint(0, m, (num_gemms - 2,))).values.tolist() dist.append(dist[-1]) # Manually add a zero @@ -588,9 +592,9 @@ def test_grouped_linear_accuracy_cutlass( ) -# ROCm grouped GEMM backends (HipKittens / CK). Kept separate from the CUTLASS test above so the -# upstream one stays untouched. The backend is picked by env var; use_cutlass only relaxes the -# tolerances, since a non-reference GEMM will not be bit-exact against the sequential Linear. +# ROCm grouped GEMM backends (HipKittens / CK / FlyDSL). Kept separate from the CUTLASS test above +# so the upstream one stays untouched. The backend is picked by env var; use_cutlass only relaxes +# the tolerances, since a non-reference GEMM will not be bit-exact against the sequential Linear. @pytest.mark.skipif(not IS_HIP_EXTENSION, reason="ROCm-only grouped GEMM backends") @pytest.mark.parametrize("dtype", param_types, ids=str) @pytest.mark.parametrize("num_gemms", [3, 6]) @@ -600,7 +604,7 @@ def test_grouped_linear_accuracy_cutlass( @pytest.mark.parametrize("fp8_model_params", all_boolean) @pytest.mark.parametrize("fuse_wgrad_accumulation", all_boolean) @pytest.mark.parametrize("delay_wgrad_compute", all_boolean) -@pytest.mark.parametrize("grouped_gemm_backend", ["hipkittens", "ck"], ids=str) +@pytest.mark.parametrize("grouped_gemm_backend", ["hipkittens", "ck", "flydsl"], ids=str) def test_grouped_linear_accuracy_rocm_backends( dtype, num_gemms, @@ -628,13 +632,35 @@ def test_grouped_linear_accuracy_rocm_backends( # Triton backends in test_grouped_linear_accuracy instead. pytest.skip("HipKittens/CK grouped GEMM backends do not support FP8 block scaling.") - monkeypatch.setenv("NVTE_USE_CUTLASS_GROUPED_GEMM", "1") + if grouped_gemm_backend == "flydsl": + # FlyDSL grouped GEMM is gfx950-only and covers MXFP8, tensor-wise FP8 + # (current/delayed scaling), and unquantized bf16/fp16. Skip configs its gate + # rejects, which would silently fall back to the default backend and give + # false "flydsl verified" coverage. + if get_device_compute_capability() != (9, 5): + pytest.skip("FlyDSL grouped GEMM is only supported on gfx950.") + if recipe is None and dtype not in (torch.bfloat16, torch.float16): + pytest.skip("FlyDSL unquantized grouped GEMM supports only bf16/fp16.") + if recipe is not None and not ( + recipe.mxfp8() or recipe.float8_current_scaling() or recipe.delayed() + ): + pytest.skip("FlyDSL grouped GEMM supports MXFP8 and tensor-wise FP8 recipes only.") + + # Select exactly one backend; leave the others' env vars unset. HipKittens and CK + # are gated behind NVTE_USE_CUTLASS_GROUPED_GEMM; FlyDSL has its own opt-in and + # must not see the cutlass/HK/CK env vars. + monkeypatch.delenv("NVTE_USE_CUTLASS_GROUPED_GEMM", raising=False) monkeypatch.delenv("NVTE_USE_HIPKITTENS_GROUPED_GEMM", raising=False) monkeypatch.delenv("NVTE_USE_CK_GROUPED_GEMM", raising=False) + monkeypatch.delenv("NVTE_USE_GROUPED_GEMM_FLYDSL", raising=False) if grouped_gemm_backend == "hipkittens": + monkeypatch.setenv("NVTE_USE_CUTLASS_GROUPED_GEMM", "1") monkeypatch.setenv("NVTE_USE_HIPKITTENS_GROUPED_GEMM", "1") - else: + elif grouped_gemm_backend == "ck": + monkeypatch.setenv("NVTE_USE_CUTLASS_GROUPED_GEMM", "1") monkeypatch.setenv("NVTE_USE_CK_GROUPED_GEMM", "1") + else: + monkeypatch.setenv("NVTE_USE_GROUPED_GEMM_FLYDSL", "1") # For the configs CK grouped GEMM is guaranteed to handle, # assert it does not silently fall back to multi-stream hipBLASLt. diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py index 130cfc250e..ed1f15a559 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -257,10 +257,12 @@ def _load_scale(scale_div, index): grouped = group or wgrad num_tiles = gd.num_tiles_bound(c_m, c_n, BLOCK_M, BLOCK_N, num_groups, const_expr(wgrad)) - # Grouped-TN forward (weight_ptrs): weights are G independently-allocated - # per-expert slabs, addressed from a per-expert base pointer read from - # ``b_data_ptrs`` (reusing gB's static layout) instead of an expert-shifted - # B row. Forward-only compile flag; dgrad/wgrad keep the packed-B fold. + # weight_ptrs: weights are G independently-allocated per-expert slabs, + # addressed from a per-expert base pointer read from ``b_data_ptrs`` (reusing + # gB's static layout) instead of an expert-shifted B row. FP8 is TN-only, so + # both grouped forward and dgrad (dgrad feeds host-transposed weights through + # this same TN kernel) set this compile-time flag; wgrad keeps the packed-B + # fold. if const_expr(weight_ptrs): b_ptr_layout = fx.get_layout(gB) @@ -1371,10 +1373,11 @@ def doGemm( work_queue_arg = torch.zeros(1, dtype=torch.int32, device=A.device) work_counter_arg = torch.zeros(1, dtype=torch.int32, device=A.device) - # Per-expert weight base pointers (grouped TN forward only): compile a kernel that - # reads each expert's weight slab from b_data_ptrs[g] instead of concatenating the - # G weights host-side. weight_ptrs is a compile-time flag keyed into the kernel - # cache; dgrad/wgrad leave it False and pass a placeholder never dereferenced. + # Per-expert weight base pointers (grouped forward and dgrad -- both run on this + # TN kernel, dgrad over host-transposed weights): compile a kernel that reads each + # expert's weight slab from b_data_ptrs[g] instead of concatenating the G weights + # host-side. weight_ptrs is a compile-time flag keyed into the kernel cache; wgrad + # leaves it False and passes a placeholder never dereferenced. weight_ptrs = b_data_ptrs is not None if weight_ptrs: b_data_ptrs_arg = b_data_ptrs.contiguous().view(-1) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 053d9b8025..c61ce22f30 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -1046,9 +1046,8 @@ def _run_mxfp8_grouped( In both, ``k`` is A's per-expert contraction (TN: K_in; NN: N_out) and ``n`` is the output width. ``m_splits`` is the Python list of per-expert token counts (already host-resident); the per-tile ``(pid_m, pid_n, expert)`` descriptor - queue is built + shape-cached from it directly -- no device->host copy of a - ``group_offs`` tensor (that copy forced a blocking ``hipDeviceSynchronize`` - that stranded the GPU idle before the kernel). Operands arrive already + queue is built + shape-cached from it directly, so no device->host copy of a + ``group_offs`` tensor (and its blocking sync) is needed. Operands arrive already MXFP8-quantized and packed -- no quantization, no BLAS ownership swap. """ if layout not in ("TN", "NN"): @@ -1108,11 +1107,12 @@ def _run_mxfp8_grouped( epilogue, bias_arg = _resolve_epilogue(bias, n) - # TN forward addresses the weight operand from per-expert base pointers supplied - # by the caller (one int64 base per expert slab); the kernel reads B slab g from - # b_data_ptrs[g] instead of an expert-shifted column into one packed B, so the - # host never concatenates the G per-expert weights. NN dgrad has no pointer array - # and keeps the packed-B contraction fold (phase 2). + # Both TN forward and NN dgrad address the weight operand from per-expert base + # pointers supplied by the caller (one int64 base per expert slab) when + # ``b_data_ptrs`` is given; the kernel reads B slab g from b_data_ptrs[g] instead + # of folding the expert into a packed B (an output-N column for TN, the + # contraction base for NN), so the host never concatenates the G per-expert + # weights. When ``b_data_ptrs`` is None the packed-B expert fold is used. mxfp8_matmul( a_data, a_scale, @@ -1254,21 +1254,28 @@ def _run_bf16_grouped( a = packed grad_output [M_total, N_out] b = packed weights [G*N_out, K_in] ``m_splits`` is the Python list of per-expert token counts (already on the - host). The descriptor queue is built + cached from it directly -- no - device->host copy of a ``group_offs`` tensor (that copy forced a blocking - ``hipDeviceSynchronize`` that stranded the GPU idle before the kernel). + host). The descriptor queue is built + cached from it directly, so no + device->host copy of a ``group_offs`` tensor (and its blocking sync) is needed. """ from .half_prec_gemm import doGemm as _half_prec_doGemm if layout not in ("TN", "NN"): - raise FlyDSLUnsupportedError(f"FlyDSL BF16 grouped GEMM supports TN/NN, got {layout}") + raise FlyDSLUnsupportedError(f"FlyDSL half-precision grouped GEMM supports TN/NN, got {layout}") + + # BF16 and FP16 share this kernel (the MFMA opcode is selected per operand dtype); + # take the input dtype from the operands so both route through unchanged. + input_dtype = a_data.dtype + if input_dtype not in (torch.bfloat16, torch.float16): + raise FlyDSLUnsupportedError( + f"FlyDSL half-precision grouped GEMM requires bf16/fp16 operands, got {input_dtype}" + ) D = _validate_or_allocate_output( D, shape=torch.Size((m_total, n)), dtype=output_dtype, device=a_data.device, - backend_name=f"BF16 grouped {layout}", + backend_name=f"half-precision grouped {layout}", ) from .half_prec_gemm import _BLOCK_M, _BLOCK_N @@ -1289,8 +1296,8 @@ def _run_bf16_grouped( m=m_total, n=n, k=k, - input_dtype=torch.bfloat16, - label="BF16 grouped", + input_dtype=input_dtype, + label="half-precision grouped", epilogue=epilogue, bias=bias_arg, group=True, @@ -1343,6 +1350,12 @@ def _run_bf16_grouped_wgrad( num_groups = len(m_splits) device = dW.device + # BF16/FP16 share the kernel; take the dtype from the operands. + input_dtype = dy_cols[0].dtype + if input_dtype not in (torch.bfloat16, torch.float16): + raise FlyDSLUnsupportedError( + f"FlyDSL half-precision grouped wgrad requires bf16/fp16 operands, got {input_dtype}" + ) if n_out % _BLOCK_M != 0: raise FlyDSLUnsupportedError(f"wgrad requires N_out divisible by {_BLOCK_M}, got {n_out}") if k_in % _BLOCK_N != 0: @@ -1366,12 +1379,10 @@ def _run_bf16_grouped_wgrad( x_data = gd.uniform_pad_concat([x_cols[e] for e in ids], k_pad) # Fast path: when this bucket's experts are exactly rows 0..G-1 in order - # (the common single-bucket / balanced case), the bucket-local expert index - # equals the global index, so the kernel can write straight into dW -- no - # bucket-local buffer, no scatter copy. This reproduces the pre-bucketing - # single-launch path byte-for-byte (and avoids a full-dW-size copy that - # otherwise regresses the balanced case ~15%). Only genuinely multi-bucket - # (skewed) launches pay the scatter, where the K savings dwarf it. + # (the single-bucket / balanced case), the bucket-local expert index equals + # the global index, so the kernel writes straight into dW -- no bucket-local + # buffer and no scatter copy. Only multi-bucket (skewed) launches allocate a + # bucket-local output and scatter it back into dW. direct = ids == list(range(num_groups)) dW_b = dW if direct else torch.empty(bsize * n_out, k_in, dtype=output_dtype, device=device) @@ -1389,8 +1400,8 @@ def _run_bf16_grouped_wgrad( m=n_out, n=k_in, k=k_pad, - input_dtype=torch.bfloat16, - label="BF16 grouped wgrad", + input_dtype=input_dtype, + label="half-precision grouped wgrad", group=False, num_groups=bsize, work_queue=work_queue, diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py index 8684195551..171284afd9 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py @@ -425,11 +425,11 @@ def kernel_gemm( wq_div = fx.logical_divide(wq_rsrc, fx.make_layout(1, 1)) desc_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32) - # Grouped-TN forward: weights arrive as G independently-allocated - # per-expert slabs, so B is addressed from a per-expert base pointer read - # from ``b_data_ptrs`` (byte-addressed, reusing gB's static layout) rather - # than an expert-shifted column into one packed B. NN/wgrad keep the - # packed-B folds; dense never enters here. + # Grouped fwd/dgrad: weights arrive as G independently-allocated per-expert + # slabs, so B is addressed from a per-expert base pointer read from + # ``b_data_ptrs`` (byte-addressed, reusing gB's static layout) rather than + # an expert-shifted offset into one packed B. wgrad keeps the packed-B + # folds; dense never enters here. b_ptr_layout = fx.get_layout(gB) # LDS broadcast scratch (one int32): thread 0 writes the atomically diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index e99dc26ae2..338cf87da5 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -64,19 +64,18 @@ def _dummy_arg(dtype: torch.dtype, device: torch.device): """One-element placeholder for an unused uniform-signature kernel arg. - Cached per (dtype, device): the kernel never dereferences these buffers, so a - single shared zero tensor is reused across calls instead of allocating+zeroing - a fresh one each launch (which added a per-call alloc + fill-kernel + launch). + Cached per (dtype, device) and shared across calls: the kernel never + dereferences these buffers, so a single zero tensor avoids a per-launch + allocate+zero of a fresh one. """ return torch.zeros(1, dtype=dtype, device=device) # Per-key cache of flyc.compile'd launchers. flydsl's JitFunction.__call__ rebuilds -# a full cache key (sig.bind + globals snapshot + arg-signature resolve) on EVERY -# launch -- ~70 us/launch of pure Python, which dominated the host cost on small -# GEMMs. flyc.compile returns a CompiledFunction whose __call__ skips all of that -# (~5 us): constexpr values are baked at compile time, so one CompiledFunction per -# (shape-class, dtype, layout) key serves every pointer/scalar/stream variation. +# a full cache key (sig.bind + globals snapshot + arg-signature resolve) on every +# launch; flyc.compile returns a CompiledFunction whose __call__ skips that, baking +# constexpr values at compile time so one CompiledFunction per (shape-class, dtype, +# layout) key serves every pointer/scalar/stream variation. # NOTE: the fast path intentionally bypasses flydsl's globals-drift safety check; # these kernel modules never mutate their captured globals between launches, so it # is safe here. Keyed identically to the underlying @lru_cache launcher builders. @@ -548,7 +547,6 @@ def kernel_gemm( work_queue: fx.Tensor = None, work_counter: fx.Tensor = None, b_data_ptrs: fx.Tensor = None, - b_scale_ptrs: fx.Tensor = None, ): # Total output tiles. Dense uses this as the grid extent; grouped uses it as # the drain bound for the persistent loop. For fwd/dgrad, c_m = M_total spans @@ -604,12 +602,12 @@ def kernel_gemm( wq_rsrc = fx.rocdl.make_buffer_tensor(work_queue, max_size=True) wq_div = fx.logical_divide(wq_rsrc, fx.make_layout(1, 1)) - # Grouped-TN forward: the weights arrive as G independently-allocated + # Grouped fwd/dgrad: the weights arrive as G independently-allocated # per-expert slabs (not one packed buffer), so B is addressed from a # per-expert base pointer read from ``b_data_ptrs`` rather than an - # expert-shifted column into a packed B. Reuse ``gB``'s static layout - # (compile-time cosize) and swap only the base per tile. NN/wgrad keep - # the packed-B expert-offset folds; dense never enters here. + # expert-shifted offset into a packed B. Reuse ``gB``'s static layout + # (compile-time cosize) and swap only the base per tile. wgrad keeps the + # packed-B expert-offset folds; dense never enters here. b_ptr_layout = fx.get_layout(gB) # LDS broadcast scratch (one int32), following the persistent-kernel @@ -1769,7 +1767,6 @@ def launch_gemm( work_queue: fx.Tensor = None, work_counter: fx.Tensor = None, b_data_ptrs: fx.Tensor = None, - b_scale_ptrs: fx.Tensor = None, stream: fx.Stream = fx.Stream(None), ): # The integration only dispatches aligned shapes; no partial-tile masking exists. @@ -1793,7 +1790,6 @@ def launch_gemm( work_queue, work_counter, b_data_ptrs, - b_scale_ptrs, value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, ).launch(grid=(grid_x, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) @@ -1817,7 +1813,6 @@ def do_gemm( work_queue: fx.Tensor = None, work_counter: fx.Tensor = None, b_data_ptrs: torch.Tensor = None, - b_scale_ptrs: torch.Tensor = None, sched: str = "worksteal", wgrad: bool = False, ): @@ -1969,11 +1964,9 @@ def do_gemm( As_arg = As.contiguous().view(-1) Bs_arg = Bs.contiguous().view(-1) C_arg = C.view(-1) - # DEFAULT keeps the kernel signature uniform with dummy 1-element buffers. - # These placeholders are never dereferenced by the kernel, so they are cached - # (per dtype+device) rather than re-allocated+zeroed every call -- each - # torch.zeros is an alloc + a fill kernel + a launch, which on small/fast GEMMs - # dominated the host launch time (4 such allocs per call -> ~4.6x CPU overhead). + # DEFAULT keeps the kernel signature uniform with dummy 1-element buffers. These + # placeholders are never dereferenced by the kernel, so a cached shared buffer + # (see _dummy_arg) is used rather than a fresh allocate+zero every call. if needs_bias: Bias_arg = bias.contiguous().view(-1) else: @@ -1993,19 +1986,15 @@ def do_gemm( work_queue_arg = _dummy_arg(torch.int32, A.device) work_counter_arg = _dummy_arg(torch.int32, A.device) - # Per-expert weight base pointers (grouped TN forward only): an int64[G] device - # array of each expert's B-data / B-scale base address, read in-kernel so the - # weight operand needs no host-side concat. Every other path (dense, NN, wgrad) - # supplies a dtype-correct placeholder to keep the flyc.compile positional ABI - # uniform; the kernel never dereferences it there. + # Per-expert weight base pointers (grouped fwd TN / dgrad NN): an int64[G] device + # array of each expert's B-data base address, read in-kernel so the weight operand + # needs no host-side concat. Dense and wgrad supply a dtype-correct placeholder to + # keep the flyc.compile positional ABI uniform; the kernel never dereferences it + # there. (The B-scale stays packed -- no per-expert scale pointers.) if b_data_ptrs is not None: b_data_ptrs_arg = b_data_ptrs.contiguous().view(-1) else: b_data_ptrs_arg = _dummy_arg(torch.int64, A.device) - if b_scale_ptrs is not None: - b_scale_ptrs_arg = b_scale_ptrs.contiguous().view(-1) - else: - b_scale_ptrs_arg = _dummy_arg(torch.int64, A.device) launch_key = ( "mxfp8_gemm", @@ -2033,7 +2022,6 @@ def do_gemm( work_queue_arg, work_counter_arg, b_data_ptrs_arg, - b_scale_ptrs_arg, stream, ) _compiled_launch( diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 2f2c813201..514be35033 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -1123,9 +1123,9 @@ def _flydsl_grouped_common_ok( """Dtype-independent gate shared by every FlyDSL grouped GEMM path. HIP extension, gfx950, the ``NVTE_USE_GROUPED_GEMM_FLYDSL=1`` opt-in, - 128-aligned N/K for the tile geometry, and the same feature-set exclusions - as the blockwise-triton path. The per-dtype gates add only the quantization - condition (BF16: none; tensor-wise FP8: current/delayed scaling; MXFP8). + 128-aligned N/K for the tile geometry, and the unsupported-feature + exclusions. The per-dtype gates add only the quantization condition + (BF16/FP16: none; tensor-wise FP8: current/delayed scaling; MXFP8). """ return ( IS_HIP_EXTENSION @@ -1161,9 +1161,9 @@ def _flydsl_grouped_kind( ): """Return which FlyDSL grouped path applies: 'mxfp8' | 'fp8' | 'bf16' | None. - One dispatch signal replacing the former MXFP8-only boolean. All kinds share - the persistent work-stealing driver; they differ only in operand packing and - (fp8/mxfp8) the scale payload. BF16 is the unquantized path (fp8 is False). + All kinds share the persistent work-stealing driver; they differ only in + operand packing and (fp8/mxfp8) the scale payload. The 'bf16' kind is the + unquantized half-precision path (bf16/fp16; fp8 is False). """ if not _GroupedLinear._flydsl_grouped_common_ok( use_bias=use_bias, @@ -1189,57 +1189,24 @@ def _flydsl_grouped_kind( return "fp8" return None - @staticmethod - def _is_flydsl_mxfp8_grouped_gemm_supported( - *, - fp8, - recipe, - use_bias, - backward_override, - cpu_offloading, - save_original_input, - debug, - unpad_output, - actual_m_splits, - in_features, - out_features, - ) -> bool: - """Back-compat shim: True only for the MXFP8 FlyDSL grouped path.""" - return ( - _GroupedLinear._flydsl_grouped_kind( - fp8=fp8, - recipe=recipe, - use_bias=use_bias, - backward_override=backward_override, - cpu_offloading=cpu_offloading, - save_original_input=save_original_input, - debug=debug, - unpad_output=unpad_output, - actual_m_splits=actual_m_splits, - in_features=in_features, - out_features=out_features, - ) - == "mxfp8" - ) - @staticmethod def _grouped_gemm_flydsl_mxfp8( *, inputmats, - weights_fp8, + weights, out, m_splits, m_splits_tensor, ): """Dispatch the FlyDSL MXFP8 grouped GEMM over already-quantized operands. - Called inline at the forward GEMM dispatch (like ``general_grouped_gemm_triton``), - riding off the quantization TE already performed -- ``inputmats`` are the - per-group MXFP8 activations from ``tex.split_quantize`` and ``weights_fp8`` the - per-expert MXFP8 weights. This first iteration does not quantize: it packs the - rowwise MXFP8 payloads into the ``[M_total, K]`` / ``[G*N, K]`` layout and - issues one packed TN launch via :func:`_run_mxfp8_grouped`, writing into - ``out``. The descriptor queue is built host-side from ``m_splits``. + Called inline at the forward GEMM dispatch. It does not quantize -- it + consumes TE's quantized operands (``inputmats`` the per-group MXFP8 + activations from ``tex.split_quantize``, ``weights`` the per-expert MXFP8 + weights): activations are packed into the ``[M_total, K]`` layout and weights + are addressed via per-expert base pointers, then one TN launch is issued via + :func:`_run_mxfp8_grouped`, writing into ``out``. The descriptor queue is + built host-side from ``m_splits``. """ from ..flydsl_kernels.gemm import _run_mxfp8_grouped from ..flydsl_kernels.gemm.grouped_gemm_utils import concat_or_view_rows @@ -1259,20 +1226,20 @@ def _grouped_gemm_flydsl_mxfp8( # per-expert representative used only for dtype/shape. The B-scale stays # packed (cheap) and keeps the in-kernel packed-column fold. b_data_ptrs = torch.tensor( - [w._rowwise_data.data_ptr() for w in weights_fp8], + [w._rowwise_data.data_ptr() for w in weights], dtype=torch.int64, device=a_data.device, ) - b_data = weights_fp8[0]._rowwise_data - b_scale = concat_or_view_rows([w._rowwise_scale_inv for w in weights_fp8]) + b_data = weights[0]._rowwise_data + b_scale = concat_or_view_rows([w._rowwise_scale_inv for w in weights]) # Preserve the native MXFP8 dtype so the wrapper reinterprets the uint8 # payload to the correct fp8 view at the kernel boundary. a_data._fp8_dtype = inputmats[0]._fp8_dtype - b_data._fp8_dtype = weights_fp8[0]._fp8_dtype + b_data._fp8_dtype = weights[0]._fp8_dtype - n = weights_fp8[0]._rowwise_data.size(0) - k = weights_fp8[0]._rowwise_data.size(-1) + n = weights[0]._rowwise_data.size(0) + k = weights[0]._rowwise_data.size(-1) m_total = a_data.size(0) _run_mxfp8_grouped( @@ -1293,7 +1260,7 @@ def _grouped_gemm_flydsl_mxfp8( def _backward_flydsl_mxfp8_grouped( *, grad_output, - weights_for_dgrad, + weights, dgrad, m_splits, m_splits_tensor, @@ -1301,11 +1268,11 @@ def _backward_flydsl_mxfp8_grouped( """FlyDSL grouped dgrad: dX = dY @ W through the packed NN kernel. Rides off backward's existing quantization -- ``grad_output`` are the - per-group MXFP8 activations (rowwise) and ``weights_for_dgrad`` the - per-expert MXFP8 weights with columnwise payloads ensured by the caller. - Packs dY as the NN kernel-A ``[M_total, N_out]`` and W (columnwise) as the - contraction-packed kernel-B ``[G*N_out, K_in]``, then issues one grouped NN - launch writing into ``dgrad`` ``[M_total, K_in]``. No quantization here. + per-group MXFP8 activations (rowwise) and ``weights`` the per-expert MXFP8 + weights with columnwise payloads ensured by the caller. Packs dY as the NN + kernel-A ``[M_total, N_out]`` and W (columnwise) as the contraction-packed + kernel-B ``[G*N_out, K_in]``, then issues one grouped NN launch writing into + ``dgrad`` ``[M_total, K_in]``. No quantization here. """ from ..flydsl_kernels.gemm import _run_mxfp8_grouped from ..flydsl_kernels.gemm.grouped_gemm_utils import concat_or_view_rows @@ -1317,15 +1284,15 @@ def _backward_flydsl_mxfp8_grouped( a_data = concat_or_view_rows([g._rowwise_data for g in grad_output]) a_scale = concat_or_view_rows([g._rowwise_scale_inv for g in grad_output]) b_data_ptrs = torch.tensor( - [w._columnwise_data.data_ptr() for w in weights_for_dgrad], + [w._columnwise_data.data_ptr() for w in weights], dtype=torch.int64, device=a_data.device, ) - b_data = weights_for_dgrad[0]._columnwise_data - b_scale = concat_or_view_rows([w._columnwise_scale_inv for w in weights_for_dgrad]) + b_data = weights[0]._columnwise_data + b_scale = concat_or_view_rows([w._columnwise_scale_inv for w in weights]) a_data._fp8_dtype = grad_output[0]._fp8_dtype - b_data._fp8_dtype = weights_for_dgrad[0]._fp8_dtype + b_data._fp8_dtype = weights[0]._fp8_dtype n_out = grad_output[0]._rowwise_data.size(-1) # per-expert contraction (N_out) k_in = dgrad.size(-1) # output width (K_in) @@ -1495,7 +1462,7 @@ def _wgrad_flydsl_bf16_grouped(*, grad_output, inputmats, wgrad_list, m_splits): wgrad_list[g].copy_(dW[g * n_out : (g + 1) * n_out]) @staticmethod - def _grouped_gemm_flydsl_fp8(*, inputmats, weights_fp8, out, m_splits, m_splits_tensor): + def _grouped_gemm_flydsl_fp8(*, inputmats, weights, out, m_splits, m_splits_tensor): """FlyDSL tensor-wise FP8 grouped forward (TN): Y_e = X_e @ W_e^T. Experts pack weights on the output-N axis ([G*N, K]); each operand carries @@ -1509,18 +1476,18 @@ def _grouped_gemm_flydsl_fp8(*, inputmats, weights_fp8, out, m_splits, m_splits_ # reads each from a per-expert base pointer, so pass an int64[G] address array # instead of concatenating the G weights. b_data is a per-expert representative. b_data_ptrs = torch.tensor( - [w._data.data_ptr() for w in weights_fp8], + [w._data.data_ptr() for w in weights], dtype=torch.int64, device=a_data.device, ) - b_data = weights_fp8[0]._data + b_data = weights[0]._data a_data._fp8_dtype = inputmats[0]._fp8_dtype - b_data._fp8_dtype = weights_fp8[0]._fp8_dtype + b_data._fp8_dtype = weights[0]._fp8_dtype a_scale = torch.cat([im._scale_inv.reshape(1) for im in inputmats], dim=0) - b_scale = torch.cat([w._scale_inv.reshape(1) for w in weights_fp8], dim=0) + b_scale = torch.cat([w._scale_inv.reshape(1) for w in weights], dim=0) - n = weights_fp8[0]._data.size(0) - k = weights_fp8[0]._data.size(-1) + n = weights[0]._data.size(0) + k = weights[0]._data.size(-1) m_total = a_data.size(0) _run_fp8_grouped( @@ -1539,7 +1506,7 @@ def _grouped_gemm_flydsl_fp8(*, inputmats, weights_fp8, out, m_splits, m_splits_ @staticmethod def _backward_flydsl_fp8_grouped( - *, grad_output, weights_for_dgrad, dgrad, m_splits, m_splits_tensor + *, grad_output, weights, dgrad, m_splits, m_splits_tensor ): """FlyDSL tensor-wise FP8 grouped dgrad: dX_e = dY_e @ W_e, via the TN kernel. @@ -1557,15 +1524,15 @@ def _backward_flydsl_fp8_grouped( # pointers through the same TN per-expert path (no concat); dY still concats. a_data = concat_or_view_rows([g._data for g in grad_output]) b_data_ptrs = torch.tensor( - [w._transpose.data_ptr() for w in weights_for_dgrad], + [w._transpose.data_ptr() for w in weights], dtype=torch.int64, device=a_data.device, ) - b_data = weights_for_dgrad[0]._transpose + b_data = weights[0]._transpose a_data._fp8_dtype = grad_output[0]._fp8_dtype - b_data._fp8_dtype = weights_for_dgrad[0]._fp8_dtype + b_data._fp8_dtype = weights[0]._fp8_dtype a_scale = torch.cat([g._scale_inv.reshape(1) for g in grad_output], dim=0) - b_scale = torch.cat([w._scale_inv.reshape(1) for w in weights_for_dgrad], dim=0) + b_scale = torch.cat([w._scale_inv.reshape(1) for w in weights], dim=0) n_out = grad_output[0]._data.size(-1) # per-expert contraction (N_out) k_in = dgrad.size(-1) # output width (K_in) @@ -1625,6 +1592,33 @@ def _wgrad_flydsl_fp8_grouped(*, grad_output, inputmats, wgrad_list, m_splits): for g in range(num_gemms): wgrad_list[g].copy_(dW[g * n_out : (g + 1) * n_out]) + @classmethod + def _flydsl_grouped_dispatch(cls, phase, kind): + """Return the FlyDSL grouped helper for ``phase`` and dtype ``kind``. + + ``phase`` is 'fwd' | 'dgrad' | 'wgrad' and ``kind`` is the value from + :meth:`_flydsl_grouped_kind` ('mxfp8' | 'fp8' | 'bf16'). All helpers of a + phase share a keyword signature, so the caller passes one kwarg set. + """ + table = { + "fwd": { + "mxfp8": cls._grouped_gemm_flydsl_mxfp8, + "fp8": cls._grouped_gemm_flydsl_fp8, + "bf16": cls._grouped_gemm_flydsl_bf16, + }, + "dgrad": { + "mxfp8": cls._backward_flydsl_mxfp8_grouped, + "fp8": cls._backward_flydsl_fp8_grouped, + "bf16": cls._backward_flydsl_bf16_grouped, + }, + "wgrad": { + "mxfp8": cls._wgrad_flydsl_mxfp8_grouped, + "fp8": cls._wgrad_flydsl_fp8_grouped, + "bf16": cls._wgrad_flydsl_bf16_grouped, + }, + } + return table[phase][kind] + @staticmethod def _forward_blockwise_fp8_triton( ctx, @@ -2039,9 +2033,10 @@ def forward( pow2_grad=blockwise_recipe.fp8_quant_bwd_grad.power_2_scale, ) - # FlyDSL MXFP8 grouped GEMM (ROCm gfx950) opt-in. Unlike the blockwise-triton - # path this does NOT quantize: it rides off TE's existing quantization and is - # dispatched inline at the forward GEMM (see below), like the triton grouped path. + # FlyDSL grouped GEMM (ROCm gfx950) opt-in, covering MXFP8, tensor-wise FP8, + # and BF16/FP16. It does not quantize -- it consumes TE's existing quantized + # operands and is dispatched inline at the forward GEMM (see below). Returns + # the dtype kind ('mxfp8' | 'fp8' | 'bf16') or None when unsupported. flydsl_grouped_kind = _GroupedLinear._flydsl_grouped_kind( fp8=fp8, recipe=blockwise_recipe, @@ -2055,9 +2050,12 @@ def forward( in_features=weights[0].size(-1), out_features=weights[0].size(0), ) - # The BF16 kernel is bfloat16-only (FP16 grouped is not wired through - # FlyDSL); drop the BF16 kind for any other activation dtype. - if flydsl_grouped_kind == "bf16" and activation_dtype != torch.bfloat16: + # The half-precision kernel handles bf16 and fp16 (the MFMA opcode is selected + # per operand dtype); drop the kind for any other (e.g. fp32) activation dtype. + if flydsl_grouped_kind == "bf16" and activation_dtype not in ( + torch.bfloat16, + torch.float16, + ): flydsl_grouped_kind = None use_grouped_gemm_flydsl = flydsl_grouped_kind is not None @@ -2284,30 +2282,13 @@ def forward( # Perform GEMM if use_grouped_gemm_flydsl: - if flydsl_grouped_kind == "mxfp8": - _GroupedLinear._grouped_gemm_flydsl_mxfp8( - inputmats=inputmats, - weights_fp8=weights_fp8, - out=out, - m_splits=m_splits, - m_splits_tensor=m_splits_tensor, - ) - elif flydsl_grouped_kind == "fp8": - _GroupedLinear._grouped_gemm_flydsl_fp8( - inputmats=inputmats, - weights_fp8=weights_fp8, - out=out, - m_splits=m_splits, - m_splits_tensor=m_splits_tensor, - ) - else: # bf16 - _GroupedLinear._grouped_gemm_flydsl_bf16( - inputmats=inputmats, - weights=weights_fp8, - out=out, - m_splits=m_splits, - m_splits_tensor=m_splits_tensor, - ) + _GroupedLinear._flydsl_grouped_dispatch("fwd", flydsl_grouped_kind)( + inputmats=inputmats, + weights=weights_fp8, + out=out, + m_splits=m_splits, + m_splits_tensor=m_splits_tensor, + ) else: if use_grouped_gemm_triton: general_grouped_gemm_func = general_grouped_gemm_triton @@ -2358,7 +2339,6 @@ def forward( # FlyDSL grouped forward routes dgrad through the grouped NN path; the # shared backward branches on this kind (wgrad still uses the default # where a dtype/feature is unsupported). - ctx.use_flydsl_mxfp8_grouped = use_grouped_gemm_flydsl ctx.flydsl_grouped_kind = flydsl_grouped_kind ctx.weight_quantizers = weight_quantizers ctx.weights_shape_1 = weights[0].shape[1] @@ -2916,27 +2896,11 @@ def backward( columnwise_usage=quantizer.columnwise_usage, ) _flydsl_kind = getattr(ctx, "flydsl_grouped_kind", None) - if _flydsl_kind == "mxfp8": + if _flydsl_kind is not None: # dgrad dX = dY @ W through the FlyDSL grouped NN kernel, riding # off the already-quantized grad_output (rowwise) and weights # (columnwise, ensured above). - _GroupedLinear._backward_flydsl_mxfp8_grouped( - grad_output=grad_output, - weights_for_dgrad=weights_for_dgrad, - dgrad=dgrad, - m_splits=ctx.m_splits, - m_splits_tensor=ctx.m_splits_tensor, - ) - elif _flydsl_kind == "fp8": - _GroupedLinear._backward_flydsl_fp8_grouped( - grad_output=grad_output, - weights_for_dgrad=weights_for_dgrad, - dgrad=dgrad, - m_splits=ctx.m_splits, - m_splits_tensor=ctx.m_splits_tensor, - ) - elif _flydsl_kind == "bf16": - _GroupedLinear._backward_flydsl_bf16_grouped( + _GroupedLinear._flydsl_grouped_dispatch("dgrad", _flydsl_kind)( grad_output=grad_output, weights=weights_for_dgrad, dgrad=dgrad, @@ -3074,9 +3038,9 @@ def backward( inputmats_dequant.append(cast_if_needed(inputmat, ctx.activation_dtype)) inputmats = inputmats_dequant - # FlyDSL grouped wgrad (NT) opt-in. Phase 1 handles the plain path - # only; fused wgrad accumulation, bias grad, and delayed wgrad still - # fall through to the default grouped GEMM below. + # FlyDSL grouped wgrad (NT) opt-in, plain path only; fused wgrad + # accumulation, bias grad, and delayed wgrad fall through to the + # default grouped GEMM below. _flydsl_kind = getattr(ctx, "flydsl_grouped_kind", None) use_flydsl_wgrad = ( _flydsl_kind is not None @@ -3087,27 +3051,12 @@ def backward( ) ) if use_flydsl_wgrad: - if _flydsl_kind == "mxfp8": - _GroupedLinear._wgrad_flydsl_mxfp8_grouped( - grad_output=grad_output, - inputmats=inputmats, - wgrad_list=wgrad_list, - m_splits=ctx.m_splits, - ) - elif _flydsl_kind == "fp8": - _GroupedLinear._wgrad_flydsl_fp8_grouped( - grad_output=grad_output, - inputmats=inputmats, - wgrad_list=wgrad_list, - m_splits=ctx.m_splits, - ) - else: # bf16 - _GroupedLinear._wgrad_flydsl_bf16_grouped( - grad_output=grad_output, - inputmats=inputmats, - wgrad_list=wgrad_list, - m_splits=ctx.m_splits, - ) + _GroupedLinear._flydsl_grouped_dispatch("wgrad", _flydsl_kind)( + grad_output=grad_output, + inputmats=inputmats, + wgrad_list=wgrad_list, + m_splits=ctx.m_splits, + ) clear_tensor_data(*inputmats) grouped_gemm_wgrad = None else: diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 02cb4f44eb..acd63fbe24 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -343,6 +343,15 @@ def get_align_size_for_quantization(recipe: Recipe) -> int: # alignment as the conservative fallback for context-free callers. if recipe.custom(): return recipe.quantization_alignment + # The FlyDSL grouped GEMM (gfx950 opt-in) tiles every expert with BLOCK_M=256 + # across all dtypes, so per-expert token counts must be 256-aligned regardless of + # recipe. + if ( + IS_HIP_EXTENSION + and get_device_compute_capability() == (9, 5) + and os.environ.get("NVTE_USE_GROUPED_GEMM_FLYDSL", "0") == "1" + ): + return 256 if recipe.mxfp8(): # HipKittens grouped GEMM requires 256-aligned expert dimensions. # HK is used by default when NVTE_USE_CUTLASS_GROUPED_GEMM=1, From 0dee39c1c39c48c1cc2d66d2918dc666fcdbbc61 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Fri, 16 Oct 2026 20:13:06 +0000 Subject: [PATCH 10/22] Cut FlyDSL grouped-GEMM host overhead across all dtypes --- .../pytorch/flydsl_kernels/gemm/fp32_gemm.py | 81 ++++++--- .../pytorch/flydsl_kernels/gemm/fp8_gemm.py | 134 +++++++++------ .../flydsl_kernels/gemm/grouped_gemm_utils.py | 11 ++ .../flydsl_kernels/gemm/half_prec_gemm.py | 117 +++++++++---- .../pytorch/module/grouped_linear.py | 158 ++++++++++++------ 5 files changed, 344 insertions(+), 157 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py index 1b7f9250e1..5523312f58 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py @@ -969,6 +969,47 @@ def _cached_launch(K: int, use_xcd_remap: bool = True, epilogue: str = "DEFAULT" return _compile_kernel(K, use_xcd_remap=use_xcd_remap, epilogue=epilogue) +@functools.lru_cache(maxsize=None) +def _dense_placeholder(device: torch.device, dtype: torch.dtype, numel: int = 1): + """One cached 1-element buffer per (device, dtype) for dense kernel-signature + placeholders. + + The FP32 binary requires the Bias/Aux epilogue tensor arguments to be present + but never dereferences them in the DEFAULT epilogue. Re-allocating them via + ``torch.zeros`` on every call is pure host overhead; since the contents are + never read, a single shared read-only buffer per (device, dtype) is safe. + """ + return torch.zeros(numel, dtype=dtype, device=device) + + +# Per-key cache of flyc.compile'd launchers. flydsl's JitFunction.__call__ rebuilds +# a full cache key (sig.bind + globals snapshot + arg-signature resolve) on every +# launch -- ~190us/call, which dominates wall time at small M where the kernel is +# short. flyc.compile returns a CompiledFunction whose __call__ skips all of that +# (ctypes pointer/scalar update + invoke, ~5us), baking constexpr values at compile +# time so one CompiledFunction per launch_key serves every pointer/scalar/stream +# variation. M/N are runtime scalar args (not baked), so the key excludes them -- +# identical to the MXFP8/FP8/half-precision paths. +# NOTE: the fast path intentionally bypasses flydsl's globals-drift safety check; +# this module never mutates its captured globals between launches, so it is safe. +_COMPILED_LAUNCH_CACHE = {} + + +def _compiled_launch(key, jit_fn, args): + """Return (and cache) a flyc.compile'd fast launcher for ``jit_fn``. + + ``key`` uniquely identifies the compiled specialization (the same tuple the + JitFunction builder's lru_cache uses). ``args`` is the full positional arg list + for the first call; used once to compile, then the cached CompiledFunction is + reused positionally for all later calls with the same ABI. + """ + cf = _COMPILED_LAUNCH_CACHE.get(key) + if cf is None: + cf = flyc.compile(jit_fn, *args) + _COMPILED_LAUNCH_CACHE[key] = cf + return cf + + def fp32_matmul( a: torch.Tensor, b: torch.Tensor, @@ -989,27 +1030,16 @@ def fp32_matmul( The optimized core streams both operands with K contiguous and consumes B as [N, K] directly, so this adapter performs no internal transpose. The wrapper is responsible for normalizing NN/NT layouts into this contract. + + Only the checks ``doGemm`` does not repeat are performed here: operand rank, + device agreement, and output contiguity. Inner-dim match, operand/output + dtype, and C shape are validated by ``doGemm`` against the runtime geometry, + so they are not duplicated in this shim. """ if a.ndim != 2 or b.ndim != 2: raise ValueError( f"FlyDSL FP32 TN expects rank-2 operands, got A{tuple(a.shape)} and B{tuple(b.shape)}" ) - - m, k = a.shape - n, kb = b.shape - if kb != k: - raise ValueError(f"Inner dimensions do not match: A{tuple(a.shape)} and B{tuple(b.shape)}") - if a.dtype != torch.float32 or b.dtype != torch.float32: - raise TypeError( - "FlyDSL FP32 GEMM expects both operands to have torch.float32 dtype, " - f"got {a.dtype} and {b.dtype}" - ) - if tuple(c.shape) != (m, n): - raise ValueError(f"C shape {tuple(c.shape)} != expected {(m, n)}") - if c.dtype != torch.float32: - raise FlyDSLUnsupportedError( - f"The current FlyDSL FP32 kernel stores torch.float32 output, got {c.dtype}" - ) if a.device != b.device or a.device != c.device: raise ValueError( f"A, B, and C must be on the same device, got {a.device}, {b.device}, and {c.device}" @@ -1090,13 +1120,24 @@ def doGemm( B_arg = B.contiguous().view(torch.uint8).view(-1) C_arg = C.view(-1) # DEFAULT keeps the kernel signature uniform with dummy 1-element buffers. + # The placeholders are never dereferenced by the kernel, so reuse one cached + # buffer per (device, dtype) instead of re-allocating each call. if needs_bias: Bias_arg = bias.contiguous().view(-1) else: - Bias_arg = torch.zeros(1, dtype=torch.float32, device=A.device) + Bias_arg = _dense_placeholder(A.device, torch.float32) if needs_aux: Aux_arg = aux.contiguous().view(-1) else: - Aux_arg = torch.zeros(1, dtype=torch.float32, device=A.device) - launch = _cached_launch(int(K_runtime), bool(use_xcd_remap), epilogue) - launch(A_arg, B_arg, C_arg, Bias_arg, Aux_arg, M_runtime, N_runtime, stream=stream) + Aux_arg = _dense_placeholder(A.device, torch.float32) + + launch_key = ("fp32_gemm", int(K_runtime), bool(use_xcd_remap), epilogue) + launch_args = (A_arg, B_arg, C_arg, Bias_arg, Aux_arg, M_runtime, N_runtime, stream) + # flyc.compile fast path: skips flydsl's per-call cache-key rebuild (sig.bind + + # globals snapshot + arg resolve), the dominant host cost at small M. M/N are + # runtime scalars, so one compiled launcher per launch_key serves all shapes. + _compiled_launch( + launch_key, + _cached_launch(int(K_runtime), bool(use_xcd_remap), epilogue), + launch_args, + )(*launch_args) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py index ed1f15a559..4b93e52bd6 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -1139,6 +1139,49 @@ def launch_gemm( return launch_gemm +@functools.lru_cache(maxsize=None) +def _dense_placeholder(device: torch.device, dtype: torch.dtype, numel: int = 1): + """One cached 1-element buffer per (device, dtype) for dense kernel-signature + placeholders. + + The dense FP8 binary requires the grouped/epilogue tensor arguments + (Bias/Aux/work_queue/work_counter/b_data_ptrs) to be present but never + dereferences them. Re-allocating them via ``torch.zeros`` on every call is + pure host overhead that dominates wall time at small M; since the contents + are never read, a single shared read-only buffer per (device, dtype) is + safe to reuse across calls. + """ + return torch.zeros(numel, dtype=dtype, device=device) + + +# Per-key cache of flyc.compile'd launchers. flydsl's JitFunction.__call__ rebuilds +# a full cache key (sig.bind + globals snapshot + arg-signature resolve) on every +# launch -- ~190us/call, which dominates wall time at small M where the kernel is +# short. flyc.compile returns a CompiledFunction whose __call__ skips all of that +# (ctypes pointer/scalar update + invoke, ~5us), baking constexpr values at compile +# time so one CompiledFunction per (shape-class, dtype, layout) key serves every +# pointer/scalar/stream variation. M/N are runtime scalar args (not baked), so the +# key excludes them -- identical to the MXFP8 path in mxfp8_gemm.py. +# NOTE: the fast path intentionally bypasses flydsl's globals-drift safety check; +# this module never mutates its captured globals between launches, so it is safe. +_COMPILED_LAUNCH_CACHE = {} + + +def _compiled_launch(key, jit_fn, args): + """Return (and cache) a flyc.compile'd fast launcher for ``jit_fn``. + + ``key`` uniquely identifies the compiled specialization (the same tuple the + JitFunction builder's lru_cache uses). ``args`` is the full positional arg list + for the first call; used once to compile, then the cached CompiledFunction is + reused positionally for all later calls with the same ABI. + """ + cf = _COMPILED_LAUNCH_CACHE.get(key) + if cf is None: + cf = flyc.compile(jit_fn, *args) + _COMPILED_LAUNCH_CACHE[key] = cf + return cf + + @functools.lru_cache(maxsize=None) def _cached_launch( K: int, @@ -1190,54 +1233,13 @@ def fp8_matmul( The wrapper is responsible for adapting TE's logical [K, N] operand into the core's row-major [N, K] representation. No operand swap or transpose is performed in this module. - """ - if not isinstance(a, torch.Tensor) or not isinstance(b, torch.Tensor): - raise TypeError("FlyDSL FP8 GEMM expects plain torch.Tensor payloads") - - if a.ndim != 2 or b.ndim != 2: - raise ValueError( - f"FlyDSL FP8 TN expects rank-2 operands, got A{tuple(a.shape)} and B{tuple(b.shape)}" - ) - - supported_fp8_dtypes = ( - torch.float8_e4m3fn, - torch.float8_e5m2, - ) - if a.dtype not in supported_fp8_dtypes or b.dtype not in supported_fp8_dtypes: - raise TypeError( - f"FlyDSL FP8 GEMM expects E4M3 or E5M2 payloads, got A={a.dtype} and B={b.dtype}" - ) - - m, k = a.shape - n, kb = b.shape - if kb != k: - raise ValueError(f"Inner dimensions do not match: A{tuple(a.shape)} and B{tuple(b.shape)}") - - for name, scale in ( - ("A_scale_inv", a_scale_inv), - ("B_scale_inv", b_scale_inv), - ): - if not isinstance(scale, torch.Tensor): - raise TypeError(f"{name} must be a torch.Tensor") - if scale.dtype != torch.float32 or scale.numel() != 1: - raise TypeError( - f"{name} must contain exactly one FP32 value, got " - f"dtype={scale.dtype}, shape={tuple(scale.shape)}" - ) - - if tuple(c.shape) != (m, n): - raise ValueError(f"C shape {tuple(c.shape)} != expected {(m, n)}") - if c.dtype not in (torch.float16, torch.bfloat16, torch.float32): - raise FlyDSLUnsupportedError( - f"FlyDSL FP8 supports only float16, bfloat16, and float32 outputs, got {c.dtype}" - ) - if not c.is_contiguous(): - raise FlyDSLUnsupportedError("FlyDSL FP8 requires contiguous output storage") - - tensors = (a, b, a_scale_inv, b_scale_inv, c) - if any(t.device != a.device for t in tensors[1:]): - raise ValueError("A, B, inverse scales, and C must be on the same device") + ``_run_fp8`` (the sole caller) has already validated operand rank, FP8 + dtypes, inverse-scale dtype/numel, and contiguity; ``doGemm`` re-validates + dtypes, inner-dim match, scale numel, and C shape/dtype against the runtime + geometry. This shim therefore just forwards, rather than repeating those + checks a third time (``doGemm`` derives M/N/K from the operands itself). + """ doGemm( a, b, @@ -1355,14 +1357,16 @@ def doGemm( A_scale_arg = A_scale_inv.contiguous().view(-1) B_scale_arg = B_scale_inv.contiguous().view(-1) # DEFAULT keeps the kernel signature uniform with dummy 1-element buffers. + # The placeholders are never dereferenced by the kernel, so reuse one cached + # buffer per (device, dtype) instead of re-allocating each call. if needs_bias: Bias_arg = bias.contiguous().view(-1) else: - Bias_arg = torch.zeros(1, dtype=torch.float32, device=A.device) + Bias_arg = _dense_placeholder(A.device, torch.float32) if needs_aux: Aux_arg = aux.view(-1) else: - Aux_arg = torch.zeros(1, dtype=C.dtype, device=A.device) + Aux_arg = _dense_placeholder(A.device, C.dtype) # Grouped runtime tensors are required by the kernel signature; the dense # binary never dereferences them but still needs placeholder buffers. @@ -1370,8 +1374,8 @@ def doGemm( work_queue_arg = work_queue.contiguous().view(-1) work_counter_arg = work_counter.contiguous().view(-1) else: - work_queue_arg = torch.zeros(1, dtype=torch.int32, device=A.device) - work_counter_arg = torch.zeros(1, dtype=torch.int32, device=A.device) + work_queue_arg = _dense_placeholder(A.device, torch.int32) + work_counter_arg = _dense_placeholder(A.device, torch.int32) # Per-expert weight base pointers (grouped forward and dgrad -- both run on this # TN kernel, dgrad over host-transposed weights): compile a kernel that reads each @@ -1382,9 +1386,10 @@ def doGemm( if weight_ptrs: b_data_ptrs_arg = b_data_ptrs.contiguous().view(-1) else: - b_data_ptrs_arg = torch.zeros(1, dtype=torch.int64, device=A.device) + b_data_ptrs_arg = _dense_placeholder(A.device, torch.int64) - launch = _cached_launch( + launch_key = ( + "fp8_gemm", int(K_runtime), A.dtype, B.dtype, @@ -1396,7 +1401,7 @@ def doGemm( wgrad, weight_ptrs, ) - launch( + launch_args = ( A_arg, B_arg, C_arg, @@ -1410,5 +1415,24 @@ def doGemm( work_queue_arg, work_counter_arg, b_data_ptrs_arg, - stream=stream, + stream, ) + # flyc.compile fast path: skips flydsl's per-call cache-key rebuild (sig.bind + + # globals snapshot + arg resolve), the dominant host cost at small M. M/N are + # runtime scalars, so one compiled launcher per launch_key serves all shapes. + _compiled_launch( + launch_key, + _cached_launch( + int(K_runtime), + A.dtype, + B.dtype, + C.dtype, + bool(use_xcd_remap), + epilogue, + group, + sched, + wgrad, + weight_ptrs, + ), + launch_args, + )(*launch_args) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py index 1b660b42f8..3e3ab949db 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py @@ -303,7 +303,18 @@ def uniform_pad_concat(tensors, rows): Used by the wgrad uniform-pad path to stack per-expert operands (or their scales) on the contraction/token axis. Padded rows are zero, so they contribute nothing to the token-axis contraction. + + When no tensor actually needs padding (every leading dim already equals + ``rows`` -- the balanced / no-skew case, where each expert's token count is a + multiple of the bucket's BLOCK_M), there is nothing to zero-fill, so this + degenerates to a plain concat. Delegate that case to ``concat_or_view_rows``, + which returns a zero-copy spanning view when the per-expert payloads are + physically adjacent (the common bulk-allocated layout), eliminating the + ``torch.cat`` copy on the wgrad hot path. """ + if all(t.size(0) == rows for t in tensors): + return concat_or_view_rows(tensors) + def _pad_rows(t): if t.size(0) == rows: return t.contiguous() diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py index 171284afd9..2cf3a98fe4 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py @@ -1412,6 +1412,48 @@ def _cached_launch( ) +@functools.lru_cache(maxsize=None) +def _dense_placeholder(device: torch.device, dtype: torch.dtype, numel: int = 1): + """One cached 1-element buffer per (device, dtype) for dense kernel-signature + placeholders. + + The dense half-precision binary requires the grouped/epilogue tensor arguments + (Bias/Aux/work_queue/work_counter/b_data_ptrs) to be present but never + dereferences them. Re-allocating them via ``torch.zeros`` on every call is pure + host overhead that dominates wall time at small M; since the contents are never + read, a single shared read-only buffer per (device, dtype) is safe to reuse. + """ + return torch.zeros(numel, dtype=dtype, device=device) + + +# Per-key cache of flyc.compile'd launchers. flydsl's JitFunction.__call__ rebuilds +# a full cache key (sig.bind + globals snapshot + arg-signature resolve) on every +# launch -- ~190us/call, which dominates wall time at small M where the kernel is +# short. flyc.compile returns a CompiledFunction whose __call__ skips all of that +# (ctypes pointer/scalar update + invoke, ~5us), baking constexpr values at compile +# time so one CompiledFunction per launch_key serves every pointer/scalar/stream +# variation. M/N are runtime scalar args (not baked), so the key excludes them -- +# identical to the MXFP8/FP8 paths. +# NOTE: the fast path intentionally bypasses flydsl's globals-drift safety check; +# this module never mutates its captured globals between launches, so it is safe. +_COMPILED_LAUNCH_CACHE = {} + + +def _compiled_launch(key, jit_fn, args): + """Return (and cache) a flyc.compile'd fast launcher for ``jit_fn``. + + ``key`` uniquely identifies the compiled specialization (the same tuple the + JitFunction builder's lru_cache uses). ``args`` is the full positional arg list + for the first call; used once to compile, then the cached CompiledFunction is + reused positionally for all later calls with the same ABI. + """ + cf = _COMPILED_LAUNCH_CACHE.get(key) + if cf is None: + cf = flyc.compile(jit_fn, *args) + _COMPILED_LAUNCH_CACHE[key] = cf + return cf + + def _half_prec_matmul( a: torch.Tensor, b: torch.Tensor, @@ -1432,17 +1474,17 @@ def _half_prec_matmul( ``input_dtype`` is the required FP16/BF16 operand dtype and ``label`` is the human-readable kernel name used in error messages. + + Only the checks ``doGemm`` does not repeat are performed here: the + layout-specific physical operand shapes, operand/output contiguity, and + device agreement. Layout validity, operand dtype, output dtype, and tiling + are validated by ``doGemm`` against the runtime geometry, so they are not + duplicated in this shim. """ - if layout not in ("TN", "NN", "NT"): - raise ValueError(f"Unsupported {label} layout: {layout}") if a.ndim != 2 or b.ndim != 2: raise ValueError( f"FlyDSL {label} expects rank-2 operands, got A{tuple(a.shape)} and B{tuple(b.shape)}" ) - if a.dtype != input_dtype or b.dtype != input_dtype: - raise TypeError( - f"FlyDSL {label} GEMM expects {input_dtype} operands, got A={a.dtype}, B={b.dtype}" - ) if not a.is_contiguous() or not b.is_contiguous(): raise FlyDSLUnsupportedError( f"FlyDSL {label} {layout} requires original contiguous row-major " @@ -1459,6 +1501,8 @@ def _half_prec_matmul( "NN": ((m, k), (k, n)), "NT": ((k, m), (k, n)), } + if layout not in expected_shapes: + raise ValueError(f"Unsupported {label} layout: {layout}") expected_a, expected_b = expected_shapes[layout] if tuple(a.shape) != expected_a or tuple(b.shape) != expected_b: raise ValueError( @@ -1469,10 +1513,6 @@ def _half_prec_matmul( if tuple(c.shape) != (m, n): raise ValueError(f"C shape {tuple(c.shape)} != expected {(m, n)}") - if c.dtype not in (torch.float16, torch.bfloat16, torch.float32): - raise TypeError( - f"FlyDSL {label} output must be float16, bfloat16, or float32, got {c.dtype}" - ) if a.device != b.device or a.device != c.device: raise ValueError( f"A, B, and C must be on the same device, got {a.device}, {b.device}, and {c.device}" @@ -1662,17 +1702,6 @@ def doGemm( if stream is None: stream = torch.cuda.current_stream() - launch = _cached_launch( - K_runtime, - C.dtype, - layout, - _MFMA_SUFFIX[input_dtype], - bool(use_xcd_remap), - epilogue, - group, - sched, - wgrad, - ) # Preserve the original validated byte-addressed G2L path. These are # metadata-only dtype/flatten views of the already-contiguous row-major # tensors selected by gemm_wrappers.py; no transpose or copy is performed. @@ -1680,14 +1709,16 @@ def doGemm( B_arg = B.view(torch.uint8).view(-1) C_arg = C.view(-1) # DEFAULT keeps the kernel signature uniform with dummy 1-element buffers. + # The placeholders are never dereferenced by the kernel, so reuse one cached + # buffer per (device, dtype) instead of re-allocating each call. if needs_bias: Bias_arg = bias.contiguous().view(-1) else: - Bias_arg = torch.zeros(1, dtype=torch.float32, device=A.device) + Bias_arg = _dense_placeholder(A.device, torch.float32) if needs_aux: Aux_arg = aux.view(-1) else: - Aux_arg = torch.zeros(1, dtype=C.dtype, device=A.device) + Aux_arg = _dense_placeholder(A.device, C.dtype) # Grouped runtime tensors are required by the kernel signature; the dense # binary never dereferences them but still needs placeholder buffers. @@ -1695,8 +1726,8 @@ def doGemm( work_queue_arg = work_queue.contiguous().view(-1) work_counter_arg = work_counter.contiguous().view(-1) else: - work_queue_arg = torch.zeros(1, dtype=torch.int32, device=A.device) - work_counter_arg = torch.zeros(1, dtype=torch.int32, device=A.device) + work_queue_arg = _dense_placeholder(A.device, torch.int32) + work_counter_arg = _dense_placeholder(A.device, torch.int32) # Per-expert weight base pointers (grouped TN only): an int64[G] device array of # each expert's weight-slab base address, read in-kernel so the weight operand @@ -1704,9 +1735,21 @@ def doGemm( if b_data_ptrs is not None: b_data_ptrs_arg = b_data_ptrs.contiguous().view(-1) else: - b_data_ptrs_arg = torch.zeros(1, dtype=torch.int64, device=A.device) + b_data_ptrs_arg = _dense_placeholder(A.device, torch.int64) - launch( + launch_key = ( + "half_prec_gemm", + K_runtime, + C.dtype, + layout, + _MFMA_SUFFIX[input_dtype], + bool(use_xcd_remap), + epilogue, + group, + sched, + wgrad, + ) + launch_args = ( A_arg, B_arg, C_arg, @@ -1718,5 +1761,23 @@ def doGemm( work_queue_arg, work_counter_arg, b_data_ptrs_arg, - stream=stream, + stream, ) + # flyc.compile fast path: skips flydsl's per-call cache-key rebuild (sig.bind + + # globals snapshot + arg resolve), the dominant host cost at small M. M/N are + # runtime scalars, so one compiled launcher per launch_key serves all shapes. + _compiled_launch( + launch_key, + _cached_launch( + K_runtime, + C.dtype, + layout, + _MFMA_SUFFIX[input_dtype], + bool(use_xcd_remap), + epilogue, + group, + sched, + wgrad, + ), + launch_args, + )(*launch_args) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 514be35033..37bf6a8a2e 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -89,6 +89,34 @@ from transformer_engine.pytorch.triton_kernels.grouped_gemm import general_grouped_gemm_triton import os + +# Cache of the per-expert weight base-pointer device tensors used by the FlyDSL +# grouped GEMM (one int64[G] array of each expert's data_ptr()). Rebuilding it +# every call -- torch.tensor([...], device=cuda) -- costs a blocking H2D copy that +# is the single largest host cost on the grouped forward/dgrad path. Expert weight +# storages are stable across calls, so key the cache on the storage pointers and +# only re-upload when one actually moves (reallocation / first use). +_B_DATA_PTRS_CACHE = {} + + +def _cached_b_data_ptrs(datas, device): + """Return (and cache) an int64[G] device tensor of ``t.data_ptr()`` for ``datas``. + + ``datas`` is the list of per-expert operand tensors (e.g. each weight's rowwise + or columnwise payload). ``data_ptr()`` reads are host-only (no sync); the + expensive H2D ``torch.tensor`` upload runs only on a cache miss. Keyed on the + pointer tuple itself, so stable weight storage (the common case) reuses one + device tensor across every call and the cache stays bounded to the distinct + weight sets actually seen. + """ + key = (tuple(t.data_ptr() for t in datas), device) + tensor = _B_DATA_PTRS_CACHE.get(key) + if tensor is None: + tensor = torch.tensor(key[0], dtype=torch.int64, device=device) + _B_DATA_PTRS_CACHE[key] = tensor + return tensor + + _NATIVE_SPLIT_QUANTIZER_TYPES = frozenset( { Float8Quantizer, @@ -1225,10 +1253,8 @@ def _grouped_gemm_flydsl_mxfp8( # weights into one packed buffer (the dominant host cost). b_data is a # per-expert representative used only for dtype/shape. The B-scale stays # packed (cheap) and keeps the in-kernel packed-column fold. - b_data_ptrs = torch.tensor( - [w._rowwise_data.data_ptr() for w in weights], - dtype=torch.int64, - device=a_data.device, + b_data_ptrs = _cached_b_data_ptrs( + [w._rowwise_data for w in weights], a_data.device ) b_data = weights[0]._rowwise_data b_scale = concat_or_view_rows([w._rowwise_scale_inv for w in weights]) @@ -1283,10 +1309,8 @@ def _backward_flydsl_mxfp8_grouped( # slabs addressed via per-expert base pointers (no concat); dY still concats. a_data = concat_or_view_rows([g._rowwise_data for g in grad_output]) a_scale = concat_or_view_rows([g._rowwise_scale_inv for g in grad_output]) - b_data_ptrs = torch.tensor( - [w._columnwise_data.data_ptr() for w in weights], - dtype=torch.int64, - device=a_data.device, + b_data_ptrs = _cached_b_data_ptrs( + [w._columnwise_data for w in weights], a_data.device ) b_data = weights[0]._columnwise_data b_scale = concat_or_view_rows([w._columnwise_scale_inv for w in weights]) @@ -1380,11 +1404,7 @@ def _grouped_gemm_flydsl_bf16(*, inputmats, weights, out, m_splits, m_splits_ten # reads each from a per-expert base pointer, so pass an int64[G] address array # instead of concatenating the G weights. b_data is a per-expert representative # (shape only; the kernel addresses weights via the pointers). - b_data_ptrs = torch.tensor( - [w.data_ptr() for w in weights], - dtype=torch.int64, - device=a_data.device, - ) + b_data_ptrs = _cached_b_data_ptrs(list(weights), a_data.device) b_data = weights[0] n = weights[0].size(0) k = weights[0].size(-1) @@ -1412,11 +1432,7 @@ def _backward_flydsl_bf16_grouped(*, grad_output, weights, dgrad, m_splits, m_sp # Weights are independent per-expert [N_out, K_in] tensors; the grouped-NN # kernel reads each from a per-expert base pointer (no concat). b_data is a # per-expert representative. - b_data_ptrs = torch.tensor( - [w.data_ptr() for w in weights], - dtype=torch.int64, - device=a_data.device, - ) + b_data_ptrs = _cached_b_data_ptrs(list(weights), a_data.device) b_data = weights[0] n_out = grad_output[0].shape[-1] # per-expert contraction (N_out) k_in = dgrad.size(-1) # output width (K_in) @@ -1475,10 +1491,8 @@ def _grouped_gemm_flydsl_fp8(*, inputmats, weights, out, m_splits, m_splits_tens # Weights are independent per-expert [N, K] FP8 tensors; the grouped-TN kernel # reads each from a per-expert base pointer, so pass an int64[G] address array # instead of concatenating the G weights. b_data is a per-expert representative. - b_data_ptrs = torch.tensor( - [w._data.data_ptr() for w in weights], - dtype=torch.int64, - device=a_data.device, + b_data_ptrs = _cached_b_data_ptrs( + [w._data for w in weights], a_data.device ) b_data = weights[0]._data a_data._fp8_dtype = inputmats[0]._fp8_dtype @@ -1523,10 +1537,8 @@ def _backward_flydsl_fp8_grouped( # backings are independent per-expert slabs, addressed via per-expert base # pointers through the same TN per-expert path (no concat); dY still concats. a_data = concat_or_view_rows([g._data for g in grad_output]) - b_data_ptrs = torch.tensor( - [w._transpose.data_ptr() for w in weights], - dtype=torch.int64, - device=a_data.device, + b_data_ptrs = _cached_b_data_ptrs( + [w._transpose for w in weights], a_data.device ) b_data = weights[0]._transpose a_data._fp8_dtype = grad_output[0]._fp8_dtype @@ -1619,6 +1631,39 @@ def _flydsl_grouped_dispatch(cls, phase, kind): } return table[phase][kind] + @classmethod + def _try_flydsl_grouped(cls, phase, kind, **kwargs): + """Run the FlyDSL grouped helper for ``phase``/``kind``; report if it ran. + + Returns True when FlyDSL handled the GEMM. A ``FlyDSLUnsupportedError`` + (e.g. a per-expert dimension the kernel tiling cannot cover) means the + config is unsupported, so the helper returns False and the caller runs the + default grouped GEMM instead -- mirroring the dense GEMM fallback in + ``cpp_extensions.gemm``. ``NVTE_FLYDSL_GEMM_WARN_FALLBACK`` warns on fallback. + """ + from ..flydsl_kernels.gemm import FlyDSLUnsupportedError + + try: + cls._flydsl_grouped_dispatch(phase, kind)(**kwargs) + return True + except FlyDSLUnsupportedError as exc: + warn_fallback = os.environ.get("NVTE_FLYDSL_GEMM_WARN_FALLBACK", "0").lower() not in ( + "", + "0", + "false", + "no", + "off", + ) + if warn_fallback: + warnings.warn( + "[FLYDSL WARNING]: FlyDSL grouped GEMM does not support this " + f"configuration ({phase}); falling back to the default backend. " + f"Reason: {exc}", + UserWarning, + stacklevel=2, + ) + return False + @staticmethod def _forward_blockwise_fp8_triton( ctx, @@ -2281,15 +2326,16 @@ def forward( use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator # Perform GEMM - if use_grouped_gemm_flydsl: - _GroupedLinear._flydsl_grouped_dispatch("fwd", flydsl_grouped_kind)( - inputmats=inputmats, - weights=weights_fp8, - out=out, - m_splits=m_splits, - m_splits_tensor=m_splits_tensor, - ) - else: + ran_flydsl_fwd = use_grouped_gemm_flydsl and _GroupedLinear._try_flydsl_grouped( + "fwd", + flydsl_grouped_kind, + inputmats=inputmats, + weights=weights_fp8, + out=out, + m_splits=m_splits, + m_splits_tensor=m_splits_tensor, + ) + if not ran_flydsl_fwd: if use_grouped_gemm_triton: general_grouped_gemm_func = general_grouped_gemm_triton kwargs = {"m_splits_tensor": m_splits_tensor} @@ -2896,18 +2942,20 @@ def backward( columnwise_usage=quantizer.columnwise_usage, ) _flydsl_kind = getattr(ctx, "flydsl_grouped_kind", None) - if _flydsl_kind is not None: - # dgrad dX = dY @ W through the FlyDSL grouped NN kernel, riding - # off the already-quantized grad_output (rowwise) and weights - # (columnwise, ensured above). - _GroupedLinear._flydsl_grouped_dispatch("dgrad", _flydsl_kind)( - grad_output=grad_output, - weights=weights_for_dgrad, - dgrad=dgrad, - m_splits=ctx.m_splits, - m_splits_tensor=ctx.m_splits_tensor, - ) - else: + # dgrad dX = dY @ W through the FlyDSL grouped NN kernel, riding off + # the already-quantized grad_output (rowwise) and weights (columnwise, + # ensured above). Falls back to the default backend when FlyDSL does + # not support this config. + ran_flydsl_dgrad = _flydsl_kind is not None and _GroupedLinear._try_flydsl_grouped( + "dgrad", + _flydsl_kind, + grad_output=grad_output, + weights=weights_for_dgrad, + dgrad=dgrad, + m_splits=ctx.m_splits, + m_splits_tensor=ctx.m_splits_tensor, + ) + if not ran_flydsl_dgrad: if ctx.use_grouped_gemm_triton: general_grouped_gemm_func = general_grouped_gemm_triton kwargs = {"m_splits_tensor": ctx.m_splits_tensor} @@ -3050,13 +3098,15 @@ def backward( ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute() ) ) - if use_flydsl_wgrad: - _GroupedLinear._flydsl_grouped_dispatch("wgrad", _flydsl_kind)( - grad_output=grad_output, - inputmats=inputmats, - wgrad_list=wgrad_list, - m_splits=ctx.m_splits, - ) + ran_flydsl_wgrad = use_flydsl_wgrad and _GroupedLinear._try_flydsl_grouped( + "wgrad", + _flydsl_kind, + grad_output=grad_output, + inputmats=inputmats, + wgrad_list=wgrad_list, + m_splits=ctx.m_splits, + ) + if ran_flydsl_wgrad: clear_tensor_data(*inputmats) grouped_gemm_wgrad = None else: From 3112c12e4535b0b3453bf7ef704581151fcc2685 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Sat, 17 Oct 2026 19:13:10 +0000 Subject: [PATCH 11/22] Hoist _SpanCAI --- .../flydsl_kernels/gemm/grouped_gemm_utils.py | 37 ++++++++++++------- 1 file changed, 24 insertions(+), 13 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py index 3e3ab949db..c2ba1dc9df 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py @@ -219,6 +219,27 @@ def build_wgrad_work_queue(num_groups, num_pid_m, num_pid_n, device): return work_queue, work_counter, num_tiles +class _SpanCAI: + """Minimal ``__cuda_array_interface__`` holder for a raw uint8 device span. + + Wraps ``total_bytes`` at ``data_ptr`` as a 1-D uint8 CAI so ``torch.as_tensor`` + can build a zero-copy view over a contiguous multi-tensor run (see + :func:`concat_or_view_rows`). Hoisted to module scope so the class object is + built once at import, not reconstructed on every call. + """ + + __slots__ = ("__cuda_array_interface__",) + + def __init__(self, total_bytes, data_ptr): + self.__cuda_array_interface__ = { + "shape": (total_bytes,), + "typestr": "|u1", + "data": (data_ptr, False), + "strides": None, + "version": 3, + } + + def concat_or_view_rows(tensors): """Concatenate per-expert tensors on dim 0, skipping the copy when possible. @@ -279,19 +300,9 @@ def concat_or_view_rows(tensors): row_bytes *= elem total_bytes = total_rows * row_bytes - class _SpanCAI: - __slots__ = ("__cuda_array_interface__",) - - def __init__(self): - self.__cuda_array_interface__ = { - "shape": (total_bytes,), - "typestr": "|u1", - "data": (first.data_ptr(), False), - "strides": None, - "version": 3, - } - - byte_view = torch.as_tensor(_SpanCAI(), device=first.device) + byte_view = torch.as_tensor( + _SpanCAI(total_bytes, first.data_ptr()), device=first.device + ) view = byte_view.view(first.dtype).view(total_rows, *tuple(trailing)) view._flydsl_span_anchor = tensors return view From d290583418d7ff474e806b555091079deb2bb802 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Sun, 18 Oct 2026 02:44:01 +0000 Subject: [PATCH 12/22] Write FlyDSL FP8 grouped wgrad output in place, skip the copy-out loop --- .../pytorch/module/grouped_linear.py | 47 +++++++++++++++++-- 1 file changed, 42 insertions(+), 5 deletions(-) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 37bf6a8a2e..5024ce0e76 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -117,6 +117,33 @@ def _cached_b_data_ptrs(datas, device): return tensor +def _packed_wgrad_view(wgrad_list, num_gemms, n_out, k_in): + """Return a [num_gemms*n_out, k_in] view over ``wgrad_list`` if it is one packed run. + + The non-fused grouped-wgrad destinations are ``[wgrad_packed[g] for g in range]`` + -- adjacent ``[n_out, k_in]`` slices of a single contiguous ``[num_gemms, n_out, + k_in]`` buffer. When that layout holds, the grouped wgrad kernel can write its + ``[num_gemms*n_out, k_in]`` output directly into that storage, eliding the + per-expert copy-out. Return ``None`` (caller allocates a scratch buffer and + scatters) whenever the destinations are not a single contiguous, row-contiguous, + correctly-shaped run -- e.g. fused main_grad accumulation into separate params. + """ + first = wgrad_list[0] + elem = first.element_size() + expected_ptr = first.data_ptr() + for w in wgrad_list: + if ( + w.dtype != first.dtype + or tuple(w.shape) != (n_out, k_in) + or not w.is_contiguous() + or w.data_ptr() != expected_ptr + ): + return None + expected_ptr += w.numel() * elem + # One contiguous run anchored at first.data_ptr(); span all rows over its storage. + return torch.as_strided(first, (num_gemms * n_out, k_in), (k_in, 1)) + + _NATIVE_SPLIT_QUANTIZER_TYPES = frozenset( { Float8Quantizer, @@ -1585,9 +1612,18 @@ def _wgrad_flydsl_fp8_grouped(*, grad_output, inputmats, wgrad_list, m_splits): n_out = dyt_cols[0].size(0) # per-expert output feature N k_in = xt_cols[0].size(0) # per-expert input feature K - dW = torch.empty( - num_gemms * n_out, k_in, dtype=wgrad_list[0].dtype, device=wgrad_list[0].device - ) + # wgrad_list[g] are normally adjacent [n_out, k_in] slices of one packed + # [num_gemms, n_out, k_in] buffer (the non-fused ``wgrad_packed`` path). When + # that holds, point the kernel's [num_gemms*n_out, k_in] output straight at + # that storage so it writes each dW_e in place -- eliding the per-expert + # copy-out scatter. Fall back to a fresh buffer + scatter only when the + # destinations are not one contiguous run (e.g. fused main_grad accumulation). + dW = _packed_wgrad_view(wgrad_list, num_gemms, n_out, k_in) + scatter = dW is None + if scatter: + dW = torch.empty( + num_gemms * n_out, k_in, dtype=wgrad_list[0].dtype, device=wgrad_list[0].device + ) _run_fp8_grouped_wgrad( dyt_cols, dy_scale, @@ -1601,8 +1637,9 @@ def _wgrad_flydsl_fp8_grouped(*, grad_output, inputmats, wgrad_list, m_splits): x_fp8_dtype=inputmats[0]._fp8_dtype, output_dtype=wgrad_list[0].dtype, ) - for g in range(num_gemms): - wgrad_list[g].copy_(dW[g * n_out : (g + 1) * n_out]) + if scatter: + for g in range(num_gemms): + wgrad_list[g].copy_(dW[g * n_out : (g + 1) * n_out]) @classmethod def _flydsl_grouped_dispatch(cls, phase, kind): From 071d251c4dddd8bbeae0633cbcb62bb72cadc966 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 20 Oct 2026 23:49:48 +0000 Subject: [PATCH 13/22] Read FlyDSL FP8 grouped operands and scales in place via per-expert pointers --- .../pytorch/flydsl_kernels/gemm/fp8_gemm.py | 96 +++++++++++++++---- .../flydsl_kernels/gemm/gemm_wrappers.py | 79 ++++++++++++--- .../flydsl_kernels/gemm/grouped_gemm_utils.py | 50 ++++++++++ .../pytorch/module/grouped_linear.py | 69 +++++++++---- 4 files changed, 245 insertions(+), 49 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py index 4b93e52bd6..d446299d6e 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -100,6 +100,8 @@ def _compile_kernel( sched: str = "worksteal", wgrad: bool = False, weight_ptrs: bool = False, + act_ptrs: bool = False, + scale_ptrs: bool = False, ): """Build the specialized kernel for compile-time K, A/B FP8 types, and output dtype. @@ -218,6 +220,7 @@ def kernel_gemm( work_queue: fx.Tensor = None, work_counter: fx.Tensor = None, b_data_ptrs: fx.Tensor = None, + a_data_ptrs: fx.Tensor = None, ): lds = fx.SharedAllocator().allocate(SharedStorage).peek() lds_a0 = (lds.a0_0, lds.a0_1) @@ -231,10 +234,14 @@ def kernel_gemm( gB = make_fp8_buffer_tensor(B, b_f8_ir_t) a_div = fx.logical_divide(gA, fx.make_layout(1, 1)) b_div = fx.logical_divide(gB, fx.make_layout(1, 1)) - a_scale_rsrc = fx.rocdl.make_buffer_tensor(A_scale_inv, max_size=True) - b_scale_rsrc = fx.rocdl.make_buffer_tensor(B_scale_inv, max_size=True) - a_scale_div = fx.logical_divide(a_scale_rsrc, fx.make_layout(1, 1)) - b_scale_div = fx.logical_divide(b_scale_rsrc, fx.make_layout(1, 1)) + # scale_ptrs (grouped only): A_scale_inv / B_scale_inv are int64[G] arrays of + # per-expert scale_inv pointers, read in place per tile instead of from a + # host-gathered length-G vector. + if const_expr(not scale_ptrs): + a_scale_rsrc = fx.rocdl.make_buffer_tensor(A_scale_inv, max_size=True) + b_scale_rsrc = fx.rocdl.make_buffer_tensor(B_scale_inv, max_size=True) + a_scale_div = fx.logical_divide(a_scale_rsrc, fx.make_layout(1, 1)) + b_scale_div = fx.logical_divide(b_scale_rsrc, fx.make_layout(1, 1)) scale_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) def _load_scale(scale_div, index): @@ -261,10 +268,17 @@ def _load_scale(scale_div, index): # addressed from a per-expert base pointer read from ``b_data_ptrs`` (reusing # gB's static layout) instead of an expert-shifted B row. FP8 is TN-only, so # both grouped forward and dgrad (dgrad feeds host-transposed weights through - # this same TN kernel) set this compile-time flag; wgrad keeps the packed-B - # fold. + # this same TN kernel) set this compile-time flag. wgrad sets it only with + # act_ptrs, for its per-expert X^T operand. if const_expr(weight_ptrs): b_ptr_layout = fx.get_layout(gB) + # act_ptrs: the per-expert activation operand (A; and in wgrad also B, via + # b_data_ptrs) stays in its own per-expert allocation instead of a host-side + # concat. The host passes each expert's base pointer pre-shifted by that + # expert's packed offset, so the packed addressing below is unchanged; only the + # buffer base moves per tile. + if const_expr(act_ptrs): + a_ptr_layout = fx.get_layout(gA) if const_expr(grouped): wq_rsrc = fx.rocdl.make_buffer_tensor(work_queue, max_size=True) @@ -290,13 +304,20 @@ def _desc_field(slot, field): return gd.read_desc_field(wq_div, desc_ld_atom, slot, field) def process_tile(pid_m, pid_n, expert): - # Per-expert scalar scale: dense reads index 0 (length-1 buffer); - # grouped reads scale[expert] from the length-num_groups vector. - if const_expr(grouped): - scale_idx = fx.Int32(expert) + # Per-expert scalar scale: dense reads index 0 (length-1 buffer); grouped + # reads scale[expert] from the length-num_groups vector, or through the + # expert's scale_inv pointer (scale_ptrs). + if const_expr(scale_ptrs): + a_scale = fx.Float32(gd.load_f32_at_ptr(A_scale_inv, expert)) + b_scale = fx.Float32(gd.load_f32_at_ptr(B_scale_inv, expert)) else: - scale_idx = fx.Int32(0) - output_scale = _load_scale(a_scale_div, scale_idx) * _load_scale(b_scale_div, scale_idx) + if const_expr(grouped): + scale_idx = fx.Int32(expert) + else: + scale_idx = fx.Int32(0) + a_scale = _load_scale(a_scale_div, scale_idx) + b_scale = _load_scale(b_scale_div, scale_idx) + output_scale = a_scale * b_scale bx_m = pid_m * BLOCK_M by_n = pid_n * BLOCK_N @@ -346,14 +367,31 @@ def process_tile(pid_m, pid_n, expert): # the global K coordinate is XOR-unswizzled for the physical LDS slot. gl_off_a = compute_global_swizzle(lane, wave_id, K, LOAD_PASSES_HALF, preshuffled=False) gl_off_b = compute_global_swizzle(lane, wave_id, K, LOAD_PASSES_HALF, preshuffled=False) - a_g2s = G2SLoader(a_div, gl_off_a, LOAD_PASSES_HALF, a_f8_ir_t, wave_id) + # Issue both per-expert base-pointer loads before either is made uniform, so a + # tile waits on one memory round trip rather than two. + if const_expr(act_ptrs): + a_base_raw = gd.read_ptr_at(a_data_ptrs, expert) + if const_expr(weight_ptrs): + b_base_raw = gd.read_ptr_at(b_data_ptrs, expert) + if const_expr(act_ptrs): + a_base_i64 = gd.make_uniform_i64(a_base_raw) + a_div_tile = gd.rebase_buffer_view(a_base_i64, a_f8_ir_t, a_ptr_layout) + else: + a_div_tile = a_div + a_g2s = G2SLoader(a_div_tile, gl_off_a, LOAD_PASSES_HALF, a_f8_ir_t, wave_id) # Grouped-TN forward B-data: rebase onto this tile's expert weight slab and # read with a LOCAL column (base selects the expert). Other paths keep the # shared packed b_div / by_n_b_idx. if const_expr(weight_ptrs): - b_base_i64 = gd.read_ptr_at(b_data_ptrs, expert) + b_base_i64 = gd.make_uniform_i64(b_base_raw) b_div_tile = gd.rebase_buffer_view(b_base_i64, b_f8_ir_t, b_ptr_layout) - b_data_col = by_n_idx + # wgrad bases are pre-shifted by the packed expert offset (see act_ptrs), + # so keep the packed column there; forward/dgrad weight bases are the + # real per-expert slabs, read with a local column. + if const_expr(wgrad): + b_data_col = by_n_b_idx + else: + b_data_col = by_n_idx else: b_div_tile = b_div b_data_col = by_n_b_idx @@ -1111,6 +1149,7 @@ def launch_gemm( work_queue: fx.Tensor = None, work_counter: fx.Tensor = None, b_data_ptrs: fx.Tensor = None, + a_data_ptrs: fx.Tensor = None, stream: fx.Stream = fx.Stream(None), ): # Dense: one workgroup per output tile. Grouped/wgrad: a fixed num-CU @@ -1133,6 +1172,7 @@ def launch_gemm( work_queue, work_counter, b_data_ptrs, + a_data_ptrs, value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, ).launch(grid=(grid_x, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) @@ -1194,6 +1234,8 @@ def _cached_launch( sched: str = "worksteal", wgrad: bool = False, weight_ptrs: bool = False, + act_ptrs: bool = False, + scale_ptrs: bool = False, ): return _compile_kernel( K, @@ -1206,6 +1248,8 @@ def _cached_launch( sched=sched, wgrad=wgrad, weight_ptrs=weight_ptrs, + act_ptrs=act_ptrs, + scale_ptrs=scale_ptrs, ) @@ -1269,6 +1313,7 @@ def doGemm( work_queue: torch.Tensor = None, work_counter: torch.Tensor = None, b_data_ptrs: torch.Tensor = None, + a_data_ptrs: torch.Tensor = None, sched: str = "worksteal", wgrad: bool = False, m: int = None, @@ -1308,10 +1353,14 @@ def doGemm( label="FP8 GEMM", ) require_launch_size("FP8 GEMM", ("A", A), ("B", B), ("C", C)) - # Dense: one scalar inverse scale per operand. Grouped: a per-expert vector. + # Dense: one scalar inverse scale per operand. Grouped: a per-expert vector, or + # (scale_ptrs) an int64[G] array of per-expert scale_inv pointers read in place. expected_scale_numel = num_groups if (group or wgrad) else 1 - assert A_scale_inv.dtype == torch.float32 and A_scale_inv.numel() == expected_scale_numel - assert B_scale_inv.dtype == torch.float32 and B_scale_inv.numel() == expected_scale_numel + scale_ptrs = A_scale_inv.dtype == torch.int64 + expected_scale_dtype = torch.int64 if scale_ptrs else torch.float32 + assert not scale_ptrs or group or wgrad, "per-expert scale pointers are grouped-only" + assert A_scale_inv.dtype == expected_scale_dtype and A_scale_inv.numel() == expected_scale_numel + assert B_scale_inv.dtype == expected_scale_dtype and B_scale_inv.numel() == expected_scale_numel # wgrad's C is packed [G*N_out, K_in] on the output-M axis; fwd/dgrad/dense C is # the unpacked (M_runtime, N_runtime). expected_c = (num_groups * M_runtime, N_runtime) if wgrad else (M_runtime, N_runtime) @@ -1387,6 +1436,12 @@ def doGemm( b_data_ptrs_arg = b_data_ptrs.contiguous().view(-1) else: b_data_ptrs_arg = _dense_placeholder(A.device, torch.int64) + # Per-expert activation base pointers (see act_ptrs in _compile_kernel). + act_ptrs = a_data_ptrs is not None + if act_ptrs: + a_data_ptrs_arg = a_data_ptrs.contiguous().view(-1) + else: + a_data_ptrs_arg = _dense_placeholder(A.device, torch.int64) launch_key = ( "fp8_gemm", @@ -1400,6 +1455,8 @@ def doGemm( sched, wgrad, weight_ptrs, + act_ptrs, + scale_ptrs, ) launch_args = ( A_arg, @@ -1415,6 +1472,7 @@ def doGemm( work_queue_arg, work_counter_arg, b_data_ptrs_arg, + a_data_ptrs_arg, stream, ) # flyc.compile fast path: skips flydsl's per-call cache-key rebuild (sig.bind + @@ -1433,6 +1491,8 @@ def doGemm( sched, wgrad, weight_ptrs, + act_ptrs, + scale_ptrs, ), launch_args, )(*launch_args) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index c61ce22f30..cfe7af9e64 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -1418,11 +1418,19 @@ def _run_bf16_grouped_wgrad( return dW +def _fp8_scale_arg(scale): + # int64 = per-expert scale_inv pointer array, read in place by the kernel; otherwise + # a length-G fp32 vector of the per-expert scales. + if scale.dtype == torch.int64: + return scale + return scale.to(torch.float32).contiguous() + + def _run_fp8_grouped( a_data, - a_scale_vec, # [num_groups] fp32 per-expert inverse scales + a_scale_vec, # [num_groups] fp32 per-expert inverse scales, or int64 scale_inv pointers b_data, - b_scale_vec, # [num_groups] fp32 per-expert inverse scales + b_scale_vec, # [num_groups] fp32 per-expert inverse scales, or int64 scale_inv pointers D, *, m_splits, @@ -1432,15 +1440,20 @@ def _run_fp8_grouped( output_dtype: torch.dtype, bias=None, b_data_ptrs=None, + a_data_ptrs=None, ): """Dispatch a packed tensor-wise FP8 grouped GEMM (TN) through the persistent kernel. Tensor-wise FP8 is TN-only; grouped dgrad/wgrad reach this same TN kernel via host-materialized transpose backings (see the backward helpers). Experts pack on B's output-N axis ([G*N, K]); each expert carries one scalar inverse scale, so the - scales arrive as length-``num_groups`` vectors the kernel indexes by ``expert``. + scales arrive as length-``num_groups`` vectors the kernel indexes by ``expert``, or + as int64 arrays of per-expert ``scale_inv`` pointers it reads in place. ``m_splits`` is the host-side per-expert token list; the descriptor queue is built - + shape-cached from it (no GPU->CPU sync). + + shape-cached from it (no GPU->CPU sync). ``a_data_ptrs`` (optional) reads A per + expert in place: ``a_data`` is then one representative expert payload and each + entry is that expert's base shifted by its packed row start (see + ``grouped_gemm_utils.shifted_expert_ptrs``). """ from .fp8_gemm import doGemm as _fp8_doGemm @@ -1479,8 +1492,8 @@ def _run_fp8_grouped( a_data, b_data, D.view(m_total, n), - a_scale_vec.to(torch.float32).contiguous(), - b_scale_vec.to(torch.float32).contiguous(), + _fp8_scale_arg(a_scale_vec), + _fp8_scale_arg(b_scale_vec), epilogue=epilogue, bias=bias_arg, group=True, @@ -1488,6 +1501,7 @@ def _run_fp8_grouped( work_queue=work_queue, work_counter=work_counter, b_data_ptrs=b_data_ptrs, + a_data_ptrs=a_data_ptrs, sched=sched, m=m_total, n=n, @@ -1498,9 +1512,9 @@ def _run_fp8_grouped( def _run_fp8_grouped_wgrad( dyt_cols, # list[G] of per-expert dY^T (host-transposed) [N_out, m_e] - dy_scale_vec, # [num_groups] fp32 per-expert inverse scales for dY + dy_scale_vec, # [num_groups] fp32 inverse scales for dY, or list[G] of per-expert scale_inv xt_cols, # list[G] of per-expert X^T (host-transposed) [K_in, m_e] - x_scale_vec, # [num_groups] fp32 per-expert inverse scales for X + x_scale_vec, # [num_groups] fp32 inverse scales for X, or list[G] of per-expert scale_inv dW, # output [G*N_out, K_in] *, m_splits, @@ -1552,21 +1566,54 @@ def _pad_cols(t, cols): num_pid_m = n_out // _BLOCK_M num_pid_n = k_in // _BLOCK_N - dy_scale_vec = dy_scale_vec.to(torch.float32).contiguous() - x_scale_vec = x_scale_vec.to(torch.float32).contiguous() + # Per-expert scale_inv lists are read in place through per-bucket pointer arrays. + scale_lists = isinstance(dy_scale_vec, (list, tuple)) + if not scale_lists: + dy_scale_vec = dy_scale_vec.to(torch.float32).contiguous() + x_scale_vec = x_scale_vec.to(torch.float32).contiguous() buckets = gd.bucketize_wgrad_tokens(m_splits, _BLOCK_M, max_buckets=_wgrad_max_buckets()) + # An expert with no tokens has dW_e = 0. Quantizing its empty input never writes + # its tensor-wise scale_inv, which can hold NaN, and a zero-padded launch would + # compute NaN * 0. Zero those rows directly and leave the expert out of the launches. + for e, m in enumerate(m_splits): + if int(m) == 0: + dW[e * n_out : (e + 1) * n_out].zero_() + buckets = [(k_pad, [e for e in ids if int(m_splits[e]) != 0]) for k_pad, ids in buckets] + buckets = [(k_pad, ids) for k_pad, ids in buckets if ids] + for k_pad, ids in buckets: bsize = len(ids) - dy_data = torch.cat([_pad_cols(dyt_cols[e], k_pad) for e in ids], dim=0) - x_data = torch.cat([_pad_cols(xt_cols[e], k_pad) for e in ids], dim=0) + dy_b = [dyt_cols[e] for e in ids] + x_b = [xt_cols[e] for e in ids] + if all(t.size(1) == k_pad and t.is_contiguous() for t in dy_b + x_b): + # No padding needed: read each expert's transpose in place. The bases are + # shifted by the expert's packed row (j*N_out for dY^T, j*K_in for X^T), + # so the kernel's packed addressing lands in each per-expert slab. + dy_data, x_data = dy_b[0], x_b[0] + row_bytes = k_pad * dy_data.element_size() + dy_ptrs = gd.shifted_expert_ptrs( + dy_b, [j * n_out for j in range(bsize)], row_bytes, device + ) + x_ptrs = gd.shifted_expert_ptrs( + x_b, [j * k_in for j in range(bsize)], row_bytes, device + ) + else: + dy_data = torch.cat([_pad_cols(t, k_pad) for t in dy_b], dim=0) + x_data = torch.cat([_pad_cols(t, k_pad) for t in x_b], dim=0) + dy_ptrs = x_ptrs = None if dy_data.dtype == torch.uint8: dy_data = reinterpret_as_fp8_tensor(dy_data, dy_fp8_dtype) if x_data.dtype == torch.uint8: x_data = reinterpret_as_fp8_tensor(x_data, x_fp8_dtype) - idx = torch.tensor(ids, device=device) - dy_scale_b = dy_scale_vec[idx] - x_scale_b = x_scale_vec[idx] + + if scale_lists: + dy_scale_b = gd.expert_ptrs([dy_scale_vec[e] for e in ids], device) + x_scale_b = gd.expert_ptrs([x_scale_vec[e] for e in ids], device) + else: + idx = torch.tensor(ids, device=device) + dy_scale_b = dy_scale_vec[idx] + x_scale_b = x_scale_vec[idx] direct = ids == list(range(num_groups)) dW_b = dW if direct else torch.empty(bsize * n_out, k_in, dtype=output_dtype, @@ -1592,6 +1639,8 @@ def _pad_cols(t, cols): m=n_out, n=k_in, k=k_pad, + a_data_ptrs=dy_ptrs, + b_data_ptrs=x_ptrs, ) if not direct: diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py index c2ba1dc9df..dcc57d33d2 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py @@ -71,6 +71,29 @@ def read_ptr_at(ptr_array, index): return llvm.LoadOp(T.i64, elt_ptr).result +def make_uniform_i64(value): + """Broadcast lane 0's i64 to the whole wave (``readfirstlane``). + + A buffer descriptor lives in SGPRs. An address loaded by :func:`read_ptr_at` is a + per-lane value, so a descriptor rebased on it can make the compiler guard each + buffer access for divergence. Use this when every lane loaded the same address + (e.g. indexed by the tile's expert), so lane 0's value is every lane's value. + """ + return fx.rocdl.readfirstlane(T.i64, value) + + +def load_f32_at_ptr(ptr_array, index): + """Load the fp32 scalar that ``ptr_array[index]`` points to. + + A plain global load needs no wave-uniform address, so unlike the buffer-descriptor + bases this skips :func:`make_uniform_i64`: nothing waits on either load until the + value is used, which keeps a per-tile scale fetch off the tile's critical path. + """ + ptr_ty = ir.Type.parse("!llvm.ptr<1>") + addr = read_ptr_at(ptr_array, index) + return llvm.LoadOp(T.f32, llvm.IntToPtrOp(ptr_ty, addr).result).result + + def rebase_buffer_view(base_i64, elem_ir_type, layout): """Rebuild a buffer-tensor logical-divide view at a runtime base address. @@ -308,6 +331,33 @@ def concat_or_view_rows(tensors): return view +@functools.lru_cache(maxsize=512) +def _ptr_array_cached(ptrs, device): + return torch.tensor(ptrs, dtype=torch.int64, device=device) + + +def shifted_expert_ptrs(tensors, packed_row_starts, row_bytes, device): + """int64[G] device array of per-expert base pointers, each shifted to its packed row. + + Lets a grouped kernel read per-expert operands in place with its packed-operand + addressing: entry ``g`` is ``tensors[g].data_ptr() - packed_row_starts[g] * + row_bytes``, so ``entry + packed_row * row_bytes`` lands on row + ``packed_row - packed_row_starts[g]`` of ``tensors[g]``, exactly where the + ``torch.cat`` of the experts would have put it. Each tensor must be row-major + with ``row_bytes`` bytes per row. The upload runs only on a cache miss; the + caching allocator usually hands back the same per-expert buffers step to step. + """ + ptrs = tuple( + t.data_ptr() - int(start) * row_bytes for t, start in zip(tensors, packed_row_starts) + ) + return _ptr_array_cached(ptrs, device) + + +def expert_ptrs(tensors, device): + """int64[G] device array of ``t.data_ptr()`` per expert (cached; see shifted_expert_ptrs).""" + return _ptr_array_cached(tuple(t.data_ptr() for t in tensors), device) + + def uniform_pad_concat(tensors, rows): """Zero-pad each tensor's leading dim up to ``rows`` and concat on dim 0. diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 1e5b45d7d6..5422b65bfe 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -118,6 +118,38 @@ def _cached_b_data_ptrs(datas, device): return tensor +def _flydsl_fp8_row_operand(datas, m_splits): + """Address per-expert FP8 row payloads in place for the FlyDSL grouped TN kernel. + + ``datas`` are the per-expert ``[m_e, K]`` rowwise payloads. Returns ``(a_data, + a_data_ptrs)``: a representative payload plus an int64[G] array of per-expert base + pointers shifted by each expert's packed row start, so the kernel reads every + expert where the ``[M_total, K]`` concat would have placed it, without the copy. + Falls back to ``(concat_or_view_rows(datas), None)`` when a payload is not a + row-major ``[m_e, K]`` block. + """ + from ..flydsl_kernels.gemm.grouped_gemm_utils import ( + concat_or_view_rows, + shifted_expert_ptrs, + ) + + k = datas[0].size(-1) + if not all( + d.dim() == 2 and d.size(0) == int(m) and d.size(1) == k and d.is_contiguous() + for d, m in zip(datas, m_splits) + ): + return concat_or_view_rows(datas), None + row_starts = [] + start = 0 + for m in m_splits: + row_starts.append(start) + start += int(m) + # Any non-empty payload serves as the kernel's A argument; its base is replaced per tile. + a_data = next((d for d in datas if d.numel() > 0), datas[0]) + a_data_ptrs = shifted_expert_ptrs(datas, row_starts, k * a_data.element_size(), a_data.device) + return a_data, a_data_ptrs + + def _packed_wgrad_view(wgrad_list, num_gemms, n_out, k_in): """Return a [num_gemms*n_out, k_in] view over ``wgrad_list`` if it is one packed run. @@ -1151,12 +1183,13 @@ def _grouped_gemm_flydsl_fp8(*, inputmats, weights, out, m_splits, m_splits_tens """FlyDSL tensor-wise FP8 grouped forward (TN): Y_e = X_e @ W_e^T. Experts pack weights on the output-N axis ([G*N, K]); each operand carries - one scalar inverse scale, gathered into length-G per-expert vectors. + one scalar inverse scale, read in place through per-expert pointers. """ from ..flydsl_kernels.gemm import _run_fp8_grouped - from ..flydsl_kernels.gemm.grouped_gemm_utils import concat_or_view_rows + from ..flydsl_kernels.gemm.grouped_gemm_utils import expert_ptrs - a_data = concat_or_view_rows([im._data for im in inputmats]) + # Activations are read in place from each expert's payload (no [M_total, K] concat). + a_data, a_data_ptrs = _flydsl_fp8_row_operand([im._data for im in inputmats], m_splits) # Weights are independent per-expert [N, K] FP8 tensors; the grouped-TN kernel # reads each from a per-expert base pointer, so pass an int64[G] address array # instead of concatenating the G weights. b_data is a per-expert representative. @@ -1166,12 +1199,13 @@ def _grouped_gemm_flydsl_fp8(*, inputmats, weights, out, m_splits, m_splits_tens b_data = weights[0]._data a_data._fp8_dtype = inputmats[0]._fp8_dtype b_data._fp8_dtype = weights[0]._fp8_dtype - a_scale = torch.cat([im._scale_inv.reshape(1) for im in inputmats], dim=0) - b_scale = torch.cat([w._scale_inv.reshape(1) for w in weights], dim=0) + # Per-expert scale_inv pointers, read in place by the kernel (no gather). + a_scale = expert_ptrs([im._scale_inv for im in inputmats], a_data.device) + b_scale = expert_ptrs([w._scale_inv for w in weights], a_data.device) n = weights[0]._data.size(0) k = weights[0]._data.size(-1) - m_total = a_data.size(0) + m_total = sum(int(m) for m in m_splits) _run_fp8_grouped( a_data, @@ -1185,6 +1219,7 @@ def _grouped_gemm_flydsl_fp8(*, inputmats, weights, out, m_splits, m_splits_tens k=k, output_dtype=out.dtype, b_data_ptrs=b_data_ptrs, + a_data_ptrs=a_data_ptrs, ) @staticmethod @@ -1199,25 +1234,25 @@ def _backward_flydsl_fp8_grouped( B packed on the output axis = K_in, i.e. per expert dY @ (W^T)^T = dY @ W. """ from ..flydsl_kernels.gemm import _run_fp8_grouped - from ..flydsl_kernels.gemm.grouped_gemm_utils import concat_or_view_rows + from ..flydsl_kernels.gemm.grouped_gemm_utils import expert_ptrs # A = dY (rowwise fp8 data); B = W columnwise transpose, the materialized - # [K_in, N_out] packed on K_in (the TN GEMM output axis). The transpose - # backings are independent per-expert slabs, addressed via per-expert base - # pointers through the same TN per-expert path (no concat); dY still concats. - a_data = concat_or_view_rows([g._data for g in grad_output]) + # [K_in, N_out] packed on K_in (the TN GEMM output axis). Both are independent + # per-expert slabs, addressed via per-expert base pointers (no concat). + a_data, a_data_ptrs = _flydsl_fp8_row_operand([g._data for g in grad_output], m_splits) b_data_ptrs = _cached_b_data_ptrs( [w._transpose for w in weights], a_data.device ) b_data = weights[0]._transpose a_data._fp8_dtype = grad_output[0]._fp8_dtype b_data._fp8_dtype = weights[0]._fp8_dtype - a_scale = torch.cat([g._scale_inv.reshape(1) for g in grad_output], dim=0) - b_scale = torch.cat([w._scale_inv.reshape(1) for w in weights], dim=0) + # Per-expert scale_inv pointers, read in place by the kernel (no gather). + a_scale = expert_ptrs([g._scale_inv for g in grad_output], a_data.device) + b_scale = expert_ptrs([w._scale_inv for w in weights], a_data.device) n_out = grad_output[0]._data.size(-1) # per-expert contraction (N_out) k_in = dgrad.size(-1) # output width (K_in) - m_total = a_data.size(0) + m_total = sum(int(m) for m in m_splits) _run_fp8_grouped( a_data, @@ -1231,6 +1266,7 @@ def _backward_flydsl_fp8_grouped( k=n_out, output_dtype=dgrad.dtype, b_data_ptrs=b_data_ptrs, + a_data_ptrs=a_data_ptrs, ) @staticmethod @@ -1246,8 +1282,9 @@ def _wgrad_flydsl_fp8_grouped(*, grad_output, inputmats, wgrad_list, m_splits): num_gemms = len(inputmats) dyt_cols = [g._transpose for g in grad_output] xt_cols = [x._transpose for x in inputmats] - dy_scale = torch.cat([g._scale_inv.reshape(1) for g in grad_output], dim=0) - x_scale = torch.cat([x._scale_inv.reshape(1) for x in inputmats], dim=0) + # Per-expert scale_inv tensors; the wgrad runner reads them in place by pointer. + dy_scale = [g._scale_inv for g in grad_output] + x_scale = [x._scale_inv for x in inputmats] # Operands may carry only the columnwise (_transpose) backing. The # transpose is [feature, m_e]: dY^T is [N_out, m_e], X^T is [K_in, m_e]. From 0c783e4ede0df487f96305d7ad6099096b85a7f4 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 21 Oct 2026 02:35:01 +0000 Subject: [PATCH 14/22] Share per-expert in-place operand reads across all FlyDSL grouped dtypes --- .../flydsl_kernels/gemm/gemm_wrappers.py | 101 ++++-- .../flydsl_kernels/gemm/grouped_gemm_utils.py | 80 ++++- .../flydsl_kernels/gemm/half_prec_gemm.py | 57 ++- .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 339 ++++++++++++++---- .../pytorch/module/grouped_linear.py | 211 +++++------ 5 files changed, 542 insertions(+), 246 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index cfe7af9e64..2676202896 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -1028,6 +1028,7 @@ def _run_mxfp8_grouped( layout: str = "TN", bias=None, b_data_ptrs=None, + a_data_ptrs=None, ): """Dispatch a packed MXFP8 grouped GEMM through the persistent kernel. @@ -1049,6 +1050,8 @@ def _run_mxfp8_grouped( queue is built + shape-cached from it directly, so no device->host copy of a ``group_offs`` tensor (and its blocking sync) is needed. Operands arrive already MXFP8-quantized and packed -- no quantization, no BLAS ownership swap. + ``a_data_ptrs`` (optional) reads A per expert in place: ``a_data`` is then one + representative expert payload (see ``grouped_gemm_utils.stacked_operand``). """ if layout not in ("TN", "NN"): raise FlyDSLUnsupportedError(f"FlyDSL MXFP8 grouped GEMM supports TN/NN, got {layout}") @@ -1066,7 +1069,8 @@ def _run_mxfp8_grouped( if b_data.dtype == torch.uint8: b_data = reinterpret_as_fp8_tensor(b_data, b_fp8_dtype) - if a_scale.dtype != torch.uint8 or b_scale.dtype != torch.uint8: + scale_dtypes = [(s[0] if isinstance(s, (list, tuple)) else s).dtype for s in (a_scale, b_scale)] + if any(d != torch.uint8 for d in scale_dtypes): raise TypeError("FlyDSL MXFP8 grouped GEMM expects raw E8M0 scales as torch.uint8") D = _validate_or_allocate_output( @@ -1127,11 +1131,29 @@ def _run_mxfp8_grouped( work_queue=work_queue, work_counter=work_counter, b_data_ptrs=b_data_ptrs, + a_data_ptrs=a_data_ptrs, + a_shape=(m_total, a_data.size(-1)) if a_data_ptrs is not None else None, sched=sched, ) return D +def _wgrad_operands(dy_b, x_b, unpadded, pad, device): + """One wgrad bucket's (dy, dy_ptrs, x, x_ptrs), shared by every dtype. + + When no expert in the bucket needs padding (``unpadded``), both operands are read + in place from their per-expert tensors (``gd.stacked_operand``), so the kernel + sees the packed layout without a host-side concat. Otherwise ``pad`` builds the + padded packed operand and the pointers are ``None``. + """ + if unpadded: + dy = gd.stacked_operand(dy_b, device) + x = gd.stacked_operand(x_b, device) + if dy is not None and x is not None: + return dy[0], dy[1], x[0], x[1] + return pad(dy_b), None, pad(x_b), None + + def _run_mxfp8_grouped_wgrad( dy_cols, # list[G] of per-expert dY columnwise fp8 data, each [m_e, N_out] dy_scales, # list[G] of per-expert dY columnwise scale, each [m_e/32, N_out] @@ -1184,14 +1206,28 @@ def _run_mxfp8_grouped_wgrad( num_pid_m = n_out // BLOCK_M num_pid_n = k_in // BLOCK_N buckets = gd.bucketize_wgrad_tokens(m_splits, BLOCK_M, max_buckets=_wgrad_max_buckets()) + buckets = gd.drop_empty_experts(buckets, m_splits, dW, n_out) for k_pad, ids in buckets: bsize = len(ids) scale_rows = k_pad // SCALE_GROUP_SIZE - dy_data = gd.uniform_pad_concat([dy_cols[e] for e in ids], k_pad) - dy_scale = gd.uniform_pad_concat([dy_scales[e] for e in ids], scale_rows) - x_data = gd.uniform_pad_concat([x_cols[e] for e in ids], k_pad) - x_scale = gd.uniform_pad_concat([x_scales[e] for e in ids], scale_rows) + dy_b = [dy_cols[e] for e in ids] + x_b = [x_cols[e] for e in ids] + dy_data, dy_ptrs, x_data, x_ptrs = _wgrad_operands( + dy_b, + x_b, + all(t.size(0) == k_pad for t in dy_b + x_b), + lambda ts: gd.uniform_pad_concat(ts, k_pad), + device, + ) + dy_scale_b = [dy_scales[e] for e in ids] + x_scale_b = [x_scales[e] for e in ids] + if dy_ptrs is not None: + # Unpadded bucket: the scale packer reads the per-expert scales in place. + dy_scale, x_scale = dy_scale_b, x_scale_b + else: + dy_scale = gd.uniform_pad_concat(dy_scale_b, scale_rows) + x_scale = gd.uniform_pad_concat(x_scale_b, scale_rows) if dy_data.dtype == torch.uint8: dy_data = reinterpret_as_fp8_tensor(dy_data, dy_fp8_dtype) @@ -1219,6 +1255,10 @@ def _run_mxfp8_grouped_wgrad( work_counter=work_counter, sched=sched, wgrad=True, + a_data_ptrs=dy_ptrs, + b_data_ptrs=x_ptrs, + a_shape=(bsize * k_pad, n_out) if dy_ptrs is not None else None, + b_shape=(bsize * k_pad, k_in) if x_ptrs is not None else None, ) if not direct: @@ -1241,6 +1281,7 @@ def _run_bf16_grouped( layout: str = "TN", bias=None, b_data_ptrs=None, + a_data_ptrs=None, ): """Dispatch a packed BF16 grouped GEMM through the persistent kernel. @@ -1305,6 +1346,7 @@ def _run_bf16_grouped( work_queue=work_queue, work_counter=work_counter, b_data_ptrs=b_data_ptrs, + a_data_ptrs=a_data_ptrs, sched=sched, ) return D @@ -1372,11 +1414,19 @@ def _run_bf16_grouped_wgrad( num_pid_m = n_out // _BLOCK_M num_pid_n = k_in // _BLOCK_N buckets = gd.bucketize_wgrad_tokens(m_splits, _BLOCK_M, max_buckets=_wgrad_max_buckets()) + buckets = gd.drop_empty_experts(buckets, m_splits, dW, n_out) for k_pad, ids in buckets: bsize = len(ids) - dy_data = gd.uniform_pad_concat([dy_cols[e] for e in ids], k_pad) - x_data = gd.uniform_pad_concat([x_cols[e] for e in ids], k_pad) + dy_b = [dy_cols[e] for e in ids] + x_b = [x_cols[e] for e in ids] + dy_data, dy_ptrs, x_data, x_ptrs = _wgrad_operands( + dy_b, + x_b, + all(t.size(0) == k_pad for t in dy_b + x_b), + lambda ts: gd.uniform_pad_concat(ts, k_pad), + device, + ) # Fast path: when this bucket's experts are exactly rows 0..G-1 in order # (the single-bucket / balanced case), the bucket-local expert index equals @@ -1408,6 +1458,8 @@ def _run_bf16_grouped_wgrad( work_counter=work_counter, sched=sched, wgrad=True, + a_data_ptrs=dy_ptrs, + b_data_ptrs=x_ptrs, ) if not direct: @@ -1453,7 +1505,7 @@ def _run_fp8_grouped( + shape-cached from it (no GPU->CPU sync). ``a_data_ptrs`` (optional) reads A per expert in place: ``a_data`` is then one representative expert payload and each entry is that expert's base shifted by its packed row start (see - ``grouped_gemm_utils.shifted_expert_ptrs``). + ``grouped_gemm_utils.stacked_operand``). """ from .fp8_gemm import doGemm as _fp8_doGemm @@ -1573,35 +1625,20 @@ def _pad_cols(t, cols): x_scale_vec = x_scale_vec.to(torch.float32).contiguous() buckets = gd.bucketize_wgrad_tokens(m_splits, _BLOCK_M, max_buckets=_wgrad_max_buckets()) - # An expert with no tokens has dW_e = 0. Quantizing its empty input never writes - # its tensor-wise scale_inv, which can hold NaN, and a zero-padded launch would - # compute NaN * 0. Zero those rows directly and leave the expert out of the launches. - for e, m in enumerate(m_splits): - if int(m) == 0: - dW[e * n_out : (e + 1) * n_out].zero_() - buckets = [(k_pad, [e for e in ids if int(m_splits[e]) != 0]) for k_pad, ids in buckets] - buckets = [(k_pad, ids) for k_pad, ids in buckets if ids] + buckets = gd.drop_empty_experts(buckets, m_splits, dW, n_out) for k_pad, ids in buckets: bsize = len(ids) dy_b = [dyt_cols[e] for e in ids] x_b = [xt_cols[e] for e in ids] - if all(t.size(1) == k_pad and t.is_contiguous() for t in dy_b + x_b): - # No padding needed: read each expert's transpose in place. The bases are - # shifted by the expert's packed row (j*N_out for dY^T, j*K_in for X^T), - # so the kernel's packed addressing lands in each per-expert slab. - dy_data, x_data = dy_b[0], x_b[0] - row_bytes = k_pad * dy_data.element_size() - dy_ptrs = gd.shifted_expert_ptrs( - dy_b, [j * n_out for j in range(bsize)], row_bytes, device - ) - x_ptrs = gd.shifted_expert_ptrs( - x_b, [j * k_in for j in range(bsize)], row_bytes, device - ) - else: - dy_data = torch.cat([_pad_cols(t, k_pad) for t in dy_b], dim=0) - x_data = torch.cat([_pad_cols(t, k_pad) for t in x_b], dim=0) - dy_ptrs = x_ptrs = None + # FP8 operands are [feature, m_e] transposes, so tokens pad on the trailing axis. + dy_data, dy_ptrs, x_data, x_ptrs = _wgrad_operands( + dy_b, + x_b, + all(t.size(1) == k_pad for t in dy_b + x_b), + lambda ts: torch.cat([_pad_cols(t, k_pad) for t in ts], dim=0), + device, + ) if dy_data.dtype == torch.uint8: dy_data = reinterpret_as_fp8_tensor(dy_data, dy_fp8_dtype) if x_data.dtype == torch.uint8: diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py index dcc57d33d2..ff3fd72d29 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py @@ -71,6 +71,11 @@ def read_ptr_at(ptr_array, index): return llvm.LoadOp(T.i64, elt_ptr).result +def load_i32_at(tensor, index): + """Load ``tensor[index]`` (int32) with a plain per-lane global load.""" + return llvm.LoadOp(T.i32, global_i32_ptr(tensor, index)).result + + def make_uniform_i64(value): """Broadcast lane 0's i64 to the whole wave (``readfirstlane``). @@ -336,26 +341,69 @@ def _ptr_array_cached(ptrs, device): return torch.tensor(ptrs, dtype=torch.int64, device=device) -def shifted_expert_ptrs(tensors, packed_row_starts, row_bytes, device): - """int64[G] device array of per-expert base pointers, each shifted to its packed row. +def expert_ptrs(tensors, device): + """int64[G] device array of each per-expert tensor's ``data_ptr()``. - Lets a grouped kernel read per-expert operands in place with its packed-operand - addressing: entry ``g`` is ``tensors[g].data_ptr() - packed_row_starts[g] * - row_bytes``, so ``entry + packed_row * row_bytes`` lands on row - ``packed_row - packed_row_starts[g]`` of ``tensors[g]``, exactly where the - ``torch.cat`` of the experts would have put it. Each tensor must be row-major - with ``row_bytes`` bytes per row. The upload runs only on a cache miss; the - caching allocator usually hands back the same per-expert buffers step to step. + Grouped kernels read per-expert operands (weights, FP8 scale_inv) through this + array instead of a host-side concat. The upload runs only on a cache miss: weights + and the caching allocator's per-expert buffers repeat step to step. """ - ptrs = tuple( - t.data_ptr() - int(start) * row_bytes for t, start in zip(tensors, packed_row_starts) - ) - return _ptr_array_cached(ptrs, device) + return _ptr_array_cached(tuple(t.data_ptr() for t in tensors), device) -def expert_ptrs(tensors, device): - """int64[G] device array of ``t.data_ptr()`` per expert (cached; see shifted_expert_ptrs).""" - return _ptr_array_cached(tuple(t.data_ptr() for t in tensors), device) +def stacked_operand(tensors, device): + """Address per-expert tensors stacked on dim 0 in place, as one packed operand. + + Returns ``(representative, ptrs)``, or ``None`` when the tensors are not + contiguous with a common trailing shape and dtype. Entry ``g`` of the int64[G] + ``ptrs`` is ``tensors[g].data_ptr()`` minus the byte offset where expert ``g`` + starts in the ``torch.cat(tensors)`` layout, so a kernel keeps its packed + addressing and only swaps the buffer base per tile. ``representative`` is a + non-empty expert tensor standing in for the operand's dtype. + """ + first = tensors[0] + trailing = first.shape[1:] + ptrs = [] + offset = 0 + for t in tensors: + if t.dtype != first.dtype or t.shape[1:] != trailing or not t.is_contiguous(): + return None + ptrs.append(t.data_ptr() - offset) + offset += t.numel() * t.element_size() + representative = next((t for t in tensors if t.numel() > 0), first) + return representative, _ptr_array_cached(tuple(ptrs), device) + + +@functools.lru_cache(maxsize=512) +def _block_expert_table_cached(blocks_per_expert, device): + table = [e for e, n in enumerate(blocks_per_expert) for _ in range(n)] + return torch.tensor(table, dtype=torch.int32, device=device) + + +def block_expert_table(extents, block, device): + """int32 table mapping each ``block``-sized slice of a stacked axis to its expert. + + ``extents`` are the per-expert lengths along the stacked axis; each must be a + multiple of ``block``, so no block straddles two experts. + """ + for n in extents: + if int(n) % block != 0: + raise ValueError(f"per-expert extent {int(n)} is not a multiple of {block}") + return _block_expert_table_cached(tuple(int(n) // block for n in extents), device) + + +def drop_empty_experts(buckets, m_splits, dW, n_out): + """Zero the ``dW`` rows of zero-token experts and drop them from the wgrad buckets. + + An expert with no tokens has ``dW_e = 0``; launching it anyway contracts a zero + padding that a non-finite scale (FP8 never writes an empty tensor's scale_inv) + would turn into NaN. + """ + for e, m in enumerate(m_splits): + if int(m) == 0: + dW[e * n_out : (e + 1) * n_out].zero_() + buckets = [(k_pad, [e for e in ids if int(m_splits[e]) != 0]) for k_pad, ids in buckets] + return [(k_pad, ids) for k_pad, ids in buckets if ids] def uniform_pad_concat(tensors, rows): diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py index 2cf3a98fe4..3c17d2696b 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py @@ -114,6 +114,7 @@ def _compile_kernel( group: bool = False, sched: str = "worksteal", wgrad: bool = False, + act_ptrs: bool = False, ): """Build one compile-time-specialized TN, NN, or NT half-precision kernel. @@ -393,6 +394,7 @@ def kernel_gemm( work_queue: fx.Tensor = None, work_counter: fx.Tensor = None, b_data_ptrs: fx.Tensor = None, + a_data_ptrs: fx.Tensor = None, ): lds = fx.SharedAllocator().allocate(SharedStorage).peek() lds_a0 = (lds.a0_0, lds.a0_1) @@ -428,9 +430,16 @@ def kernel_gemm( # Grouped fwd/dgrad: weights arrive as G independently-allocated per-expert # slabs, so B is addressed from a per-expert base pointer read from # ``b_data_ptrs`` (byte-addressed, reusing gB's static layout) rather than - # an expert-shifted offset into one packed B. wgrad keeps the packed-B - # folds; dense never enters here. + # an expert-shifted offset into one packed B. wgrad reads B by pointer only + # with act_ptrs; dense never enters here. b_ptr_layout = fx.get_layout(gB) + # act_ptrs: the per-expert activation operand (A; and in wgrad also B, via + # b_data_ptrs) stays in its own per-expert allocation instead of a host-side + # concat. The host passes each expert's base pointer pre-shifted by that + # expert's packed offset, so the packed addressing below is unchanged; only + # the buffer base moves per tile. + if const_expr(act_ptrs): + a_ptr_layout = fx.get_layout(gA) # LDS broadcast scratch (one int32): thread 0 writes the atomically # claimed slot, a barrier broadcasts it to the whole workgroup. The @@ -504,8 +513,19 @@ def process_tile(pid_m, pid_n, expert): gl_off_a = _a_global_offsets(lane, wave_id, c_m) gl_off_b = _b_global_offsets(lane, wave_id, c_n) + # Issue both per-expert base-pointer loads before either is made uniform, so a + # tile waits on one memory round trip rather than two. + if const_expr(act_ptrs): + a_base_raw = gd.read_ptr_at(a_data_ptrs, expert) + if const_expr(group or act_ptrs): + b_base_raw = gd.read_ptr_at(b_data_ptrs, expert) + if const_expr(act_ptrs): + a_base_i64 = gd.make_uniform_i64(a_base_raw) + a_div_tile = gd.rebase_buffer_view(a_base_i64, fx.Uint8.ir_type, a_ptr_layout) + else: + a_div_tile = a_div a_g2s = G2SLoader( - a_div, + a_div_tile, gl_off_a, LOAD_PASSES_HALF, fx.Uint8.ir_type, @@ -515,12 +535,19 @@ def process_tile(pid_m, pid_n, expert): # the per-expert base already selects the expert, so in-tile addressing is # LOCAL. TN reads a local column (by_n_idx); NN (dgrad, transpose-read) # packs experts on the contraction axis, so the base replaces the expert*K - # contraction shift (b_k_off -> 0). Dense/wgrad keep the shared packed view. + # contraction shift (b_k_off -> 0). wgrad bases are pre-shifted (act_ptrs), + # so they keep the packed column and contraction shift. Dense keeps the + # shared packed view. if const_expr(group): - b_base_i64 = gd.read_ptr_at(b_data_ptrs, expert) + b_base_i64 = gd.make_uniform_i64(b_base_raw) b_div_tile = gd.rebase_buffer_view(b_base_i64, fx.Uint8.ir_type, b_ptr_layout) b_data_col = by_n_idx b_k_off = fx.Index(0) + elif const_expr(act_ptrs): + b_base_i64 = gd.make_uniform_i64(b_base_raw) + b_div_tile = gd.rebase_buffer_view(b_base_i64, fx.Uint8.ir_type, b_ptr_layout) + b_data_col = by_n_b_idx + b_k_off = b_expert_k_off else: b_div_tile = b_div b_data_col = by_n_b_idx @@ -1361,6 +1388,7 @@ def launch_gemm( work_queue: fx.Tensor = None, work_counter: fx.Tensor = None, b_data_ptrs: fx.Tensor = None, + a_data_ptrs: fx.Tensor = None, stream: fx.Stream = fx.Stream(None), ): # Dense: one workgroup per output tile. Grouped/wgrad: a fixed num-CU @@ -1381,6 +1409,7 @@ def launch_gemm( work_queue, work_counter, b_data_ptrs, + a_data_ptrs, value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, ).launch(grid=(grid_x, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) @@ -1398,6 +1427,7 @@ def _cached_launch( group: bool = False, sched: str = "worksteal", wgrad: bool = False, + act_ptrs: bool = False, ): return _compile_kernel( K, @@ -1409,6 +1439,7 @@ def _cached_launch( group=group, sched=sched, wgrad=wgrad, + act_ptrs=act_ptrs, ) @@ -1622,6 +1653,7 @@ def doGemm( work_queue: torch.Tensor = None, work_counter: torch.Tensor = None, b_data_ptrs: torch.Tensor = None, + a_data_ptrs: torch.Tensor = None, sched: str = "worksteal", wgrad: bool = False, ): @@ -1729,13 +1761,19 @@ def doGemm( work_queue_arg = _dense_placeholder(A.device, torch.int32) work_counter_arg = _dense_placeholder(A.device, torch.int32) - # Per-expert weight base pointers (grouped TN only): an int64[G] device array of - # each expert's weight-slab base address, read in-kernel so the weight operand - # needs no host-side concat. Dense/NN/wgrad pass a placeholder never read there. + # Per-expert weight base pointers (grouped TN/NN; wgrad's pre-shifted X with + # act_ptrs): an int64[G] device array read in-kernel so the operand needs no + # host-side concat. Dense passes a placeholder that is never read. if b_data_ptrs is not None: b_data_ptrs_arg = b_data_ptrs.contiguous().view(-1) else: b_data_ptrs_arg = _dense_placeholder(A.device, torch.int64) + # Per-expert activation base pointers (see act_ptrs in _compile_kernel). + act_ptrs = a_data_ptrs is not None + if act_ptrs: + a_data_ptrs_arg = a_data_ptrs.contiguous().view(-1) + else: + a_data_ptrs_arg = _dense_placeholder(A.device, torch.int64) launch_key = ( "half_prec_gemm", @@ -1748,6 +1786,7 @@ def doGemm( group, sched, wgrad, + act_ptrs, ) launch_args = ( A_arg, @@ -1761,6 +1800,7 @@ def doGemm( work_queue_arg, work_counter_arg, b_data_ptrs_arg, + a_data_ptrs_arg, stream, ) # flyc.compile fast path: skips flydsl's per-call cache-key rebuild (sig.bind + @@ -1778,6 +1818,7 @@ def doGemm( group, sched, wgrad, + act_ptrs, ), launch_args, )(*launch_args) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index 338cf87da5..80fa15d877 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -116,12 +116,19 @@ def _compile_mx32_scale_pack_kernel( source_colwise: bool, stride0: int, stride1: int, + per_expert: bool = False, ): """Build one fused scale packing kernel. One GPU thread produces one final ``uint32`` word in the GEMM-consumed ``[K/128, dim]`` layout. There is no intermediate ``scale_iter`` tensor and no eager PyTorch shift/index/OR kernels. + + ``per_expert`` packs the concatenation of per-expert scale tensors without + building it: the source is an int64[G] array of per-expert base pointers, each + shifted to its expert's offset in the concatenated layout (so the offsets below + are unchanged), plus an int32 table mapping each 64-row tile (rowwise) or K128 + block (columnwise) to its expert. """ if dim % 64 != 0: raise FlyDSLUnsupportedError(f"Scale outer dimension={dim} must be a multiple of 64") @@ -150,25 +157,8 @@ def _source_offset(source_k32, source_row): # Logical source is [dim, K/32], with arbitrary positive strides. return source_row * fx.Index(stride0) + source_k32 * fx.Index(stride1) - @flyc.kernel(known_block_size=[_SCALE_PACK_THREADS, 1, 1]) - def kernel_pack_mx32_scales(src: fx.Tensor, dst: fx.Tensor): - # Address both buffers as flat 1-D arrays so a linear slice coordinate - # maps to base+off, matching the legacy buffer_load semantics. Building - # buffer tensors directly from the 2-D src/dst would make logical_divide - # walk the layout column-major and corrupt every strided access. - src_flat = fx.rocdl.make_buffer_tensor( - fx.Tensor(fx.make_view(fx.get_iter(src), fx.make_layout(dim * qk, 1))), - max_size=True, - ) - dst_flat = fx.rocdl.make_buffer_tensor( - fx.Tensor(fx.make_view(fx.get_iter(dst), fx.make_layout(total_words, 1))), - max_size=True, - ) - src_div = fx.logical_divide(src_flat, fx.make_layout(1, 1)) - dst_div = fx.logical_divide(dst_flat, fx.make_layout(1, 1)) - scale_load_atom = fx.make_copy_atom(fx.rocdl.BufferCopy8b(), fx.Uint8) - scale_store_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32) - + def _word_coords(): + # This thread's output word and the source coordinates of its four bytes. linear = fx.Index(fx.block_idx.x) * fx.Index(_SCALE_PACK_THREADS) + fx.Index( gpu.thread_id("x") ) @@ -179,26 +169,80 @@ def kernel_pack_mx32_scales(src: fx.Tensor, dst: fx.Tensor): k_subgroup = (dst_row // fx.Index(16)) % fx.Index(4) tile = dst_row // fx.Index(64) source_k32 = k128 * fx.Index(4) + k_subgroup + return linear, k128, tile, row_within_16, source_k32 - def load_scale_byte(group): - source_row = tile * fx.Index(64) + fx.Index(group * 16) + row_within_16 - off = _source_offset(source_k32, source_row) - reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Uint8) - fx.copy(scale_load_atom, fx.slice(src_div, (None, fx.Int32(off))), reg) - value_u8 = fx.memref_load_vec(reg)[0] - # The byte load widens via sext, so mask to the low 8 bits to - # preserve the raw E8M0 byte for scales >= 0x80 before packing. - return fx.Int32(value_u8) & fx.Int32(0xFF) - - b0 = load_scale_byte(0) - b1 = load_scale_byte(1) - b2 = load_scale_byte(2) - b3 = load_scale_byte(3) + def _load_scale_byte(src_div, coords, group): + _, _, tile, row_within_16, source_k32 = coords + source_row = tile * fx.Index(64) + fx.Index(group * 16) + row_within_16 + off = _source_offset(source_k32, source_row) + scale_load_atom = fx.make_copy_atom(fx.rocdl.BufferCopy8b(), fx.Uint8) + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Uint8) + fx.copy(scale_load_atom, fx.slice(src_div, (None, fx.Int32(off))), reg) + value_u8 = fx.memref_load_vec(reg)[0] + # The byte load widens via sext, so mask to the low 8 bits to + # preserve the raw E8M0 byte for scales >= 0x80 before packing. + return fx.Int32(value_u8) & fx.Int32(0xFF) + + def _store_word(dst, src_div, coords): + dst_flat = fx.rocdl.make_buffer_tensor( + fx.Tensor(fx.make_view(fx.get_iter(dst), fx.make_layout(total_words, 1))), + max_size=True, + ) + dst_div = fx.logical_divide(dst_flat, fx.make_layout(1, 1)) + scale_store_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32) + b0 = _load_scale_byte(src_div, coords, 0) + b1 = _load_scale_byte(src_div, coords, 1) + b2 = _load_scale_byte(src_div, coords, 2) + b3 = _load_scale_byte(src_div, coords, 3) packed = b0 | (b1 << 8) | (b2 << 16) | (b3 << 24) reg_i32 = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Int32) fx.memref_store_vec(Vec.filled(1, packed, fx.Int32), reg_i32) + linear = coords[0] fx.copy(scale_store_atom, reg_i32, fx.slice(dst_div, (None, fx.Int32(linear)))) + if per_expert: + + @flyc.kernel(known_block_size=[_SCALE_PACK_THREADS, 1, 1]) + def kernel_pack_mx32_scales(src_ptrs: fx.Tensor, expert_table: fx.Tensor, dst: fx.Tensor): + coords = _word_coords() + _, k128, tile, _, _ = coords + # A wave's 64 consecutive words share one 64-row tile (dim % 64 == 0) and one + # K128 block, so one expert: make its base wave-uniform and rebase a buffer + # view, keeping the buffer loads of the packed path. + block = k128 if source_colwise else tile + expert = fx.Int32(gd.load_i32_at(expert_table, block)) + base = gd.make_uniform_i64(gd.read_ptr_at(src_ptrs, expert)) + src_div = gd.rebase_buffer_view(base, fx.Uint8.ir_type, fx.make_layout(dim * qk, 1)) + _store_word(dst, src_div, coords) + + @flyc.jit + def launch_pack_mx32_scales( + src_ptrs: fx.Tensor, + expert_table: fx.Tensor, + dst: fx.Tensor, + stream: fx.Stream = fx.Stream(None), + ): + kernel_pack_mx32_scales(src_ptrs, expert_table, dst).launch( + grid=(total_words // _SCALE_PACK_THREADS, 1, 1), + block=(_SCALE_PACK_THREADS, 1, 1), + stream=stream, + ) + + return launch_pack_mx32_scales + + @flyc.kernel(known_block_size=[_SCALE_PACK_THREADS, 1, 1]) + def kernel_pack_mx32_scales(src: fx.Tensor, dst: fx.Tensor): + # Address both buffers as flat 1-D arrays so a linear slice coordinate + # maps to base+off, matching the legacy buffer_load semantics. Building + # buffer tensors directly from the 2-D src/dst would make logical_divide + # walk the layout column-major and corrupt every strided access. + src_flat = fx.rocdl.make_buffer_tensor( + fx.Tensor(fx.make_view(fx.get_iter(src), fx.make_layout(dim * qk, 1))), + max_size=True, + ) + src_div = fx.logical_divide(src_flat, fx.make_layout(1, 1)) + _store_word(dst, src_div, _word_coords()) + @flyc.jit def launch_pack_mx32_scales( src: fx.Tensor, @@ -221,9 +265,10 @@ def _cached_mx32_scale_pack_launch( source_colwise: bool, stride0: int, stride1: int, + per_expert: bool = False, ): """Cache orientation-and-stride-specialized fused pack binaries.""" - return _compile_mx32_scale_pack_kernel(dim, qk, source_colwise, stride0, stride1) + return _compile_mx32_scale_pack_kernel(dim, qk, source_colwise, stride0, stride1, per_expert) def pack_mx32_scales_for_hk( @@ -279,6 +324,80 @@ def pack_mx32_scales_for_hk( return packed +def pack_mx32_scales_for_hk_experts( + scales, + *, + source_colwise: bool = False, + stream=None, +) -> torch.Tensor: + """:func:`pack_mx32_scales_for_hk` of ``torch.cat(scales)``, without the concat. + + ``scales`` are per-expert uint8 E8M0 tensors stacked on dim 0: rowwise + ``[rows_e, K/32]`` or columnwise ``[K32_e, dim]``. The packing kernel reads each + expert in place through ``grouped_gemm_utils.stacked_operand`` pointers. Falls + back to packing a real concat when the tensors cannot be addressed in place. + """ + stacked = gd.stacked_operand(scales, scales[0].device) + if stacked is None or any(t.dim() != 2 for t in scales): + return pack_mx32_scales_for_hk( + torch.cat(scales, dim=0), source_colwise=source_colwise, stream=stream + ) + representative, src_ptrs = stacked + if representative.dtype != torch.uint8: + raise TypeError(f"MXFP8 scales must be torch.uint8 E8M0 bytes, got {representative.dtype}") + + total = sum(t.size(0) for t in scales) + inner = representative.size(1) + if source_colwise: + qk, dim = total, inner + # Experts stack on K/32; each thread's bytes share one K128 block (4 rows). + block = 4 + else: + dim, qk = total, inner + # Experts stack on dim; each thread's bytes share one 64-row tile. + block = 64 + expert_table = gd.block_expert_table( + [t.size(0) for t in scales], block, representative.device + ) + + packed = torch.empty((qk // 4, dim), dtype=torch.int32, device=representative.device) + if stream is None: + stream = torch.cuda.current_stream(representative.device) + # The shifted pointers address the contiguous concatenation, so its strides apply. + stride0, stride1 = inner, 1 + pack_key = ("mx32_scale_pack_experts", dim, qk, bool(source_colwise), stride0, stride1) + pack_args = (src_ptrs, expert_table, packed, stream) + _compiled_launch( + pack_key, + _cached_mx32_scale_pack_launch( + dim, qk, bool(source_colwise), stride0, stride1, per_expert=True + ), + pack_args, + )(*pack_args) + return packed + + +def _scale_shape(scale): + # A scale operand is one tensor or a list of per-expert tensors stacked on dim 0. + if isinstance(scale, (list, tuple)): + return (sum(t.size(0) for t in scale), scale[0].size(1)) + return tuple(scale.shape) + + +def _pack_scales(scale, *, source_colwise, stream): + if isinstance(scale, (list, tuple)): + return pack_mx32_scales_for_hk_experts(scale, source_colwise=source_colwise, stream=stream) + return pack_mx32_scales_for_hk(scale, source_colwise=source_colwise, stream=stream) + + +def _scale_dtype(scale): + return scale[0].dtype if isinstance(scale, (list, tuple)) else scale.dtype + + +def _scale_device(scale): + return scale[0].device if isinstance(scale, (list, tuple)) else scale.device + + def _compile_kernel( K: int, a_fp8_dtype: torch.dtype, @@ -289,6 +408,7 @@ def _compile_kernel( group: bool = False, sched: str = "worksteal", wgrad: bool = False, + act_ptrs: bool = False, ): """Build one compile-time-specialized TN, NN, or NT kernel. @@ -547,6 +667,7 @@ def kernel_gemm( work_queue: fx.Tensor = None, work_counter: fx.Tensor = None, b_data_ptrs: fx.Tensor = None, + a_data_ptrs: fx.Tensor = None, ): # Total output tiles. Dense uses this as the grid extent; grouped uses it as # the drain bound for the persistent loop. For fwd/dgrad, c_m = M_total spans @@ -606,9 +727,16 @@ def kernel_gemm( # per-expert slabs (not one packed buffer), so B is addressed from a # per-expert base pointer read from ``b_data_ptrs`` rather than an # expert-shifted offset into a packed B. Reuse ``gB``'s static layout - # (compile-time cosize) and swap only the base per tile. wgrad keeps the - # packed-B expert-offset folds; dense never enters here. + # (compile-time cosize) and swap only the base per tile. wgrad reads B by + # pointer only with act_ptrs; dense never enters here. b_ptr_layout = fx.get_layout(gB) + # act_ptrs: the per-expert activation operand (A; and in wgrad also B, via + # b_data_ptrs) stays in its own per-expert allocation instead of a host-side + # concat. The host passes each expert's base pointer pre-shifted by that + # expert's packed offset, so the packed addressing below is unchanged; only + # the buffer base moves per tile. + if const_expr(act_ptrs): + a_ptr_layout = fx.get_layout(gA) # LDS broadcast scratch (one int32), following the persistent-kernel # idiom in mega_moe_stage1: recast the LDS slot to an int32 iterator and @@ -751,8 +879,19 @@ def process_tile(pid_m, pid_n, expert): LOAD_PASSES_HALF, preshuffled=False, ) + # Issue both per-expert base-pointer loads before either is made uniform, so a + # tile waits on one memory round trip rather than two. + if const_expr(act_ptrs): + a_base_raw = gd.read_ptr_at(a_data_ptrs, expert) + if const_expr(group or act_ptrs): + b_base_raw = gd.read_ptr_at(b_data_ptrs, expert) + if const_expr(act_ptrs): + a_base_i64 = gd.make_uniform_i64(a_base_raw) + a_div_tile = gd.rebase_buffer_view(a_base_i64, a_f8_ir_t, a_ptr_layout) + else: + a_div_tile = a_div a_g2s = G2SLoader( - a_div, + a_div_tile, gl_off_a, LOAD_PASSES_HALF, a_f8_ir_t, @@ -763,12 +902,18 @@ def process_tile(pid_m, pid_n, expert): # addressing uses LOCAL coords. TN reads with a local column (by_n_idx); # NN (dgrad, transpose-read) packs experts on the contraction axis, so the # per-expert base replaces the expert*K contraction shift (b_k_off -> 0). - # Dense/wgrad keep the shared packed ``b_div`` and its expert folds. + # wgrad bases are pre-shifted (act_ptrs), so they keep the packed column and + # contraction shift. Dense keeps the shared packed ``b_div``. if const_expr(group): - b_base_i64 = gd.read_ptr_at(b_data_ptrs, expert) + b_base_i64 = gd.make_uniform_i64(b_base_raw) b_div_tile = gd.rebase_buffer_view(b_base_i64, b_f8_ir_t, b_ptr_layout) b_data_col = by_n_idx b_k_off = fx.Index(0) + elif const_expr(act_ptrs): + b_base_i64 = gd.make_uniform_i64(b_base_raw) + b_div_tile = gd.rebase_buffer_view(b_base_i64, b_f8_ir_t, b_ptr_layout) + b_data_col = by_n_b_idx + b_k_off = b_expert_k_off else: b_div_tile = b_div b_data_col = by_n_b_idx @@ -1767,6 +1912,7 @@ def launch_gemm( work_queue: fx.Tensor = None, work_counter: fx.Tensor = None, b_data_ptrs: fx.Tensor = None, + a_data_ptrs: fx.Tensor = None, stream: fx.Stream = fx.Stream(None), ): # The integration only dispatches aligned shapes; no partial-tile masking exists. @@ -1790,6 +1936,7 @@ def launch_gemm( work_queue, work_counter, b_data_ptrs, + a_data_ptrs, value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, ).launch(grid=(grid_x, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) @@ -1813,6 +1960,9 @@ def do_gemm( work_queue: fx.Tensor = None, work_counter: fx.Tensor = None, b_data_ptrs: torch.Tensor = None, + a_data_ptrs: torch.Tensor = None, + a_shape=None, + b_shape=None, sched: str = "worksteal", wgrad: bool = False, ): @@ -1820,16 +1970,23 @@ def do_gemm( The GELU epilogues require a caller-allocated ``aux`` output (M x N, same dtype as C), filled in place with the pre-activation values for backward. + + With ``a_data_ptrs`` the per-expert A operand is read in place (see act_ptrs in + ``_compile_kernel``): ``A`` is then one representative expert and ``a_shape`` the + packed operand shape. ``b_shape`` likewise gives wgrad's packed B shape when B + is read through ``b_data_ptrs``. """ + a_rows, a_cols = a_shape if a_shape is not None else A.shape + b_rows, b_cols = b_shape if b_shape is not None else B.shape if layout == "TN": - M_runtime, K_runtime = A.shape - N_runtime, Kb_runtime = B.shape + M_runtime, K_runtime = a_rows, a_cols + N_runtime, Kb_runtime = b_rows, b_cols elif layout == "NN": - M_runtime, K_runtime = A.shape - Kb_runtime, N_runtime = B.shape + M_runtime, K_runtime = a_rows, a_cols + Kb_runtime, N_runtime = b_rows, b_cols elif layout == "NT": - K_runtime, M_runtime = A.shape - Kb_runtime, N_runtime = B.shape + K_runtime, M_runtime = a_rows, a_cols + Kb_runtime, N_runtime = b_rows, b_cols else: raise ValueError(f"Unsupported MXFP8 kernel layout: {layout}") @@ -1986,15 +2143,20 @@ def do_gemm( work_queue_arg = _dummy_arg(torch.int32, A.device) work_counter_arg = _dummy_arg(torch.int32, A.device) - # Per-expert weight base pointers (grouped fwd TN / dgrad NN): an int64[G] device - # array of each expert's B-data base address, read in-kernel so the weight operand - # needs no host-side concat. Dense and wgrad supply a dtype-correct placeholder to - # keep the flyc.compile positional ABI uniform; the kernel never dereferences it - # there. (The B-scale stays packed -- no per-expert scale pointers.) + # Per-expert B base pointers (grouped fwd TN / dgrad NN weights; wgrad's pre-shifted + # X with act_ptrs): an int64[G] device array read in-kernel so the operand needs no + # host-side concat. Dense supplies a dtype-correct placeholder to keep the + # flyc.compile positional ABI uniform; the kernel never dereferences it there. if b_data_ptrs is not None: b_data_ptrs_arg = b_data_ptrs.contiguous().view(-1) else: b_data_ptrs_arg = _dummy_arg(torch.int64, A.device) + # Per-expert activation base pointers (see act_ptrs in _compile_kernel). + act_ptrs = a_data_ptrs is not None + if act_ptrs: + a_data_ptrs_arg = a_data_ptrs.contiguous().view(-1) + else: + a_data_ptrs_arg = _dummy_arg(torch.int64, A.device) launch_key = ( "mxfp8_gemm", @@ -2007,6 +2169,7 @@ def do_gemm( group, sched, wgrad, + act_ptrs, ) launch_args = ( A_arg, @@ -2022,12 +2185,13 @@ def do_gemm( work_queue_arg, work_counter_arg, b_data_ptrs_arg, + a_data_ptrs_arg, stream, ) _compiled_launch( launch_key, _cached_launch( - K_runtime, A.dtype, B.dtype, C.dtype, layout, epilogue, group, sched, wgrad + K_runtime, A.dtype, B.dtype, C.dtype, layout, epilogue, group, sched, wgrad, act_ptrs ), launch_args, )(*launch_args) @@ -2044,6 +2208,7 @@ def _cached_launch( group: bool = False, sched: str = "worksteal", wgrad: bool = False, + act_ptrs: bool = False, ): """Cache independent TN/NN/NT binaries with no runtime layout argument.""" return _compile_kernel( @@ -2056,6 +2221,7 @@ def _cached_launch( group, sched, wgrad, + act_ptrs, ) @@ -2087,9 +2253,9 @@ def _validate_common_payloads( def _mxfp8_matmul_grouped_wgrad( a, # dY packed [G*M_pad, N_out], columnwise fp8 - a_scale, # [G*M_pad/32, N_out] raw E8M0 uint8 + a_scale, # [G*M_pad/32, N_out] raw E8M0 uint8, or list[G] of per-expert scales b, # X packed [G*M_pad, K_in], columnwise fp8 - b_scale, # [G*M_pad/32, K_in] raw E8M0 uint8 + b_scale, # [G*M_pad/32, K_in] raw E8M0 uint8, or list[G] of per-expert scales D, # dW packed [G*N_out, K_in] *, stream, @@ -2097,16 +2263,22 @@ def _mxfp8_matmul_grouped_wgrad( work_queue, work_counter, sched, + a_data_ptrs=None, + b_data_ptrs=None, + a_shape=None, + b_shape=None, ): """Grouped wgrad (NT) launch: dW = dY^T @ X, contracting over padded tokens. Both operands transpose-read (columnwise) and contract over the token axis; the kernel's compile-time contraction is the per-expert padded token count M_pad = (G*M_pad)/G. Experts are packed on the contraction axis for a/b and on - the output-M axis for D. Operands arrive already packed + MXFP8-quantized. + the output-M axis for D. Operands arrive already packed + MXFP8-quantized, or + (``a_data_ptrs``/``b_data_ptrs``) as one representative expert each plus + per-expert pointers, with ``a_shape``/``b_shape`` the packed shapes. """ - gm, n_out = a.shape # gm = G*M_pad, n_out = per-expert N_out - gm_b, k_in = b.shape + gm, n_out = a_shape if a_shape is not None else a.shape # gm = G*M_pad + gm_b, k_in = b_shape if b_shape is not None else b.shape if gm != gm_b: raise ValueError(f"wgrad a/b token dims differ: {gm} vs {gm_b}") if gm % num_groups != 0: @@ -2121,13 +2293,13 @@ def _mxfp8_matmul_grouped_wgrad( raise ValueError( f"wgrad D shape {tuple(D.shape)} != expected {(num_groups * n_out, k_in)}" ) - if a_scale.dtype != torch.uint8 or b_scale.dtype != torch.uint8: + if _scale_dtype(a_scale) != torch.uint8 or _scale_dtype(b_scale) != torch.uint8: raise TypeError("FlyDSL MXFP8 wgrad expects raw E8M0 scales as torch.uint8") # NT transpose-read: both scales are packed from their columnwise (K-major) # source into the kernel's [K/128, dim] int32 layout. - a_scale_hk = pack_mx32_scales_for_hk(a_scale, source_colwise=True, stream=stream) - b_scale_hk = pack_mx32_scales_for_hk(b_scale, source_colwise=True, stream=stream) + a_scale_hk = _pack_scales(a_scale, source_colwise=True, stream=stream) + b_scale_hk = _pack_scales(b_scale, source_colwise=True, stream=stream) do_gemm( a, @@ -2140,6 +2312,10 @@ def _mxfp8_matmul_grouped_wgrad( num_groups=num_groups, work_queue=work_queue, work_counter=work_counter, + b_data_ptrs=b_data_ptrs, + a_data_ptrs=a_data_ptrs, + a_shape=a_shape, + b_shape=b_shape, sched=sched, wgrad=True, ) @@ -2163,6 +2339,9 @@ def mxfp8_matmul( work_queue: fx.Tensor = None, work_counter: fx.Tensor = None, b_data_ptrs: torch.Tensor = None, + a_data_ptrs: torch.Tensor = None, + a_shape=None, + b_shape=None, sched: str = "worksteal", wgrad: bool = False, ): @@ -2181,6 +2360,11 @@ def mxfp8_matmul( TN consumes the selected rowwise payloads directly. NN and NT preserve K-major payloads and use ``ds_read_b64_tr_b8`` inside their compile-time-specialized kernels. + + ``a_data_ptrs`` reads the grouped A operand per expert in place: ``a`` is then one + representative expert and ``a_shape`` the packed shape (see ``do_gemm``). A scale + may be a list of per-expert tensors stacked on dim 0, packed in place + (:func:`pack_mx32_scales_for_hk_experts`). """ if wgrad: @@ -2195,6 +2379,10 @@ def mxfp8_matmul( work_queue=work_queue, work_counter=work_counter, sched=sched, + a_data_ptrs=a_data_ptrs, + b_data_ptrs=b_data_ptrs, + a_shape=a_shape, + b_shape=b_shape, ) if group and layout not in ("TN", "NN"): @@ -2205,14 +2393,15 @@ def mxfp8_matmul( _validate_common_payloads(a, b, D, layout=layout) + a_rows, a_cols = a_shape if a_shape is not None else a.shape if layout == "TN": - m, k = a.shape + m, k = a_rows, a_cols n, kb = b.shape elif layout == "NN": - m, k = a.shape + m, k = a_rows, a_cols kb, n = b.shape else: - k, m = a.shape + k, m = a_rows, a_cols kb, n = b.shape # Grouped GEMM is G INDEPENDENT matmuls stacked into one launch, not a dense @@ -2304,17 +2493,17 @@ def mxfp8_matmul( kb_packed = kb * num_groups if (group and b_data_ptrs is not None) else kb expected_b_scale = (kb_packed // SCALE_GROUP_SIZE, n) - if tuple(a_scale.shape) != expected_a_scale: + if _scale_shape(a_scale) != expected_a_scale: raise ValueError( - f"a_scale shape {tuple(a_scale.shape)} != expected {expected_a_scale} for {layout}" + f"a_scale shape {_scale_shape(a_scale)} != expected {expected_a_scale} for {layout}" ) - if tuple(b_scale.shape) != expected_b_scale: + if _scale_shape(b_scale) != expected_b_scale: raise ValueError( - f"b_scale shape {tuple(b_scale.shape)} != expected {expected_b_scale} for {layout}" + f"b_scale shape {_scale_shape(b_scale)} != expected {expected_b_scale} for {layout}" ) - if a_scale.dtype != torch.uint8 or b_scale.dtype != torch.uint8: + if _scale_dtype(a_scale) != torch.uint8 or _scale_dtype(b_scale) != torch.uint8: raise TypeError("FlyDSL MXFP8 expects raw E8M0 scales as torch.uint8") - if a_scale.device != a.device or b_scale.device != a.device: + if _scale_device(a_scale) != a.device or _scale_device(b_scale) != a.device: raise ValueError("A, B, scales, and D must be on the same device") if layout == "TN": @@ -2322,12 +2511,12 @@ def mxfp8_matmul( # a [M,K], b [N,K] a_kernel = a b_kernel = b - a_scale_hk = pack_mx32_scales_for_hk( + a_scale_hk = _pack_scales( a_scale, source_colwise=False, stream=stream, ) - b_scale_hk = pack_mx32_scales_for_hk( + b_scale_hk = _pack_scales( b_scale, source_colwise=False, stream=stream, @@ -2335,12 +2524,12 @@ def mxfp8_matmul( elif layout == "NN": a_kernel = a b_kernel = b - a_scale_hk = pack_mx32_scales_for_hk( + a_scale_hk = _pack_scales( a_scale, source_colwise=False, stream=stream, ) - b_scale_hk = pack_mx32_scales_for_hk( + b_scale_hk = _pack_scales( b_scale, source_colwise=True, stream=stream, @@ -2348,12 +2537,12 @@ def mxfp8_matmul( else: a_kernel = a b_kernel = b - a_scale_hk = pack_mx32_scales_for_hk( + a_scale_hk = _pack_scales( a_scale, source_colwise=True, stream=stream, ) - b_scale_hk = pack_mx32_scales_for_hk( + b_scale_hk = _pack_scales( b_scale, source_colwise=True, stream=stream, @@ -2382,6 +2571,8 @@ def mxfp8_matmul( work_queue=work_queue, work_counter=work_counter, b_data_ptrs=b_data_ptrs, + a_data_ptrs=a_data_ptrs, + a_shape=a_shape, sched=sched, ) return D diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 5422b65bfe..1887058fc8 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -91,63 +91,22 @@ import os -# Cache of the per-expert weight base-pointer device tensors used by the FlyDSL -# grouped GEMM (one int64[G] array of each expert's data_ptr()). Rebuilding it -# every call -- torch.tensor([...], device=cuda) -- costs a blocking H2D copy that -# is the single largest host cost on the grouped forward/dgrad path. Expert weight -# storages are stable across calls, so key the cache on the storage pointers and -# only re-upload when one actually moves (reallocation / first use). -_B_DATA_PTRS_CACHE = {} - - -def _cached_b_data_ptrs(datas, device): - """Return (and cache) an int64[G] device tensor of ``t.data_ptr()`` for ``datas``. - - ``datas`` is the list of per-expert operand tensors (e.g. each weight's rowwise - or columnwise payload). ``data_ptr()`` reads are host-only (no sync); the - expensive H2D ``torch.tensor`` upload runs only on a cache miss. Keyed on the - pointer tuple itself, so stable weight storage (the common case) reuses one - device tensor across every call and the cache stays bounded to the distinct - weight sets actually seen. +def _flydsl_row_operand(datas, m_splits): + """Per-expert ``[m_e, K]`` row payloads as one FlyDSL grouped-kernel operand. + + Returns ``(a_data, a_data_ptrs)``: each expert's rows are read in place through + base pointers shifted to their packed row (``grouped_gemm_utils.stacked_operand``), + so the kernel sees the ``[M_total, K]`` layout without a concat. Falls back to + ``(concat_or_view_rows(datas), None)`` when the payloads are not row-major + ``[m_e, K]`` blocks. """ - key = (tuple(t.data_ptr() for t in datas), device) - tensor = _B_DATA_PTRS_CACHE.get(key) - if tensor is None: - tensor = torch.tensor(key[0], dtype=torch.int64, device=device) - _B_DATA_PTRS_CACHE[key] = tensor - return tensor - - -def _flydsl_fp8_row_operand(datas, m_splits): - """Address per-expert FP8 row payloads in place for the FlyDSL grouped TN kernel. - - ``datas`` are the per-expert ``[m_e, K]`` rowwise payloads. Returns ``(a_data, - a_data_ptrs)``: a representative payload plus an int64[G] array of per-expert base - pointers shifted by each expert's packed row start, so the kernel reads every - expert where the ``[M_total, K]`` concat would have placed it, without the copy. - Falls back to ``(concat_or_view_rows(datas), None)`` when a payload is not a - row-major ``[m_e, K]`` block. - """ - from ..flydsl_kernels.gemm.grouped_gemm_utils import ( - concat_or_view_rows, - shifted_expert_ptrs, - ) - - k = datas[0].size(-1) - if not all( - d.dim() == 2 and d.size(0) == int(m) and d.size(1) == k and d.is_contiguous() - for d, m in zip(datas, m_splits) - ): - return concat_or_view_rows(datas), None - row_starts = [] - start = 0 - for m in m_splits: - row_starts.append(start) - start += int(m) - # Any non-empty payload serves as the kernel's A argument; its base is replaced per tile. - a_data = next((d for d in datas if d.numel() > 0), datas[0]) - a_data_ptrs = shifted_expert_ptrs(datas, row_starts, k * a_data.element_size(), a_data.device) - return a_data, a_data_ptrs + from ..flydsl_kernels.gemm.grouped_gemm_utils import concat_or_view_rows, stacked_operand + + if all(d.dim() == 2 and d.size(0) == int(m) for d, m in zip(datas, m_splits)): + stacked = stacked_operand(datas, datas[0].device) + if stacked is not None: + return stacked + return concat_or_view_rows(datas), None def _packed_wgrad_view(wgrad_list, num_gemms, n_out, k_in): @@ -932,33 +891,31 @@ def _grouped_gemm_flydsl_mxfp8( Called inline at the forward GEMM dispatch. It does not quantize -- it consumes TE's quantized operands (``inputmats`` the per-group MXFP8 activations from ``tex.split_quantize``, ``weights`` the per-expert MXFP8 - weights): activations are packed into the ``[M_total, K]`` layout and weights - are addressed via per-expert base pointers, then one TN launch is issued via + weights): activations, weights, and their scales are read per expert in place + through per-expert base pointers (no concat), then one TN launch is issued via :func:`_run_mxfp8_grouped`, writing into ``out``. The descriptor queue is built host-side from ``m_splits``. """ from ..flydsl_kernels.gemm import _run_mxfp8_grouped - from ..flydsl_kernels.gemm.grouped_gemm_utils import concat_or_view_rows - - # TN packed operands: activations concatenated along the token (M) axis, - # weights along the expert*N axis. Each MXFP8 operand carries its rowwise - # fp8 payload and E8M0 scale. Activations are bulk-allocated by - # split_quantize (per-expert payloads physically adjacent), so pack them - # with a zero-copy spanning view when contiguous; weights come from - # separate parameter storages and fall back to a real concat. - a_data = concat_or_view_rows([im._rowwise_data for im in inputmats]) - a_scale = concat_or_view_rows([im._rowwise_scale_inv for im in inputmats]) + from ..flydsl_kernels.gemm.grouped_gemm_utils import expert_ptrs + + # TN operands: activations are read in place from each expert's rows (no + # [M_total, K] concat). Each MXFP8 operand carries its rowwise fp8 payload and + # E8M0 scale; the scales are passed per expert and packed in place. + a_data, a_data_ptrs = _flydsl_row_operand( + [im._rowwise_data for im in inputmats], m_splits + ) + a_scale = [im._rowwise_scale_inv for im in inputmats] # Weights are independently-allocated per-expert slabs. The grouped-TN # kernel reads each expert's weight from a per-expert base pointer, so the # host passes an int64[G] address array instead of concatenating the G - # weights into one packed buffer (the dominant host cost). b_data is a - # per-expert representative used only for dtype/shape. The B-scale stays - # packed (cheap) and keeps the in-kernel packed-column fold. - b_data_ptrs = _cached_b_data_ptrs( + # weights into one packed buffer. b_data is a per-expert representative used + # only for dtype/shape. + b_data_ptrs = expert_ptrs( [w._rowwise_data for w in weights], a_data.device ) b_data = weights[0]._rowwise_data - b_scale = concat_or_view_rows([w._rowwise_scale_inv for w in weights]) + b_scale = [w._rowwise_scale_inv for w in weights] # Preserve the native MXFP8 dtype so the wrapper reinterprets the uint8 # payload to the correct fp8 view at the kernel boundary. @@ -967,7 +924,7 @@ def _grouped_gemm_flydsl_mxfp8( n = weights[0]._rowwise_data.size(0) k = weights[0]._rowwise_data.size(-1) - m_total = a_data.size(0) + m_total = sum(int(m) for m in m_splits) _run_mxfp8_grouped( a_data, @@ -981,6 +938,7 @@ def _grouped_gemm_flydsl_mxfp8( k=k, output_dtype=out.dtype, b_data_ptrs=b_data_ptrs, + a_data_ptrs=a_data_ptrs, ) @staticmethod @@ -996,32 +954,35 @@ def _backward_flydsl_mxfp8_grouped( Rides off backward's existing quantization -- ``grad_output`` are the per-group MXFP8 activations (rowwise) and ``weights`` the per-expert MXFP8 - weights with columnwise payloads ensured by the caller. Packs dY as the NN - kernel-A ``[M_total, N_out]`` and W (columnwise) as the contraction-packed - kernel-B ``[G*N_out, K_in]``, then issues one grouped NN launch writing into - ``dgrad`` ``[M_total, K_in]``. No quantization here. + weights with columnwise payloads ensured by the caller. dY is the NN kernel-A + ``[M_total, N_out]`` and W (columnwise) the contraction-packed kernel-B + ``[G*N_out, K_in]``, both read per expert in place through per-expert base + pointers; one grouped NN launch writes into ``dgrad`` ``[M_total, K_in]``. No + quantization here. """ from ..flydsl_kernels.gemm import _run_mxfp8_grouped - from ..flydsl_kernels.gemm.grouped_gemm_utils import concat_or_view_rows - - # NN kernel-A = dY, rowwise (contraction over N_out): concat on the token - # (M) axis. NN kernel-B = W, columnwise (transpose-read, contraction on the - # N_out axis). Weights are independent per-expert [N_out, K_in] columnwise - # slabs addressed via per-expert base pointers (no concat); dY still concats. - a_data = concat_or_view_rows([g._rowwise_data for g in grad_output]) - a_scale = concat_or_view_rows([g._rowwise_scale_inv for g in grad_output]) - b_data_ptrs = _cached_b_data_ptrs( + from ..flydsl_kernels.gemm.grouped_gemm_utils import expert_ptrs + + # NN kernel-A = dY, rowwise (contraction over N_out), read in place from each + # expert's rows. NN kernel-B = W, columnwise (transpose-read, contraction on + # the N_out axis): independent per-expert [N_out, K_in] columnwise slabs + # addressed via per-expert base pointers. Neither is concatenated. + a_data, a_data_ptrs = _flydsl_row_operand( + [g._rowwise_data for g in grad_output], m_splits + ) + a_scale = [g._rowwise_scale_inv for g in grad_output] + b_data_ptrs = expert_ptrs( [w._columnwise_data for w in weights], a_data.device ) b_data = weights[0]._columnwise_data - b_scale = concat_or_view_rows([w._columnwise_scale_inv for w in weights]) + b_scale = [w._columnwise_scale_inv for w in weights] a_data._fp8_dtype = grad_output[0]._fp8_dtype b_data._fp8_dtype = weights[0]._fp8_dtype n_out = grad_output[0]._rowwise_data.size(-1) # per-expert contraction (N_out) k_in = dgrad.size(-1) # output width (K_in) - m_total = a_data.size(0) + m_total = sum(int(m) for m in m_splits) _run_mxfp8_grouped( a_data, @@ -1036,6 +997,7 @@ def _backward_flydsl_mxfp8_grouped( output_dtype=dgrad.dtype, layout="NN", b_data_ptrs=b_data_ptrs, + a_data_ptrs=a_data_ptrs, ) @staticmethod @@ -1065,11 +1027,14 @@ def _wgrad_flydsl_mxfp8_grouped( n_out = grad_output[0]._columnwise_data.size(-1) # per-expert output feature N k_in = inputmats[0]._columnwise_data.size(-1) # per-expert input feature K - # Pack the per-expert [N_out, K_in] outputs into one [G*N_out, K_in] buffer; - # wgrad_list[g] are the per-expert slices the caller expects written. - dW = torch.empty( - num_gemms * n_out, k_in, dtype=wgrad_list[0].dtype, device=wgrad_list[0].device - ) + # Write straight into the packed wgrad storage when wgrad_list is one contiguous + # run; otherwise launch into scratch and scatter. + dW = _packed_wgrad_view(wgrad_list, num_gemms, n_out, k_in) + scatter = dW is None + if scatter: + dW = torch.empty( + num_gemms * n_out, k_in, dtype=wgrad_list[0].dtype, device=wgrad_list[0].device + ) _run_mxfp8_grouped_wgrad( dy_cols, @@ -1085,31 +1050,34 @@ def _wgrad_flydsl_mxfp8_grouped( output_dtype=wgrad_list[0].dtype, ) - # Scatter the packed result back into the caller's per-expert wgrad tensors. - for g in range(num_gemms): - wgrad_list[g].copy_(dW[g * n_out : (g + 1) * n_out]) + if scatter: + for g in range(num_gemms): + wgrad_list[g].copy_(dW[g * n_out : (g + 1) * n_out]) @staticmethod def _grouped_gemm_flydsl_bf16(*, inputmats, weights, out, m_splits, m_splits_tensor): """FlyDSL BF16 grouped forward (TN): Y_e = X_e @ W_e^T, one packed launch. Unquantized analog of :meth:`_grouped_gemm_flydsl_mxfp8`: activations are the - per-group BF16 splits (concat on the token axis), weights the per-expert BF16 - tensors (concat on the output-N axis). No scales. + per-group BF16 splits and weights the per-expert BF16 tensors, both read in + place through per-expert base pointers. No scales. """ from ..flydsl_kernels.gemm import _run_bf16_grouped - from ..flydsl_kernels.gemm.grouped_gemm_utils import concat_or_view_rows + from ..flydsl_kernels.gemm.grouped_gemm_utils import expert_ptrs - a_data = concat_or_view_rows([im.reshape(-1, im.shape[-1]) for im in inputmats]) + # Activations are read in place from each expert's rows (no [M_total, K] concat). + a_data, a_data_ptrs = _flydsl_row_operand( + [im.reshape(-1, im.shape[-1]) for im in inputmats], m_splits + ) # Weights are independent per-expert [N_out, K] tensors; the grouped-TN kernel # reads each from a per-expert base pointer, so pass an int64[G] address array # instead of concatenating the G weights. b_data is a per-expert representative # (shape only; the kernel addresses weights via the pointers). - b_data_ptrs = _cached_b_data_ptrs(list(weights), a_data.device) + b_data_ptrs = expert_ptrs(list(weights), a_data.device) b_data = weights[0] n = weights[0].size(0) k = weights[0].size(-1) - m_total = a_data.size(0) + m_total = sum(int(m) for m in m_splits) _run_bf16_grouped( a_data, @@ -1121,23 +1089,27 @@ def _grouped_gemm_flydsl_bf16(*, inputmats, weights, out, m_splits, m_splits_ten k=k, output_dtype=out.dtype, b_data_ptrs=b_data_ptrs, + a_data_ptrs=a_data_ptrs, ) @staticmethod def _backward_flydsl_bf16_grouped(*, grad_output, weights, dgrad, m_splits, m_splits_tensor): """FlyDSL BF16 grouped dgrad (NN): dX_e = dY_e @ W_e, one packed launch.""" from ..flydsl_kernels.gemm import _run_bf16_grouped - from ..flydsl_kernels.gemm.grouped_gemm_utils import concat_or_view_rows + from ..flydsl_kernels.gemm.grouped_gemm_utils import expert_ptrs - a_data = concat_or_view_rows([g.reshape(-1, g.shape[-1]) for g in grad_output]) + # dY is read in place from each expert's rows (no [M_total, N_out] concat). + a_data, a_data_ptrs = _flydsl_row_operand( + [g.reshape(-1, g.shape[-1]) for g in grad_output], m_splits + ) # Weights are independent per-expert [N_out, K_in] tensors; the grouped-NN # kernel reads each from a per-expert base pointer (no concat). b_data is a # per-expert representative. - b_data_ptrs = _cached_b_data_ptrs(list(weights), a_data.device) + b_data_ptrs = expert_ptrs(list(weights), a_data.device) b_data = weights[0] n_out = grad_output[0].shape[-1] # per-expert contraction (N_out) k_in = dgrad.size(-1) # output width (K_in) - m_total = a_data.size(0) + m_total = sum(int(m) for m in m_splits) _run_bf16_grouped( a_data, @@ -1150,6 +1122,7 @@ def _backward_flydsl_bf16_grouped(*, grad_output, weights, dgrad, m_splits, m_sp output_dtype=dgrad.dtype, layout="NN", b_data_ptrs=b_data_ptrs, + a_data_ptrs=a_data_ptrs, ) @staticmethod @@ -1163,9 +1136,14 @@ def _wgrad_flydsl_bf16_grouped(*, grad_output, inputmats, wgrad_list, m_splits): n_out = dy_cols[0].shape[-1] k_in = x_cols[0].shape[-1] - dW = torch.empty( - num_gemms * n_out, k_in, dtype=wgrad_list[0].dtype, device=wgrad_list[0].device - ) + # Write straight into the packed wgrad storage when wgrad_list is one contiguous + # run; otherwise launch into scratch and scatter. + dW = _packed_wgrad_view(wgrad_list, num_gemms, n_out, k_in) + scatter = dW is None + if scatter: + dW = torch.empty( + num_gemms * n_out, k_in, dtype=wgrad_list[0].dtype, device=wgrad_list[0].device + ) _run_bf16_grouped_wgrad( dy_cols, x_cols, @@ -1175,8 +1153,9 @@ def _wgrad_flydsl_bf16_grouped(*, grad_output, inputmats, wgrad_list, m_splits): k_in=k_in, output_dtype=wgrad_list[0].dtype, ) - for g in range(num_gemms): - wgrad_list[g].copy_(dW[g * n_out : (g + 1) * n_out]) + if scatter: + for g in range(num_gemms): + wgrad_list[g].copy_(dW[g * n_out : (g + 1) * n_out]) @staticmethod def _grouped_gemm_flydsl_fp8(*, inputmats, weights, out, m_splits, m_splits_tensor): @@ -1189,11 +1168,11 @@ def _grouped_gemm_flydsl_fp8(*, inputmats, weights, out, m_splits, m_splits_tens from ..flydsl_kernels.gemm.grouped_gemm_utils import expert_ptrs # Activations are read in place from each expert's payload (no [M_total, K] concat). - a_data, a_data_ptrs = _flydsl_fp8_row_operand([im._data for im in inputmats], m_splits) + a_data, a_data_ptrs = _flydsl_row_operand([im._data for im in inputmats], m_splits) # Weights are independent per-expert [N, K] FP8 tensors; the grouped-TN kernel # reads each from a per-expert base pointer, so pass an int64[G] address array # instead of concatenating the G weights. b_data is a per-expert representative. - b_data_ptrs = _cached_b_data_ptrs( + b_data_ptrs = expert_ptrs( [w._data for w in weights], a_data.device ) b_data = weights[0]._data @@ -1239,8 +1218,8 @@ def _backward_flydsl_fp8_grouped( # A = dY (rowwise fp8 data); B = W columnwise transpose, the materialized # [K_in, N_out] packed on K_in (the TN GEMM output axis). Both are independent # per-expert slabs, addressed via per-expert base pointers (no concat). - a_data, a_data_ptrs = _flydsl_fp8_row_operand([g._data for g in grad_output], m_splits) - b_data_ptrs = _cached_b_data_ptrs( + a_data, a_data_ptrs = _flydsl_row_operand([g._data for g in grad_output], m_splits) + b_data_ptrs = expert_ptrs( [w._transpose for w in weights], a_data.device ) b_data = weights[0]._transpose From d14a87f32190c7f21332bc4a2250dfd6134d5193 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 21 Oct 2026 12:58:59 +0000 Subject: [PATCH 15/22] Make FlyDSL grouped GEMM dispatch non-blocking and cheaper on the host --- .../pytorch/flydsl_kernels/gemm/fp32_gemm.py | 3 + .../pytorch/flydsl_kernels/gemm/fp8_gemm.py | 54 +++-- .../flydsl_kernels/gemm/gemm_wrappers.py | 115 +++++----- .../flydsl_kernels/gemm/grouped_gemm_utils.py | 209 +++++++++++------- .../flydsl_kernels/gemm/half_prec_gemm.py | 39 +++- .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 174 +++++++-------- .../pytorch/module/grouped_linear.py | 109 +++++---- transformer_engine/pytorch/quantization.py | 26 ++- 8 files changed, 415 insertions(+), 314 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py index 5523312f58..68e95fe471 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py @@ -1003,6 +1003,9 @@ def _compiled_launch(key, jit_fn, args): for the first call; used once to compile, then the cached CompiledFunction is reused positionally for all later calls with the same ABI. """ + # A compiled launcher is built for the current device; key on it so a launch on + # another device compiles its own instead of reusing this one. + key = (key, torch.cuda.current_device()) cf = _COMPILED_LAUNCH_CACHE.get(key) if cf is None: cf = flyc.compile(jit_fn, *args) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py index d446299d6e..8e0d686d0c 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -221,6 +221,7 @@ def kernel_gemm( work_counter: fx.Tensor = None, b_data_ptrs: fx.Tensor = None, a_data_ptrs: fx.Tensor = None, + num_work: fx.Int32 = 0, ): lds = fx.SharedAllocator().allocate(SharedStorage).peek() lds_a0 = (lds.a0_0, lds.a0_1) @@ -258,11 +259,12 @@ def _load_scale(scale_div, index): num_blocks_m = c_m // BLOCK_M num_blocks_n = c_n // BLOCK_N - # Total output tiles. Dense uses this as the grid extent; grouped uses it as - # the persistent-loop drain bound. ``grouped`` gates the shared descriptor - # queue / work-steal driver; const_expr folds it out of the dense binary. + # ``grouped`` gates the shared descriptor queue / work-steal driver; + # const_expr folds it out of the dense binary. grouped = group or wgrad - num_tiles = gd.num_tiles_bound(c_m, c_n, BLOCK_M, BLOCK_N, num_groups, const_expr(wgrad)) + # Grouped drain bound: the exact number of descriptors the host wrote to + # ``work_queue`` (``num_work``), so no slot past the queue end is ever read. + num_tiles = num_work # weight_ptrs: weights are G independently-allocated per-expert slabs, # addressed from a per-expert base pointer read from ``b_data_ptrs`` (reusing @@ -437,7 +439,11 @@ def load_bias(col): gAux = fx.rocdl.make_buffer_tensor(Aux, max_size=True) aux_div = fx.logical_divide(gAux, fx.make_layout(1, 1)) aux_store_atom = fx.make_copy_atom( - fx.rocdl.BufferCopy32b() if output_element_bytes == 4 else fx.rocdl.BufferCopy16b(), + ( + fx.rocdl.BufferCopy32b() + if output_element_bytes == 4 + else fx.rocdl.BufferCopy16b() + ), output_fx_dtype, ) @@ -547,9 +553,8 @@ def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): # wgrad, where A = dY^T is packed [G*N_out, m_pad] on the M-row); 0 # for dense and grouped fwd/dgrad (A is per-token, absolute rows). global_base = ( - (bx_m_idx + a_expert_m_off + fx.Index(subtile * (BLOCK_M // 2))) * fx.Index(K) - + k_base - ) + bx_m_idx + a_expert_m_off + fx.Index(subtile * (BLOCK_M // 2)) + ) * fx.Index(K) + k_base a_g2s.load_one(lds_a[subtile], fx.Int32(global_base), pass_in_subtile) def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): @@ -557,10 +562,8 @@ def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): # view (b_g2s) with a LOCAL column (b_data_col). Other paths: by_n_b_idx # shifts the B row to the expert block (grouped dgrad weights, grouped # wgrad X^T); reduces to dense by_n for the dense path. - global_base = ( - (b_data_col + fx.Index(subtile * (BLOCK_N // 2))) * fx.Index(K) - + k_base - ) + b_row = b_data_col + fx.Index(subtile * (BLOCK_N // 2)) + global_base = b_row * fx.Index(K) + k_base b_g2s.load_one(lds_b[subtile], fx.Int32(global_base), pass_in_subtile) def stage_a_subtile(k_base, subtile, lds_a): @@ -594,7 +597,9 @@ def load_b_frag(lds_b, local_row, half): def _acc_idx(subtile_id, mi, ni): return ( - subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + mi * MFMA_N_PER_SUBTILE + ni + subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + + mi * MFMA_N_PER_SUBTILE + + ni ) def pinned_mfma(acc_idx, a_frag, b_frag): @@ -713,7 +718,9 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): def load_b_subtile_ni_regs(lds_b, sn, ni): subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - b_row_addr = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 + b_row_addr = ( + subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 + ) return load_b_frag(lds_b, b_row_addr, sn) def load_b_subtile_regs(lds_b, sn): @@ -726,7 +733,9 @@ def load_b_subtile_regs(lds_b, sn): def load_a_subtile_mi_half(lds_a, sm, mi, half): subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - a_row_addr = subtile_m_idx * fx.Index(SUBTILE_M) + fx.Index(mi * MFMA_M) + lane_mod_16 + a_row_addr = ( + subtile_m_idx * fx.Index(SUBTILE_M) + fx.Index(mi * MFMA_M) + lane_mod_16 + ) half_row = a_row_addr - fx.Index(sm * (BLOCK_M // 2)) row_byte_base = half_row * fx.Index(BLOCK_K) return load_frag_half_at_byte_base(lds_a[sm], row_byte_base, half) @@ -1150,6 +1159,7 @@ def launch_gemm( work_counter: fx.Tensor = None, b_data_ptrs: fx.Tensor = None, a_data_ptrs: fx.Tensor = None, + num_work: fx.Int32 = 0, stream: fx.Stream = fx.Stream(None), ): # Dense: one workgroup per output tile. Grouped/wgrad: a fixed num-CU @@ -1173,6 +1183,7 @@ def launch_gemm( work_counter, b_data_ptrs, a_data_ptrs, + num_work, value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, ).launch(grid=(grid_x, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) @@ -1215,6 +1226,9 @@ def _compiled_launch(key, jit_fn, args): for the first call; used once to compile, then the cached CompiledFunction is reused positionally for all later calls with the same ABI. """ + # A compiled launcher is built for the current device; key on it so a launch on + # another device compiles its own instead of reusing this one. + key = (key, torch.cuda.current_device()) cf = _COMPILED_LAUNCH_CACHE.get(key) if cf is None: cf = flyc.compile(jit_fn, *args) @@ -1314,6 +1328,7 @@ def doGemm( work_counter: torch.Tensor = None, b_data_ptrs: torch.Tensor = None, a_data_ptrs: torch.Tensor = None, + num_work: int = 0, sched: str = "worksteal", wgrad: bool = False, m: int = None, @@ -1398,7 +1413,7 @@ def doGemm( raise ValueError(f"FP8 epilogue {epilogue} does not accept an aux tensor") if stream is None: - stream = torch.cuda.current_stream() + stream = gd.current_stream_handle(A.device) A_arg = A.view(torch.uint8).contiguous().view(-1) B_arg = B.view(torch.uint8).contiguous().view(-1) @@ -1420,6 +1435,12 @@ def doGemm( # Grouped runtime tensors are required by the kernel signature; the dense # binary never dereferences them but still needs placeholder buffers. if group or wgrad: + # num_work is the kernel's drain bound: it must be the queue's descriptor count. + if work_queue.numel() != 3 * int(num_work): + raise ValueError( + f"num_work={num_work} does not match work_queue of {work_queue.numel()} " + "int32 (3 per descriptor)" + ) work_queue_arg = work_queue.contiguous().view(-1) work_counter_arg = work_counter.contiguous().view(-1) else: @@ -1473,6 +1494,7 @@ def doGemm( work_counter_arg, b_data_ptrs_arg, a_data_ptrs_arg, + int(num_work), stream, ) # flyc.compile fast path: skips flydsl's per-call cache-key rebuild (sig.bind + diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 2676202896..9e94590a47 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -1032,24 +1032,26 @@ def _run_mxfp8_grouped( ): """Dispatch a packed MXFP8 grouped GEMM through the persistent kernel. - Two layouts, both single packed launches spanning all groups: + Two layouts, both single launches spanning all groups. Weights are read per + expert through the required ``b_data_ptrs`` (``b_data`` is one representative); + scale shapes are those of the per-expert scales stacked on dim 0: TN (forward, C = A @ B^T): - a = packed activations [M_total, K] scales [M_total, K/32] - b = packed weights [G*N, K] scales [G*N, K/32] - output width n = N_out (per expert); experts pack on B's output-N axis. + a = activations [M_total, K] scales [M_total, K/32] + b = per-expert weights [N, K] scales [G*N, K/32] + output width n = N_out (per expert); experts stack on B's output-N axis. NN (dgrad, dX = dY @ W): - a = packed grad_output [M_total, N_out] scales [M_total, N_out/32] - b = packed weights [G*N_out, K_in] scales [G*N_out/32, K_in] - output width n = K_in; experts pack on B's contraction (N_out) axis. + a = grad_output [M_total, N_out] scales [M_total, N_out/32] + b = per-expert weights [N_out, K_in] scales [G*N_out/32, K_in] + output width n = K_in; experts stack on B's contraction (N_out) axis. In both, ``k`` is A's per-expert contraction (TN: K_in; NN: N_out) and ``n`` is the output width. ``m_splits`` is the Python list of per-expert token counts (already host-resident); the per-tile ``(pid_m, pid_n, expert)`` descriptor - queue is built + shape-cached from it directly, so no device->host copy of a - ``group_offs`` tensor (and its blocking sync) is needed. Operands arrive already - MXFP8-quantized and packed -- no quantization, no BLAS ownership swap. + queue is built from it on the host and uploaded without blocking, so no + device->host copy of a ``group_offs`` tensor (and its blocking sync) is needed. + Operands arrive already MXFP8-quantized -- no quantization, no BLAS ownership swap. ``a_data_ptrs`` (optional) reads A per expert in place: ``a_data`` is then one representative expert payload (see ``grouped_gemm_utils.stacked_operand``). """ @@ -1097,7 +1099,7 @@ def _run_mxfp8_grouped( # token count must be BLOCK_M-aligned (enforced upstream by the tiling check), # so tiles never straddle a group boundary. Built + cached from the host-side # m_splits list (no GPU->CPU sync). - work_queue, work_counter, num_tiles = gd.build_fwd_dgrad_work_queue( + work_queue, num_tiles = gd.build_fwd_dgrad_work_queue( m_splits, num_pid_n, BLOCK_M, a_data.device ) @@ -1108,15 +1110,14 @@ def _run_mxfp8_grouped( # the launch grid, one grid-stride sweep already balances well, so the static # path avoids atomic contention. NVTE_FLYDSL_GROUPED_SCHED overrides for tests. sched = _select_grouped_sched(num_tiles, a_data.device) + work_counter = gd.work_counter(sched, a_data.device) epilogue, bias_arg = _resolve_epilogue(bias, n) # Both TN forward and NN dgrad address the weight operand from per-expert base - # pointers supplied by the caller (one int64 base per expert slab) when - # ``b_data_ptrs`` is given; the kernel reads B slab g from b_data_ptrs[g] instead - # of folding the expert into a packed B (an output-N column for TN, the - # contraction base for NN), so the host never concatenates the G per-expert - # weights. When ``b_data_ptrs`` is None the packed-B expert fold is used. + # pointers supplied by the caller (one int64 base per expert slab, required): the + # kernel reads B slab g from b_data_ptrs[g], so the host never concatenates the G + # per-expert weights. mxfp8_matmul( a_data, a_scale, @@ -1132,6 +1133,7 @@ def _run_mxfp8_grouped( work_counter=work_counter, b_data_ptrs=b_data_ptrs, a_data_ptrs=a_data_ptrs, + num_work=num_tiles, a_shape=(m_total, a_data.size(-1)) if a_data_ptrs is not None else None, sched=sched, ) @@ -1155,13 +1157,13 @@ def _wgrad_operands(dy_b, x_b, unpadded, pad, device): def _run_mxfp8_grouped_wgrad( - dy_cols, # list[G] of per-expert dY columnwise fp8 data, each [m_e, N_out] - dy_scales, # list[G] of per-expert dY columnwise scale, each [m_e/32, N_out] - x_cols, # list[G] of per-expert X columnwise fp8 data, each [m_e, K_in] - x_scales, # list[G] of per-expert X columnwise scale, each [m_e/32, K_in] - dW, # output [G*N_out, K_in] + dy_cols, # list[G] of per-expert dY columnwise fp8 data, each [m_e, N_out] + dy_scales, # list[G] of per-expert dY columnwise scale, each [m_e/32, N_out] + x_cols, # list[G] of per-expert X columnwise fp8 data, each [m_e, K_in] + x_scales, # list[G] of per-expert X columnwise scale, each [m_e/32, K_in] + dW, # output [G*N_out, K_in] *, - m_splits, # list[int] per-expert token counts + m_splits, # list[int] per-expert token counts n_out, k_in, dy_fp8_dtype, @@ -1235,13 +1237,11 @@ def _run_mxfp8_grouped_wgrad( x_data = reinterpret_as_fp8_tensor(x_data, x_fp8_dtype) direct = ids == list(range(num_groups)) - dW_b = dW if direct else torch.empty(bsize * n_out, k_in, dtype=output_dtype, - device=device) + dW_b = dW if direct else torch.empty(bsize * n_out, k_in, dtype=output_dtype, device=device) - work_queue, work_counter, num_tiles = gd.build_wgrad_work_queue( - bsize, num_pid_m, num_pid_n, device - ) + work_queue, num_tiles = gd.build_wgrad_work_queue(bsize, num_pid_m, num_pid_n, device) sched = _select_grouped_sched(num_tiles, device) + work_counter = gd.work_counter(sched, device) mxfp8_matmul( dy_data, @@ -1256,6 +1256,7 @@ def _run_mxfp8_grouped_wgrad( sched=sched, wgrad=True, a_data_ptrs=dy_ptrs, + num_work=num_tiles, b_data_ptrs=x_ptrs, a_shape=(bsize * k_pad, n_out) if dy_ptrs is not None else None, b_shape=(bsize * k_pad, k_in) if x_ptrs is not None else None, @@ -1285,14 +1286,14 @@ def _run_bf16_grouped( ): """Dispatch a packed BF16 grouped GEMM through the persistent kernel. - Scaleless analog of :func:`_run_mxfp8_grouped`; the packed operand contracts - match the MXFP8 TN/NN layouts exactly (experts packed on B's output-N axis for - TN, on B's contraction axis for NN), only without E8M0 scales: + Scaleless analog of :func:`_run_mxfp8_grouped`, with the same TN/NN operand + contracts, only without E8M0 scales. Weights are read per expert through the + required ``b_data_ptrs``; ``b_data`` is one representative expert: TN (forward, C = A @ B^T): - a = packed activations [M_total, K] b = packed weights [G*N, K] + a = activations [M_total, K] b = per-expert weights [N, K] NN (dgrad, dX = dY @ W): - a = packed grad_output [M_total, N_out] b = packed weights [G*N_out, K_in] + a = grad_output [M_total, N_out] b = per-expert weights [N_out, K_in] ``m_splits`` is the Python list of per-expert token counts (already on the host). The descriptor queue is built + cached from it directly, so no @@ -1301,7 +1302,9 @@ def _run_bf16_grouped( from .half_prec_gemm import doGemm as _half_prec_doGemm if layout not in ("TN", "NN"): - raise FlyDSLUnsupportedError(f"FlyDSL half-precision grouped GEMM supports TN/NN, got {layout}") + raise FlyDSLUnsupportedError( + f"FlyDSL half-precision grouped GEMM supports TN/NN, got {layout}" + ) # BF16 and FP16 share this kernel (the MFMA opcode is selected per operand dtype); # take the input dtype from the operands so both route through unchanged. @@ -1323,10 +1326,11 @@ def _run_bf16_grouped( num_groups = len(m_splits) num_pid_n = n // _BLOCK_N - work_queue, work_counter, num_tiles = gd.build_fwd_dgrad_work_queue( + work_queue, num_tiles = gd.build_fwd_dgrad_work_queue( m_splits, num_pid_n, _BLOCK_M, a_data.device ) sched = _select_grouped_sched(num_tiles, a_data.device) + work_counter = gd.work_counter(sched, a_data.device) epilogue, bias_arg = _resolve_epilogue(bias, n) _half_prec_doGemm( @@ -1347,6 +1351,7 @@ def _run_bf16_grouped( work_counter=work_counter, b_data_ptrs=b_data_ptrs, a_data_ptrs=a_data_ptrs, + num_work=num_tiles, sched=sched, ) return D @@ -1361,9 +1366,9 @@ def _wgrad_max_buckets(): def _run_bf16_grouped_wgrad( - dy_cols, # list[G] of per-expert dY, each [m_e, N_out] - x_cols, # list[G] of per-expert X, each [m_e, K_in] - dW, # output [G*N_out, K_in] + dy_cols, # list[G] of per-expert dY, each [m_e, N_out] + x_cols, # list[G] of per-expert X, each [m_e, K_in] + dW, # output [G*N_out, K_in] *, m_splits, n_out, @@ -1434,13 +1439,11 @@ def _run_bf16_grouped_wgrad( # buffer and no scatter copy. Only multi-bucket (skewed) launches allocate a # bucket-local output and scatter it back into dW. direct = ids == list(range(num_groups)) - dW_b = dW if direct else torch.empty(bsize * n_out, k_in, dtype=output_dtype, - device=device) + dW_b = dW if direct else torch.empty(bsize * n_out, k_in, dtype=output_dtype, device=device) - work_queue, work_counter, num_tiles = gd.build_wgrad_work_queue( - bsize, num_pid_m, num_pid_n, device - ) + work_queue, num_tiles = gd.build_wgrad_work_queue(bsize, num_pid_m, num_pid_n, device) sched = _select_grouped_sched(num_tiles, device) + work_counter = gd.work_counter(sched, device) _half_prec_doGemm( dy_data, @@ -1459,6 +1462,7 @@ def _run_bf16_grouped_wgrad( sched=sched, wgrad=True, a_data_ptrs=dy_ptrs, + num_work=num_tiles, b_data_ptrs=x_ptrs, ) @@ -1480,9 +1484,9 @@ def _fp8_scale_arg(scale): def _run_fp8_grouped( a_data, - a_scale_vec, # [num_groups] fp32 per-expert inverse scales, or int64 scale_inv pointers + a_scale_vec, # [num_groups] fp32 per-expert inverse scales, or int64 scale_inv pointers b_data, - b_scale_vec, # [num_groups] fp32 per-expert inverse scales, or int64 scale_inv pointers + b_scale_vec, # [num_groups] fp32 per-expert inverse scales, or int64 scale_inv pointers D, *, m_splits, @@ -1534,10 +1538,11 @@ def _run_fp8_grouped( from .fp8_gemm import _BLOCK_M, _BLOCK_N num_pid_n = n // _BLOCK_N - work_queue, work_counter, num_tiles = gd.build_fwd_dgrad_work_queue( + work_queue, num_tiles = gd.build_fwd_dgrad_work_queue( m_splits, num_pid_n, _BLOCK_M, a_data.device ) sched = _select_grouped_sched(num_tiles, a_data.device) + work_counter = gd.work_counter(sched, a_data.device) epilogue, bias_arg = _resolve_epilogue(bias, n) _fp8_doGemm( @@ -1554,6 +1559,7 @@ def _run_fp8_grouped( work_counter=work_counter, b_data_ptrs=b_data_ptrs, a_data_ptrs=a_data_ptrs, + num_work=num_tiles, sched=sched, m=m_total, n=n, @@ -1563,11 +1569,11 @@ def _run_fp8_grouped( def _run_fp8_grouped_wgrad( - dyt_cols, # list[G] of per-expert dY^T (host-transposed) [N_out, m_e] - dy_scale_vec, # [num_groups] fp32 inverse scales for dY, or list[G] of per-expert scale_inv - xt_cols, # list[G] of per-expert X^T (host-transposed) [K_in, m_e] - x_scale_vec, # [num_groups] fp32 inverse scales for X, or list[G] of per-expert scale_inv - dW, # output [G*N_out, K_in] + dyt_cols, # list[G] of per-expert dY^T (host-transposed) [N_out, m_e] + dy_scale_vec, # [num_groups] fp32 inverse scales for dY, or list[G] of per-expert scale_inv + xt_cols, # list[G] of per-expert X^T (host-transposed) [K_in, m_e] + x_scale_vec, # [num_groups] fp32 inverse scales for X, or list[G] of per-expert scale_inv + dW, # output [G*N_out, K_in] *, m_splits, n_out, @@ -1653,13 +1659,11 @@ def _pad_cols(t, cols): x_scale_b = x_scale_vec[idx] direct = ids == list(range(num_groups)) - dW_b = dW if direct else torch.empty(bsize * n_out, k_in, dtype=output_dtype, - device=device) + dW_b = dW if direct else torch.empty(bsize * n_out, k_in, dtype=output_dtype, device=device) - work_queue, work_counter, num_tiles = gd.build_wgrad_work_queue( - bsize, num_pid_m, num_pid_n, device - ) + work_queue, num_tiles = gd.build_wgrad_work_queue(bsize, num_pid_m, num_pid_n, device) sched = _select_grouped_sched(num_tiles, device) + work_counter = gd.work_counter(sched, device) _fp8_doGemm( dy_data, @@ -1677,6 +1681,7 @@ def _pad_cols(t, cols): n=k_in, k=k_pad, a_data_ptrs=dy_ptrs, + num_work=num_tiles, b_data_ptrs=x_ptrs, ) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py index ff3fd72d29..8623227950 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py @@ -15,7 +15,7 @@ ``kernel_gemm``. What's shared instead: * **Device-side IR emitters** (``global_i32_ptr``, ``make_claim_broadcast``, - ``read_desc_field``, ``num_tiles_bound``) -- control-flow-free functions that + ``read_desc_field``) -- control-flow-free functions that append IR at the current insertion point, so they emit identically whether defined inline or imported. * **Host-side builders** (descriptor-queue construction, schedule selection, @@ -26,6 +26,7 @@ import functools import os +import numpy as np import torch import flydsl.expr as fx @@ -34,6 +35,8 @@ from flydsl._mlir.dialects import llvm from flydsl.expr.typing import T +from .exceptions import FlyDSLUnsupportedError + # --------------------------------------------------------------------------- # Device-side IR emitters (called while tracing a kernel; pure, no control flow) @@ -138,20 +141,6 @@ def read_desc_field(wq_div, desc_ld_atom, slot, field): return fx.memref_load_vec(reg)[0] -def num_tiles_bound(c_m, c_n, block_m, block_n, num_groups, wgrad): - """Drain bound for the persistent loop. - - For dense/grouped fwd/dgrad, ``c_m`` already spans all experts, so the tile - count is complete. For wgrad, ``c_m = N_out`` is the PER-EXPERT output-row - count (the packed ``[G*N_out, K_in]`` output is addressed via - ``store_m_expert_off``), so multiply by ``num_groups``. - """ - num_tiles = (c_m // block_m) * (c_n // block_n) - if wgrad: - num_tiles = num_tiles * num_groups - return num_tiles - - # --------------------------------------------------------------------------- # Host-side builders (plain torch; dtype-independent) # --------------------------------------------------------------------------- @@ -176,75 +165,126 @@ def select_grouped_sched(num_tiles: int, device) -> str: raise ValueError( f"NVTE_FLYDSL_GROUPED_SCHED must be auto|static|worksteal, got {override!r}" ) - num_cus = torch.cuda.get_device_properties(device).multi_processor_count - return "static" if num_tiles <= _GROUPED_SCHED_ROUNDS * num_cus else "worksteal" + return "static" if num_tiles <= _GROUPED_SCHED_ROUNDS * _num_cus(device) else "worksteal" + + +@functools.lru_cache(maxsize=None) +def _num_cus(device): + return torch.cuda.get_device_properties(device).multi_processor_count + +# Per-call metadata caches. Keys include the current stream: a cached device tensor is +# filled by a non-blocking copy on the stream that first built it, so only launches on +# that stream are ordered after the copy. Pointer arrays and block tables are tiny +# (<= a few KB each), so their cache is large enough to hold a training step's worth of +# rotating per-expert buffers (each step adds several arrays); the work queues can be +# much larger, so their cache stays small. +_SMALL_CACHE_SIZE = 512 +_QUEUE_CACHE_SIZE = 64 -@functools.lru_cache(maxsize=512) -def _fwd_dgrad_queue_cached(m_splits, num_pid_n, block_m, device): - """Cached flat ``(pid_m, pid_n, expert)`` work queue for a fwd/dgrad shape. - Keyed by the per-expert token counts + tiling, so repeated calls with the same - ``m_splits`` (every microbatch in training, every bench iteration) reuse one - already-uploaded queue instead of rebuilding + re-H2D-copying it each forward. - Built from the Python ``m_splits`` tuple directly -- NO device->host copy of - ``group_offs`` (that copy forced a ``hipDeviceSynchronize`` that stranded the - GPU idle for hundreds of microseconds before the kernel). +def current_stream_handle(device=None): + """Raw handle of the current stream on ``device`` (default: the current device). + + FlyDSL launchers accept an int stream, and this is ~20x cheaper than building a + torch ``Stream`` object per launch. """ - descriptors = [] - off = 0 - for g, m in enumerate(m_splits): - m_start_blk = off // block_m - m_end_blk = (off + m) // block_m - for pid_m in range(m_start_blk, m_end_blk): - for pid_n in range(num_pid_n): - descriptors.append((pid_m, pid_n, g)) - off += m - num_tiles = len(descriptors) - work_queue = torch.tensor(descriptors, dtype=torch.int32, device=device).reshape(-1) - return work_queue, num_tiles + index = None if device is None else device.index + if index is None: + index = torch.cuda.current_device() + return torch._C._cuda_getCurrentRawStream(index) + + +def upload_to_device(host, device): + """Copy a host CPU tensor to ``device`` without blocking the host. + + ``torch.tensor(..., device="cuda")`` waits for all queued GPU work before copying, + which stalls the host on every call. Staging through pinned memory and copying + ``non_blocking`` does not; the caching host allocator keeps the pinned block alive + until the copy has run. The staging block comes from ``torch.empty(pin_memory=True)`` + (the caching host allocator) rather than ``Tensor.pin_memory()``, which is ~3x slower + per call while GPU work is queued. + """ + staging = torch.empty(host.shape, dtype=host.dtype, pin_memory=True) + staging.copy_(host) + return staging.to(device, non_blocking=True) + + +@functools.lru_cache(maxsize=_QUEUE_CACHE_SIZE) +def _fwd_dgrad_queue_cached(m_splits, num_pid_n, block_m, device, stream): + """Flat ``(pid_m, pid_n, expert)`` work queue for a fwd/dgrad shape. + + Every expert's token count is a ``block_m`` multiple, so the expert row blocks tile + ``[0, M_total)`` in order: the queue is each row block repeated over the output + columns, tagged with its expert. Built with numpy from the host ``m_splits`` (no + device->host copy of ``group_offs``) and uploaded without blocking the host. + """ + blocks = np.asarray(m_splits, dtype=np.int64) // block_m + num_blocks = int(blocks.sum()) + queue = np.empty((num_blocks, num_pid_n, 3), dtype=np.int32) + queue[:, :, 0] = np.arange(num_blocks, dtype=np.int32)[:, None] + queue[:, :, 1] = np.arange(num_pid_n, dtype=np.int32)[None, :] + queue[:, :, 2] = np.repeat(np.arange(len(m_splits), dtype=np.int32), blocks)[:, None] + work_queue = upload_to_device(torch.from_numpy(queue.reshape(-1)), device) + return work_queue, num_blocks * num_pid_n def build_fwd_dgrad_work_queue(m_splits, num_pid_n, block_m, device): - """Return ``(work_queue, work_counter, num_tiles)`` for a fwd/dgrad launch. + """Return ``(work_queue, num_tiles)`` for a fwd/dgrad launch. - The descriptor queue is cached (shape-only, data-independent) so only the first - call with a given ``m_splits`` pays the build+upload; the zeroed atomic counter - is cheap and allocated fresh each call. ``m_splits`` is a Python list/tuple of - per-expert token counts -- no GPU->CPU sync. + Raises ``FlyDSLUnsupportedError`` when an expert's token count is not a ``block_m`` + multiple: the kernel tiles each expert in whole ``block_m`` row blocks, so a + misaligned split would silently mix experts within a tile. ``num_tiles`` is the + exact queue length, which the kernel uses as its drain bound. ``m_splits`` is a + Python list/tuple of per-expert token counts -- no GPU->CPU sync. """ - work_queue, num_tiles = _fwd_dgrad_queue_cached( - tuple(int(m) for m in m_splits), int(num_pid_n), int(block_m), device + m_splits = tuple(int(m) for m in m_splits) + for g, m in enumerate(m_splits): + if m % block_m != 0: + raise FlyDSLUnsupportedError( + f"m_splits[{g}]={m} is not a multiple of BLOCK_M={block_m}" + ) + return _fwd_dgrad_queue_cached( + m_splits, int(num_pid_n), int(block_m), device, current_stream_handle(device) ) - work_counter = torch.zeros(1, dtype=torch.int32, device=device) - return work_queue, work_counter, num_tiles -@functools.lru_cache(maxsize=512) -def _wgrad_queue_cached(num_groups, num_pid_m, num_pid_n, device): - descriptors = [ - (pid_m, pid_n, g) - for g in range(num_groups) - for pid_m in range(num_pid_m) - for pid_n in range(num_pid_n) - ] - num_tiles = len(descriptors) - work_queue = torch.tensor(descriptors, dtype=torch.int32, device=device).reshape(-1) - return work_queue, num_tiles +@functools.lru_cache(maxsize=_QUEUE_CACHE_SIZE) +def _wgrad_queue_cached(num_groups, num_pid_m, num_pid_n, device, stream): + queue = np.empty((num_groups, num_pid_m, num_pid_n, 3), dtype=np.int32) + queue[..., 0] = np.arange(num_pid_m, dtype=np.int32)[None, :, None] + queue[..., 1] = np.arange(num_pid_n, dtype=np.int32)[None, None, :] + queue[..., 2] = np.arange(num_groups, dtype=np.int32)[:, None, None] + work_queue = upload_to_device(torch.from_numpy(queue.reshape(-1)), device) + return work_queue, num_groups * num_pid_m * num_pid_n def build_wgrad_work_queue(num_groups, num_pid_m, num_pid_n, device): - """Return cached ``(work_queue, work_counter, num_tiles)`` for a wgrad launch. + """Return ``(work_queue, num_tiles)`` for a wgrad launch. The wgrad descriptor queue depends only on ``(num_groups, num_pid_m, num_pid_n)`` - (per-expert tile geometry), so it is shape-cached and reused across calls -- no - per-call Python rebuild or H2D re-upload. Fresh zeroed atomic counter each call. + (per-expert tile geometry), so it is shape-cached and reused across calls; + ``num_tiles`` is the kernel's drain bound. """ - work_queue, num_tiles = _wgrad_queue_cached( - int(num_groups), int(num_pid_m), int(num_pid_n), device + return _wgrad_queue_cached( + int(num_groups), int(num_pid_m), int(num_pid_n), device, current_stream_handle(device) ) - work_counter = torch.zeros(1, dtype=torch.int32, device=device) - return work_queue, work_counter, num_tiles + + +@functools.lru_cache(maxsize=None) +def _unused_counter(device): + return torch.zeros(1, dtype=torch.int32, device=device) + + +def work_counter(sched, device): + """The atomic tile counter a grouped launch needs for schedule ``sched``. + + Work-stealing claims tiles from a freshly zeroed counter. The static schedule never + touches it, so it gets a shared placeholder instead of a per-launch zero fill. + """ + if sched == "worksteal": + return torch.zeros(1, dtype=torch.int32, device=device) + return _unused_counter(device) class _SpanCAI: @@ -328,27 +368,27 @@ def concat_or_view_rows(tensors): row_bytes *= elem total_bytes = total_rows * row_bytes - byte_view = torch.as_tensor( - _SpanCAI(total_bytes, first.data_ptr()), device=first.device - ) + byte_view = torch.as_tensor(_SpanCAI(total_bytes, first.data_ptr()), device=first.device) view = byte_view.view(first.dtype).view(total_rows, *tuple(trailing)) view._flydsl_span_anchor = tensors return view -@functools.lru_cache(maxsize=512) -def _ptr_array_cached(ptrs, device): - return torch.tensor(ptrs, dtype=torch.int64, device=device) +@functools.lru_cache(maxsize=_SMALL_CACHE_SIZE) +def _ptr_array_cached(ptrs, device, stream): + return upload_to_device(torch.tensor(ptrs, dtype=torch.int64), device) def expert_ptrs(tensors, device): """int64[G] device array of each per-expert tensor's ``data_ptr()``. Grouped kernels read per-expert operands (weights, FP8 scale_inv) through this - array instead of a host-side concat. The upload runs only on a cache miss: weights - and the caching allocator's per-expert buffers repeat step to step. + array instead of a host-side concat. Repeated pointer sets reuse the uploaded array; + a new set costs one non-blocking upload. """ - return _ptr_array_cached(tuple(t.data_ptr() for t in tensors), device) + return _ptr_array_cached( + tuple(t.data_ptr() for t in tensors), device, current_stream_handle(device) + ) def stacked_operand(tensors, device): @@ -362,22 +402,23 @@ def stacked_operand(tensors, device): non-empty expert tensor standing in for the operand's dtype. """ first = tensors[0] + dtype = first.dtype trailing = first.shape[1:] ptrs = [] offset = 0 for t in tensors: - if t.dtype != first.dtype or t.shape[1:] != trailing or not t.is_contiguous(): + if t.dtype is not dtype or t.shape[1:] != trailing or not t.is_contiguous(): return None ptrs.append(t.data_ptr() - offset) - offset += t.numel() * t.element_size() + offset += t.nbytes representative = next((t for t in tensors if t.numel() > 0), first) - return representative, _ptr_array_cached(tuple(ptrs), device) + return representative, _ptr_array_cached(tuple(ptrs), device, current_stream_handle(device)) -@functools.lru_cache(maxsize=512) -def _block_expert_table_cached(blocks_per_expert, device): - table = [e for e, n in enumerate(blocks_per_expert) for _ in range(n)] - return torch.tensor(table, dtype=torch.int32, device=device) +@functools.lru_cache(maxsize=_SMALL_CACHE_SIZE) +def _block_expert_table_cached(blocks_per_expert, device, stream): + table = np.repeat(np.arange(len(blocks_per_expert), dtype=np.int32), blocks_per_expert) + return upload_to_device(torch.from_numpy(table), device) def block_expert_table(extents, block, device): @@ -389,7 +430,9 @@ def block_expert_table(extents, block, device): for n in extents: if int(n) % block != 0: raise ValueError(f"per-expert extent {int(n)} is not a multiple of {block}") - return _block_expert_table_cached(tuple(int(n) // block for n in extents), device) + return _block_expert_table_cached( + tuple(int(n) // block for n in extents), device, current_stream_handle(device) + ) def drop_empty_experts(buckets, m_splits, dW, n_out): @@ -450,6 +493,7 @@ def bucketize_wgrad_tokens(m_splits, block_m, max_buckets=None): unchanged. Returns a list of ``(k_pad, [global_expert_ids])`` tuples, ordered by ascending ``k_pad``. """ + def _round(m): return max(((int(m) + block_m - 1) // block_m) * block_m, 512) @@ -463,6 +507,7 @@ def _round(m): # Simple policy: collapse the sorted distinct values into `max_buckets` # contiguous groups, each group promoted to its group-max. import math + per = math.ceil(len(distinct) / max_buckets) ladder = {} for i in range(0, len(distinct), per): diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py index 3c17d2696b..317c10c768 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py @@ -395,6 +395,7 @@ def kernel_gemm( work_counter: fx.Tensor = None, b_data_ptrs: fx.Tensor = None, a_data_ptrs: fx.Tensor = None, + num_work: fx.Int32 = 0, ): lds = fx.SharedAllocator().allocate(SharedStorage).peek() lds_a0 = (lds.a0_0, lds.a0_1) @@ -415,12 +416,13 @@ def kernel_gemm( num_blocks_m = c_m // BLOCK_M num_blocks_n = c_n // BLOCK_N - # Total output tiles. Dense uses this as the grid extent; grouped uses it as - # the persistent-loop drain bound. ``grouped`` gates the shared descriptor - # queue / work-steal driver (group fwd/dgrad or the wgrad specialization); - # const_expr folds it all out of the dense binary. + # ``grouped`` gates the shared descriptor queue / work-steal driver (group + # fwd/dgrad or the wgrad specialization); const_expr folds it all out of the + # dense binary. grouped = group or wgrad - num_tiles = gd.num_tiles_bound(c_m, c_n, BLOCK_M, BLOCK_N, num_groups, const_expr(wgrad)) + # Grouped drain bound: the exact number of descriptors the host wrote to + # ``work_queue`` (``num_work``), so no slot past the queue end is ever read. + num_tiles = num_work if const_expr(grouped): wq_rsrc = fx.rocdl.make_buffer_tensor(work_queue, max_size=True) @@ -600,7 +602,11 @@ def load_bias(col): gAux = fx.rocdl.make_buffer_tensor(Aux, max_size=True) aux_div = fx.logical_divide(gAux, fx.make_layout(1, 1)) aux_store_atom = fx.make_copy_atom( - fx.rocdl.BufferCopy32b() if output_element_bytes == 4 else fx.rocdl.BufferCopy16b(), + ( + fx.rocdl.BufferCopy32b() + if output_element_bytes == 4 + else fx.rocdl.BufferCopy16b() + ), output_fx_dtype, ) @@ -792,7 +798,9 @@ def load_transposed_frag(lds_page, local_x_tile): def _acc_idx(subtile_id, mi, ni): return ( - subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + mi * MFMA_N_PER_SUBTILE + ni + subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + + mi * MFMA_N_PER_SUBTILE + + ni ) def _k32_frag(full_frag, k32): @@ -1389,6 +1397,7 @@ def launch_gemm( work_counter: fx.Tensor = None, b_data_ptrs: fx.Tensor = None, a_data_ptrs: fx.Tensor = None, + num_work: fx.Int32 = 0, stream: fx.Stream = fx.Stream(None), ): # Dense: one workgroup per output tile. Grouped/wgrad: a fixed num-CU @@ -1410,6 +1419,7 @@ def launch_gemm( work_counter, b_data_ptrs, a_data_ptrs, + num_work, value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, ).launch(grid=(grid_x, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) @@ -1478,6 +1488,9 @@ def _compiled_launch(key, jit_fn, args): for the first call; used once to compile, then the cached CompiledFunction is reused positionally for all later calls with the same ABI. """ + # A compiled launcher is built for the current device; key on it so a launch on + # another device compiles its own instead of reusing this one. + key = (key, torch.cuda.current_device()) cf = _COMPILED_LAUNCH_CACHE.get(key) if cf is None: cf = flyc.compile(jit_fn, *args) @@ -1654,6 +1667,7 @@ def doGemm( work_counter: torch.Tensor = None, b_data_ptrs: torch.Tensor = None, a_data_ptrs: torch.Tensor = None, + num_work: int = 0, sched: str = "worksteal", wgrad: bool = False, ): @@ -1671,6 +1685,8 @@ def doGemm( """ if layout not in ("TN", "NN", "NT"): raise ValueError(f"Unsupported {label} layout: {layout}") + if group and b_data_ptrs is None: + raise FlyDSLUnsupportedError(f"grouped {label} GEMM requires per-expert weight pointers") M_runtime = int(m) N_runtime = int(n) @@ -1732,7 +1748,7 @@ def doGemm( raise ValueError(f"{label} epilogue {epilogue} does not accept an aux tensor") if stream is None: - stream = torch.cuda.current_stream() + stream = gd.current_stream_handle(A.device) # Preserve the original validated byte-addressed G2L path. These are # metadata-only dtype/flatten views of the already-contiguous row-major @@ -1755,6 +1771,12 @@ def doGemm( # Grouped runtime tensors are required by the kernel signature; the dense # binary never dereferences them but still needs placeholder buffers. if group or wgrad: + # num_work is the kernel's drain bound: it must be the queue's descriptor count. + if work_queue.numel() != 3 * int(num_work): + raise ValueError( + f"num_work={num_work} does not match work_queue of {work_queue.numel()} " + "int32 (3 per descriptor)" + ) work_queue_arg = work_queue.contiguous().view(-1) work_counter_arg = work_counter.contiguous().view(-1) else: @@ -1801,6 +1823,7 @@ def doGemm( work_counter_arg, b_data_ptrs_arg, a_data_ptrs_arg, + int(num_work), stream, ) # flyc.compile fast path: skips flydsl's per-call cache-key rebuild (sig.bind + diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index 80fa15d877..71f2ffecc4 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -90,6 +90,9 @@ def _compiled_launch(key, jit_fn, args): first call; it is used once to compile, then the cached CompiledFunction is reused positionally for all later calls with the same ABI. """ + # A compiled launcher is built for the current device; key on it so a launch on + # another device compiles its own instead of reusing this one. + key = (key, torch.cuda.current_device()) cf = _COMPILED_LAUNCH_CACHE.get(key) if cf is None: cf = flyc.compile(jit_fn, *args) @@ -311,7 +314,7 @@ def pack_mx32_scales_for_hk( device=scales_u8.device, ) if stream is None: - stream = torch.cuda.current_stream(scales_u8.device) + stream = gd.current_stream_handle(scales_u8.device) stride0, stride1 = (int(x) for x in scales_u8.stride()) pack_key = ("mx32_scale_pack", dim, qk, bool(source_colwise), stride0, stride1) @@ -337,8 +340,9 @@ def pack_mx32_scales_for_hk_experts( expert in place through ``grouped_gemm_utils.stacked_operand`` pointers. Falls back to packing a real concat when the tensors cannot be addressed in place. """ - stacked = gd.stacked_operand(scales, scales[0].device) - if stacked is None or any(t.dim() != 2 for t in scales): + # stacked_operand requires one common trailing shape, so rank 2 holds for all. + stacked = gd.stacked_operand(scales, scales[0].device) if scales[0].dim() == 2 else None + if stacked is None: return pack_mx32_scales_for_hk( torch.cat(scales, dim=0), source_colwise=source_colwise, stream=stream ) @@ -346,8 +350,9 @@ def pack_mx32_scales_for_hk_experts( if representative.dtype != torch.uint8: raise TypeError(f"MXFP8 scales must be torch.uint8 E8M0 bytes, got {representative.dtype}") - total = sum(t.size(0) for t in scales) - inner = representative.size(1) + extents = [t.shape[0] for t in scales] + total = sum(extents) + inner = representative.shape[1] if source_colwise: qk, dim = total, inner # Experts stack on K/32; each thread's bytes share one K128 block (4 rows). @@ -356,13 +361,11 @@ def pack_mx32_scales_for_hk_experts( dim, qk = total, inner # Experts stack on dim; each thread's bytes share one 64-row tile. block = 64 - expert_table = gd.block_expert_table( - [t.size(0) for t in scales], block, representative.device - ) + expert_table = gd.block_expert_table(extents, block, representative.device) packed = torch.empty((qk // 4, dim), dtype=torch.int32, device=representative.device) if stream is None: - stream = torch.cuda.current_stream(representative.device) + stream = gd.current_stream_handle(representative.device) # The shifted pointers address the contiguous concatenation, so its strides apply. stride0, stride1 = inner, 1 pack_key = ("mx32_scale_pack_experts", dim, qk, bool(source_colwise), stride0, stride1) @@ -668,13 +671,11 @@ def kernel_gemm( work_counter: fx.Tensor = None, b_data_ptrs: fx.Tensor = None, a_data_ptrs: fx.Tensor = None, + num_work: fx.Int32 = 0, ): - # Total output tiles. Dense uses this as the grid extent; grouped uses it as - # the drain bound for the persistent loop. For fwd/dgrad, c_m = M_total spans - # all experts, so this is already the total. For wgrad, c_m = N_out is the - # PER-EXPERT output-row count (the packed output [G*N_out, K_in] is addressed - # via store_m_expert_off), so multiply by num_groups for the full count. - num_tiles = gd.num_tiles_bound(c_m, c_n, BLOCK_M, BLOCK_N, num_groups, const_expr(wgrad)) + # Grouped drain bound: the exact number of descriptors the host wrote to + # ``work_queue`` (``num_work``), so no slot past the queue end is ever read. + num_tiles = num_work lds = fx.SharedAllocator().allocate(SharedStorage).peek() lds_a0 = (lds.a0_0, lds.a0_1) @@ -938,9 +939,7 @@ def process_tile(pid_m, pid_n, expert): # store_m_expert_off (grouped wgrad only) shifts the output row into this # tile's expert block of the packed [G*N_out, K_in] weight-gradient; the # a/b operands used local (per-expert) pid_m, so the store re-packs here. - c_tile_base_elems = ( - (bx_m_idx + store_m_expert_off) * fx.Index(c_n) + by_n_idx - ) + c_tile_base_elems = (bx_m_idx + store_m_expert_off) * fx.Index(c_n) + by_n_idx gC = fx.rocdl.make_buffer_tensor(C, max_size=True) c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) c_store_atom = fx.make_copy_atom( @@ -971,7 +970,11 @@ def load_bias(col): gAux = fx.rocdl.make_buffer_tensor(Aux, max_size=True) aux_div = fx.logical_divide(gAux, fx.make_layout(1, 1)) aux_store_atom = fx.make_copy_atom( - fx.rocdl.BufferCopy32b() if output_element_bytes == 4 else fx.rocdl.BufferCopy16b(), + ( + fx.rocdl.BufferCopy32b() + if output_element_bytes == 4 + else fx.rocdl.BufferCopy16b() + ), output_fx_dtype, ) @@ -1137,9 +1140,7 @@ def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): # every other path. a_g2s.load_one( lds_a[subtile], - fx.Int32( - _a_global_base(k_base + a_expert_k_off, subtile, c_m, bx_m_idx) - ), + fx.Int32(_a_global_base(k_base + a_expert_k_off, subtile, c_m, bx_m_idx)), pass_in_subtile, ) @@ -1150,9 +1151,7 @@ def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): # packed column by_n_b_idx and the expert contraction shift b_k_off. b_g2s.load_one( lds_b[subtile], - fx.Int32( - _b_global_base(k_base + b_k_off, subtile, c_n, b_data_col) - ), + fx.Int32(_b_global_base(k_base + b_k_off, subtile, c_n, b_data_col)), pass_in_subtile, ) @@ -1214,7 +1213,9 @@ def load_transposed_frag(lds_page, local_x_tile): def _acc_idx(subtile_id, mi, ni): return ( - subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + mi * MFMA_N_PER_SUBTILE + ni + subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + + mi * MFMA_N_PER_SUBTILE + + ni ) def pinned_mfma(acc_idx, a_frag, b_frag, a_scale, b_scale, mi, ni): @@ -1527,7 +1528,9 @@ def hk_one_k_with_refill( # LDS reads consume the already-ready next page, not the page being refilled. rocdl.sched_barrier(0) barrier( - vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE + LOAD_PASSES_SCALES, + vmcnt=2 * LOAD_PASSES_A_SUBTILE + + 2 * LOAD_PASSES_B_SUBTILE + + LOAD_PASSES_SCALES, lgkmcnt=0, ) rocdl.sched_barrier(0) @@ -1913,6 +1916,7 @@ def launch_gemm( work_counter: fx.Tensor = None, b_data_ptrs: fx.Tensor = None, a_data_ptrs: fx.Tensor = None, + num_work: fx.Int32 = 0, stream: fx.Stream = fx.Stream(None), ): # The integration only dispatches aligned shapes; no partial-tile masking exists. @@ -1937,6 +1941,7 @@ def launch_gemm( work_counter, b_data_ptrs, a_data_ptrs, + num_work, value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, ).launch(grid=(grid_x, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) @@ -1961,6 +1966,7 @@ def do_gemm( work_counter: fx.Tensor = None, b_data_ptrs: torch.Tensor = None, a_data_ptrs: torch.Tensor = None, + num_work: int = 0, a_shape=None, b_shape=None, sched: str = "worksteal", @@ -1990,25 +1996,20 @@ def do_gemm( else: raise ValueError(f"Unsupported MXFP8 kernel layout: {layout}") - # Grouped NN packs B along the contraction axis, so the packed B's contraction - # dim is G x A's (Kb_runtime = G*N_out vs K_runtime = N_out). With per-expert - # weight pointers B is a per-expert [N_out, K_in] representative, so its - # contraction already equals A's. Every other case has them equal. - if group and layout == "NN" and b_data_ptrs is None: - assert Kb_runtime == num_groups * K_runtime, ( - f"grouped NN B.K={Kb_runtime} != num_groups*A.K={num_groups * K_runtime}" - ) - else: - assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" + # Grouped launches read B from per-expert weight pointers, so B is a per-expert + # representative whose contraction equals A's (as in the dense case). + if group and b_data_ptrs is None: + raise FlyDSLUnsupportedError("grouped MXFP8 GEMM requires per-expert weight pointers") + assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" # Grouped wgrad (NT): both operands (dY, X) are packed on the token/contraction # axis to G*M_pad, but the kernel's compile-time contraction is the PER-EXPERT # padded token count M_pad. M_runtime (=N_out) and N_runtime (=K_in) are already # per-expert; only K_runtime must be de-packed. if wgrad: - assert K_runtime % num_groups == 0, ( - f"grouped wgrad contraction {K_runtime} not divisible by num_groups={num_groups}" - ) + assert ( + K_runtime % num_groups == 0 + ), f"grouped wgrad contraction {K_runtime} not divisible by num_groups={num_groups}" K_runtime = K_runtime // num_groups # Grouped fwd/dgrad launch over the per-expert output width = C's last dim (TN: @@ -2048,13 +2049,11 @@ def do_gemm( expected_bs = (packed_k128, N_runtime) else: expected_as = (K_runtime // _BLOCK_K, M_runtime) - # TN packs B along N (column count = G*N_out); NN packs B along the - # contraction axis (row count on K128 = G*N_out/128). The B-scale stays packed - # even when the B data is addressed per-expert, so recover the packed - # contraction (Kb_runtime is per-expert N_out when pointers are used). + # TN packs the B-scale along N (column count = G*N_out); NN packs it along the + # contraction axis (row count on K128 = G*N_out/128). Kb_runtime is the + # per-expert N_out, since the B data is addressed per expert. if group and layout == "NN": - kb_packed = Kb_runtime * num_groups if b_data_ptrs is not None else Kb_runtime - expected_bs = (kb_packed // _BLOCK_K, N_runtime) + expected_bs = ((Kb_runtime * num_groups) // _BLOCK_K, N_runtime) else: n_packed = N_runtime * num_groups if group else N_runtime expected_bs = (K_runtime // _BLOCK_K, n_packed) @@ -2107,7 +2106,7 @@ def do_gemm( raise ValueError(f"MXFP8 epilogue {epilogue} does not accept an aux tensor") if stream is None: - stream = torch.cuda.current_stream() + stream = gd.current_stream_handle(A.device) # A non-contiguous output cannot receive the kernel's writes: the flat # descriptor below would target a throwaway ``.contiguous()`` copy. Reject @@ -2137,6 +2136,12 @@ def do_gemm( # binary never dereferences them but still needs placeholder buffers. wgrad # uses the same persistent descriptor queue as the grouped fwd/dgrad path. if group or wgrad: + # num_work is the kernel's drain bound: it must be the queue's descriptor count. + if work_queue.numel() != 3 * int(num_work): + raise ValueError( + f"num_work={num_work} does not match work_queue of {work_queue.numel()} " + "int32 (3 per descriptor)" + ) work_queue_arg = work_queue.contiguous().view(-1) work_counter_arg = work_counter.contiguous().view(-1) else: @@ -2186,6 +2191,7 @@ def do_gemm( work_counter_arg, b_data_ptrs_arg, a_data_ptrs_arg, + int(num_work), stream, ) _compiled_launch( @@ -2252,11 +2258,11 @@ def _validate_common_payloads( def _mxfp8_matmul_grouped_wgrad( - a, # dY packed [G*M_pad, N_out], columnwise fp8 - a_scale, # [G*M_pad/32, N_out] raw E8M0 uint8, or list[G] of per-expert scales - b, # X packed [G*M_pad, K_in], columnwise fp8 - b_scale, # [G*M_pad/32, K_in] raw E8M0 uint8, or list[G] of per-expert scales - D, # dW packed [G*N_out, K_in] + a, # dY packed [G*M_pad, N_out], columnwise fp8 + a_scale, # [G*M_pad/32, N_out] raw E8M0 uint8, or list[G] of per-expert scales + b, # X packed [G*M_pad, K_in], columnwise fp8 + b_scale, # [G*M_pad/32, K_in] raw E8M0 uint8, or list[G] of per-expert scales + D, # dW packed [G*N_out, K_in] *, stream, num_groups, @@ -2265,6 +2271,7 @@ def _mxfp8_matmul_grouped_wgrad( sched, a_data_ptrs=None, b_data_ptrs=None, + num_work=0, a_shape=None, b_shape=None, ): @@ -2286,13 +2293,10 @@ def _mxfp8_matmul_grouped_wgrad( m_pad = gm // num_groups if m_pad % SCALE_GROUP_SIZE != 0: raise FlyDSLUnsupportedError( - f"wgrad padded per-expert tokens M_pad={m_pad} must be divisible by " - f"{SCALE_GROUP_SIZE}" + f"wgrad padded per-expert tokens M_pad={m_pad} must be divisible by {SCALE_GROUP_SIZE}" ) if tuple(D.shape) != (num_groups * n_out, k_in): - raise ValueError( - f"wgrad D shape {tuple(D.shape)} != expected {(num_groups * n_out, k_in)}" - ) + raise ValueError(f"wgrad D shape {tuple(D.shape)} != expected {(num_groups * n_out, k_in)}") if _scale_dtype(a_scale) != torch.uint8 or _scale_dtype(b_scale) != torch.uint8: raise TypeError("FlyDSL MXFP8 wgrad expects raw E8M0 scales as torch.uint8") @@ -2316,6 +2320,7 @@ def _mxfp8_matmul_grouped_wgrad( a_data_ptrs=a_data_ptrs, a_shape=a_shape, b_shape=b_shape, + num_work=num_work, sched=sched, wgrad=True, ) @@ -2340,6 +2345,7 @@ def mxfp8_matmul( work_counter: fx.Tensor = None, b_data_ptrs: torch.Tensor = None, a_data_ptrs: torch.Tensor = None, + num_work: int = 0, a_shape=None, b_shape=None, sched: str = "worksteal", @@ -2383,6 +2389,7 @@ def mxfp8_matmul( b_data_ptrs=b_data_ptrs, a_shape=a_shape, b_shape=b_shape, + num_work=num_work, ) if group and layout not in ("TN", "NN"): @@ -2420,42 +2427,18 @@ def mxfp8_matmul( # [G*N_out, K_in] on that CONTRACTION axis. So B's first dim kb = G*N_out # is G x A's per-expert contraction k = N_out (kb != k by design), and the # output width n = K_in is NOT packed. D = [M_total, K_in] = (m, n). + if group and b_data_ptrs is None: + raise FlyDSLUnsupportedError("grouped MXFP8 GEMM requires per-expert weight pointers") if group and layout == "TN": - # Per-expert weight pointers: B is a per-expert [N_out, K] representative, - # so b.shape[0] already IS N_out. Without pointers B is the packed [G*N_out, - # K] and N_out unpacks. Either way n_packed = G*N_out for the B-scale, which - # stays packed. - if b_data_ptrs is not None: - n_out = n - else: - if n % num_groups != 0: - raise ValueError( - f"Grouped MXFP8 TN packed N={n} is not divisible by num_groups={num_groups}" - ) - n_out = n // num_groups + # B is a per-expert [N_out, K] representative read through b_data_ptrs, so + # b.shape[0] already IS N_out; the B-scale still spans G*N_out rows. + n_out = n expected_kb = k # contraction unpacked; B rows already per-K elif group and layout == "NN": - # Per-expert weight pointers: B is a per-expert [N_out, K_in] representative, - # so b.shape[0] IS the per-expert contraction N_out (= k). Without pointers B - # is the packed [G*N_out, K_in] and the contraction unpacks. - if b_data_ptrs is not None: - if kb != k: - raise ValueError( - f"Grouped MXFP8 NN per-expert B contraction {kb} != A contraction {k}" - ) - else: - if kb % num_groups != 0: - raise ValueError( - f"Grouped MXFP8 NN packed contraction={kb} not divisible by " - f"num_groups={num_groups}" - ) - if kb // num_groups != k: - raise ValueError( - f"Grouped MXFP8 NN per-expert contraction {kb // num_groups} " - f"!= A contraction {k}" - ) + # B is a per-expert [N_out, K_in] representative read through b_data_ptrs, so + # b.shape[0] IS the per-expert contraction N_out (= k). n_out = n # output width (K_in) is not packed - expected_kb = kb # checked above (per-expert or packed) + expected_kb = k else: n_out = n expected_kb = k @@ -2480,17 +2463,13 @@ def mxfp8_matmul( expected_a_scale = (m, k // SCALE_GROUP_SIZE) if layout == "TN": - # Grouped TN packs the B scale along N, spanning G*N_out rows. With per-expert - # weight pointers B is per-expert so n == N_out; scale rows are n_out*num_groups. - # Without pointers (packed B) n already is G*N_out. - b_scale_rows = n_out * num_groups if (group and b_data_ptrs is not None) else n + # Grouped TN packs the B scale along N, spanning G*N_out rows (n == N_out). + b_scale_rows = n_out * num_groups if group else n expected_b_scale = (b_scale_rows, k // SCALE_GROUP_SIZE) else: # NN/NT B scale is [contraction/32, N]. Grouped NN packs the scale along the - # contraction axis (G*N_out rows). With per-expert weight pointers kb is the - # per-expert N_out, so the packed scale spans kb*num_groups; without pointers - # kb already is G*N_out. - kb_packed = kb * num_groups if (group and b_data_ptrs is not None) else kb + # contraction axis: kb is the per-expert N_out, so it spans kb*num_groups rows. + kb_packed = kb * num_groups if group else kb expected_b_scale = (kb_packed // SCALE_GROUP_SIZE, n) if _scale_shape(a_scale) != expected_a_scale: @@ -2573,6 +2552,7 @@ def mxfp8_matmul( b_data_ptrs=b_data_ptrs, a_data_ptrs=a_data_ptrs, a_shape=a_shape, + num_work=num_work, sched=sched, ) return D diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 1887058fc8..e2b1027ab8 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -91,6 +91,25 @@ import os +def _flydsl_rows(t): + """``t`` as 2-D rows; GroupedLinear's per-expert tensors already are, so skip the view.""" + return t if t.dim() == 2 else t.reshape(-1, t.shape[-1]) + + +def _flydsl_fp8_transposes(tensors): + """The materialized FP8 transposes of ``tensors`` for the FlyDSL TN-only kernel. + + Raises ``FlyDSLUnsupportedError`` (so the caller falls back to the default grouped + GEMM) when a transpose is missing or marked stale by ``_transpose_invalid``. + """ + from ..flydsl_kernels.gemm import FlyDSLUnsupportedError + + for t in tensors: + if t._transpose is None or t._transpose_invalid: + raise FlyDSLUnsupportedError("FP8 operand has no valid transpose") + return [t._transpose for t in tensors] + + def _flydsl_row_operand(datas, m_splits): """Per-expert ``[m_e, K]`` row payloads as one FlyDSL grouped-kernel operand. @@ -799,6 +818,7 @@ def _is_blockwise_fp8_triton_grouped_gemm_supported( def _flydsl_grouped_common_ok( *, use_bias, + use_grouped_gemm_triton, backward_override, cpu_offloading, save_original_input, @@ -823,6 +843,7 @@ def _flydsl_grouped_common_ok( and out_features % 128 == 0 and not ( use_bias + or use_grouped_gemm_triton or backward_override or cpu_offloading or save_original_input @@ -838,6 +859,7 @@ def _flydsl_grouped_kind( fp8, recipe, use_bias, + use_grouped_gemm_triton, backward_override, cpu_offloading, save_original_input, @@ -855,6 +877,7 @@ def _flydsl_grouped_kind( """ if not _GroupedLinear._flydsl_grouped_common_ok( use_bias=use_bias, + use_grouped_gemm_triton=use_grouped_gemm_triton, backward_override=backward_override, cpu_offloading=cpu_offloading, save_original_input=save_original_input, @@ -902,18 +925,14 @@ def _grouped_gemm_flydsl_mxfp8( # TN operands: activations are read in place from each expert's rows (no # [M_total, K] concat). Each MXFP8 operand carries its rowwise fp8 payload and # E8M0 scale; the scales are passed per expert and packed in place. - a_data, a_data_ptrs = _flydsl_row_operand( - [im._rowwise_data for im in inputmats], m_splits - ) + a_data, a_data_ptrs = _flydsl_row_operand([im._rowwise_data for im in inputmats], m_splits) a_scale = [im._rowwise_scale_inv for im in inputmats] # Weights are independently-allocated per-expert slabs. The grouped-TN # kernel reads each expert's weight from a per-expert base pointer, so the # host passes an int64[G] address array instead of concatenating the G # weights into one packed buffer. b_data is a per-expert representative used # only for dtype/shape. - b_data_ptrs = expert_ptrs( - [w._rowwise_data for w in weights], a_data.device - ) + b_data_ptrs = expert_ptrs([w._rowwise_data for w in weights], a_data.device) b_data = weights[0]._rowwise_data b_scale = [w._rowwise_scale_inv for w in weights] @@ -967,13 +986,9 @@ def _backward_flydsl_mxfp8_grouped( # expert's rows. NN kernel-B = W, columnwise (transpose-read, contraction on # the N_out axis): independent per-expert [N_out, K_in] columnwise slabs # addressed via per-expert base pointers. Neither is concatenated. - a_data, a_data_ptrs = _flydsl_row_operand( - [g._rowwise_data for g in grad_output], m_splits - ) + a_data, a_data_ptrs = _flydsl_row_operand([g._rowwise_data for g in grad_output], m_splits) a_scale = [g._rowwise_scale_inv for g in grad_output] - b_data_ptrs = expert_ptrs( - [w._columnwise_data for w in weights], a_data.device - ) + b_data_ptrs = expert_ptrs([w._columnwise_data for w in weights], a_data.device) b_data = weights[0]._columnwise_data b_scale = [w._columnwise_scale_inv for w in weights] @@ -1066,9 +1081,7 @@ def _grouped_gemm_flydsl_bf16(*, inputmats, weights, out, m_splits, m_splits_ten from ..flydsl_kernels.gemm.grouped_gemm_utils import expert_ptrs # Activations are read in place from each expert's rows (no [M_total, K] concat). - a_data, a_data_ptrs = _flydsl_row_operand( - [im.reshape(-1, im.shape[-1]) for im in inputmats], m_splits - ) + a_data, a_data_ptrs = _flydsl_row_operand([_flydsl_rows(im) for im in inputmats], m_splits) # Weights are independent per-expert [N_out, K] tensors; the grouped-TN kernel # reads each from a per-expert base pointer, so pass an int64[G] address array # instead of concatenating the G weights. b_data is a per-expert representative @@ -1099,9 +1112,7 @@ def _backward_flydsl_bf16_grouped(*, grad_output, weights, dgrad, m_splits, m_sp from ..flydsl_kernels.gemm.grouped_gemm_utils import expert_ptrs # dY is read in place from each expert's rows (no [M_total, N_out] concat). - a_data, a_data_ptrs = _flydsl_row_operand( - [g.reshape(-1, g.shape[-1]) for g in grad_output], m_splits - ) + a_data, a_data_ptrs = _flydsl_row_operand([_flydsl_rows(g) for g in grad_output], m_splits) # Weights are independent per-expert [N_out, K_in] tensors; the grouped-NN # kernel reads each from a per-expert base pointer (no concat). b_data is a # per-expert representative. @@ -1131,8 +1142,8 @@ def _wgrad_flydsl_bf16_grouped(*, grad_output, inputmats, wgrad_list, m_splits): from ..flydsl_kernels.gemm import _run_bf16_grouped_wgrad num_gemms = len(inputmats) - dy_cols = [g.reshape(-1, g.shape[-1]) for g in grad_output] - x_cols = [x.reshape(-1, x.shape[-1]) for x in inputmats] + dy_cols = [_flydsl_rows(g) for g in grad_output] + x_cols = [_flydsl_rows(x) for x in inputmats] n_out = dy_cols[0].shape[-1] k_in = x_cols[0].shape[-1] @@ -1172,9 +1183,7 @@ def _grouped_gemm_flydsl_fp8(*, inputmats, weights, out, m_splits, m_splits_tens # Weights are independent per-expert [N, K] FP8 tensors; the grouped-TN kernel # reads each from a per-expert base pointer, so pass an int64[G] address array # instead of concatenating the G weights. b_data is a per-expert representative. - b_data_ptrs = expert_ptrs( - [w._data for w in weights], a_data.device - ) + b_data_ptrs = expert_ptrs([w._data for w in weights], a_data.device) b_data = weights[0]._data a_data._fp8_dtype = inputmats[0]._fp8_dtype b_data._fp8_dtype = weights[0]._fp8_dtype @@ -1202,9 +1211,7 @@ def _grouped_gemm_flydsl_fp8(*, inputmats, weights, out, m_splits, m_splits_tens ) @staticmethod - def _backward_flydsl_fp8_grouped( - *, grad_output, weights, dgrad, m_splits, m_splits_tensor - ): + def _backward_flydsl_fp8_grouped(*, grad_output, weights, dgrad, m_splits, m_splits_tensor): """FlyDSL tensor-wise FP8 grouped dgrad: dX_e = dY_e @ W_e, via the TN kernel. FP8 is TN-only, so dgrad reaches the same TN kernel using materialized @@ -1219,10 +1226,9 @@ def _backward_flydsl_fp8_grouped( # [K_in, N_out] packed on K_in (the TN GEMM output axis). Both are independent # per-expert slabs, addressed via per-expert base pointers (no concat). a_data, a_data_ptrs = _flydsl_row_operand([g._data for g in grad_output], m_splits) - b_data_ptrs = expert_ptrs( - [w._transpose for w in weights], a_data.device - ) - b_data = weights[0]._transpose + w_transposes = _flydsl_fp8_transposes(weights) + b_data_ptrs = expert_ptrs(w_transposes, a_data.device) + b_data = w_transposes[0] a_data._fp8_dtype = grad_output[0]._fp8_dtype b_data._fp8_dtype = weights[0]._fp8_dtype # Per-expert scale_inv pointers, read in place by the kernel (no gather). @@ -1259,8 +1265,8 @@ def _wgrad_flydsl_fp8_grouped(*, grad_output, inputmats, wgrad_list, m_splits): from ..flydsl_kernels.gemm import _run_fp8_grouped_wgrad num_gemms = len(inputmats) - dyt_cols = [g._transpose for g in grad_output] - xt_cols = [x._transpose for x in inputmats] + dyt_cols = _flydsl_fp8_transposes(grad_output) + xt_cols = _flydsl_fp8_transposes(inputmats) # Per-expert scale_inv tensors; the wgrad runner reads them in place by pointer. dy_scale = [g._scale_inv for g in grad_output] x_scale = [x._scale_inv for x in inputmats] @@ -1779,6 +1785,7 @@ def forward( fp8=fp8, recipe=blockwise_recipe, use_bias=use_bias, + use_grouped_gemm_triton=use_grouped_gemm_triton, backward_override=backward_override, cpu_offloading=cpu_offloading, save_original_input=save_original_input, @@ -1945,12 +1952,19 @@ def forward( # tensors (like scales), but bulk allocation shares storage across all tensors, # so if scales can't be offloaded, nothing in the group can be offloaded. fused_padding_kwargs = {} - if actual_m_splits is not None and IS_HIP_EXTENSION \ - and inp_view.shape[0] == sum(actual_m_splits): + if ( + actual_m_splits is not None + and IS_HIP_EXTENSION + and inp_view.shape[0] == sum(actual_m_splits) + ): fused_padding_kwargs["valid_split_sections"] = actual_m_splits inputmats = tex.split_quantize( - inp_view, m_splits, input_quantizers, - disable_bulk_allocation=cpu_offloading, **fused_padding_kwargs) + inp_view, + m_splits, + input_quantizers, + disable_bulk_allocation=cpu_offloading, + **fused_padding_kwargs, + ) elif debug: inputmats = DebugQuantizer.multi_tensor_quantize( inp_view, input_quantizers, m_splits, activation_dtype @@ -2672,14 +2686,21 @@ def backward( **kwargs, ) - if ctx.actual_m_splits is not None and ctx.actual_m_splits != ctx.m_splits \ - and not ctx.output_unpadded: + if ( + ctx.actual_m_splits is not None + and ctx.actual_m_splits != ctx.m_splits + and not ctx.output_unpadded + ): dgrad_unpadded = torch.empty( (sum(ctx.actual_m_splits), dgrad.shape[-1]), - dtype=dgrad.dtype, device=dgrad.device, + dtype=dgrad.dtype, + device=dgrad.device, ) tex.fused_multi_row_unpadding( - dgrad, dgrad_unpadded, ctx.m_splits, ctx.actual_m_splits, + dgrad, + dgrad_unpadded, + ctx.m_splits, + ctx.actual_m_splits, ) dgrad = dgrad_unpadded @@ -2789,9 +2810,7 @@ def backward( _flydsl_kind is not None and not ctx.fuse_wgrad_accumulation and not ctx.use_bias # bias-grad-in-wgrad-GEMM not supported yet - and not ( - ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute() - ) + and not (ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute()) ) ran_flydsl_wgrad = use_flydsl_wgrad and _GroupedLinear._try_flydsl_grouped( "wgrad", @@ -2835,9 +2854,7 @@ def backward( [inputmats, grad_output, wgrad_list], grouped_gemm_wgrad ) else: - _, grad_biases_, _ = grouped_gemm_wgrad( - inputmats, grad_output, wgrad_list - ) + _, grad_biases_, _ = grouped_gemm_wgrad(inputmats, grad_output, wgrad_list) for i in range(ctx.num_gemms): if grad_biases[i] is None: diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 4fc1abe062..3ff2c7dc1b 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -177,7 +177,7 @@ def _compute_mxfp8_support() -> Tuple[bool, str]: if os.getenv("NVTE_ROCM_ENABLE_MXFP8", "0") == "0": return False, "MXFP8 support is not enabled." gpu_arch = get_device_compute_capability() - if gpu_arch in ((9, 5),): #TODO: enable for gfx1250 when GEMM is available + if gpu_arch in ((9, 5),): # TODO: enable for gfx1250 when GEMM is available return True, "" return False, "Device arch gfx95x or newer is required for MXFP8 execution." if get_device_compute_capability() >= (12, 0): @@ -343,13 +343,14 @@ def get_align_size_for_quantization(recipe: Recipe) -> int: # alignment as the conservative fallback for context-free callers. if recipe.custom(): return recipe.quantization_alignment - # The FlyDSL grouped GEMM (gfx950 opt-in) tiles every expert with BLOCK_M=256 - # across all dtypes, so per-expert token counts must be 256-aligned regardless of - # recipe. + # The FlyDSL grouped GEMM (gfx950 opt-in) tiles every expert with BLOCK_M=256, so + # per-expert token counts must be 256-aligned for the recipes it handles (MXFP8 and + # tensor-wise FP8); other recipes fall back and keep their own alignment. if ( IS_HIP_EXTENSION and get_device_compute_capability() == (9, 5) and os.environ.get("NVTE_USE_GROUPED_GEMM_FLYDSL", "0") == "1" + and (recipe.mxfp8() or recipe.float8_current_scaling() or recipe.delayed()) ): return 256 if recipe.mxfp8(): @@ -360,11 +361,7 @@ def get_align_size_for_quantization(recipe: Recipe) -> int: use_ck = os.environ.get("NVTE_USE_CK_GROUPED_GEMM", "0") == "1" use_cutlass = os.environ.get("NVTE_USE_CUTLASS_GROUPED_GEMM", "0") == "1" use_hipkittens = use_hk or (use_cutlass and not use_ck) - if ( - IS_HIP_EXTENSION - and get_device_compute_capability() == (9, 5) - and use_hipkittens - ): + if IS_HIP_EXTENSION and get_device_compute_capability() == (9, 5) and use_hipkittens: return 256 return 32 if recipe.nvfp4(): @@ -470,6 +467,7 @@ def is_fp8_block_scaling_available(return_reason: bool = False) -> Union[bool, T return check_fp8_block_scaling_support() return check_fp8_block_scaling_support()[0] + def is_nvfp4_available(return_reason: bool = False) -> Union[bool, Tuple[bool, str]]: """ Determine if support is available for the NVFP4 recipe. @@ -487,6 +485,7 @@ def is_nvfp4_available(return_reason: bool = False) -> Union[bool, Tuple[bool, s return check_nvfp4_support() return check_nvfp4_support()[0] + def is_mxfp4_available(return_reason: bool = False) -> Union[bool, Tuple[bool, str]]: """ Determine if support is available for the MXFP4 recipe. @@ -503,6 +502,7 @@ def is_mxfp4_available(return_reason: bool = False) -> Union[bool, Tuple[bool, s return check_mxfp4_support() return check_mxfp4_support()[0] + @dataclass(slots=True) class FP8GlobalState: """Mutable process-global FP8 state stored on an instance. @@ -923,7 +923,13 @@ def autocast_exit(cls, enabled: bool, _graph: bool) -> None: # Reduce only the non-FP8 weight modules here. # FP8 weight modules are reduced at the end of the optimizer # step after the weight amax is populated. - if not cls.SKIP_FP8_REDUCTION_FOR_FSDP2 and enabled and qstate.autocast_depth == 0 and not _graph and torch.is_grad_enabled(): + if ( + not cls.SKIP_FP8_REDUCTION_FOR_FSDP2 + and enabled + and qstate.autocast_depth == 0 + and not _graph + and torch.is_grad_enabled() + ): # delayed scaling only function, for other recipes (current scaling with any granularity), # this is noop for other recipes because cls.global_amax_buffer is empty list cls.reduce_and_update_fp8_tensors(forward=True) From bcab563b2e84c79202a87fd82026116a79d3949a Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 21 Oct 2026 15:14:15 +0000 Subject: [PATCH 16/22] remove redundant checks, clean up some code comments --- .../flydsl_kernels/gemm/gemm_wrappers.py | 88 +++++-------------- .../flydsl_kernels/gemm/half_prec_gemm.py | 6 -- .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 84 +++++------------- .../pytorch/module/grouped_linear.py | 45 +++++----- 4 files changed, 64 insertions(+), 159 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 9e94590a47..44e8a0d117 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -1009,7 +1009,6 @@ def _run_mxfp8( return D, aux -# Schedule selection now lives in the shared grouped driver (dtype-independent). _select_grouped_sched = gd.select_grouped_sched @@ -1055,15 +1054,8 @@ def _run_mxfp8_grouped( ``a_data_ptrs`` (optional) reads A per expert in place: ``a_data`` is then one representative expert payload (see ``grouped_gemm_utils.stacked_operand``). """ - if layout not in ("TN", "NN"): - raise FlyDSLUnsupportedError(f"FlyDSL MXFP8 grouped GEMM supports TN/NN, got {layout}") - if k % 32 != 0: - raise FlyDSLUnsupportedError(f"K={k} must be divisible by MXFP8 scale group size 32") - if n % 128 != 0: - raise FlyDSLUnsupportedError( - f"FlyDSL MXFP8 grouped GEMM requires output width N divisible by 128, got N={n}" - ) - + # Layout, K/32, tile alignment and scale dtype are validated by mxfp8_matmul / + # do_gemm, which every launch goes through. a_fp8_dtype = getattr(a_data, "_fp8_dtype", None) b_fp8_dtype = getattr(b_data, "_fp8_dtype", None) if a_data.dtype == torch.uint8: @@ -1071,10 +1063,6 @@ def _run_mxfp8_grouped( if b_data.dtype == torch.uint8: b_data = reinterpret_as_fp8_tensor(b_data, b_fp8_dtype) - scale_dtypes = [(s[0] if isinstance(s, (list, tuple)) else s).dtype for s in (a_scale, b_scale)] - if any(d != torch.uint8 for d in scale_dtypes): - raise TypeError("FlyDSL MXFP8 grouped GEMM expects raw E8M0 scales as torch.uint8") - D = _validate_or_allocate_output( D, shape=torch.Size((m_total, n)), @@ -1084,40 +1072,32 @@ def _run_mxfp8_grouped( ) # Persistent-kernel work distribution: - # work_queue [num_tiles, 3] int32 -- one precomputed (pid_m, pid_n, expert) - # descriptor per output tile. The host bakes the tile->expert - # mapping here so the kernel needs no in-kernel search, and may - # permute the rows (XCD-aware, group-balanced) freely. - # work_counter [1] int32 -- shared atomic slot dispenser, zeroed per launch. - # A workgroup claims slot = atomicAdd(work_counter, 1), reads work_queue[slot], - # and repeats until slot >= num_tiles. + # work_queue one (pid_m, pid_n, expert) int32 descriptor per output tile, so + # the kernel needs no tile-to-expert search. + # work_counter atomic slot counter, used only by the work-stealing schedule. + # Workgroups take queue slots until slot >= num_work (the queue length). from .mxfp8_gemm import BLOCK_M, BLOCK_N num_groups = len(m_splits) num_pid_n = n // BLOCK_N - # Per-expert M-block ranges over the packed [M_total, N] output. Each expert's - # token count must be BLOCK_M-aligned (enforced upstream by the tiling check), - # so tiles never straddle a group boundary. Built + cached from the host-side - # m_splits list (no GPU->CPU sync). + # Tiles cover each expert's rows of the packed [M_total, N] output. Expert token + # counts must be BLOCK_M multiples (the queue builder raises otherwise), so no + # tile spans two experts. Built from the host-side m_splits (no GPU->CPU sync). work_queue, num_tiles = gd.build_fwd_dgrad_work_queue( m_splits, num_pid_n, BLOCK_M, a_data.device ) - # Schedule selection (compile-time specialization on the kernel side): - # "worksteal": persistent atomic tile-claiming, best for imbalanced experts. - # "static": grid-stride slot ownership, no atomic; deadlock-proof. - # Heuristic mirrors the HK reference: when the tile count is small relative to - # the launch grid, one grid-stride sweep already balances well, so the static - # path avoids atomic contention. NVTE_FLYDSL_GROUPED_SCHED overrides for tests. + # Schedule (a compile-time kernel variant): "static" walks the queue with a grid + # stride and no atomics; "worksteal" claims tiles atomically, which balances + # uneven experts. Static is chosen when there are at most a few tiles per CU. + # NVTE_FLYDSL_GROUPED_SCHED overrides the choice. sched = _select_grouped_sched(num_tiles, a_data.device) work_counter = gd.work_counter(sched, a_data.device) epilogue, bias_arg = _resolve_epilogue(bias, n) - # Both TN forward and NN dgrad address the weight operand from per-expert base - # pointers supplied by the caller (one int64 base per expert slab, required): the - # kernel reads B slab g from b_data_ptrs[g], so the host never concatenates the G - # per-expert weights. + # TN forward and NN dgrad read expert g's weight from b_data_ptrs[g] (required), + # so the weights are never concatenated. mxfp8_matmul( a_data, a_scale, @@ -1185,14 +1165,6 @@ def _run_mxfp8_grouped_wgrad( num_groups = len(m_splits) device = dW.device - if n_out % BLOCK_M != 0: - raise FlyDSLUnsupportedError( - f"wgrad requires N_out divisible by BLOCK_M={BLOCK_M}, got {n_out}" - ) - if k_in % BLOCK_N != 0: - raise FlyDSLUnsupportedError( - f"wgrad requires K_in divisible by BLOCK_N={BLOCK_N}, got {k_in}" - ) dW = _validate_or_allocate_output( dW, @@ -1202,9 +1174,9 @@ def _run_mxfp8_grouped_wgrad( backend_name="MXFP8 grouped wgrad NT", ) - # Descriptor queue: one (pid_m, pid_n, expert) per output tile of every expert's - # [N_out, K_in] weight-gradient block. pid_m/pid_n are LOCAL (per-expert); the - # kernel re-packs the output row via expert*N_out. + # One (pid_m, pid_n, expert) descriptor per tile of each expert's [N_out, K_in] + # weight gradient. pid_m/pid_n are per-expert; the kernel offsets the output row + # by expert*N_out. num_pid_m = n_out // BLOCK_M num_pid_n = k_in // BLOCK_N buckets = gd.bucketize_wgrad_tokens(m_splits, BLOCK_M, max_buckets=_wgrad_max_buckets()) @@ -1303,7 +1275,8 @@ def _run_bf16_grouped( if layout not in ("TN", "NN"): raise FlyDSLUnsupportedError( - f"FlyDSL half-precision grouped GEMM supports TN/NN, got {layout}" + f"FlyDSL half-precision grouped fwd/dgrad supports TN/NN, got {layout} " + "(grouped NT wgrad goes through _run_bf16_grouped_wgrad)" ) # BF16 and FP16 share this kernel (the MFMA opcode is selected per operand dtype); @@ -1403,10 +1376,6 @@ def _run_bf16_grouped_wgrad( raise FlyDSLUnsupportedError( f"FlyDSL half-precision grouped wgrad requires bf16/fp16 operands, got {input_dtype}" ) - if n_out % _BLOCK_M != 0: - raise FlyDSLUnsupportedError(f"wgrad requires N_out divisible by {_BLOCK_M}, got {n_out}") - if k_in % _BLOCK_N != 0: - raise FlyDSLUnsupportedError(f"wgrad requires K_in divisible by {_BLOCK_N}, got {k_in}") dW = _validate_or_allocate_output( dW, @@ -1433,11 +1402,9 @@ def _run_bf16_grouped_wgrad( device, ) - # Fast path: when this bucket's experts are exactly rows 0..G-1 in order - # (the single-bucket / balanced case), the bucket-local expert index equals - # the global index, so the kernel writes straight into dW -- no bucket-local - # buffer and no scatter copy. Only multi-bucket (skewed) launches allocate a - # bucket-local output and scatter it back into dW. + # When the bucket holds all experts in order (the balanced case), bucket-local + # and global expert indices match, so the kernel writes dW directly. Otherwise + # it writes a bucket-local buffer that is copied into dW below. direct = ids == list(range(num_groups)) dW_b = dW if direct else torch.empty(bsize * n_out, k_in, dtype=output_dtype, device=device) @@ -1514,11 +1481,6 @@ def _run_fp8_grouped( from .fp8_gemm import doGemm as _fp8_doGemm num_groups = len(m_splits) - if a_scale_vec.numel() != num_groups or b_scale_vec.numel() != num_groups: - raise ValueError( - f"FP8 grouped expects length-{num_groups} scale vectors, got " - f"A={a_scale_vec.numel()}, B={b_scale_vec.numel()}" - ) a_fp8_dtype = getattr(a_data, "_fp8_dtype", None) b_fp8_dtype = getattr(b_data, "_fp8_dtype", None) @@ -1601,10 +1563,6 @@ def _run_fp8_grouped_wgrad( num_groups = len(m_splits) device = dW.device - if n_out % _BLOCK_M != 0: - raise FlyDSLUnsupportedError(f"wgrad requires N_out divisible by {_BLOCK_M}, got {n_out}") - if k_in % _BLOCK_N != 0: - raise FlyDSLUnsupportedError(f"wgrad requires K_in divisible by {_BLOCK_N}, got {k_in}") dW = _validate_or_allocate_output( dW, diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py index 317c10c768..8d7181273a 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py @@ -1710,12 +1710,6 @@ def doGemm( ) require_launch_size(f"{label} GEMM", ("A", A), ("B", B), ("C", C)) - # wgrad's C is packed [G*N_out, K_in] on the output-M axis; fwd/dgrad/dense C is - # the unpacked (M_runtime, N_runtime). - expected_c = (num_groups * M_runtime, N_runtime) if wgrad else (M_runtime, N_runtime) - if tuple(C.shape) != expected_c: - raise ValueError(f"C shape {tuple(C.shape)} != expected {expected_c}") - if epilogue not in ("DEFAULT", "BIAS", "GELU_AUX", "GELU_AUX_BIAS"): raise ValueError(f"Unsupported {label} epilogue: {epilogue}") needs_bias = epilogue in ("BIAS", "GELU_AUX_BIAS") diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index 71f2ffecc4..90fdc4a968 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -1996,26 +1996,17 @@ def do_gemm( else: raise ValueError(f"Unsupported MXFP8 kernel layout: {layout}") - # Grouped launches read B from per-expert weight pointers, so B is a per-expert - # representative whose contraction equals A's (as in the dense case). - if group and b_data_ptrs is None: - raise FlyDSLUnsupportedError("grouped MXFP8 GEMM requires per-expert weight pointers") - assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" - - # Grouped wgrad (NT): both operands (dY, X) are packed on the token/contraction - # axis to G*M_pad, but the kernel's compile-time contraction is the PER-EXPERT - # padded token count M_pad. M_runtime (=N_out) and N_runtime (=K_in) are already - # per-expert; only K_runtime must be de-packed. + # Grouped wgrad (NT): dY and X stack the G experts on the contraction (token) + # axis as G*M_pad, while the kernel's compile-time contraction is the per-expert + # M_pad. M_runtime (N_out) and N_runtime (K_in) are already per-expert. if wgrad: assert ( K_runtime % num_groups == 0 ), f"grouped wgrad contraction {K_runtime} not divisible by num_groups={num_groups}" K_runtime = K_runtime // num_groups - # Grouped fwd/dgrad launch over the per-expert output width = C's last dim (TN: - # N_out; NN: K_in); the kernel re-derives the packed stride from num_groups. - # wgrad keeps M_runtime/N_runtime per-expert (C is packed on the output-M axis, - # handled in-kernel via store_m_expert_off). + # Grouped fwd/dgrad: N is the per-expert output width, C's last dim (TN: N_out, + # NN: K_in). wgrad keeps M/N per-expert; its C stacks experts on the M axis. if group: N_runtime = C.shape[-1] supported_fp8_dtypes = (torch.float8_e4m3fn, torch.float8_e5m2) @@ -2061,10 +2052,6 @@ def do_gemm( assert Bs.dtype == torch.int32, f"Bs dtype {Bs.dtype} != torch.int32 packed scales" assert tuple(As.shape) == expected_as, f"As shape {tuple(As.shape)} != {expected_as}" assert tuple(Bs.shape) == expected_bs, f"Bs shape {tuple(Bs.shape)} != {expected_bs}" - # wgrad's C is packed [G*N_out, K_in] on the output-M axis; fwd/dgrad C is the - # unpacked (M_runtime, N_runtime). - expected_c = (num_groups * M_runtime, N_runtime) if wgrad else (M_runtime, N_runtime) - assert tuple(C.shape) == expected_c, f"C shape {tuple(C.shape)} != {expected_c}" if C.dtype not in (torch.float16, torch.bfloat16, torch.float32): raise TypeError( f"C dtype must be torch.float16, torch.bfloat16, or torch.float32, got {C.dtype}" @@ -2114,7 +2101,7 @@ def do_gemm( if not C.is_contiguous(): raise FlyDSLUnsupportedError("FlyDSL MXFP8 requires contiguous output storage") - # Preserve the exact flat descriptor contract used by the passing kernels. + # The kernel addresses every operand through a flat 1-D buffer descriptor. A_arg = A.view(torch.uint8).contiguous().view(-1) B_arg = B.view(torch.uint8).contiguous().view(-1) As_arg = As.contiguous().view(-1) @@ -2148,10 +2135,9 @@ def do_gemm( work_queue_arg = _dummy_arg(torch.int32, A.device) work_counter_arg = _dummy_arg(torch.int32, A.device) - # Per-expert B base pointers (grouped fwd TN / dgrad NN weights; wgrad's pre-shifted - # X with act_ptrs): an int64[G] device array read in-kernel so the operand needs no - # host-side concat. Dense supplies a dtype-correct placeholder to keep the - # flyc.compile positional ABI uniform; the kernel never dereferences it there. + # Per-expert B base pointers: the weights in grouped fwd/dgrad, or wgrad's + # pre-shifted X with act_ptrs. Dense passes an unused placeholder so the + # positional kernel signature stays the same. if b_data_ptrs is not None: b_data_ptrs_arg = b_data_ptrs.contiguous().view(-1) else: @@ -2284,8 +2270,8 @@ def _mxfp8_matmul_grouped_wgrad( (``a_data_ptrs``/``b_data_ptrs``) as one representative expert each plus per-expert pointers, with ``a_shape``/``b_shape`` the packed shapes. """ - gm, n_out = a_shape if a_shape is not None else a.shape # gm = G*M_pad - gm_b, k_in = b_shape if b_shape is not None else b.shape + gm = (a_shape if a_shape is not None else a.shape)[0] # gm = G*M_pad + gm_b = (b_shape if b_shape is not None else b.shape)[0] if gm != gm_b: raise ValueError(f"wgrad a/b token dims differ: {gm} vs {gm_b}") if gm % num_groups != 0: @@ -2295,8 +2281,6 @@ def _mxfp8_matmul_grouped_wgrad( raise FlyDSLUnsupportedError( f"wgrad padded per-expert tokens M_pad={m_pad} must be divisible by {SCALE_GROUP_SIZE}" ) - if tuple(D.shape) != (num_groups * n_out, k_in): - raise ValueError(f"wgrad D shape {tuple(D.shape)} != expected {(num_groups * n_out, k_in)}") if _scale_dtype(a_scale) != torch.uint8 or _scale_dtype(b_scale) != torch.uint8: raise TypeError("FlyDSL MXFP8 wgrad expects raw E8M0 scales as torch.uint8") @@ -2393,7 +2377,10 @@ def mxfp8_matmul( ) if group and layout not in ("TN", "NN"): - raise ValueError("FlyDSL MXFP8 Group GEMM only supports the TN and NN layouts") + raise ValueError( + "FlyDSL MXFP8 grouped fwd/dgrad supports only the TN and NN layouts " + "(grouped NT wgrad uses wgrad=True)" + ) if layout not in ("TN", "NN", "NT"): raise ValueError(f"Unsupported MXFP8 kernel layout: {layout}") @@ -2411,47 +2398,18 @@ def mxfp8_matmul( k, m = a_rows, a_cols kb, n = b.shape - # Grouped GEMM is G INDEPENDENT matmuls stacked into one launch, not a dense - # A @ B_packed. Where the packed B axis lands, and hence which operand shape - # unpacks, depends on the layout: - # - # TN forward (C = A @ B^T): B is packed [G*N_out, K] on the OUTPUT-N axis. - # The kernel-K contraction (K_in) matches A's K, so kb == k. The output - # width unpacks: n = G*N_out -> n_out = N_out. D = [M_total, N_out]. - # (Per expert g: A_g[m_g,K] @ W_g^T[K,N_out] = D_g[m_g,N_out]; stacking the - # rows gives [M_total, N_out] -- the packed G*N_out rows are storage only, - # never the output width; a dense A@B_packed^T would give [M_total,G*N_out], - # G x too much compute, which grouped GEMM avoids.) - # - # NN dgrad (dX = dY @ W): the kernel CONTRACTION is N_out, and B=W is packed - # [G*N_out, K_in] on that CONTRACTION axis. So B's first dim kb = G*N_out - # is G x A's per-expert contraction k = N_out (kb != k by design), and the - # output width n = K_in is NOT packed. D = [M_total, K_in] = (m, n). + # Grouped fwd (TN) and dgrad (NN) read B per expert through b_data_ptrs, so b is + # one expert's weight: its shape gives the per-expert output width (TN: N_out, + # NN: K_in) and contraction, exactly as in the dense case. Only the B scales are + # stacked over the G experts (checked below). if group and b_data_ptrs is None: raise FlyDSLUnsupportedError("grouped MXFP8 GEMM requires per-expert weight pointers") - if group and layout == "TN": - # B is a per-expert [N_out, K] representative read through b_data_ptrs, so - # b.shape[0] already IS N_out; the B-scale still spans G*N_out rows. - n_out = n - expected_kb = k # contraction unpacked; B rows already per-K - elif group and layout == "NN": - # B is a per-expert [N_out, K_in] representative read through b_data_ptrs, so - # b.shape[0] IS the per-expert contraction N_out (= k). - n_out = n # output width (K_in) is not packed - expected_kb = k - else: - n_out = n - expected_kb = k - - if kb != expected_kb: + n_out = n + if kb != k: raise ValueError( f"Incompatible MXFP8 {layout} operands: A{tuple(a.shape)} and B{tuple(b.shape)}" ) - # D is the full stacked output [M_total, n_out]; expert g writes its - # [m_g, n_out] slice into rows [group_offs[g], group_offs[g+1]). - if tuple(D.shape) != (m, n_out): - raise ValueError(f"D shape {tuple(D.shape)} != expected {(m, n_out)}") if k % SCALE_GROUP_SIZE != 0: raise FlyDSLUnsupportedError( f"K={k} must be divisible by MXFP8 scale group size {SCALE_GROUP_SIZE}" diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index e2b1027ab8..7195b12863 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -869,12 +869,7 @@ def _flydsl_grouped_kind( in_features, out_features, ): - """Return which FlyDSL grouped path applies: 'mxfp8' | 'fp8' | 'bf16' | None. - - All kinds share the persistent work-stealing driver; they differ only in - operand packing and (fp8/mxfp8) the scale payload. The 'bf16' kind is the - unquantized half-precision path (bf16/fp16; fp8 is False). - """ + """Return which FlyDSL grouped path applies: 'mxfp8' | 'fp8' | 'bf16' | None.""" if not _GroupedLinear._flydsl_grouped_common_ok( use_bias=use_bias, use_grouped_gemm_triton=use_grouped_gemm_triton, @@ -889,8 +884,8 @@ def _flydsl_grouped_kind( ): return None if not fp8: - # Unquantized BF16 grouped GEMM (no recipe). FP16 is not routed here; - # the activation dtype is checked at the dispatch site. + # Unquantized grouped GEMM. The 'bf16' kind covers BF16 and FP16 (one + # half-precision kernel); the caller drops it for other activation dtypes. return "bf16" if recipe is None: return None @@ -901,7 +896,7 @@ def _flydsl_grouped_kind( return None @staticmethod - def _grouped_gemm_flydsl_mxfp8( + def _fwd_flydsl_mxfp8_grouped( *, inputmats, weights, @@ -961,7 +956,7 @@ def _grouped_gemm_flydsl_mxfp8( ) @staticmethod - def _backward_flydsl_mxfp8_grouped( + def _dgrad_flydsl_mxfp8_grouped( *, grad_output, weights, @@ -1070,10 +1065,10 @@ def _wgrad_flydsl_mxfp8_grouped( wgrad_list[g].copy_(dW[g * n_out : (g + 1) * n_out]) @staticmethod - def _grouped_gemm_flydsl_bf16(*, inputmats, weights, out, m_splits, m_splits_tensor): + def _fwd_flydsl_bf16_grouped(*, inputmats, weights, out, m_splits, m_splits_tensor): """FlyDSL BF16 grouped forward (TN): Y_e = X_e @ W_e^T, one packed launch. - Unquantized analog of :meth:`_grouped_gemm_flydsl_mxfp8`: activations are the + Unquantized analog of :meth:`_fwd_flydsl_mxfp8_grouped`: activations are the per-group BF16 splits and weights the per-expert BF16 tensors, both read in place through per-expert base pointers. No scales. """ @@ -1106,7 +1101,7 @@ def _grouped_gemm_flydsl_bf16(*, inputmats, weights, out, m_splits, m_splits_ten ) @staticmethod - def _backward_flydsl_bf16_grouped(*, grad_output, weights, dgrad, m_splits, m_splits_tensor): + def _dgrad_flydsl_bf16_grouped(*, grad_output, weights, dgrad, m_splits, m_splits_tensor): """FlyDSL BF16 grouped dgrad (NN): dX_e = dY_e @ W_e, one packed launch.""" from ..flydsl_kernels.gemm import _run_bf16_grouped from ..flydsl_kernels.gemm.grouped_gemm_utils import expert_ptrs @@ -1169,7 +1164,7 @@ def _wgrad_flydsl_bf16_grouped(*, grad_output, inputmats, wgrad_list, m_splits): wgrad_list[g].copy_(dW[g * n_out : (g + 1) * n_out]) @staticmethod - def _grouped_gemm_flydsl_fp8(*, inputmats, weights, out, m_splits, m_splits_tensor): + def _fwd_flydsl_fp8_grouped(*, inputmats, weights, out, m_splits, m_splits_tensor): """FlyDSL tensor-wise FP8 grouped forward (TN): Y_e = X_e @ W_e^T. Experts pack weights on the output-N axis ([G*N, K]); each operand carries @@ -1211,7 +1206,7 @@ def _grouped_gemm_flydsl_fp8(*, inputmats, weights, out, m_splits, m_splits_tens ) @staticmethod - def _backward_flydsl_fp8_grouped(*, grad_output, weights, dgrad, m_splits, m_splits_tensor): + def _dgrad_flydsl_fp8_grouped(*, grad_output, weights, dgrad, m_splits, m_splits_tensor): """FlyDSL tensor-wise FP8 grouped dgrad: dX_e = dY_e @ W_e, via the TN kernel. FP8 is TN-only, so dgrad reaches the same TN kernel using materialized @@ -1315,14 +1310,14 @@ def _flydsl_grouped_dispatch(cls, phase, kind): """ table = { "fwd": { - "mxfp8": cls._grouped_gemm_flydsl_mxfp8, - "fp8": cls._grouped_gemm_flydsl_fp8, - "bf16": cls._grouped_gemm_flydsl_bf16, + "mxfp8": cls._fwd_flydsl_mxfp8_grouped, + "fp8": cls._fwd_flydsl_fp8_grouped, + "bf16": cls._fwd_flydsl_bf16_grouped, }, "dgrad": { - "mxfp8": cls._backward_flydsl_mxfp8_grouped, - "fp8": cls._backward_flydsl_fp8_grouped, - "bf16": cls._backward_flydsl_bf16_grouped, + "mxfp8": cls._dgrad_flydsl_mxfp8_grouped, + "fp8": cls._dgrad_flydsl_fp8_grouped, + "bf16": cls._dgrad_flydsl_bf16_grouped, }, "wgrad": { "mxfp8": cls._wgrad_flydsl_mxfp8_grouped, @@ -1358,8 +1353,8 @@ def _try_flydsl_grouped(cls, phase, kind, **kwargs): if warn_fallback: warnings.warn( "[FLYDSL WARNING]: FlyDSL grouped GEMM does not support this " - f"configuration ({phase}); falling back to the default backend. " - f"Reason: {exc}", + f"configuration (kind={kind}, phase={phase}); falling back to the " + f"default backend. Reason: {exc}", UserWarning, stacklevel=2, ) @@ -1795,8 +1790,8 @@ def forward( in_features=weights[0].size(-1), out_features=weights[0].size(0), ) - # The half-precision kernel handles bf16 and fp16 (the MFMA opcode is selected - # per operand dtype); drop the kind for any other (e.g. fp32) activation dtype. + # The half-precision kernel handles bf16 and fp16. + # Drop the kind for any other (e.g. fp32) activation dtype. if flydsl_grouped_kind == "bf16" and activation_dtype not in ( torch.bfloat16, torch.float16, From ff139a49b9f0d69fce5b2b754e19b5c9cf37e3c5 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 21 Oct 2026 15:26:38 +0000 Subject: [PATCH 17/22] address bot comments --- .../flydsl_kernels/gemm/half_prec_gemm.py | 34 +++++--------- .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 45 ++++++++----------- 2 files changed, 29 insertions(+), 50 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py index 8d7181273a..e3b3ea6aef 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py @@ -475,28 +475,16 @@ def process_tile(pid_m, pid_n, expert): bx_m_idx = fx.Index(bx_m) by_n_idx = fx.Index(by_n) - # Per-expert offsets fold the expert index into the packed operands / - # output (const_expr-zeroed for dense, so the dense binary is unchanged): - # TN fwd -- experts on B's output-N axis (by_n_b_idx = expert*c_n + n) - # NN dgrad-- experts on B's contraction axis (b_expert_k_off = expert*K) - # NT wgrad-- both operands on the contraction (token) axis - # (a/b_expert_k_off = expert*K); output packed on M (store_m_expert_off) + # Per-expert offsets into the packed operands / output. Only wgrad (NT) needs + # them: both operands are packed on the contraction (token) axis + # (a/b_expert_k_off = expert*K) and the output on M (store_m_expert_off). + # Grouped fwd/dgrad read B per expert through b_data_ptrs (see b_div_tile + # below), so they use local coordinates and zero offsets, like dense. if const_expr(wgrad): by_n_b_idx = by_n_idx a_expert_k_off = fx.Index(expert * fx.Int32(K)) b_expert_k_off = fx.Index(expert * fx.Int32(K)) store_m_expert_off = fx.Index(expert * fx.Int32(c_m)) - elif const_expr(group): - if const_expr(b_transpose_read): - by_n_b_idx = by_n_idx - a_expert_k_off = fx.Index(0) - b_expert_k_off = fx.Index(expert * fx.Int32(K)) - store_m_expert_off = fx.Index(0) - else: - by_n_b_idx = fx.Index(expert * fx.Int32(c_n) + by_n) - a_expert_k_off = fx.Index(0) - b_expert_k_off = fx.Index(0) - store_m_expert_off = fx.Index(0) else: by_n_b_idx = by_n_idx a_expert_k_off = fx.Index(0) @@ -533,13 +521,11 @@ def process_tile(pid_m, pid_n, expert): fx.Uint8.ir_type, wave_id, ) - # Grouped B-data: rebase onto this tile's expert weight slab (byte buffer); - # the per-expert base already selects the expert, so in-tile addressing is - # LOCAL. TN reads a local column (by_n_idx); NN (dgrad, transpose-read) - # packs experts on the contraction axis, so the base replaces the expert*K - # contraction shift (b_k_off -> 0). wgrad bases are pre-shifted (act_ptrs), - # so they keep the packed column and contraction shift. Dense keeps the - # shared packed view. + # Grouped B-data: rebase onto this tile's expert weight (byte buffer); the + # per-expert base already selects the expert, so in-tile addressing is + # local (column by_n_idx, no contraction shift) for both TN and NN. + # wgrad bases are pre-shifted (act_ptrs), so they keep the packed column + # and contraction shift. Dense keeps the shared packed view. if const_expr(group): b_base_i64 = gd.make_uniform_i64(b_base_raw) b_div_tile = gd.rebase_buffer_view(b_base_i64, fx.Uint8.ir_type, b_ptr_layout) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index 90fdc4a968..de7dffd4fe 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -777,23 +777,19 @@ def process_tile(pid_m, pid_n, expert): bx_m_idx = fx.Index(bx_m) by_n_idx = fx.Index(by_n) - # Per-expert offsets into the packed B operand / B-scale. Where the - # expert index lands depends on the layout, because the packed weight - # stacks experts on B's *first* logical axis, which is a different - # kernel axis for TN vs the transpose-read (NN/NT) case: + # Per-expert offsets. Grouped fwd/dgrad read B data per expert through + # b_data_ptrs (see b_div_tile below), so B-data addressing is local and + # b_expert_k_off is 0. Only the B scales stay stacked over the G experts, + # on B's first logical axis, which is a different kernel axis per layout: # - # TN (forward, C = A @ B^T): B is packed [G*N, K], experts on the - # OUTPUT-N axis. Expert offset folds into the B column base - # (by_n_b_idx = expert*N + n) and the B-scale row stride grows to - # G*N (c_n_bs). The B-scale k128 term is untouched. + # TN (forward, C = A @ B^T): B-scale stacks experts on the OUTPUT-N + # axis, so the scale column is by_n_b_idx = expert*N + n and the + # scale row stride grows to G*N (c_n_bs). # # NN (dgrad, C = A @ B where A=dY[M,N_out], B=W[N_out,K_in]): the - # kernel contraction axis IS N_out, and B is packed [G*N_out, K_in] - # with experts on that CONTRACTION axis. Expert offset instead folds - # into the contraction (k_base) term as expert*N_out elements - # (b_expert_k_off), and into the B-scale's k128 (contraction-128) - # term as expert*(N_out/128) (bs_expert_k128_off). The B column base - # and B-scale stride stay per-expert (c_n = K_in, unchanged). + # kernel contraction axis is N_out, and the B-scale stacks experts + # on it, so the scale k128 term shifts by expert*(N_out/128) + # (bs_expert_k128_off). The scale column and stride stay per-expert. # # Dense (group is False) zeroes every expert term. # @@ -817,14 +813,13 @@ def process_tile(pid_m, pid_n, expert): store_m_expert_off = fx.Index(expert * fx.Int32(c_m)) elif const_expr(group): if const_expr(b_transpose_read): - # NN: expert on the contraction axis. The kernel's compile-time - # ``K`` IS the contraction dim (= N_out for dgrad), so the packed - # weight [G*N_out, K_in] shifts by expert*K elements along k_base, - # and its scale [G*N_out/128, K_in] by expert*NUM_K_TILES. + # NN: the B-scale [G*N_out/128, K_in] stacks experts on the + # contraction axis (compile-time K = N_out), so it shifts by + # expert*NUM_K_TILES. B data is read per expert (no shift). by_n_b_idx = by_n_idx c_n_bs = c_n a_expert_k_off = fx.Index(0) - b_expert_k_off = fx.Index(expert * fx.Int32(K)) + b_expert_k_off = fx.Index(0) as_expert_k128_off = fx.Int32(0) bs_expert_k128_off = expert * fx.Int32(NUM_K_TILES) store_m_expert_off = fx.Index(0) @@ -898,13 +893,11 @@ def process_tile(pid_m, pid_n, expert): a_f8_ir_t, wave_id, ) - # Grouped B-data source: rebase onto this tile's expert weight slab (the - # per-expert base pointer already selects the expert), so the in-tile - # addressing uses LOCAL coords. TN reads with a local column (by_n_idx); - # NN (dgrad, transpose-read) packs experts on the contraction axis, so the - # per-expert base replaces the expert*K contraction shift (b_k_off -> 0). - # wgrad bases are pre-shifted (act_ptrs), so they keep the packed column and - # contraction shift. Dense keeps the shared packed ``b_div``. + # Grouped B-data source: rebase onto this tile's expert weight (the + # per-expert base pointer already selects the expert), so in-tile + # addressing is local (column by_n_idx, no contraction shift) for both + # TN and NN. wgrad bases are pre-shifted (act_ptrs), so they keep the + # packed column and contraction shift. Dense keeps the shared ``b_div``. if const_expr(group): b_base_i64 = gd.make_uniform_i64(b_base_raw) b_div_tile = gd.rebase_buffer_view(b_base_i64, b_f8_ir_t, b_ptr_layout) From 6635016cf71a331ac2416c329f9bb7b2875300a8 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 21 Oct 2026 18:08:09 +0000 Subject: [PATCH 18/22] address bot comments --- tests/pytorch/test_grouped_linear.py | 50 ++++++++++++++----- .../pytorch/flydsl_kernels/gemm/fp8_gemm.py | 12 ++++- .../flydsl_kernels/gemm/gemm_common_utils.py | 14 +++--- .../flydsl_kernels/gemm/half_prec_gemm.py | 25 +++++++++- .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 18 ++++++- 5 files changed, 96 insertions(+), 23 deletions(-) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 4f8944ce8f..909fbe5214 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -4,8 +4,10 @@ # # See LICENSE for license information. +import contextlib import os import random +import warnings from typing import Dict, List, Optional, Sequence import pytest @@ -672,21 +674,43 @@ def test_grouped_linear_accuracy_rocm_backends( if expect_ck_no_fallback: monkeypatch.setenv("NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK", "1") - test_grouped_linear_accuracy( - dtype, - num_gemms, - bs, - model, - recipe, - fp8_model_params, - fuse_wgrad_accumulation, - False, - delay_wgrad_compute, - use_triton=False, - parallel_mode=None, - use_cutlass=True, + # FlyDSL falls back to the default grouped GEMM silently on any config it rejects + # past the gate skipped above, which would compare the default backend against + # itself. Enable its fallback warning and assert none was emitted. + expect_flydsl_ran = grouped_gemm_backend == "flydsl" + if expect_flydsl_ran: + monkeypatch.setenv("NVTE_FLYDSL_GEMM_WARN_FALLBACK", "1") + capture = ( + warnings.catch_warnings(record=True) if expect_flydsl_ran else contextlib.nullcontext() ) + with capture as caught: + if expect_flydsl_ran: + warnings.simplefilter("always") + test_grouped_linear_accuracy( + dtype, + num_gemms, + bs, + model, + recipe, + fp8_model_params, + fuse_wgrad_accumulation, + False, + delay_wgrad_compute, + use_triton=False, + parallel_mode=None, + use_cutlass=True, + ) + + if expect_flydsl_ran: + fallbacks = sorted({str(w.message) for w in caught if "[FLYDSL WARNING]" in str(w.message)}) + assert not fallbacks, ( + "FlyDSL grouped GEMM fell back to the default backend, so it never ran " + f"(dtype={dtype}, recipe={recipe_id(recipe)}, fp8_model_params={fp8_model_params}, " + f"fuse_wgrad={fuse_wgrad_accumulation}, delay_wgrad={delay_wgrad_compute}):\n" + + "\n".join(fallbacks) + ) + if expect_ck_no_fallback: # NVTE_WARN writes to std::cerr; capfd captures file-descriptor-level # output, including C/C++ stderr. diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py index 8e0d686d0c..cc1c3ddd8b 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -1367,7 +1367,17 @@ def doGemm( block_k=_BLOCK_K, label="FP8 GEMM", ) - require_launch_size("FP8 GEMM", ("A", A), ("B", B), ("C", C)) + # With per-expert pointers, A (and wgrad's B) is one representative expert, but + # the kernel addresses the whole packed operand through 32-bit offsets: grouped A + # is [M_total, K]; wgrad's A and B stack the G experts' padded tokens ([G*K, M], + # [G*K, N]). Grouped fwd/dgrad weights use local offsets, so B's own size holds. + a_extent = b_extent = () + if a_data_ptrs is not None: + a_numel = num_groups * K_runtime * M_runtime if wgrad else M_runtime * K_runtime + a_extent = (a_numel * A.element_size(),) + if wgrad and b_data_ptrs is not None: + b_extent = (num_groups * K_runtime * N_runtime * B.element_size(),) + require_launch_size("FP8 GEMM", ("A", A, *a_extent), ("B", B, *b_extent), ("C", C)) # Dense: one scalar inverse scale per operand. Grouped: a per-expert vector, or # (scale_ptrs) an int64[G] array of per-expert scale_inv pointers read in place. expected_scale_numel = num_groups if (group or wgrad) else 1 diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py index 36136f08a9..d5e68e534d 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py @@ -63,15 +63,17 @@ def require_block_tiling(m, n, k, *, block_m, block_n, block_k, label, min_k_til def require_launch_size(label, *tensors): """Reject operands whose byte size exceeds the int32 launch-argument limit. - ``tensors`` is an iterable of ``(name, tensor)`` pairs. Raises - ``FlyDSLUnsupportedError`` (so ``general_gemm`` falls back to the C++ - backend) for the first operand at or above ``_MAX_LAUNCH_BYTES``; skips - ``None`` tensors (e.g. an absent bias/aux). + ``tensors`` is an iterable of ``(name, tensor)`` or ``(name, tensor, nbytes)`` + entries. ``nbytes`` overrides the tensor's own size: a grouped operand read + through per-expert pointers is passed as one representative expert, while the + kernel addresses the whole packed operand. Raises ``FlyDSLUnsupportedError`` + (so ``general_gemm`` falls back to the C++ backend) for the first operand at or + above ``_MAX_LAUNCH_BYTES``; skips ``None`` tensors (e.g. an absent bias/aux). """ - for name, t in tensors: + for name, t, *override in tensors: if t is None: continue - nbytes = t.numel() * t.element_size() + nbytes = override[0] if override else t.numel() * t.element_size() if nbytes > _MAX_LAUNCH_BYTES: raise FlyDSLUnsupportedError( f"FlyDSL {label} operand {name} is {nbytes} bytes, which exceeds " diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py index e3b3ea6aef..e26dbf20c8 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py @@ -1694,7 +1694,23 @@ def doGemm( block_k=_BLOCK_K, label=f"{label} GEMM", ) - require_launch_size(f"{label} GEMM", ("A", A), ("B", B), ("C", C)) + # With per-expert pointers, A (and wgrad's B) is one representative expert, but + # the kernel addresses the whole packed operand through 32-bit offsets: grouped A + # is [M_total, K]; wgrad's A and B stack the G experts' padded tokens ([G*K, M], + # [G*K, N]). Grouped fwd/dgrad weights use local offsets, so B's own size holds. + a_extent = b_extent = () + if a_data_ptrs is not None: + a_numel = num_groups * K_runtime * M_runtime if wgrad else M_runtime * K_runtime + a_extent = (a_numel * A.element_size(),) + if wgrad and b_data_ptrs is not None: + b_extent = (num_groups * K_runtime * N_runtime * B.element_size(),) + require_launch_size(f"{label} GEMM", ("A", A, *a_extent), ("B", B, *b_extent), ("C", C)) + + # wgrad's C is packed [G*N_out, K_in] on the output-M axis; fwd/dgrad/dense C is + # the unpacked (M_runtime, N_runtime). + expected_c = (num_groups * M_runtime, N_runtime) if wgrad else (M_runtime, N_runtime) + if tuple(C.shape) != expected_c: + raise ValueError(f"C shape {tuple(C.shape)} != expected {expected_c}") if epilogue not in ("DEFAULT", "BIAS", "GELU_AUX", "GELU_AUX_BIAS"): raise ValueError(f"Unsupported {label} epilogue: {epilogue}") @@ -1770,6 +1786,13 @@ def doGemm( b_data_ptrs_arg = b_data_ptrs.contiguous().view(-1) else: b_data_ptrs_arg = _dense_placeholder(A.device, torch.int64) + # act_ptrs also selects how the kernel reads B: outside grouped fwd/dgrad, B is + # read through b_data_ptrs only when act_ptrs is set. So activation pointers are + # grouped-only, and wgrad takes the dY and X pointer arrays together or not at all. + if a_data_ptrs is not None and not (group or wgrad): + raise ValueError(f"{label}: per-expert activation pointers require a grouped launch") + if wgrad and (a_data_ptrs is None) != (b_data_ptrs is None): + raise ValueError(f"{label} wgrad: a_data_ptrs and b_data_ptrs must be given together") # Per-expert activation base pointers (see act_ptrs in _compile_kernel). act_ptrs = a_data_ptrs is not None if act_ptrs: diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index de7dffd4fe..817c8a2bc1 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -2015,10 +2015,13 @@ def do_gemm( block_k=_BLOCK_K, label=f"MXFP8 {layout} GEMM", ) + # With per-expert pointers, A (and wgrad's B) is one representative expert, but + # the kernel addresses the packed operand (a_shape/b_shape) through 32-bit + # offsets, so the bound is on the packed extent. require_launch_size( f"MXFP8 {layout} GEMM", - ("A", A), - ("B", B), + ("A", A, a_rows * a_cols * A.element_size()), + ("B", B, b_rows * b_cols * B.element_size()), ("As", As), ("Bs", Bs), ("C", C), @@ -2045,6 +2048,10 @@ def do_gemm( assert Bs.dtype == torch.int32, f"Bs dtype {Bs.dtype} != torch.int32 packed scales" assert tuple(As.shape) == expected_as, f"As shape {tuple(As.shape)} != {expected_as}" assert tuple(Bs.shape) == expected_bs, f"Bs shape {tuple(Bs.shape)} != {expected_bs}" + # wgrad's C is packed [G*N_out, K_in] on the output-M axis; fwd/dgrad C is the + # unpacked (M_runtime, N_runtime). + expected_c = (num_groups * M_runtime, N_runtime) if wgrad else (M_runtime, N_runtime) + assert tuple(C.shape) == expected_c, f"C shape {tuple(C.shape)} != {expected_c}" if C.dtype not in (torch.float16, torch.bfloat16, torch.float32): raise TypeError( f"C dtype must be torch.float16, torch.bfloat16, or torch.float32, got {C.dtype}" @@ -2135,6 +2142,13 @@ def do_gemm( b_data_ptrs_arg = b_data_ptrs.contiguous().view(-1) else: b_data_ptrs_arg = _dummy_arg(torch.int64, A.device) + # act_ptrs also selects how the kernel reads B: outside grouped fwd/dgrad, B is + # read through b_data_ptrs only when act_ptrs is set. So activation pointers are + # grouped-only, and wgrad takes the dY and X pointer arrays together or not at all. + if a_data_ptrs is not None and not (group or wgrad): + raise ValueError("MXFP8: per-expert activation pointers require a grouped launch") + if wgrad and (a_data_ptrs is None) != (b_data_ptrs is None): + raise ValueError("MXFP8 wgrad: a_data_ptrs and b_data_ptrs must be given together") # Per-expert activation base pointers (see act_ptrs in _compile_kernel). act_ptrs = a_data_ptrs is not None if act_ptrs: From e169938d4197626d78b05d6614382397288af265 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 21 Oct 2026 18:30:19 +0000 Subject: [PATCH 19/22] address bot comments --- tests/pytorch/test_grouped_linear.py | 73 +++++++++++-------- .../pytorch/flydsl_kernels/gemm/fp8_gemm.py | 24 +++--- .../flydsl_kernels/gemm/gemm_wrappers.py | 7 ++ .../flydsl_kernels/gemm/half_prec_gemm.py | 26 ++++--- 4 files changed, 77 insertions(+), 53 deletions(-) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 909fbe5214..68458039aa 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -4,10 +4,8 @@ # # See LICENSE for license information. -import contextlib import os import random -import warnings from typing import Dict, List, Optional, Sequence import pytest @@ -675,40 +673,51 @@ def test_grouped_linear_accuracy_rocm_backends( monkeypatch.setenv("NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK", "1") # FlyDSL falls back to the default grouped GEMM silently on any config it rejects - # past the gate skipped above, which would compare the default backend against - # itself. Enable its fallback warning and assert none was emitted. - expect_flydsl_ran = grouped_gemm_backend == "flydsl" - if expect_flydsl_ran: + # past the gate skipped above, and wgrad is not attempted at all with fused or + # delayed wgrad, so a missing fallback warning does not show that FlyDSL ran. + # Record each FlyDSL dispatch and assert on the phases that actually ran; the + # fallback warning stays on so a failure shows the reason. + flydsl_dispatches = [] + if grouped_gemm_backend == "flydsl": monkeypatch.setenv("NVTE_FLYDSL_GEMM_WARN_FALLBACK", "1") - capture = ( - warnings.catch_warnings(record=True) if expect_flydsl_ran else contextlib.nullcontext() - ) + try_flydsl_grouped = _GroupedLinear._try_flydsl_grouped - with capture as caught: - if expect_flydsl_ran: - warnings.simplefilter("always") - test_grouped_linear_accuracy( - dtype, - num_gemms, - bs, - model, - recipe, - fp8_model_params, - fuse_wgrad_accumulation, - False, - delay_wgrad_compute, - use_triton=False, - parallel_mode=None, - use_cutlass=True, + def recording_try_flydsl_grouped(cls, phase, kind, **kwargs): + ran = try_flydsl_grouped(phase, kind, **kwargs) + flydsl_dispatches.append((phase, ran)) + return ran + + monkeypatch.setattr( + _GroupedLinear, "_try_flydsl_grouped", classmethod(recording_try_flydsl_grouped) ) - if expect_flydsl_ran: - fallbacks = sorted({str(w.message) for w in caught if "[FLYDSL WARNING]" in str(w.message)}) - assert not fallbacks, ( - "FlyDSL grouped GEMM fell back to the default backend, so it never ran " - f"(dtype={dtype}, recipe={recipe_id(recipe)}, fp8_model_params={fp8_model_params}, " - f"fuse_wgrad={fuse_wgrad_accumulation}, delay_wgrad={delay_wgrad_compute}):\n" - + "\n".join(fallbacks) + test_grouped_linear_accuracy( + dtype, + num_gemms, + bs, + model, + recipe, + fp8_model_params, + fuse_wgrad_accumulation, + False, + delay_wgrad_compute, + use_triton=False, + parallel_mode=None, + use_cutlass=True, + ) + + if grouped_gemm_backend == "flydsl": + expected = {"fwd", "dgrad"} + if not fuse_wgrad_accumulation and not delay_wgrad_compute: + expected.add("wgrad") + ran = {phase for phase, ok in flydsl_dispatches if ok} + fell_back = sorted({phase for phase, ok in flydsl_dispatches if not ok}) + assert not fell_back and expected <= ran, ( + f"FlyDSL grouped GEMM did not run every expected phase: expected {sorted(expected)}, " + f"ran {sorted(ran)}, fell back in {fell_back} (dtype={dtype}, " + f"recipe={recipe_id(recipe)}, fp8_model_params={fp8_model_params}, " + f"fuse_wgrad={fuse_wgrad_accumulation}, delay_wgrad={delay_wgrad_compute}). " + "See the [FLYDSL WARNING] for the fallback reason." ) if expect_ck_no_fallback: diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py index cc1c3ddd8b..7e73babd09 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -1334,12 +1334,16 @@ def doGemm( m: int = None, n: int = None, k: int = None, + a_shape=None, + b_shape=None, ): """Launch tensor-wise FP8 GEMM with TE-style inverse input scales. Dense derives M/N/K from operand shapes. Grouped/wgrad pass explicit per-expert ``m``/``n``/``k`` (the packed operands do not reveal the per-expert geometry), a length-``num_groups`` scale vector for each operand, and the descriptor queue. + With ``a_data_ptrs`` (``b_data_ptrs`` in wgrad) the operand is one representative + expert and ``a_shape`` (``b_shape``) the packed operand shape. """ if group or wgrad: M_runtime, N_runtime, K_runtime = int(m), int(n), int(k) @@ -1368,16 +1372,16 @@ def doGemm( label="FP8 GEMM", ) # With per-expert pointers, A (and wgrad's B) is one representative expert, but - # the kernel addresses the whole packed operand through 32-bit offsets: grouped A - # is [M_total, K]; wgrad's A and B stack the G experts' padded tokens ([G*K, M], - # [G*K, N]). Grouped fwd/dgrad weights use local offsets, so B's own size holds. - a_extent = b_extent = () - if a_data_ptrs is not None: - a_numel = num_groups * K_runtime * M_runtime if wgrad else M_runtime * K_runtime - a_extent = (a_numel * A.element_size(),) - if wgrad and b_data_ptrs is not None: - b_extent = (num_groups * K_runtime * N_runtime * B.element_size(),) - require_launch_size("FP8 GEMM", ("A", A, *a_extent), ("B", B, *b_extent), ("C", C)) + # the kernel addresses the packed operand (a_shape/b_shape) through 32-bit + # offsets, so the bound is on the packed extent. + a_rows, a_cols = a_shape if a_shape is not None else A.shape + b_rows, b_cols = b_shape if b_shape is not None else B.shape + require_launch_size( + "FP8 GEMM", + ("A", A, a_rows * a_cols * A.element_size()), + ("B", B, b_rows * b_cols * B.element_size()), + ("C", C), + ) # Dense: one scalar inverse scale per operand. Grouped: a per-expert vector, or # (scale_ptrs) an int64[G] array of per-expert scale_inv pointers read in place. expected_scale_numel = num_groups if (group or wgrad) else 1 diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 44e8a0d117..9988008692 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -1326,6 +1326,7 @@ def _run_bf16_grouped( a_data_ptrs=a_data_ptrs, num_work=num_tiles, sched=sched, + a_shape=(m_total, a_data.size(-1)) if a_data_ptrs is not None else None, ) return D @@ -1431,6 +1432,8 @@ def _run_bf16_grouped_wgrad( a_data_ptrs=dy_ptrs, num_work=num_tiles, b_data_ptrs=x_ptrs, + a_shape=(bsize * k_pad, n_out) if dy_ptrs is not None else None, + b_shape=(bsize * k_pad, k_in) if x_ptrs is not None else None, ) if not direct: @@ -1526,6 +1529,7 @@ def _run_fp8_grouped( m=m_total, n=n, k=k, + a_shape=(m_total, a_data.size(-1)) if a_data_ptrs is not None else None, ) return D @@ -1641,6 +1645,9 @@ def _pad_cols(t, cols): a_data_ptrs=dy_ptrs, num_work=num_tiles, b_data_ptrs=x_ptrs, + # TN over transposes: the bucket's experts stack on the feature axis. + a_shape=(bsize * n_out, k_pad) if dy_ptrs is not None else None, + b_shape=(bsize * k_in, k_pad) if x_ptrs is not None else None, ) if not direct: diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py index e26dbf20c8..5e2f7ded7b 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py @@ -1656,6 +1656,8 @@ def doGemm( num_work: int = 0, sched: str = "worksteal", wgrad: bool = False, + a_shape=None, + b_shape=None, ): """Launch one cached K/output/layout-specialized half-precision core. @@ -1667,7 +1669,9 @@ def doGemm( NT: A backing [K,M], B backing [K,N] NN/NT orientation is implemented by compile-time global addressing and - ``ds_read_b64_tr_b16`` only. + ``ds_read_b64_tr_b16`` only. With ``a_data_ptrs`` (``b_data_ptrs`` in wgrad) the + operand is one representative expert and ``a_shape`` (``b_shape``) the packed + operand shape. """ if layout not in ("TN", "NN", "NT"): raise ValueError(f"Unsupported {label} layout: {layout}") @@ -1695,16 +1699,16 @@ def doGemm( label=f"{label} GEMM", ) # With per-expert pointers, A (and wgrad's B) is one representative expert, but - # the kernel addresses the whole packed operand through 32-bit offsets: grouped A - # is [M_total, K]; wgrad's A and B stack the G experts' padded tokens ([G*K, M], - # [G*K, N]). Grouped fwd/dgrad weights use local offsets, so B's own size holds. - a_extent = b_extent = () - if a_data_ptrs is not None: - a_numel = num_groups * K_runtime * M_runtime if wgrad else M_runtime * K_runtime - a_extent = (a_numel * A.element_size(),) - if wgrad and b_data_ptrs is not None: - b_extent = (num_groups * K_runtime * N_runtime * B.element_size(),) - require_launch_size(f"{label} GEMM", ("A", A, *a_extent), ("B", B, *b_extent), ("C", C)) + # the kernel addresses the packed operand (a_shape/b_shape) through 32-bit + # offsets, so the bound is on the packed extent. + a_rows, a_cols = a_shape if a_shape is not None else A.shape + b_rows, b_cols = b_shape if b_shape is not None else B.shape + require_launch_size( + f"{label} GEMM", + ("A", A, a_rows * a_cols * A.element_size()), + ("B", B, b_rows * b_cols * B.element_size()), + ("C", C), + ) # wgrad's C is packed [G*N_out, K_in] on the output-M axis; fwd/dgrad/dense C is # the unpacked (M_runtime, N_runtime). From 9efa0d8b98182ae5a5a7e75b267ce6e987480f51 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 21 Oct 2026 23:05:35 +0000 Subject: [PATCH 20/22] remove formatting-only changes from quantization.py --- transformer_engine/pytorch/quantization.py | 19 +++++++------------ 1 file changed, 7 insertions(+), 12 deletions(-) diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 3ff2c7dc1b..e608f4a9ca 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -177,7 +177,7 @@ def _compute_mxfp8_support() -> Tuple[bool, str]: if os.getenv("NVTE_ROCM_ENABLE_MXFP8", "0") == "0": return False, "MXFP8 support is not enabled." gpu_arch = get_device_compute_capability() - if gpu_arch in ((9, 5),): # TODO: enable for gfx1250 when GEMM is available + if gpu_arch in ((9, 5),): #TODO: enable for gfx1250 when GEMM is available return True, "" return False, "Device arch gfx95x or newer is required for MXFP8 execution." if get_device_compute_capability() >= (12, 0): @@ -361,7 +361,11 @@ def get_align_size_for_quantization(recipe: Recipe) -> int: use_ck = os.environ.get("NVTE_USE_CK_GROUPED_GEMM", "0") == "1" use_cutlass = os.environ.get("NVTE_USE_CUTLASS_GROUPED_GEMM", "0") == "1" use_hipkittens = use_hk or (use_cutlass and not use_ck) - if IS_HIP_EXTENSION and get_device_compute_capability() == (9, 5) and use_hipkittens: + if ( + IS_HIP_EXTENSION + and get_device_compute_capability() == (9, 5) + and use_hipkittens + ): return 256 return 32 if recipe.nvfp4(): @@ -467,7 +471,6 @@ def is_fp8_block_scaling_available(return_reason: bool = False) -> Union[bool, T return check_fp8_block_scaling_support() return check_fp8_block_scaling_support()[0] - def is_nvfp4_available(return_reason: bool = False) -> Union[bool, Tuple[bool, str]]: """ Determine if support is available for the NVFP4 recipe. @@ -485,7 +488,6 @@ def is_nvfp4_available(return_reason: bool = False) -> Union[bool, Tuple[bool, s return check_nvfp4_support() return check_nvfp4_support()[0] - def is_mxfp4_available(return_reason: bool = False) -> Union[bool, Tuple[bool, str]]: """ Determine if support is available for the MXFP4 recipe. @@ -502,7 +504,6 @@ def is_mxfp4_available(return_reason: bool = False) -> Union[bool, Tuple[bool, s return check_mxfp4_support() return check_mxfp4_support()[0] - @dataclass(slots=True) class FP8GlobalState: """Mutable process-global FP8 state stored on an instance. @@ -923,13 +924,7 @@ def autocast_exit(cls, enabled: bool, _graph: bool) -> None: # Reduce only the non-FP8 weight modules here. # FP8 weight modules are reduced at the end of the optimizer # step after the weight amax is populated. - if ( - not cls.SKIP_FP8_REDUCTION_FOR_FSDP2 - and enabled - and qstate.autocast_depth == 0 - and not _graph - and torch.is_grad_enabled() - ): + if not cls.SKIP_FP8_REDUCTION_FOR_FSDP2 and enabled and qstate.autocast_depth == 0 and not _graph and torch.is_grad_enabled(): # delayed scaling only function, for other recipes (current scaling with any granularity), # this is noop for other recipes because cls.global_amax_buffer is empty list cls.reduce_and_update_fp8_tensors(forward=True) From 2114ec4bdee5c3b8a7270752491b79f82a92220b Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 21 Oct 2026 23:56:06 +0000 Subject: [PATCH 21/22] Replace deprecated fx.Index with fx.Int64 in FlyDSL GEMM kernels --- .../pytorch/flydsl_kernels/gemm/fp32_gemm.py | 62 +++---- .../pytorch/flydsl_kernels/gemm/fp8_gemm.py | 78 ++++----- .../flydsl_kernels/gemm/grouped_gemm_utils.py | 4 +- .../flydsl_kernels/gemm/half_prec_gemm.py | 104 ++++++------ .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 156 +++++++++--------- 5 files changed, 202 insertions(+), 202 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py index 68e95fe471..12c85c7663 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py @@ -200,8 +200,8 @@ def kernel_gemm( # The flattened/XCD-swizzled block coordinates are i32, while global # address arithmetic below is expressed in MLIR index type. Convert # once here and use these index-typed tile bases for every address. - bx_m_idx = fx.Index(bx_m) - by_n_idx = fx.Index(by_n) + bx_m_idx = fx.Int64(bx_m) + by_n_idx = fx.Int64(by_n) # Keep wave/lane arithmetic in i32. compute_global_swizzle() combines # these values with i32 constants, so Index-typed coordinates would make @@ -231,7 +231,7 @@ def kernel_gemm( # Per-CTA tile base in C elements, folded into each store's linear # coordinate below (matching the scale-load addressing on this build; # add_offset on a dynamic Index is unsupported here). - c_tile_base_elems = bx_m_idx * fx.Index(c_n) + by_n_idx + c_tile_base_elems = bx_m_idx * fx.Int64(c_n) + by_n_idx gC = fx.rocdl.make_buffer_tensor(C, max_size=True) c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) c_store_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) @@ -366,15 +366,15 @@ def hot_loop_scheduler_q_prefetch_4n(): def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): # One pass writes 256 threads * 16 B = 4 KiB. Four passes fill one # 128x64 half-page (16 KiB). Each half has its own LDS base. - global_base = (bx_m_idx + fx.Index(subtile * (BLOCK_M // 2))) * fx.Index( + global_base = (bx_m_idx + fx.Int64(subtile * (BLOCK_M // 2))) * fx.Int64( K * ELEM_BYTES - ) + k_base * fx.Index(ELEM_BYTES) + ) + k_base * fx.Int64(ELEM_BYTES) a_g2s.load_one(lds_a[subtile], fx.Int32(global_base), pass_in_subtile) def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): - global_base = (by_n_idx + fx.Index(subtile * (BLOCK_N // 2))) * fx.Index( + global_base = (by_n_idx + fx.Int64(subtile * (BLOCK_N // 2))) * fx.Int64( K * ELEM_BYTES - ) + k_base * fx.Index(ELEM_BYTES) + ) + k_base * fx.Int64(ELEM_BYTES) b_g2s.load_one(lds_b[subtile], fx.Int32(global_base), pass_in_subtile) def stage_a_subtile(k_base, subtile, lds_a): @@ -403,8 +403,8 @@ def load_frag_at_byte_base(lds_page, row_byte_base): def load_b_frag(lds_b, local_row, half): # B is [N, K]. Each 128-row half-page has a local row origin of 0. - half_row = local_row - fx.Index(half * (BLOCK_N // 2)) - return load_frag_at_byte_base(lds_b[half], half_row * fx.Index(BLOCK_K * ELEM_BYTES)) + half_row = local_row - fx.Int64(half * (BLOCK_N // 2)) + return load_frag_at_byte_base(lds_b[half], half_row * fx.Int64(BLOCK_K * ELEM_BYTES)) def _acc_idx(subtile_id, mi, ni): return ( @@ -491,10 +491,10 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): mi = local_idx // MFMA_N_PER_SUBTILE ni = local_idx % MFMA_N_PER_SUBTILE - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - row_base = subtile_m_idx * SUBTILE_M + fx.Index(mi * MFMA_M) + lane_div_16 * 4 - col = subtile_n_idx * SUBTILE_N + fx.Index(ni * MFMA_N) + lane_mod_16 + subtile_m_idx = reg_subtile_m_idx0 + fx.Int64(sm * 2) + subtile_n_idx = reg_subtile_n_idx0 + fx.Int64(sn * 2) + row_base = subtile_m_idx * SUBTILE_M + fx.Int64(mi * MFMA_M) + lane_div_16 * 4 + col = subtile_n_idx * SUBTILE_N + fx.Int64(ni * MFMA_N) + lane_mod_16 # Bias depends only on the output-feature (N) coordinate, so read it # once per column (global index by_n_idx + col) and reuse across the @@ -503,8 +503,8 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): bias_value = load_bias(by_n_idx + col) for ii in range_constexpr(4): - row = row_base + fx.Index(ii) - c_idx = c_tile_base_elems + row * fx.Index(c_n) + col + row = row_base + fx.Int64(ii) + c_idx = c_tile_base_elems + row * fx.Int64(c_n) + col # Epilogue stages run on the fp32 accumulator; output is fp32 # so there is no dtype narrowing: @@ -546,8 +546,8 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): zero_pinned_accumulators() def load_b_subtile_ni_regs(lds_b, sn, ni): - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - b_row_addr = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 + subtile_n_idx = reg_subtile_n_idx0 + fx.Int64(sn * 2) + b_row_addr = subtile_n_idx * fx.Int64(SUBTILE_N) + fx.Int64(ni * MFMA_N) + lane_mod_16 return load_b_frag(lds_b, b_row_addr, sn) def load_b_subtile_regs(lds_b, sn): @@ -559,10 +559,10 @@ def load_b_subtile_regs(lds_b, sn): ) def load_a_subtile_mi_half(lds_a, sm, mi, half): - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - a_row_addr = subtile_m_idx * fx.Index(SUBTILE_M) + fx.Index(mi * MFMA_M) + lane_mod_16 - half_row = a_row_addr - fx.Index(sm * (BLOCK_M // 2)) - row_byte_base = half_row * fx.Index(BLOCK_K * ELEM_BYTES) + subtile_m_idx = reg_subtile_m_idx0 + fx.Int64(sm * 2) + a_row_addr = subtile_m_idx * fx.Int64(SUBTILE_M) + fx.Int64(mi * MFMA_M) + lane_mod_16 + half_row = a_row_addr - fx.Int64(sm * (BLOCK_M // 2)) + row_byte_base = half_row * fx.Int64(BLOCK_K * ELEM_BYTES) return load_frag_half_at_byte_base(lds_a[sm], row_byte_base, half) def load_a_subtile_mi_regs(lds_a, sm, mi): @@ -605,7 +605,7 @@ def hk_one_k_with_refill( b13 = load_b_subtile_ni_regs(cur_b, 1, 3) # Refill the current ping-pong page with K+2, alternating A and B passes. - k_refill = fx.Index((k128 + 2) * BLOCK_K) + k_refill = fx.Int64((k128 + 2) * BLOCK_K) # Q0: interleave the current tile's A-bottom LDS reads with K+2 # refills and Q0 compute. Compute is K16-half-major across all 16 @@ -864,15 +864,15 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): store_acc_vector_for_logical_idx(drain_acc_idx, acc) # Prologue: stage K0/K1 data into ping-pong LDS pages. - stage_a_subtile(fx.Index(0), 0, lds_a0) - stage_b_subtile(fx.Index(0), 0, lds_b0) - stage_b_subtile(fx.Index(0), 1, lds_b0) - stage_a_subtile(fx.Index(0), 1, lds_a0) - - stage_a_subtile(fx.Index(BLOCK_K), 0, lds_a1) - stage_b_subtile(fx.Index(BLOCK_K), 0, lds_b1) - stage_b_subtile(fx.Index(BLOCK_K), 1, lds_b1) - stage_a_subtile(fx.Index(BLOCK_K), 1, lds_a1) + stage_a_subtile(fx.Int64(0), 0, lds_a0) + stage_b_subtile(fx.Int64(0), 0, lds_b0) + stage_b_subtile(fx.Int64(0), 1, lds_b0) + stage_a_subtile(fx.Int64(0), 1, lds_a0) + + stage_a_subtile(fx.Int64(BLOCK_K), 0, lds_a1) + stage_b_subtile(fx.Int64(BLOCK_K), 0, lds_b1) + stage_b_subtile(fx.Int64(BLOCK_K), 1, lds_b1) + stage_a_subtile(fx.Int64(BLOCK_K), 1, lds_a1) rocdl.sched_barrier(0) barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 4 * LOAD_PASSES_B_SUBTILE) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py index 7e73babd09..0ba1ae79ca 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -327,8 +327,8 @@ def process_tile(pid_m, pid_n, expert): # The flattened/XCD-swizzled block coordinates are i32, while global # address arithmetic below is expressed in MLIR index type. Convert # once here and use these index-typed tile bases for every address. - bx_m_idx = fx.Index(bx_m) - by_n_idx = fx.Index(by_n) + bx_m_idx = fx.Int64(bx_m) + by_n_idx = fx.Int64(by_n) # Per-expert offsets fold the expert index into the packed operands / # output (const_expr-zeroed for dense). FP8 is TN-NORMAL-read, so the @@ -346,17 +346,17 @@ def process_tile(pid_m, pid_n, expert): # (by_n_b_idx = expert*c_n + n). Output dW packed [G*N_out, K_in] on # the M-row (store_m_expert_off = expert*c_m). if const_expr(wgrad): - by_n_b_idx = fx.Index(expert * fx.Int32(c_n) + by_n) - a_expert_m_off = fx.Index(expert * fx.Int32(c_m)) - store_m_expert_off = fx.Index(expert * fx.Int32(c_m)) + by_n_b_idx = fx.Int64(expert * fx.Int32(c_n) + by_n) + a_expert_m_off = fx.Int64(expert * fx.Int32(c_m)) + store_m_expert_off = fx.Int64(expert * fx.Int32(c_m)) elif const_expr(group): - by_n_b_idx = fx.Index(expert * fx.Int32(c_n) + by_n) - a_expert_m_off = fx.Index(0) - store_m_expert_off = fx.Index(0) + by_n_b_idx = fx.Int64(expert * fx.Int32(c_n) + by_n) + a_expert_m_off = fx.Int64(0) + store_m_expert_off = fx.Int64(0) else: by_n_b_idx = by_n_idx - a_expert_m_off = fx.Index(0) - store_m_expert_off = fx.Index(0) + a_expert_m_off = fx.Int64(0) + store_m_expert_off = fx.Int64(0) # Keep wave/lane arithmetic in i32. compute_global_swizzle() combines # these values with i32 constants, so Index-typed coordinates would make @@ -410,7 +410,7 @@ def process_tile(pid_m, pid_n, expert): # add_offset on a dynamic Index is unsupported here). # store_m_expert_off (grouped wgrad only) repacks the output row to # this expert's block; 0 for dense and grouped fwd/dgrad. - c_tile_base_elems = (bx_m_idx + store_m_expert_off) * fx.Index(c_n) + by_n_idx + c_tile_base_elems = (bx_m_idx + store_m_expert_off) * fx.Int64(c_n) + by_n_idx gC = fx.rocdl.make_buffer_tensor(C, max_size=True) c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) c_store_atom = fx.make_copy_atom( @@ -553,8 +553,8 @@ def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): # wgrad, where A = dY^T is packed [G*N_out, m_pad] on the M-row); 0 # for dense and grouped fwd/dgrad (A is per-token, absolute rows). global_base = ( - bx_m_idx + a_expert_m_off + fx.Index(subtile * (BLOCK_M // 2)) - ) * fx.Index(K) + k_base + bx_m_idx + a_expert_m_off + fx.Int64(subtile * (BLOCK_M // 2)) + ) * fx.Int64(K) + k_base a_g2s.load_one(lds_a[subtile], fx.Int32(global_base), pass_in_subtile) def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): @@ -562,8 +562,8 @@ def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): # view (b_g2s) with a LOCAL column (b_data_col). Other paths: by_n_b_idx # shifts the B row to the expert block (grouped dgrad weights, grouped # wgrad X^T); reduces to dense by_n for the dense path. - b_row = b_data_col + fx.Index(subtile * (BLOCK_N // 2)) - global_base = b_row * fx.Index(K) + k_base + b_row = b_data_col + fx.Int64(subtile * (BLOCK_N // 2)) + global_base = b_row * fx.Int64(K) + k_base b_g2s.load_one(lds_b[subtile], fx.Int32(global_base), pass_in_subtile) def stage_a_subtile(k_base, subtile, lds_a): @@ -592,8 +592,8 @@ def load_frag_at_byte_base(lds_page, row_byte_base): def load_b_frag(lds_b, local_row, half): # B is [N, K]. Each 128-row half-page has a local row origin of 0. - half_row = local_row - fx.Index(half * (BLOCK_N // 2)) - return load_frag_at_byte_base(lds_b[half], half_row * fx.Index(BLOCK_K)) + half_row = local_row - fx.Int64(half * (BLOCK_N // 2)) + return load_frag_at_byte_base(lds_b[half], half_row * fx.Int64(BLOCK_K)) def _acc_idx(subtile_id, mi, ni): return ( @@ -657,10 +657,10 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): mi = local_idx // MFMA_N_PER_SUBTILE ni = local_idx % MFMA_N_PER_SUBTILE - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - row_base = subtile_m_idx * SUBTILE_M + fx.Index(mi * MFMA_M) + lane_div_16 * 4 - col = subtile_n_idx * SUBTILE_N + fx.Index(ni * MFMA_N) + lane_mod_16 + subtile_m_idx = reg_subtile_m_idx0 + fx.Int64(sm * 2) + subtile_n_idx = reg_subtile_n_idx0 + fx.Int64(sn * 2) + row_base = subtile_m_idx * SUBTILE_M + fx.Int64(mi * MFMA_M) + lane_div_16 * 4 + col = subtile_n_idx * SUBTILE_N + fx.Int64(ni * MFMA_N) + lane_mod_16 # Bias depends only on the output-feature (N) coordinate, so read it # once per column (global index by_n_idx + col) and reuse across the @@ -669,8 +669,8 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): bias_value = load_bias(by_n_idx + col) for ii in range_constexpr(4): - row = row_base + fx.Index(ii) - c_idx = c_tile_base_elems + row * fx.Index(c_n) + col + row = row_base + fx.Int64(ii) + c_idx = c_tile_base_elems + row * fx.Int64(c_n) + col # Epilogue stages run on the fp32 accumulator, in order, before # the output-dtype narrowing. Bias is added after the tensor-wise @@ -717,9 +717,9 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): zero_pinned_accumulators() def load_b_subtile_ni_regs(lds_b, sn, ni): - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + subtile_n_idx = reg_subtile_n_idx0 + fx.Int64(sn * 2) b_row_addr = ( - subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 + subtile_n_idx * fx.Int64(SUBTILE_N) + fx.Int64(ni * MFMA_N) + lane_mod_16 ) return load_b_frag(lds_b, b_row_addr, sn) @@ -732,12 +732,12 @@ def load_b_subtile_regs(lds_b, sn): ) def load_a_subtile_mi_half(lds_a, sm, mi, half): - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + subtile_m_idx = reg_subtile_m_idx0 + fx.Int64(sm * 2) a_row_addr = ( - subtile_m_idx * fx.Index(SUBTILE_M) + fx.Index(mi * MFMA_M) + lane_mod_16 + subtile_m_idx * fx.Int64(SUBTILE_M) + fx.Int64(mi * MFMA_M) + lane_mod_16 ) - half_row = a_row_addr - fx.Index(sm * (BLOCK_M // 2)) - row_byte_base = half_row * fx.Index(BLOCK_K) + half_row = a_row_addr - fx.Int64(sm * (BLOCK_M // 2)) + row_byte_base = half_row * fx.Int64(BLOCK_K) return load_frag_half_at_byte_base(lds_a[sm], row_byte_base, half) def load_a_subtile_mi_regs(lds_a, sm, mi): @@ -780,7 +780,7 @@ def hk_one_k_with_refill( b13 = load_b_subtile_ni_regs(cur_b, 1, 3) # Refill the current ping-pong page with K+2, alternating A and B passes. - k_refill = fx.Index((k128 + 2) * BLOCK_K) + k_refill = fx.Int64((k128 + 2) * BLOCK_K) # Q0: interleave the current tile's A-bottom LDS reads with K+2 # refills and Q0 compute. Each complete A-bottom fragment is assembled @@ -1038,15 +1038,15 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): store_acc_vector_for_logical_idx(drain_acc_idx, acc) # Prologue: stage K0/K1 data into ping-pong LDS pages. - stage_a_subtile(fx.Index(0), 0, lds_a0) - stage_b_subtile(fx.Index(0), 0, lds_b0) - stage_b_subtile(fx.Index(0), 1, lds_b0) - stage_a_subtile(fx.Index(0), 1, lds_a0) - - stage_a_subtile(fx.Index(BLOCK_K), 0, lds_a1) - stage_b_subtile(fx.Index(BLOCK_K), 0, lds_b1) - stage_b_subtile(fx.Index(BLOCK_K), 1, lds_b1) - stage_a_subtile(fx.Index(BLOCK_K), 1, lds_a1) + stage_a_subtile(fx.Int64(0), 0, lds_a0) + stage_b_subtile(fx.Int64(0), 0, lds_b0) + stage_b_subtile(fx.Int64(0), 1, lds_b0) + stage_a_subtile(fx.Int64(0), 1, lds_a0) + + stage_a_subtile(fx.Int64(BLOCK_K), 0, lds_a1) + stage_b_subtile(fx.Int64(BLOCK_K), 0, lds_b1) + stage_b_subtile(fx.Int64(BLOCK_K), 1, lds_b1) + stage_a_subtile(fx.Int64(BLOCK_K), 1, lds_a1) rocdl.sched_barrier(0) barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 4 * LOAD_PASSES_B_SUBTILE) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py index 8623227950..b078fe7fa5 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py @@ -52,7 +52,7 @@ def global_i32_ptr(tensor, index): ptr_ty = ir.Type.parse("!llvm.ptr<1>") base = _fly.extract_aligned_pointer_as_index(ptr_ty, tensor) base = llvm.PtrToIntOp(T.i64, base).result - byte_off = fx.as_ir_value(fx.Int64(fx.Index(index) * fx.Index(4))) + byte_off = fx.as_ir_value(fx.Int64(fx.Int64(index) * fx.Int64(4))) addr = llvm.AddOp(base, byte_off, llvm.IntegerOverflowFlags(0)).result return llvm.IntToPtrOp(ptr_ty, addr).result @@ -68,7 +68,7 @@ def read_ptr_at(ptr_array, index): ptr_ty = ir.Type.parse("!llvm.ptr<1>") base = _fly.extract_aligned_pointer_as_index(ptr_ty, ptr_array) base = llvm.PtrToIntOp(T.i64, base).result - byte_off = fx.as_ir_value(fx.Int64(fx.Index(index) * fx.Index(8))) + byte_off = fx.as_ir_value(fx.Int64(fx.Int64(index) * fx.Int64(8))) addr = llvm.AddOp(base, byte_off, llvm.IntegerOverflowFlags(0)).result elt_ptr = llvm.IntToPtrOp(ptr_ty, addr).result return llvm.LoadOp(T.i64, elt_ptr).result diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py index 5e2f7ded7b..922f58d85a 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py @@ -209,9 +209,9 @@ def _a_leading_dim_bytes(c_m): return c_m * ELEM_BYTES def _a_global_base_bytes(k_base, subtile, c_m, bx_m_idx): - return k_base * fx.Index(c_m * ELEM_BYTES) + ( - bx_m_idx + fx.Index(subtile * (BLOCK_M // 2)) - ) * fx.Index(ELEM_BYTES) + return k_base * fx.Int64(c_m * ELEM_BYTES) + ( + bx_m_idx + fx.Int64(subtile * (BLOCK_M // 2)) + ) * fx.Int64(ELEM_BYTES) def _load_a_half( load_transposed_frag_half, @@ -224,11 +224,11 @@ def _load_a_half( lane_mod_16, ): del load_frag_half_at_byte_base, lane_mod_16 - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + subtile_m_idx = reg_subtile_m_idx0 + fx.Int64(sm * 2) local_m_tile = ( - subtile_m_idx * fx.Index(SUBTILE_M) - + fx.Index(mi * MFMA_M) - - fx.Index(sm * (BLOCK_M // 2)) + subtile_m_idx * fx.Int64(SUBTILE_M) + + fx.Int64(mi * MFMA_M) + - fx.Int64(sm * (BLOCK_M // 2)) ) return load_transposed_frag_half(lds_a[sm], local_m_tile, half) @@ -240,9 +240,9 @@ def _a_leading_dim_bytes(c_m): def _a_global_base_bytes(k_base, subtile, c_m, bx_m_idx): del c_m - return (bx_m_idx + fx.Index(subtile * (BLOCK_M // 2))) * fx.Index( + return (bx_m_idx + fx.Int64(subtile * (BLOCK_M // 2))) * fx.Int64( K * ELEM_BYTES - ) + k_base * fx.Index(ELEM_BYTES) + ) + k_base * fx.Int64(ELEM_BYTES) def _load_a_half( load_transposed_frag_half, @@ -255,12 +255,12 @@ def _load_a_half( lane_mod_16, ): del load_transposed_frag_half - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - a_row_addr = subtile_m_idx * fx.Index(SUBTILE_M) + fx.Index(mi * MFMA_M) + lane_mod_16 - half_row = a_row_addr - fx.Index(sm * (BLOCK_M // 2)) + subtile_m_idx = reg_subtile_m_idx0 + fx.Int64(sm * 2) + a_row_addr = subtile_m_idx * fx.Int64(SUBTILE_M) + fx.Int64(mi * MFMA_M) + lane_mod_16 + half_row = a_row_addr - fx.Int64(sm * (BLOCK_M // 2)) return load_frag_half_at_byte_base( lds_a[sm], - half_row * fx.Index(BLOCK_K * ELEM_BYTES), + half_row * fx.Int64(BLOCK_K * ELEM_BYTES), half, ) @@ -270,9 +270,9 @@ def _b_leading_dim_bytes(c_n): return c_n * ELEM_BYTES def _b_global_base_bytes(k_base, subtile, c_n, by_n_idx): - return k_base * fx.Index(c_n * ELEM_BYTES) + ( - by_n_idx + fx.Index(subtile * (BLOCK_N // 2)) - ) * fx.Index(ELEM_BYTES) + return k_base * fx.Int64(c_n * ELEM_BYTES) + ( + by_n_idx + fx.Int64(subtile * (BLOCK_N // 2)) + ) * fx.Int64(ELEM_BYTES) def _load_b_ni( load_transposed_frag, @@ -284,11 +284,11 @@ def _load_b_ni( lane_mod_16, ): del load_normal_b_frag, lane_mod_16 - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + subtile_n_idx = reg_subtile_n_idx0 + fx.Int64(sn * 2) local_n_tile = ( - subtile_n_idx * fx.Index(SUBTILE_N) - + fx.Index(ni * MFMA_N) - - fx.Index(sn * (BLOCK_N // 2)) + subtile_n_idx * fx.Int64(SUBTILE_N) + + fx.Int64(ni * MFMA_N) + - fx.Int64(sn * (BLOCK_N // 2)) ) return load_transposed_frag(lds_b[sn], local_n_tile) @@ -300,9 +300,9 @@ def _b_leading_dim_bytes(c_n): def _b_global_base_bytes(k_base, subtile, c_n, by_n_idx): del c_n - return (by_n_idx + fx.Index(subtile * (BLOCK_N // 2))) * fx.Index( + return (by_n_idx + fx.Int64(subtile * (BLOCK_N // 2))) * fx.Int64( K * ELEM_BYTES - ) + k_base * fx.Index(ELEM_BYTES) + ) + k_base * fx.Int64(ELEM_BYTES) def _load_b_ni( load_transposed_frag, @@ -314,8 +314,8 @@ def _load_b_ni( lane_mod_16, ): del load_transposed_frag - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - b_row_addr = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 + subtile_n_idx = reg_subtile_n_idx0 + fx.Int64(sn * 2) + b_row_addr = subtile_n_idx * fx.Int64(SUBTILE_N) + fx.Int64(ni * MFMA_N) + lane_mod_16 return load_normal_b_frag(lds_b, b_row_addr, sn) # Resolve global staging maps before FlyDSL captures ``kernel_gemm``. @@ -472,8 +472,8 @@ def process_tile(pid_m, pid_n, expert): # The flattened/XCD-swizzled block coordinates are i32, while global # address arithmetic below is expressed in MLIR index type. Convert # once here and use these index-typed tile bases for every address. - bx_m_idx = fx.Index(bx_m) - by_n_idx = fx.Index(by_n) + bx_m_idx = fx.Int64(bx_m) + by_n_idx = fx.Int64(by_n) # Per-expert offsets into the packed operands / output. Only wgrad (NT) needs # them: both operands are packed on the contraction (token) axis @@ -482,14 +482,14 @@ def process_tile(pid_m, pid_n, expert): # below), so they use local coordinates and zero offsets, like dense. if const_expr(wgrad): by_n_b_idx = by_n_idx - a_expert_k_off = fx.Index(expert * fx.Int32(K)) - b_expert_k_off = fx.Index(expert * fx.Int32(K)) - store_m_expert_off = fx.Index(expert * fx.Int32(c_m)) + a_expert_k_off = fx.Int64(expert * fx.Int32(K)) + b_expert_k_off = fx.Int64(expert * fx.Int32(K)) + store_m_expert_off = fx.Int64(expert * fx.Int32(c_m)) else: by_n_b_idx = by_n_idx - a_expert_k_off = fx.Index(0) - b_expert_k_off = fx.Index(0) - store_m_expert_off = fx.Index(0) + a_expert_k_off = fx.Int64(0) + b_expert_k_off = fx.Int64(0) + store_m_expert_off = fx.Int64(0) # Keep wave/lane arithmetic in i32. compute_global_swizzle() combines # these values with i32 constants, so Index-typed coordinates would make @@ -530,7 +530,7 @@ def process_tile(pid_m, pid_n, expert): b_base_i64 = gd.make_uniform_i64(b_base_raw) b_div_tile = gd.rebase_buffer_view(b_base_i64, fx.Uint8.ir_type, b_ptr_layout) b_data_col = by_n_idx - b_k_off = fx.Index(0) + b_k_off = fx.Int64(0) elif const_expr(act_ptrs): b_base_i64 = gd.make_uniform_i64(b_base_raw) b_div_tile = gd.rebase_buffer_view(b_base_i64, fx.Uint8.ir_type, b_ptr_layout) @@ -559,7 +559,7 @@ def process_tile(pid_m, pid_n, expert): # add_offset on a dynamic Index is unsupported here). store_m_expert_off # (grouped wgrad only) repacks the output row to this expert's block; # it is 0 for dense and for grouped fwd/dgrad. - c_tile_base_elems = (bx_m_idx + store_m_expert_off) * fx.Index(c_n) + by_n_idx + c_tile_base_elems = (bx_m_idx + store_m_expert_off) * fx.Int64(c_n) + by_n_idx gC = fx.rocdl.make_buffer_tensor(C, max_size=True) c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) c_store_atom = fx.make_copy_atom( @@ -743,10 +743,10 @@ def load_frag_at_byte_base(lds_page, row_byte_base): def load_b_frag(lds_b, local_row, half): # B is [N, K]. Each 128-row half-page has a local row origin of 0. - half_row = local_row - fx.Index(half * (BLOCK_N // 2)) + half_row = local_row - fx.Int64(half * (BLOCK_N // 2)) return load_frag_at_byte_base( lds_b[half], - half_row * fx.Index(BLOCK_K * ELEM_BYTES), + half_row * fx.Int64(BLOCK_K * ELEM_BYTES), ) def load_transposed_frag_half(lds_page, local_x_tile, half): @@ -870,10 +870,10 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): mi = local_idx // MFMA_N_PER_SUBTILE ni = local_idx % MFMA_N_PER_SUBTILE - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - row_base = subtile_m_idx * SUBTILE_M + fx.Index(mi * MFMA_M) + lane_div_16 * 4 - col = subtile_n_idx * SUBTILE_N + fx.Index(ni * MFMA_N) + lane_mod_16 + subtile_m_idx = reg_subtile_m_idx0 + fx.Int64(sm * 2) + subtile_n_idx = reg_subtile_n_idx0 + fx.Int64(sn * 2) + row_base = subtile_m_idx * SUBTILE_M + fx.Int64(mi * MFMA_M) + lane_div_16 * 4 + col = subtile_n_idx * SUBTILE_N + fx.Int64(ni * MFMA_N) + lane_mod_16 # Bias depends only on the output-feature (N) coordinate, so read it # once per column (global index by_n_idx + col) and reuse across the @@ -882,8 +882,8 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): bias_value = load_bias(by_n_idx + col) for ii in range_constexpr(4): - row = row_base + fx.Index(ii) - c_idx = c_tile_base_elems + row * fx.Index(c_n) + col + row = row_base + fx.Int64(ii) + c_idx = c_tile_base_elems + row * fx.Int64(c_n) + col # Epilogue stages run on the fp32 accumulator, in order, before # the output-dtype narrowing: @@ -1000,7 +1000,7 @@ def hk_one_k_with_refill( b13 = load_b_subtile_ni_regs(cur_b, 1, 3) # Refill the current ping-pong page with K+2, alternating A and B passes. - k_refill = fx.Index((k128 + 2) * BLOCK_K) + k_refill = fx.Int64((k128 + 2) * BLOCK_K) # Q0: interleave the current tile's A-bottom LDS reads with K+2 # refills and Q0 compute. Compute is K32-major across all 16 @@ -1259,15 +1259,15 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): store_acc_vector_for_logical_idx(drain_acc_idx, acc) # Prologue: stage K0/K1 data into ping-pong LDS pages. - stage_a_subtile(fx.Index(0), 0, lds_a0) - stage_b_subtile(fx.Index(0), 0, lds_b0) - stage_b_subtile(fx.Index(0), 1, lds_b0) - stage_a_subtile(fx.Index(0), 1, lds_a0) - - stage_a_subtile(fx.Index(BLOCK_K), 0, lds_a1) - stage_b_subtile(fx.Index(BLOCK_K), 0, lds_b1) - stage_b_subtile(fx.Index(BLOCK_K), 1, lds_b1) - stage_a_subtile(fx.Index(BLOCK_K), 1, lds_a1) + stage_a_subtile(fx.Int64(0), 0, lds_a0) + stage_b_subtile(fx.Int64(0), 0, lds_b0) + stage_b_subtile(fx.Int64(0), 1, lds_b0) + stage_a_subtile(fx.Int64(0), 1, lds_a0) + + stage_a_subtile(fx.Int64(BLOCK_K), 0, lds_a1) + stage_b_subtile(fx.Int64(BLOCK_K), 0, lds_b1) + stage_b_subtile(fx.Int64(BLOCK_K), 1, lds_b1) + stage_a_subtile(fx.Int64(BLOCK_K), 1, lds_a1) rocdl.sched_barrier(0) barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 4 * LOAD_PASSES_B_SUBTILE) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index 817c8a2bc1..c41b7b6aa2 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -152,31 +152,31 @@ def _compile_mx32_scale_pack_kernel( def _source_offset(source_k32, source_row): # Logical source is [K/32, dim], but the underlying TE tensor may # be a non-contiguous view. Strides are in uint8 elements. - return source_k32 * fx.Index(stride0) + source_row * fx.Index(stride1) + return source_k32 * fx.Int64(stride0) + source_row * fx.Int64(stride1) else: def _source_offset(source_k32, source_row): # Logical source is [dim, K/32], with arbitrary positive strides. - return source_row * fx.Index(stride0) + source_k32 * fx.Index(stride1) + return source_row * fx.Int64(stride0) + source_k32 * fx.Int64(stride1) def _word_coords(): # This thread's output word and the source coordinates of its four bytes. - linear = fx.Index(fx.block_idx.x) * fx.Index(_SCALE_PACK_THREADS) + fx.Index( + linear = fx.Int64(fx.block_idx.x) * fx.Int64(_SCALE_PACK_THREADS) + fx.Int64( gpu.thread_id("x") ) - k128 = linear // fx.Index(dim) - dst_row = linear % fx.Index(dim) + k128 = linear // fx.Int64(dim) + dst_row = linear % fx.Int64(dim) - row_within_16 = dst_row % fx.Index(16) - k_subgroup = (dst_row // fx.Index(16)) % fx.Index(4) - tile = dst_row // fx.Index(64) - source_k32 = k128 * fx.Index(4) + k_subgroup + row_within_16 = dst_row % fx.Int64(16) + k_subgroup = (dst_row // fx.Int64(16)) % fx.Int64(4) + tile = dst_row // fx.Int64(64) + source_k32 = k128 * fx.Int64(4) + k_subgroup return linear, k128, tile, row_within_16, source_k32 def _load_scale_byte(src_div, coords, group): _, _, tile, row_within_16, source_k32 = coords - source_row = tile * fx.Index(64) + fx.Index(group * 16) + row_within_16 + source_row = tile * fx.Int64(64) + fx.Int64(group * 16) + row_within_16 off = _source_offset(source_k32, source_row) scale_load_atom = fx.make_copy_atom(fx.rocdl.BufferCopy8b(), fx.Uint8) reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Uint8) @@ -519,7 +519,7 @@ def _a_leading_dim(c_m): return c_m def _a_global_base(k_base, subtile, c_m, bx_m_idx): - return k_base * fx.Index(c_m) + bx_m_idx + fx.Index(subtile * (BLOCK_M // 2)) + return k_base * fx.Int64(c_m) + bx_m_idx + fx.Int64(subtile * (BLOCK_M // 2)) def _load_a_half( load_transposed_frag_half, @@ -531,11 +531,11 @@ def _load_a_half( reg_subtile_m_idx0, lane_mod_16, ): - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + subtile_m_idx = reg_subtile_m_idx0 + fx.Int64(sm * 2) local_m_tile = ( - subtile_m_idx * fx.Index(SUBTILE_M) - + fx.Index(mi * MFMA_M) - - fx.Index(sm * (BLOCK_M // 2)) + subtile_m_idx * fx.Int64(SUBTILE_M) + + fx.Int64(mi * MFMA_M) + - fx.Int64(sm * (BLOCK_M // 2)) ) return load_transposed_frag_half( lds_a[sm], @@ -551,7 +551,7 @@ def _a_leading_dim(c_m): def _a_global_base(k_base, subtile, c_m, bx_m_idx): del c_m - return (bx_m_idx + fx.Index(subtile * (BLOCK_M // 2))) * fx.Index(K) + k_base + return (bx_m_idx + fx.Int64(subtile * (BLOCK_M // 2))) * fx.Int64(K) + k_base def _load_a_half( load_transposed_frag_half, @@ -564,10 +564,10 @@ def _load_a_half( lane_mod_16, ): del load_transposed_frag_half - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - a_row_addr = subtile_m_idx * fx.Index(SUBTILE_M) + fx.Index(mi * MFMA_M) + lane_mod_16 - half_row = a_row_addr - fx.Index(sm * (BLOCK_M // 2)) - row_byte_base = half_row * fx.Index(BLOCK_K) + subtile_m_idx = reg_subtile_m_idx0 + fx.Int64(sm * 2) + a_row_addr = subtile_m_idx * fx.Int64(SUBTILE_M) + fx.Int64(mi * MFMA_M) + lane_mod_16 + half_row = a_row_addr - fx.Int64(sm * (BLOCK_M // 2)) + row_byte_base = half_row * fx.Int64(BLOCK_K) return load_frag_half_at_byte_base( lds_a[sm], row_byte_base, @@ -580,7 +580,7 @@ def _b_leading_dim(c_n): return c_n def _b_global_base(k_base, subtile, c_n, by_n_idx): - return k_base * fx.Index(c_n) + by_n_idx + fx.Index(subtile * (BLOCK_N // 2)) + return k_base * fx.Int64(c_n) + by_n_idx + fx.Int64(subtile * (BLOCK_N // 2)) def _load_b_ni( load_transposed_frag, @@ -592,11 +592,11 @@ def _load_b_ni( lane_mod_16, ): del load_normal_b_frag, lane_mod_16 - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + subtile_n_idx = reg_subtile_n_idx0 + fx.Int64(sn * 2) local_n_tile = ( - subtile_n_idx * fx.Index(SUBTILE_N) - + fx.Index(ni * MFMA_N) - - fx.Index(sn * (BLOCK_N // 2)) + subtile_n_idx * fx.Int64(SUBTILE_N) + + fx.Int64(ni * MFMA_N) + - fx.Int64(sn * (BLOCK_N // 2)) ) return load_transposed_frag(lds_b[sn], local_n_tile) @@ -608,7 +608,7 @@ def _b_leading_dim(c_n): def _b_global_base(k_base, subtile, c_n, by_n_idx): del c_n - return (by_n_idx + fx.Index(subtile * (BLOCK_N // 2))) * fx.Index(K) + k_base + return (by_n_idx + fx.Int64(subtile * (BLOCK_N // 2))) * fx.Int64(K) + k_base def _load_b_ni( load_transposed_frag, @@ -620,8 +620,8 @@ def _load_b_ni( lane_mod_16, ): del load_transposed_frag - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - b_row_addr = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 + subtile_n_idx = reg_subtile_n_idx0 + fx.Int64(sn * 2) + b_row_addr = subtile_n_idx * fx.Int64(SUBTILE_N) + fx.Int64(ni * MFMA_N) + lane_mod_16 return load_normal_b_frag(lds_b, b_row_addr, sn) if (not a_transpose_read) or (not b_transpose_read): @@ -774,8 +774,8 @@ def process_tile(pid_m, pid_n, expert): # The flattened/XCD-swizzled block coordinates are i32, while global # address arithmetic below is expressed in MLIR index type. Convert # once here and use these index-typed tile bases for every address. - bx_m_idx = fx.Index(bx_m) - by_n_idx = fx.Index(by_n) + bx_m_idx = fx.Int64(bx_m) + by_n_idx = fx.Int64(by_n) # Per-expert offsets. Grouped fwd/dgrad read B data per expert through # b_data_ptrs (see b_div_tile below), so B-data addressing is local and @@ -806,11 +806,11 @@ def process_tile(pid_m, pid_n, expert): if const_expr(wgrad): by_n_b_idx = by_n_idx c_n_bs = c_n - a_expert_k_off = fx.Index(expert * fx.Int32(K)) - b_expert_k_off = fx.Index(expert * fx.Int32(K)) + a_expert_k_off = fx.Int64(expert * fx.Int32(K)) + b_expert_k_off = fx.Int64(expert * fx.Int32(K)) as_expert_k128_off = expert * fx.Int32(NUM_K_TILES) bs_expert_k128_off = expert * fx.Int32(NUM_K_TILES) - store_m_expert_off = fx.Index(expert * fx.Int32(c_m)) + store_m_expert_off = fx.Int64(expert * fx.Int32(c_m)) elif const_expr(group): if const_expr(b_transpose_read): # NN: the B-scale [G*N_out/128, K_in] stacks experts on the @@ -818,11 +818,11 @@ def process_tile(pid_m, pid_n, expert): # expert*NUM_K_TILES. B data is read per expert (no shift). by_n_b_idx = by_n_idx c_n_bs = c_n - a_expert_k_off = fx.Index(0) - b_expert_k_off = fx.Index(0) + a_expert_k_off = fx.Int64(0) + b_expert_k_off = fx.Int64(0) as_expert_k128_off = fx.Int32(0) bs_expert_k128_off = expert * fx.Int32(NUM_K_TILES) - store_m_expert_off = fx.Index(0) + store_m_expert_off = fx.Int64(0) else: # TN: weights are per-expert slabs addressed via a base-pointer # swap (b_div_tile below), so B's column coord stays LOCAL @@ -831,21 +831,21 @@ def process_tile(pid_m, pid_n, expert): # grows to G*N (c_n_bs) and column folds in expert (handled at # by_n_b_idx for the scale read; B-data uses by_n_idx directly # via the per-expert base). - by_n_b_idx = fx.Index(expert * fx.Int32(c_n) + by_n) + by_n_b_idx = fx.Int64(expert * fx.Int32(c_n) + by_n) c_n_bs = num_groups * c_n - a_expert_k_off = fx.Index(0) - b_expert_k_off = fx.Index(0) + a_expert_k_off = fx.Int64(0) + b_expert_k_off = fx.Int64(0) as_expert_k128_off = fx.Int32(0) bs_expert_k128_off = fx.Int32(0) - store_m_expert_off = fx.Index(0) + store_m_expert_off = fx.Int64(0) else: by_n_b_idx = by_n_idx c_n_bs = c_n - a_expert_k_off = fx.Index(0) - b_expert_k_off = fx.Index(0) + a_expert_k_off = fx.Int64(0) + b_expert_k_off = fx.Int64(0) as_expert_k128_off = fx.Int32(0) bs_expert_k128_off = fx.Int32(0) - store_m_expert_off = fx.Index(0) + store_m_expert_off = fx.Int64(0) wave_id = tx_i32 // fx.Int32(WARP_SIZE) lane = tx_i32 % fx.Int32(WARP_SIZE) @@ -902,7 +902,7 @@ def process_tile(pid_m, pid_n, expert): b_base_i64 = gd.make_uniform_i64(b_base_raw) b_div_tile = gd.rebase_buffer_view(b_base_i64, b_f8_ir_t, b_ptr_layout) b_data_col = by_n_idx - b_k_off = fx.Index(0) + b_k_off = fx.Int64(0) elif const_expr(act_ptrs): b_base_i64 = gd.make_uniform_i64(b_base_raw) b_div_tile = gd.rebase_buffer_view(b_base_i64, b_f8_ir_t, b_ptr_layout) @@ -932,7 +932,7 @@ def process_tile(pid_m, pid_n, expert): # store_m_expert_off (grouped wgrad only) shifts the output row into this # tile's expert block of the packed [G*N_out, K_in] weight-gradient; the # a/b operands used local (per-expert) pid_m, so the store re-packs here. - c_tile_base_elems = (bx_m_idx + store_m_expert_off) * fx.Index(c_n) + by_n_idx + c_tile_base_elems = (bx_m_idx + store_m_expert_off) * fx.Int64(c_n) + by_n_idx gC = fx.rocdl.make_buffer_tensor(C, max_size=True) c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) c_store_atom = fx.make_copy_atom( @@ -1054,10 +1054,10 @@ def read_physical_accumulator_slot(slot_idx): # bytes for this lane's 64-row A/B half. The MFMA instruction selects # the byte via op_sel/op_sel_hi, so there is intentionally no hot-loop # byte extraction and no 0x01010101 broadcast here. - c_m_idx = fx.Index(c_m) - c_n_idx = fx.Index(c_n) + c_m_idx = fx.Int64(c_m) + c_n_idx = fx.Int64(c_n) # B-scale row stride: c_n dense, packed G*N grouped (const_expr-folded). - c_n_bs_idx = fx.Index(c_n_bs) + c_n_bs_idx = fx.Int64(c_n_bs) def hot_loop_scheduler_q_refill_2n(): # Steady-state Q1 schedule: eight chunks of one K+2 VMEM/LDS @@ -1092,7 +1092,7 @@ def load_a_scale_row(k128, row): # as_expert_k128_off shifts the contraction-128 base to this tile's # expert A-scale block for grouped wgrad (both operands packed on the # token/contraction axis); 0 for every other path. - off = (k128 + fx.Index(as_expert_k128_off)) * c_m_idx + bx_m_idx + row + off = (k128 + fx.Int64(as_expert_k128_off)) * c_m_idx + bx_m_idx + row reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Int32) fx.copy(scale_ld_atom, fx.slice(as_div, (None, fx.Int32(off))), reg) return fx.memref_load_vec(reg)[0] @@ -1100,20 +1100,20 @@ def load_a_scale_row(k128, row): def load_b_scale_row(k128, row): # bs_expert_k128_off shifts the contraction-128 base to this tile's # expert scale block for grouped NN/NT; 0 for dense and grouped TN. - off = (k128 + fx.Index(bs_expert_k128_off)) * c_n_bs_idx + by_n_b_idx + row + off = (k128 + fx.Int64(bs_expert_k128_off)) * c_n_bs_idx + by_n_b_idx + row reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Int32) fx.copy(scale_ld_atom, fx.slice(bs_div, (None, fx.Int32(off))), reg) return fx.memref_load_vec(reg)[0] def load_a_scale_subtile(k128, sm): - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - a_row = subtile_m_idx * fx.Index(SUBTILE_M) + fx.Index(lane) + subtile_m_idx = reg_subtile_m_idx0 + fx.Int64(sm * 2) + a_row = subtile_m_idx * fx.Int64(SUBTILE_M) + fx.Int64(lane) a_scale = load_a_scale_row(k128, a_row) return (a_scale, a_scale, a_scale, a_scale) def load_b_scale_subtile(k128, sn): - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - b_row = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(lane) + subtile_n_idx = reg_subtile_n_idx0 + fx.Int64(sn * 2) + b_row = subtile_n_idx * fx.Int64(SUBTILE_N) + fx.Int64(lane) b_scale = load_b_scale_row(k128, b_row) return (b_scale, b_scale, b_scale, b_scale) @@ -1174,10 +1174,10 @@ def load_frag_at_byte_base(lds_page, row_byte_base): def load_normal_b_frag(lds_b, local_row, half): # Physical [N,K] page, ordinary TN-style fixed-row read. - half_row = local_row - fx.Index(half * (BLOCK_N // 2)) + half_row = local_row - fx.Int64(half * (BLOCK_N // 2)) return load_frag_at_byte_base( lds_b[half], - half_row * fx.Index(BLOCK_K), + half_row * fx.Int64(BLOCK_K), ) def load_transposed_frag_half(lds_page, local_x_tile, half): @@ -1284,10 +1284,10 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): mi = local_idx // MFMA_N_PER_SUBTILE ni = local_idx % MFMA_N_PER_SUBTILE - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - row_base = subtile_m_idx * SUBTILE_M + fx.Index(mi * MFMA_M) + lane_div_16 * 4 - col = subtile_n_idx * SUBTILE_N + fx.Index(ni * MFMA_N) + lane_mod_16 + subtile_m_idx = reg_subtile_m_idx0 + fx.Int64(sm * 2) + subtile_n_idx = reg_subtile_n_idx0 + fx.Int64(sn * 2) + row_base = subtile_m_idx * SUBTILE_M + fx.Int64(mi * MFMA_M) + lane_div_16 * 4 + col = subtile_n_idx * SUBTILE_N + fx.Int64(ni * MFMA_N) + lane_mod_16 # Bias depends only on the output-feature (N) coordinate, so read it # once per column and reuse across the four M rows below. Index by @@ -1298,8 +1298,8 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): bias_value = load_bias(by_n_idx + col) for ii in range_constexpr(4): - row = row_base + fx.Index(ii) - c_idx = c_tile_base_elems + row * fx.Index(c_n) + col + row = row_base + fx.Int64(ii) + c_idx = c_tile_base_elems + row * fx.Int64(c_n) + col # Epilogue stages run on the fp32 accumulator, in order, before # the output-dtype narrowing: @@ -1345,7 +1345,7 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): zero_pinned_accumulators() def load_b_subtile_ni_regs(lds_b, scale_tile, sn, ni): - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + subtile_n_idx = reg_subtile_n_idx0 + fx.Int64(sn * 2) b_scales = scale_tile[2] if sn == 0 else scale_tile[3] b_ni = _load_b_ni( @@ -1367,7 +1367,7 @@ def load_b_subtile_regs(lds_b, scale_tile, sn): return b0, b1, b2, b3, bs0, bs1, bs2, bs3 def load_a_subtile_mi_half(lds_a, sm, mi, half): - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + subtile_m_idx = reg_subtile_m_idx0 + fx.Int64(sm * 2) return _load_a_half( load_transposed_frag_half, @@ -1422,7 +1422,7 @@ def hk_one_k_with_refill( # Immediately issue MFMA-ready K+2 scale loads. # They are returned for the next iteration without any in-kernel # byte extraction or broadcast. - refill_scales = load_scale_tile(fx.Index(k128 + 2)) + refill_scales = load_scale_tile(fx.Int64(k128 + 2)) next_scales_ready = prev_refill_scales # A-top and B-left are both carried as complete 64-row register tiles, # so their LDS half-pages can be refilled immediately. @@ -1435,7 +1435,7 @@ def hk_one_k_with_refill( b13, bs13 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 3) # Refill the current ping-pong page with K+2, alternating A and B passes. - k_refill = fx.Index((k128 + 2) * BLOCK_K) + k_refill = fx.Int64((k128 + 2) * BLOCK_K) # Q0: interleave the current tile's A-bottom LDS reads with K+2 # refills and Q0 compute. Each complete A-bottom fragment is assembled @@ -1750,18 +1750,18 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs, cur_scales): # and load_scale_tile returns the current wave's scale operands in VGPRs. # Load scales first, so that they become the oldest VMEM ops. - scales0 = load_scale_tile(fx.Index(0)) - scales1 = load_scale_tile(fx.Index(1)) - - stage_a_subtile(fx.Index(0), 0, lds_a0) - stage_b_subtile(fx.Index(0), 0, lds_b0) - stage_b_subtile(fx.Index(0), 1, lds_b0) - stage_a_subtile(fx.Index(0), 1, lds_a0) - - stage_a_subtile(fx.Index(BLOCK_K), 0, lds_a1) - stage_b_subtile(fx.Index(BLOCK_K), 0, lds_b1) - stage_b_subtile(fx.Index(BLOCK_K), 1, lds_b1) - stage_a_subtile(fx.Index(BLOCK_K), 1, lds_a1) + scales0 = load_scale_tile(fx.Int64(0)) + scales1 = load_scale_tile(fx.Int64(1)) + + stage_a_subtile(fx.Int64(0), 0, lds_a0) + stage_b_subtile(fx.Int64(0), 0, lds_b0) + stage_b_subtile(fx.Int64(0), 1, lds_b0) + stage_a_subtile(fx.Int64(0), 1, lds_a0) + + stage_a_subtile(fx.Int64(BLOCK_K), 0, lds_a1) + stage_b_subtile(fx.Int64(BLOCK_K), 0, lds_b1) + stage_b_subtile(fx.Int64(BLOCK_K), 1, lds_b1) + stage_a_subtile(fx.Int64(BLOCK_K), 1, lds_a1) rocdl.sched_barrier(0) barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 4 * LOAD_PASSES_B_SUBTILE) From 87b343cc4eec3e6b714e04b6ab6f2dafa4789c15 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 28 Oct 2026 00:12:00 +0000 Subject: [PATCH 22/22] address reviewer comments --- tests/pytorch/test_grouped_linear.py | 146 ++- .../pytorch/flydsl_kernels/gemm/__init__.py | 19 +- .../pytorch/flydsl_kernels/gemm/fp32_gemm.py | 62 +- .../pytorch/flydsl_kernels/gemm/fp8_gemm.py | 103 +- .../flydsl_kernels/gemm/gemm_common_utils.py | 60 +- .../flydsl_kernels/gemm/gemm_wrappers.py | 972 ++++++++---------- .../flydsl_kernels/gemm/grouped_gemm_utils.py | 15 +- .../flydsl_kernels/gemm/half_prec_gemm.py | 112 +- .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 131 +-- .../pytorch/module/grouped_linear.py | 848 ++------------- transformer_engine/pytorch/quantization.py | 23 +- 11 files changed, 881 insertions(+), 1610 deletions(-) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 63fb61bee5..0fb36a0b7d 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -656,20 +656,17 @@ def test_grouped_linear_accuracy_rocm_backends( ): pytest.skip("FlyDSL grouped GEMM supports MXFP8 and tensor-wise FP8 recipes only.") - # Select exactly one backend; leave the others' env vars unset. HipKittens and CK - # are gated behind NVTE_USE_CUTLASS_GROUPED_GEMM; FlyDSL has its own opt-in and - # must not see the cutlass/HK/CK env vars. - monkeypatch.delenv("NVTE_USE_CUTLASS_GROUPED_GEMM", raising=False) + monkeypatch.setenv("NVTE_USE_CUTLASS_GROUPED_GEMM", "1") monkeypatch.delenv("NVTE_USE_HIPKITTENS_GROUPED_GEMM", raising=False) monkeypatch.delenv("NVTE_USE_CK_GROUPED_GEMM", raising=False) monkeypatch.delenv("NVTE_USE_GROUPED_GEMM_FLYDSL", raising=False) if grouped_gemm_backend == "hipkittens": - monkeypatch.setenv("NVTE_USE_CUTLASS_GROUPED_GEMM", "1") monkeypatch.setenv("NVTE_USE_HIPKITTENS_GROUPED_GEMM", "1") elif grouped_gemm_backend == "ck": - monkeypatch.setenv("NVTE_USE_CUTLASS_GROUPED_GEMM", "1") monkeypatch.setenv("NVTE_USE_CK_GROUPED_GEMM", "1") else: + # FlyDSL has its own opt-in; keep the cutlass/HK/CK grouped GEMM path disabled. + monkeypatch.delenv("NVTE_USE_CUTLASS_GROUPED_GEMM", raising=False) monkeypatch.setenv("NVTE_USE_GROUPED_GEMM_FLYDSL", "1") # For the configs CK grouped GEMM is guaranteed to handle, @@ -689,17 +686,24 @@ def test_grouped_linear_accuracy_rocm_backends( # fallback warning stays on so a failure shows the reason. flydsl_dispatches = [] if grouped_gemm_backend == "flydsl": + from transformer_engine.pytorch.flydsl_kernels.gemm import ( + FlyDSLUnsupportedError, + gemm_wrappers, + ) + monkeypatch.setenv("NVTE_FLYDSL_GEMM_WARN_FALLBACK", "1") - try_flydsl_grouped = _GroupedLinear._try_flydsl_grouped + grouped_gemm_flydsl = gemm_wrappers._grouped_gemm_flydsl + phases = {"TN": "fwd", "NN": "dgrad", "NT": "wgrad"} - def recording_try_flydsl_grouped(cls, phase, kind, **kwargs): - ran = try_flydsl_grouped(phase, kind, **kwargs) - flydsl_dispatches.append((phase, ran)) - return ran + def recording_grouped_gemm_flydsl(*args, layout, **kwargs): + try: + grouped_gemm_flydsl(*args, layout=layout, **kwargs) + except FlyDSLUnsupportedError: + flydsl_dispatches.append((phases[layout], False)) + raise + flydsl_dispatches.append((phases[layout], True)) - monkeypatch.setattr( - _GroupedLinear, "_try_flydsl_grouped", classmethod(recording_try_flydsl_grouped) - ) + monkeypatch.setattr(gemm_wrappers, "_grouped_gemm_flydsl", recording_grouped_gemm_flydsl) test_grouped_linear_accuracy( dtype, @@ -741,6 +745,120 @@ def recording_try_flydsl_grouped(cls, phase, kind, **kwargs): ) +@pytest.mark.skipif(not IS_HIP_EXTENSION, reason="ROCm-only grouped GEMM backend") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize( + "recipe", [recipe.Float8CurrentScaling(), recipe.DelayedScaling()], ids=recipe_id +) +@pytest.mark.parametrize("fp8_model_params", all_boolean) +def test_grouped_linear_flydsl_fp8_scaling(recipe, fp8_model_params, monkeypatch): + """FlyDSL tensor-wise FP8 grouped GEMM vs. sequential Linear over several steps. + + Delayed scaling quantizes its first step with unit scales (there is no amax history + yet), so the single-step accuracy test cannot catch a per-expert scale bug for it. + Run a few steps with uneven experts and require the compared step to hand FlyDSL + non-unit scales that differ across experts, for every phase. + """ + if get_device_compute_capability() != (9, 5): + pytest.skip("FlyDSL grouped GEMM is only supported on gfx950.") + pytest.importorskip("flydsl", reason="FlyDSL package is not installed.") + from transformer_engine.pytorch.flydsl_kernels.gemm import ( + FlyDSLUnsupportedError, + gemm_wrappers, + ) + + num_gemms, num_steps, dtype = 3, 3, torch.bfloat16 + config = model_configs["126m"] + monkeypatch.setenv("NVTE_USE_GROUPED_GEMM_FLYDSL", "1") + monkeypatch.setenv("NVTE_FLYDSL_GEMM_WARN_FALLBACK", "1") + + # Record each FlyDSL grouped GEMM: its phase, whether it ran, and the operands' + # per-tensor inverse scales. + phases = {"TN": "fwd", "NN": "dgrad", "NT": "wgrad"} + dispatches = [] + grouped_gemm_flydsl = gemm_wrappers._grouped_gemm_flydsl + + def recording_grouped_gemm_flydsl(A, B, out, *, layout, **kwargs): + scales = {float(t._scale_inv) for t in list(A) + list(B)} + try: + grouped_gemm_flydsl(A, B, out, layout=layout, **kwargs) + except FlyDSLUnsupportedError: + dispatches.append((phases[layout], False, scales)) + raise + dispatches.append((phases[layout], True, scales)) + + monkeypatch.setattr(gemm_wrappers, "_grouped_gemm_flydsl", recording_grouped_gemm_flydsl) + + with quantized_model_init(enabled=fp8_model_params, recipe=recipe): + grouped_linear = GroupedLinear( + num_gemms, + config.hidden_size, + 4 * config.hidden_size, + bias=False, + params_dtype=dtype, + device="cuda", + ) + sequential_linear = torch.nn.ModuleList( + [ + Linear( + config.hidden_size, + 4 * config.hidden_size, + bias=False, + params_dtype=dtype, + device="cuda", + ) + for _ in range(num_gemms) + ] + ) + with torch.no_grad(): + for i in range(num_gemms): + sequential_linear[i].weight = Parameter(getattr(grouped_linear, f"weight{i}").clone()) + + # Uneven, FlyDSL-aligned experts so each carries its own scales. + align = get_align_size_for_quantization(recipe) + tokens = config.max_seqlen_q + m_splits = [tokens // 4, tokens // 2, tokens // 4] + assert sum(m_splits) == tokens and all(m % align == 0 for m in m_splits) + reset_rng_states() + inputs = [ + torch.randn(tokens, config.hidden_size, dtype=dtype, device="cuda") + for _ in range(num_steps) + ] + + def run(block): + FP8GlobalStateManager.reset() + for inp in inputs: + dispatches.clear() + for p in block.parameters(): + p.grad = None + inp = inp.clone().requires_grad_() + with autocast(enabled=True, recipe=recipe): + if isinstance(block, GroupedLinear): + out = block(inp, m_splits) + else: + out = torch.cat([block[i](x) for i, x in enumerate(torch.split(inp, m_splits))]) + out.sum().backward() + torch.cuda.synchronize() + return [out, inp.grad] + [p.grad for p in block.parameters()] + + outputs_ref = run(sequential_linear) + outputs = run(grouped_linear) + + ran = {phase for phase, ok, _ in dispatches if ok} + fell_back = sorted({phase for phase, ok, _ in dispatches if not ok}) + assert not fell_back and ran == {"fwd", "dgrad", "wgrad"}, ( + f"FlyDSL grouped GEMM ran {sorted(ran)}, fell back in {fell_back}. " + "See the [FLYDSL WARNING] for the fallback reason." + ) + for phase, _, scales in dispatches: + assert scales != {1.0} and len(scales) > 1, ( + f"{phase}: the compared step used scales {sorted(scales)}; it must exercise " + "non-unit, per-expert FP8 scales." + ) + for o, o_ref in zip(outputs, outputs_ref): + torch.testing.assert_close(o, o_ref, rtol=8e-3, atol=1e-3) + + @pytest.mark.skipif(not IS_HIP_EXTENSION, reason="ROCm-only grouped GEMM backend") @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) def test_grouped_linear_cuda_graph_hipkittens_mxfp8(monkeypatch, capfd): diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py b/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py index 2b330ae131..6e6390ec65 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py @@ -2,7 +2,7 @@ # # See LICENSE for license information. -"""FlyDSL GEMM kernels (dense, non-grouped) for BF16/FP16/FP32/FP8/MXFP8.""" +"""FlyDSL GEMM kernels: dense BF16/FP16/FP32/FP8/MXFP8, grouped BF16/FP16/FP8/MXFP8.""" # Supported flydsl release series. Single source of truth for the version # requirement, since flydsl is user-installed (not a declared dependency). A @@ -41,23 +41,10 @@ def _check_flydsl_version() -> None: _check_flydsl_version() from .exceptions import FlyDSLUnsupportedError -from .gemm_wrappers import ( - te_generic_gemm_flydsl, - _run_mxfp8_grouped, - _run_mxfp8_grouped_wgrad, - _run_bf16_grouped, - _run_bf16_grouped_wgrad, - _run_fp8_grouped, - _run_fp8_grouped_wgrad, -) +from .gemm_wrappers import general_grouped_gemm_flydsl, te_generic_gemm_flydsl __all__ = [ "FlyDSLUnsupportedError", + "general_grouped_gemm_flydsl", "te_generic_gemm_flydsl", - "_run_mxfp8_grouped", - "_run_mxfp8_grouped_wgrad", - "_run_bf16_grouped", - "_run_bf16_grouped_wgrad", - "_run_fp8_grouped", - "_run_fp8_grouped_wgrad", ] diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py index 12c85c7663..9afda422c0 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py @@ -27,7 +27,12 @@ # Transformer Engine-local FlyDSL utilities. from .exceptions import FlyDSLUnsupportedError -from .gemm_common_utils import require_block_tiling, require_launch_size +from .gemm_common_utils import ( + compiled_launch, + placeholder_arg, + require_block_tiling, + require_launch_size, +) from .fp16_gemm_utils import ( G2SLoader, S2RLoader, @@ -198,8 +203,9 @@ def kernel_gemm( by_n = pid_n * BLOCK_N # The flattened/XCD-swizzled block coordinates are i32, while global - # address arithmetic below is expressed in MLIR index type. Convert - # once here and use these index-typed tile bases for every address. + # address arithmetic below is 64-bit (fx.Int64, which FlyDSL prefers over the + # deprecated fx.Index). Convert once here and use these i64 tile bases for + # every address, so large operands cannot overflow i32 offsets. bx_m_idx = fx.Int64(bx_m) by_n_idx = fx.Int64(by_n) @@ -969,50 +975,6 @@ def _cached_launch(K: int, use_xcd_remap: bool = True, epilogue: str = "DEFAULT" return _compile_kernel(K, use_xcd_remap=use_xcd_remap, epilogue=epilogue) -@functools.lru_cache(maxsize=None) -def _dense_placeholder(device: torch.device, dtype: torch.dtype, numel: int = 1): - """One cached 1-element buffer per (device, dtype) for dense kernel-signature - placeholders. - - The FP32 binary requires the Bias/Aux epilogue tensor arguments to be present - but never dereferences them in the DEFAULT epilogue. Re-allocating them via - ``torch.zeros`` on every call is pure host overhead; since the contents are - never read, a single shared read-only buffer per (device, dtype) is safe. - """ - return torch.zeros(numel, dtype=dtype, device=device) - - -# Per-key cache of flyc.compile'd launchers. flydsl's JitFunction.__call__ rebuilds -# a full cache key (sig.bind + globals snapshot + arg-signature resolve) on every -# launch -- ~190us/call, which dominates wall time at small M where the kernel is -# short. flyc.compile returns a CompiledFunction whose __call__ skips all of that -# (ctypes pointer/scalar update + invoke, ~5us), baking constexpr values at compile -# time so one CompiledFunction per launch_key serves every pointer/scalar/stream -# variation. M/N are runtime scalar args (not baked), so the key excludes them -- -# identical to the MXFP8/FP8/half-precision paths. -# NOTE: the fast path intentionally bypasses flydsl's globals-drift safety check; -# this module never mutates its captured globals between launches, so it is safe. -_COMPILED_LAUNCH_CACHE = {} - - -def _compiled_launch(key, jit_fn, args): - """Return (and cache) a flyc.compile'd fast launcher for ``jit_fn``. - - ``key`` uniquely identifies the compiled specialization (the same tuple the - JitFunction builder's lru_cache uses). ``args`` is the full positional arg list - for the first call; used once to compile, then the cached CompiledFunction is - reused positionally for all later calls with the same ABI. - """ - # A compiled launcher is built for the current device; key on it so a launch on - # another device compiles its own instead of reusing this one. - key = (key, torch.cuda.current_device()) - cf = _COMPILED_LAUNCH_CACHE.get(key) - if cf is None: - cf = flyc.compile(jit_fn, *args) - _COMPILED_LAUNCH_CACHE[key] = cf - return cf - - def fp32_matmul( a: torch.Tensor, b: torch.Tensor, @@ -1128,18 +1090,18 @@ def doGemm( if needs_bias: Bias_arg = bias.contiguous().view(-1) else: - Bias_arg = _dense_placeholder(A.device, torch.float32) + Bias_arg = placeholder_arg(A.device, torch.float32) if needs_aux: Aux_arg = aux.contiguous().view(-1) else: - Aux_arg = _dense_placeholder(A.device, torch.float32) + Aux_arg = placeholder_arg(A.device, torch.float32) launch_key = ("fp32_gemm", int(K_runtime), bool(use_xcd_remap), epilogue) launch_args = (A_arg, B_arg, C_arg, Bias_arg, Aux_arg, M_runtime, N_runtime, stream) # flyc.compile fast path: skips flydsl's per-call cache-key rebuild (sig.bind + # globals snapshot + arg resolve), the dominant host cost at small M. M/N are # runtime scalars, so one compiled launcher per launch_key serves all shapes. - _compiled_launch( + compiled_launch( launch_key, _cached_launch(int(K_runtime), bool(use_xcd_remap), epilogue), launch_args, diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py index 0ba1ae79ca..95468104d4 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -25,9 +25,15 @@ from flydsl.expr.typing import T from flydsl.expr.typing import Vector as Vec -from . import grouped_gemm_utils as gd +from . import grouped_gemm_utils as grouped_utils from .exceptions import FlyDSLUnsupportedError -from .gemm_common_utils import get_num_cus, require_block_tiling, require_launch_size +from .gemm_common_utils import ( + compiled_launch, + get_num_cus, + placeholder_arg, + require_block_tiling, + require_launch_size, +) # Transformer Engine-local FlyDSL utilities. from .fp8_gemm_utils import ( @@ -260,7 +266,7 @@ def _load_scale(scale_div, index): num_blocks_n = c_n // BLOCK_N # ``grouped`` gates the shared descriptor queue / work-steal driver; - # const_expr folds it out of the dense binary. + # const_expr folds it out of the dense compiled kernel. grouped = group or wgrad # Grouped drain bound: the exact number of descriptors the host wrote to # ``work_queue`` (``num_work``), so no slot past the queue end is ever read. @@ -286,13 +292,13 @@ def _load_scale(scale_div, index): wq_rsrc = fx.rocdl.make_buffer_tensor(work_queue, max_size=True) wq_div = fx.logical_divide(wq_rsrc, fx.make_layout(1, 1)) desc_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32) - claim_iter, claim_view = gd.make_claim_broadcast(lds.claim_slot) + claim_iter, claim_view = grouped_utils.make_claim_broadcast(lds.claim_slot) def _claim_next_slot(): if tx_i32 == fx.Int32(0): old = llvm.AtomicRMWOp( llvm.AtomicBinOp.add, - gd.global_i32_ptr(work_counter, fx.Int32(0)), + grouped_utils.global_i32_ptr(work_counter, fx.Int32(0)), fx.as_ir_value(fx.Int32(1)), llvm.AtomicOrdering.monotonic, syncscope="agent", @@ -303,15 +309,15 @@ def _claim_next_slot(): return fx.Vector(claim_view.load())[0] def _desc_field(slot, field): - return gd.read_desc_field(wq_div, desc_ld_atom, slot, field) + return grouped_utils.read_desc_field(wq_div, desc_ld_atom, slot, field) def process_tile(pid_m, pid_n, expert): # Per-expert scalar scale: dense reads index 0 (length-1 buffer); grouped # reads scale[expert] from the length-num_groups vector, or through the # expert's scale_inv pointer (scale_ptrs). if const_expr(scale_ptrs): - a_scale = fx.Float32(gd.load_f32_at_ptr(A_scale_inv, expert)) - b_scale = fx.Float32(gd.load_f32_at_ptr(B_scale_inv, expert)) + a_scale = fx.Float32(grouped_utils.load_f32_at_ptr(A_scale_inv, expert)) + b_scale = fx.Float32(grouped_utils.load_f32_at_ptr(B_scale_inv, expert)) else: if const_expr(grouped): scale_idx = fx.Int32(expert) @@ -325,8 +331,9 @@ def process_tile(pid_m, pid_n, expert): by_n = pid_n * BLOCK_N # The flattened/XCD-swizzled block coordinates are i32, while global - # address arithmetic below is expressed in MLIR index type. Convert - # once here and use these index-typed tile bases for every address. + # address arithmetic below is 64-bit (fx.Int64, which FlyDSL prefers over the + # deprecated fx.Index). Convert once here and use these i64 tile bases for + # every address, so large operands cannot overflow i32 offsets. bx_m_idx = fx.Int64(bx_m) by_n_idx = fx.Int64(by_n) @@ -372,12 +379,12 @@ def process_tile(pid_m, pid_n, expert): # Issue both per-expert base-pointer loads before either is made uniform, so a # tile waits on one memory round trip rather than two. if const_expr(act_ptrs): - a_base_raw = gd.read_ptr_at(a_data_ptrs, expert) + a_base_raw = grouped_utils.read_ptr_at(a_data_ptrs, expert) if const_expr(weight_ptrs): - b_base_raw = gd.read_ptr_at(b_data_ptrs, expert) + b_base_raw = grouped_utils.read_ptr_at(b_data_ptrs, expert) if const_expr(act_ptrs): - a_base_i64 = gd.make_uniform_i64(a_base_raw) - a_div_tile = gd.rebase_buffer_view(a_base_i64, a_f8_ir_t, a_ptr_layout) + a_base_i64 = grouped_utils.make_uniform_i64(a_base_raw) + a_div_tile = grouped_utils.rebase_buffer_view(a_base_i64, a_f8_ir_t, a_ptr_layout) else: a_div_tile = a_div a_g2s = G2SLoader(a_div_tile, gl_off_a, LOAD_PASSES_HALF, a_f8_ir_t, wave_id) @@ -385,8 +392,8 @@ def process_tile(pid_m, pid_n, expert): # read with a LOCAL column (base selects the expert). Other paths keep the # shared packed b_div / by_n_b_idx. if const_expr(weight_ptrs): - b_base_i64 = gd.make_uniform_i64(b_base_raw) - b_div_tile = gd.rebase_buffer_view(b_base_i64, b_f8_ir_t, b_ptr_layout) + b_base_i64 = grouped_utils.make_uniform_i64(b_base_raw) + b_div_tile = grouped_utils.rebase_buffer_view(b_base_i64, b_f8_ir_t, b_ptr_layout) # wgrad bases are pre-shifted by the packed expert offset (see act_ptrs), # so keep the packed column there; forward/dgrad weight bases are the # real per-expert slabs, read with a local column. @@ -1190,52 +1197,6 @@ def launch_gemm( return launch_gemm -@functools.lru_cache(maxsize=None) -def _dense_placeholder(device: torch.device, dtype: torch.dtype, numel: int = 1): - """One cached 1-element buffer per (device, dtype) for dense kernel-signature - placeholders. - - The dense FP8 binary requires the grouped/epilogue tensor arguments - (Bias/Aux/work_queue/work_counter/b_data_ptrs) to be present but never - dereferences them. Re-allocating them via ``torch.zeros`` on every call is - pure host overhead that dominates wall time at small M; since the contents - are never read, a single shared read-only buffer per (device, dtype) is - safe to reuse across calls. - """ - return torch.zeros(numel, dtype=dtype, device=device) - - -# Per-key cache of flyc.compile'd launchers. flydsl's JitFunction.__call__ rebuilds -# a full cache key (sig.bind + globals snapshot + arg-signature resolve) on every -# launch -- ~190us/call, which dominates wall time at small M where the kernel is -# short. flyc.compile returns a CompiledFunction whose __call__ skips all of that -# (ctypes pointer/scalar update + invoke, ~5us), baking constexpr values at compile -# time so one CompiledFunction per (shape-class, dtype, layout) key serves every -# pointer/scalar/stream variation. M/N are runtime scalar args (not baked), so the -# key excludes them -- identical to the MXFP8 path in mxfp8_gemm.py. -# NOTE: the fast path intentionally bypasses flydsl's globals-drift safety check; -# this module never mutates its captured globals between launches, so it is safe. -_COMPILED_LAUNCH_CACHE = {} - - -def _compiled_launch(key, jit_fn, args): - """Return (and cache) a flyc.compile'd fast launcher for ``jit_fn``. - - ``key`` uniquely identifies the compiled specialization (the same tuple the - JitFunction builder's lru_cache uses). ``args`` is the full positional arg list - for the first call; used once to compile, then the cached CompiledFunction is - reused positionally for all later calls with the same ABI. - """ - # A compiled launcher is built for the current device; key on it so a launch on - # another device compiles its own instead of reusing this one. - key = (key, torch.cuda.current_device()) - cf = _COMPILED_LAUNCH_CACHE.get(key) - if cf is None: - cf = flyc.compile(jit_fn, *args) - _COMPILED_LAUNCH_CACHE[key] = cf - return cf - - @functools.lru_cache(maxsize=None) def _cached_launch( K: int, @@ -1427,7 +1388,7 @@ def doGemm( raise ValueError(f"FP8 epilogue {epilogue} does not accept an aux tensor") if stream is None: - stream = gd.current_stream_handle(A.device) + stream = grouped_utils.current_stream_handle(A.device) A_arg = A.view(torch.uint8).contiguous().view(-1) B_arg = B.view(torch.uint8).contiguous().view(-1) @@ -1440,14 +1401,14 @@ def doGemm( if needs_bias: Bias_arg = bias.contiguous().view(-1) else: - Bias_arg = _dense_placeholder(A.device, torch.float32) + Bias_arg = placeholder_arg(A.device, torch.float32) if needs_aux: Aux_arg = aux.view(-1) else: - Aux_arg = _dense_placeholder(A.device, C.dtype) + Aux_arg = placeholder_arg(A.device, C.dtype) # Grouped runtime tensors are required by the kernel signature; the dense - # binary never dereferences them but still needs placeholder buffers. + # compiled kernel never dereferences them but still needs placeholder buffers. if group or wgrad: # num_work is the kernel's drain bound: it must be the queue's descriptor count. if work_queue.numel() != 3 * int(num_work): @@ -1458,8 +1419,8 @@ def doGemm( work_queue_arg = work_queue.contiguous().view(-1) work_counter_arg = work_counter.contiguous().view(-1) else: - work_queue_arg = _dense_placeholder(A.device, torch.int32) - work_counter_arg = _dense_placeholder(A.device, torch.int32) + work_queue_arg = placeholder_arg(A.device, torch.int32) + work_counter_arg = placeholder_arg(A.device, torch.int32) # Per-expert weight base pointers (grouped forward and dgrad -- both run on this # TN kernel, dgrad over host-transposed weights): compile a kernel that reads each @@ -1470,13 +1431,13 @@ def doGemm( if weight_ptrs: b_data_ptrs_arg = b_data_ptrs.contiguous().view(-1) else: - b_data_ptrs_arg = _dense_placeholder(A.device, torch.int64) + b_data_ptrs_arg = placeholder_arg(A.device, torch.int64) # Per-expert activation base pointers (see act_ptrs in _compile_kernel). act_ptrs = a_data_ptrs is not None if act_ptrs: a_data_ptrs_arg = a_data_ptrs.contiguous().view(-1) else: - a_data_ptrs_arg = _dense_placeholder(A.device, torch.int64) + a_data_ptrs_arg = placeholder_arg(A.device, torch.int64) launch_key = ( "fp8_gemm", @@ -1514,7 +1475,7 @@ def doGemm( # flyc.compile fast path: skips flydsl's per-call cache-key rebuild (sig.bind + # globals snapshot + arg resolve), the dominant host cost at small M. M/N are # runtime scalars, so one compiled launcher per launch_key serves all shapes. - _compiled_launch( + compiled_launch( launch_key, _cached_launch( int(K_runtime), diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py index d5e68e534d..a5996a37f6 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py @@ -12,16 +12,26 @@ G2S/S2R loaders, buffer-tensor makers) stays in the per-dtype modules. """ +import functools + import torch +import flydsl.compiler as flyc import flydsl.expr as fx from flydsl.expr import arith, const_expr, range_constexpr, rocdl from .exceptions import FlyDSLUnsupportedError -def get_num_cus() -> int: - return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count +def get_num_cus(device=None) -> int: + """Number of CUs on ``device`` (default: the current device), cached per device index.""" + index = None if device is None else torch.device(device).index + return _num_cus_on(torch.cuda.current_device() if index is None else index) + + +@functools.lru_cache(maxsize=None) +def _num_cus_on(index: int) -> int: + return torch.cuda.get_device_properties(index).multi_processor_count def require_block_tiling(m, n, k, *, block_m, block_n, block_k, label, min_k_tiles=4): @@ -81,6 +91,52 @@ def require_launch_size(label, *tensors): ) +@functools.lru_cache(maxsize=None) +def placeholder_arg(device: torch.device, dtype: torch.dtype) -> torch.Tensor: + """One cached 1-element buffer per (device, dtype) for an unused kernel argument. + + The FlyDSL kernels keep a uniform signature, so a dense launch still passes the + grouped/epilogue tensors (Bias/Aux/work_queue/work_counter/data-pointer arrays) + and a static-schedule grouped launch still passes the work counter. The kernel + never reads or writes those buffers, so one shared zero tensor per (device, dtype) + avoids a per-launch allocate+zero. + """ + return torch.zeros(1, dtype=dtype, device=device) + + +# Per-key cache of flyc.compile'd launchers, shared by every FlyDSL GEMM module. +# flydsl's JitFunction.__call__ rebuilds a full cache key (sig.bind + globals snapshot +# + arg-signature resolve) on every launch -- ~190us/call, which dominates wall time +# at small M where the kernel is short. flyc.compile returns a CompiledFunction whose +# __call__ skips all of that (ctypes pointer/scalar update + invoke, ~5us), baking +# constexpr values at compile time so one CompiledFunction per key serves every +# pointer/scalar/stream variation. M/N are runtime scalar args (not baked), so keys +# exclude them. Each caller's key starts with a module-unique tag (e.g. "fp8_gemm", +# "mx32_scale_pack"), so specializations from different modules never collide. +# NOTE: the fast path intentionally bypasses flydsl's globals-drift safety check; +# these kernel modules never mutate their captured globals between launches, so it +# is safe here. +_COMPILED_LAUNCH_CACHE = {} + + +def compiled_launch(key, jit_fn, args): + """Return (and cache) a flyc.compile'd fast launcher for ``jit_fn``. + + ``key`` uniquely identifies the compiled specialization (the same tuple the + JitFunction builder's lru_cache uses). ``args`` is the full positional arg list + for the first call; it is used once to compile, then the cached CompiledFunction + is reused positionally for all later calls with the same ABI. + """ + # A compiled launcher is built for the current device; key on it so a launch on + # another device compiles its own instead of reusing this one. + key = (key, torch.cuda.current_device()) + cf = _COMPILED_LAUNCH_CACHE.get(key) + if cf is None: + cf = flyc.compile(jit_fn, *args) + _COMPILED_LAUNCH_CACHE[key] = cf + return cf + + def cdiv(numer: int, denom: int) -> int: return (numer + denom - 1) // denom diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 9988008692..b99445d26b 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -5,17 +5,23 @@ """TE entry points for the FlyDSL GEMM backend.""" import os +import warnings +from dataclasses import dataclass +from typing import Callable import torch import transformer_engine_torch as tex -from . import grouped_gemm_utils as gd +from . import grouped_gemm_utils as grouped_utils from .exceptions import FlyDSLUnsupportedError -from .half_prec_gemm import bf16_matmul, fp16_matmul +from .half_prec_gemm import _BLOCK_M as HALF_PREC_BLOCK_M, _BLOCK_N as HALF_PREC_BLOCK_N +from .half_prec_gemm import bf16_matmul, doGemm as half_prec_doGemm, fp16_matmul from .fp32_gemm import fp32_matmul -from .fp8_gemm import fp8_matmul -from .mxfp8_gemm import mxfp8_matmul +from .fp8_gemm import _BLOCK_M as FP8_BLOCK_M, _BLOCK_N as FP8_BLOCK_N +from .fp8_gemm import doGemm as fp8_doGemm, fp8_matmul +from .mxfp8_gemm import BLOCK_M as MXFP8_BLOCK_M, BLOCK_N as MXFP8_BLOCK_N +from .mxfp8_gemm import SCALE_GROUP_SIZE as MXFP8_SCALE_GROUP_SIZE, mxfp8_matmul # TODO: Some backend-independent GEMM wrapper utilities overlap with the # Triton GEMM backend (PR#667), including operand classification, logical @@ -1009,328 +1015,250 @@ def _run_mxfp8( return D, aux -_select_grouped_sched = gd.select_grouped_sched +# --------------------------------------------------------------------------------------- +# Grouped (MoE) GEMM: one persistent launch spans every expert. +# +# general_grouped_gemm_flydsl is a drop-in for cpp_extensions.general_grouped_gemm. The +# layout selects the phase (TN forward, NN dgrad, NT wgrad) and the operands' storage +# selects the kernel family. A _GroupedKernel describes the only parts that differ per +# family -- how an operand's payload and scale are read and how the kernel is launched -- +# while output validation, the descriptor queue, schedule choice, and the wgrad +# K-bucketing and scatter are shared by every family. +# --------------------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class _GroupedKernel: + """How one FlyDSL kernel family reads grouped operands and launches.""" + + name: str + block_m: int + block_n: int + # tensor -> (payload, scale) whose contraction runs along the payload's last dim. + rows: Callable + # tensor -> (payload, scale) read the other way: dgrad's B and both wgrad operands. + cols: Callable + # (per-expert scales, device, k_pad) -> the kernel's scale argument. ``k_pad`` is + # the wgrad bucket's padded token count when the operands were padded, else None. + scale_arg: Callable + # (a, a_scale, b, b_scale, D, *, layout, m, n, k, **grouped) -> None + launch: Callable + # Kernel layout for dgrad: NN, or TN when cols() yields a materialized transpose. + dgrad_layout: str + # Axis of the wgrad operands (as returned by cols) that holds the tokens. + wgrad_token_axis: int + + +def _mxfp8_scale_arg(scales, device, k_pad): + # Per-expert E8M0 scales, packed by the kernel wrapper; a padded wgrad bucket pads + # them to the bucket's token count as well. + if k_pad is None: + return list(scales) + return grouped_utils.uniform_pad_concat(list(scales), k_pad // MXFP8_SCALE_GROUP_SIZE) + + +def _launch_mxfp8(a, a_scale, b, b_scale, D, *, layout, m, n, k, **grouped): + mxfp8_matmul(a, a_scale, b, b_scale, D, layout=layout, **grouped) + + +def _fp8_transpose(t): + """The materialized transpose of FP8 operand ``t`` (the TN-only kernel reads it).""" + if t._transpose is None or t._transpose_invalid: + raise FlyDSLUnsupportedError("FP8 operand has no valid transpose") + return t._transpose + + +def _launch_fp8(a, a_scale, b, b_scale, D, *, layout, m, n, k, **grouped): + # Tensor-wise FP8 is TN-only; dgrad and wgrad reach it through transposed operands. + fp8_doGemm(a, b, D, a_scale, b_scale, m=m, n=n, k=k, **grouped) + + +def _as_rows(t): + """``t`` as 2-D rows; GroupedLinear's per-expert tensors already are, so skip the view.""" + return t if t.dim() == 2 else t.reshape(-1, t.shape[-1]) + + +def _launch_half(a, a_scale, b, b_scale, D, *, layout, m, n, k, **grouped): + # BF16 and FP16 share this kernel; the MFMA opcode is selected per operand dtype. + half_prec_doGemm( + a, + b, + D, + layout=layout, + m=m, + n=n, + k=k, + input_dtype=a.dtype, + label="half-precision grouped", + **grouped, + ) -def _run_mxfp8_grouped( - a_data, - a_scale, - b_data, - b_scale, - D, - *, - m_splits, - m_total, - n, - k, - output_dtype: torch.dtype, - layout: str = "TN", - bias=None, - b_data_ptrs=None, - a_data_ptrs=None, -): - """Dispatch a packed MXFP8 grouped GEMM through the persistent kernel. - - Two layouts, both single launches spanning all groups. Weights are read per - expert through the required ``b_data_ptrs`` (``b_data`` is one representative); - scale shapes are those of the per-expert scales stacked on dim 0: - - TN (forward, C = A @ B^T): - a = activations [M_total, K] scales [M_total, K/32] - b = per-expert weights [N, K] scales [G*N, K/32] - output width n = N_out (per expert); experts stack on B's output-N axis. - - NN (dgrad, dX = dY @ W): - a = grad_output [M_total, N_out] scales [M_total, N_out/32] - b = per-expert weights [N_out, K_in] scales [G*N_out/32, K_in] - output width n = K_in; experts stack on B's contraction (N_out) axis. - - In both, ``k`` is A's per-expert contraction (TN: K_in; NN: N_out) and ``n`` is - the output width. ``m_splits`` is the Python list of per-expert token counts - (already host-resident); the per-tile ``(pid_m, pid_n, expert)`` descriptor - queue is built from it on the host and uploaded without blocking, so no - device->host copy of a ``group_offs`` tensor (and its blocking sync) is needed. - Operands arrive already MXFP8-quantized -- no quantization, no BLAS ownership swap. - ``a_data_ptrs`` (optional) reads A per expert in place: ``a_data`` is then one - representative expert payload (see ``grouped_gemm_utils.stacked_operand``). - """ - # Layout, K/32, tile alignment and scale dtype are validated by mxfp8_matmul / - # do_gemm, which every launch goes through. - a_fp8_dtype = getattr(a_data, "_fp8_dtype", None) - b_fp8_dtype = getattr(b_data, "_fp8_dtype", None) - if a_data.dtype == torch.uint8: - a_data = reinterpret_as_fp8_tensor(a_data, a_fp8_dtype) - if b_data.dtype == torch.uint8: - b_data = reinterpret_as_fp8_tensor(b_data, b_fp8_dtype) +_MXFP8_GROUPED = _GroupedKernel( + name="MXFP8", + block_m=MXFP8_BLOCK_M, + block_n=MXFP8_BLOCK_N, + rows=lambda t: (t._rowwise_data, t._rowwise_scale_inv), + cols=lambda t: (t._columnwise_data, t._columnwise_scale_inv), + scale_arg=_mxfp8_scale_arg, + launch=_launch_mxfp8, + dgrad_layout="NN", + wgrad_token_axis=0, +) + +# One scalar inverse scale per expert, read in place through a pointer array. +_FP8_GROUPED = _GroupedKernel( + name="FP8", + block_m=FP8_BLOCK_M, + block_n=FP8_BLOCK_N, + rows=lambda t: (t._data, t._scale_inv), + cols=lambda t: (_fp8_transpose(t), t._scale_inv), + scale_arg=lambda scales, device, k_pad: grouped_utils.expert_ptrs(list(scales), device), + launch=_launch_fp8, + dgrad_layout="TN", + wgrad_token_axis=1, +) + +_HALF_GROUPED = _GroupedKernel( + name="half-precision", + block_m=HALF_PREC_BLOCK_M, + block_n=HALF_PREC_BLOCK_N, + rows=lambda t: (_as_rows(t), None), + cols=lambda t: (_as_rows(t), None), + scale_arg=lambda scales, device, k_pad: None, + launch=_launch_half, + dgrad_layout="NN", + wgrad_token_axis=0, +) + + +def _grouped_kernel_for(tensors): + """The kernel family for ``tensors``, which must all share one storage type.""" + from ...quantized_tensor import QuantizedTensorStorage + from ...tensor.storage import Float8TensorStorage, MXFP8TensorStorage + + if all(isinstance(t, MXFP8TensorStorage) for t in tensors): + return _MXFP8_GROUPED + if all(isinstance(t, Float8TensorStorage) for t in tensors): + return _FP8_GROUPED + if any(isinstance(t, QuantizedTensorStorage) for t in tensors): + raise FlyDSLUnsupportedError( + "FlyDSL grouped GEMM supports only MXFP8, tensor-wise FP8, and BF16/FP16 " + f"operands, got {sorted({type(t).__name__ for t in tensors})}" + ) + dtypes = {t.dtype for t in tensors} + if len(dtypes) == 1 and dtypes <= {torch.bfloat16, torch.float16}: + return _HALF_GROUPED + raise FlyDSLUnsupportedError( + f"FlyDSL half-precision grouped GEMM requires bf16/fp16 operands, got {dtypes}" + ) + + +def _as_kernel_dtype(data, t): + """Reinterpret TE's uint8 FP8 payload ``data`` of operand ``t`` as its FP8 dtype.""" + if data.dtype == torch.uint8: + return reinterpret_as_fp8_tensor(data, t._fp8_dtype) + return data + +def _row_operand(datas, m_splits): + """Per-expert ``[m_e, K]`` row payloads as one grouped-kernel operand. + + Returns ``(data, data_ptrs)``: each expert's rows are read in place through base + pointers shifted to their packed row (``grouped_gemm_utils.stacked_operand``), so + the kernel sees the ``[M_total, K]`` layout without a concat. Falls back to + ``(concat_or_view_rows(datas), None)`` when the payloads are not row-major + ``[m_e, K]`` blocks. + """ + if all(d.dim() == 2 and d.size(0) == int(m) for d, m in zip(datas, m_splits)): + stacked = grouped_utils.stacked_operand(datas, datas[0].device) + if stacked is not None: + return stacked + return grouped_utils.concat_or_view_rows(datas), None + + +def _run_grouped(kernel, a_tensors, b_tensors, D, *, m_splits, layout): + """One packed forward or dgrad launch: ``D[rows of e] = A_e @ op(B_e)`` per expert. + + ``layout`` is the GEMM's TE layout. TN (forward): A = activations ``[M_total, K]``, + B = per-expert weights ``[N, K]``. NN (dgrad): A = grad_output ``[M_total, N_out]``, + B = per-expert weights ``[N_out, K_in]``, read through ``kernel.cols`` and launched + with ``kernel.dgrad_layout``. A is read per expert in place; the weights are never + concatenated -- the kernel reads expert ``e``'s from ``b_data_ptrs[e]`` (``b_data`` + is one representative, for dtype and shape). The per-tile ``(pid_m, pid_n, expert)`` + descriptor queue is built on the host from ``m_splits``, so there is no + device-to-host sync. Operands arrive already quantized. + """ + if layout == "NN": + b_read, layout = kernel.cols, kernel.dgrad_layout + else: + b_read = kernel.rows + a_datas, a_scales = zip(*(kernel.rows(t) for t in a_tensors)) + b_datas, b_scales = zip(*(b_read(t) for t in b_tensors)) + a_data, a_data_ptrs = _row_operand(list(a_datas), m_splits) + device = a_data.device + b_data_ptrs = grouped_utils.expert_ptrs(list(b_datas), device) + a_data = _as_kernel_dtype(a_data, a_tensors[0]) + b_data = _as_kernel_dtype(b_datas[0], b_tensors[0]) + + # TN: B is [N, K]; NN: B is [K, N]. + n = b_data.size(0) if layout == "TN" else b_data.size(-1) + k = a_data.size(-1) + m_total = sum(int(m) for m in m_splits) D = _validate_or_allocate_output( D, shape=torch.Size((m_total, n)), - dtype=output_dtype, - device=a_data.device, - backend_name=f"MXFP8 grouped {layout}", + dtype=D.dtype, + device=device, + backend_name=f"{kernel.name} grouped {layout}", ) - # Persistent-kernel work distribution: - # work_queue one (pid_m, pid_n, expert) int32 descriptor per output tile, so - # the kernel needs no tile-to-expert search. - # work_counter atomic slot counter, used only by the work-stealing schedule. - # Workgroups take queue slots until slot >= num_work (the queue length). - from .mxfp8_gemm import BLOCK_M, BLOCK_N - - num_groups = len(m_splits) - num_pid_n = n // BLOCK_N - # Tiles cover each expert's rows of the packed [M_total, N] output. Expert token - # counts must be BLOCK_M multiples (the queue builder raises otherwise), so no - # tile spans two experts. Built from the host-side m_splits (no GPU->CPU sync). - work_queue, num_tiles = gd.build_fwd_dgrad_work_queue( - m_splits, num_pid_n, BLOCK_M, a_data.device + # Tiles cover each expert's rows of the packed output; expert token counts must be + # BLOCK_M multiples (the queue builder raises otherwise), so no tile spans two + # experts. "static" walks the queue with a grid stride; "worksteal" claims tiles + # atomically to balance uneven experts (NVTE_FLYDSL_GROUPED_SCHED overrides). + work_queue, num_tiles = grouped_utils.build_fwd_dgrad_work_queue( + m_splits, n // kernel.block_n, kernel.block_m, device ) - - # Schedule (a compile-time kernel variant): "static" walks the queue with a grid - # stride and no atomics; "worksteal" claims tiles atomically, which balances - # uneven experts. Static is chosen when there are at most a few tiles per CU. - # NVTE_FLYDSL_GROUPED_SCHED overrides the choice. - sched = _select_grouped_sched(num_tiles, a_data.device) - work_counter = gd.work_counter(sched, a_data.device) - - epilogue, bias_arg = _resolve_epilogue(bias, n) - - # TN forward and NN dgrad read expert g's weight from b_data_ptrs[g] (required), - # so the weights are never concatenated. - mxfp8_matmul( + sched = grouped_utils.select_grouped_sched(num_tiles, device) + kernel.launch( a_data, - a_scale, + kernel.scale_arg(a_scales, device, None), b_data, - b_scale, + kernel.scale_arg(b_scales, device, None), D.view(m_total, n), layout=layout, - epilogue=epilogue, - bias=bias_arg, + m=m_total, + n=n, + k=k, group=True, - num_groups=num_groups, + num_groups=len(m_splits), work_queue=work_queue, - work_counter=work_counter, + work_counter=grouped_utils.work_counter(sched, device), + sched=sched, + num_work=num_tiles, b_data_ptrs=b_data_ptrs, a_data_ptrs=a_data_ptrs, - num_work=num_tiles, a_shape=(m_total, a_data.size(-1)) if a_data_ptrs is not None else None, - sched=sched, ) - return D def _wgrad_operands(dy_b, x_b, unpadded, pad, device): """One wgrad bucket's (dy, dy_ptrs, x, x_ptrs), shared by every dtype. When no expert in the bucket needs padding (``unpadded``), both operands are read - in place from their per-expert tensors (``gd.stacked_operand``), so the kernel + in place from their per-expert tensors (``grouped_utils.stacked_operand``), so the kernel sees the packed layout without a host-side concat. Otherwise ``pad`` builds the padded packed operand and the pointers are ``None``. """ if unpadded: - dy = gd.stacked_operand(dy_b, device) - x = gd.stacked_operand(x_b, device) + dy = grouped_utils.stacked_operand(dy_b, device) + x = grouped_utils.stacked_operand(x_b, device) if dy is not None and x is not None: return dy[0], dy[1], x[0], x[1] return pad(dy_b), None, pad(x_b), None -def _run_mxfp8_grouped_wgrad( - dy_cols, # list[G] of per-expert dY columnwise fp8 data, each [m_e, N_out] - dy_scales, # list[G] of per-expert dY columnwise scale, each [m_e/32, N_out] - x_cols, # list[G] of per-expert X columnwise fp8 data, each [m_e, K_in] - x_scales, # list[G] of per-expert X columnwise scale, each [m_e/32, K_in] - dW, # output [G*N_out, K_in] - *, - m_splits, # list[int] per-expert token counts - n_out, - k_in, - dy_fp8_dtype, - x_fp8_dtype, - output_dtype, -): - """Grouped MXFP8 wgrad (NT): dW_e = dY_e^T @ X_e, via K-bucketed launches. - - K-bucketing (see :func:`_run_bf16_grouped_wgrad`): experts are grouped by their - own padded token count (rounded up to BLOCK_M, floor 512) rather than padded to - a single global max, so one large expert no longer forces every small expert to - contract over its token count. Each bucket is one fixed-K launch; MXFP8 also - packs each operand's E8M0 scale (``scale_rows = k_pad / SCALE_GROUP_SIZE``) to - the same bucket K. Per-bucket output goes to a bucket-local buffer then scatters - to global dW rows (fast direct-write when the bucket is all experts in order). - The 4-wave MFMA pipeline is untouched -- K stays compile-time per launch. - """ - from .mxfp8_gemm import BLOCK_M, BLOCK_N, SCALE_GROUP_SIZE - - num_groups = len(m_splits) - device = dW.device - - dW = _validate_or_allocate_output( - dW, - shape=torch.Size((num_groups * n_out, k_in)), - dtype=output_dtype, - device=device, - backend_name="MXFP8 grouped wgrad NT", - ) - - # One (pid_m, pid_n, expert) descriptor per tile of each expert's [N_out, K_in] - # weight gradient. pid_m/pid_n are per-expert; the kernel offsets the output row - # by expert*N_out. - num_pid_m = n_out // BLOCK_M - num_pid_n = k_in // BLOCK_N - buckets = gd.bucketize_wgrad_tokens(m_splits, BLOCK_M, max_buckets=_wgrad_max_buckets()) - buckets = gd.drop_empty_experts(buckets, m_splits, dW, n_out) - - for k_pad, ids in buckets: - bsize = len(ids) - scale_rows = k_pad // SCALE_GROUP_SIZE - dy_b = [dy_cols[e] for e in ids] - x_b = [x_cols[e] for e in ids] - dy_data, dy_ptrs, x_data, x_ptrs = _wgrad_operands( - dy_b, - x_b, - all(t.size(0) == k_pad for t in dy_b + x_b), - lambda ts: gd.uniform_pad_concat(ts, k_pad), - device, - ) - dy_scale_b = [dy_scales[e] for e in ids] - x_scale_b = [x_scales[e] for e in ids] - if dy_ptrs is not None: - # Unpadded bucket: the scale packer reads the per-expert scales in place. - dy_scale, x_scale = dy_scale_b, x_scale_b - else: - dy_scale = gd.uniform_pad_concat(dy_scale_b, scale_rows) - x_scale = gd.uniform_pad_concat(x_scale_b, scale_rows) - - if dy_data.dtype == torch.uint8: - dy_data = reinterpret_as_fp8_tensor(dy_data, dy_fp8_dtype) - if x_data.dtype == torch.uint8: - x_data = reinterpret_as_fp8_tensor(x_data, x_fp8_dtype) - - direct = ids == list(range(num_groups)) - dW_b = dW if direct else torch.empty(bsize * n_out, k_in, dtype=output_dtype, device=device) - - work_queue, num_tiles = gd.build_wgrad_work_queue(bsize, num_pid_m, num_pid_n, device) - sched = _select_grouped_sched(num_tiles, device) - work_counter = gd.work_counter(sched, device) - - mxfp8_matmul( - dy_data, - dy_scale, - x_data, - x_scale, - dW_b, - group=False, - num_groups=bsize, - work_queue=work_queue, - work_counter=work_counter, - sched=sched, - wgrad=True, - a_data_ptrs=dy_ptrs, - num_work=num_tiles, - b_data_ptrs=x_ptrs, - a_shape=(bsize * k_pad, n_out) if dy_ptrs is not None else None, - b_shape=(bsize * k_pad, k_in) if x_ptrs is not None else None, - ) - - if not direct: - for j, e in enumerate(ids): - dW[e * n_out : (e + 1) * n_out].copy_(dW_b[j * n_out : (j + 1) * n_out]) - - return dW - - -def _run_bf16_grouped( - a_data, - b_data, - D, - *, - m_splits, - m_total, - n, - k, - output_dtype: torch.dtype, - layout: str = "TN", - bias=None, - b_data_ptrs=None, - a_data_ptrs=None, -): - """Dispatch a packed BF16 grouped GEMM through the persistent kernel. - - Scaleless analog of :func:`_run_mxfp8_grouped`, with the same TN/NN operand - contracts, only without E8M0 scales. Weights are read per expert through the - required ``b_data_ptrs``; ``b_data`` is one representative expert: - - TN (forward, C = A @ B^T): - a = activations [M_total, K] b = per-expert weights [N, K] - NN (dgrad, dX = dY @ W): - a = grad_output [M_total, N_out] b = per-expert weights [N_out, K_in] - - ``m_splits`` is the Python list of per-expert token counts (already on the - host). The descriptor queue is built + cached from it directly, so no - device->host copy of a ``group_offs`` tensor (and its blocking sync) is needed. - """ - from .half_prec_gemm import doGemm as _half_prec_doGemm - - if layout not in ("TN", "NN"): - raise FlyDSLUnsupportedError( - f"FlyDSL half-precision grouped fwd/dgrad supports TN/NN, got {layout} " - "(grouped NT wgrad goes through _run_bf16_grouped_wgrad)" - ) - - # BF16 and FP16 share this kernel (the MFMA opcode is selected per operand dtype); - # take the input dtype from the operands so both route through unchanged. - input_dtype = a_data.dtype - if input_dtype not in (torch.bfloat16, torch.float16): - raise FlyDSLUnsupportedError( - f"FlyDSL half-precision grouped GEMM requires bf16/fp16 operands, got {input_dtype}" - ) - - D = _validate_or_allocate_output( - D, - shape=torch.Size((m_total, n)), - dtype=output_dtype, - device=a_data.device, - backend_name=f"half-precision grouped {layout}", - ) - - from .half_prec_gemm import _BLOCK_M, _BLOCK_N - - num_groups = len(m_splits) - num_pid_n = n // _BLOCK_N - work_queue, num_tiles = gd.build_fwd_dgrad_work_queue( - m_splits, num_pid_n, _BLOCK_M, a_data.device - ) - sched = _select_grouped_sched(num_tiles, a_data.device) - work_counter = gd.work_counter(sched, a_data.device) - - epilogue, bias_arg = _resolve_epilogue(bias, n) - _half_prec_doGemm( - a_data, - b_data, - D.view(m_total, n), - layout=layout, - m=m_total, - n=n, - k=k, - input_dtype=input_dtype, - label="half-precision grouped", - epilogue=epilogue, - bias=bias_arg, - group=True, - num_groups=num_groups, - work_queue=work_queue, - work_counter=work_counter, - b_data_ptrs=b_data_ptrs, - a_data_ptrs=a_data_ptrs, - num_work=num_tiles, - sched=sched, - a_shape=(m_total, a_data.size(-1)) if a_data_ptrs is not None else None, - ) - return D - - def _wgrad_max_buckets(): """Launch-count cap for the K-bucketed wgrad (env override, default uncapped).""" v = os.getenv("NVTE_FLYDSL_WGRAD_MAX_BUCKETS", "") @@ -1339,322 +1267,260 @@ def _wgrad_max_buckets(): return int(v) -def _run_bf16_grouped_wgrad( - dy_cols, # list[G] of per-expert dY, each [m_e, N_out] - x_cols, # list[G] of per-expert X, each [m_e, K_in] - dW, # output [G*N_out, K_in] - *, - m_splits, - n_out, - k_in, - output_dtype, -): - """Grouped BF16 wgrad (NT): dW_e = dY_e^T @ X_e, via K-bucketed launches. - - Scaleless. The kernel's contraction K (the per-expert token count) is baked in - at compile time, so experts are grouped into **K-buckets** by their own padded - token count (``bucketize_wgrad_tokens``): each expert rounds up to its own - BLOCK_M multiple (floor 512), and experts sharing a rounded K share one launch. - This avoids the old global-max pad -- a single large expert no longer forces - every small expert to contract over its token count (which wasted work and - OOM'd the int32 launch-arg limit under skew). Balanced token counts collapse to - a single bucket, reproducing the prior single-launch path bit-for-bit. - - Per bucket ``b`` with experts ``E_b`` and padded K ``k_pad``: pack those - experts' dY/X to ``[len(E_b)*k_pad, *]`` (bucket-local, so the kernel's - ``expert*K`` operand addressing uses the bucket-local index), launch into a - bucket-local output ``[len(E_b)*N_out, K_in]``, then copy each expert's - ``[N_out, K_in]`` block into its GLOBAL row of ``dW``. The hand-scheduled 4-wave - MFMA pipeline is untouched -- K stays compile-time-constant per launch. +def _pad_tokens(tensors, k_pad, token_axis): + """Zero-pad each expert's tokens to ``k_pad`` and stack the bucket's experts.""" + if token_axis == 0: + return grouped_utils.uniform_pad_concat(tensors, k_pad) + padded = [] + for t in tensors: + # [feature, m_e]: tokens are the trailing axis. + if t.size(1) == k_pad: + padded.append(t.contiguous()) + else: + out = t.new_zeros((t.size(0), k_pad)) + out[:, : t.size(1)].copy_(t) + padded.append(out) + return torch.cat(padded, dim=0) + + +def _run_grouped_wgrad(kernel, dy_tensors, x_tensors, dW, *, m_splits): + """Grouped wgrad: ``dW_e = dY_e^T @ X_e`` per expert, via K-bucketed launches. + + Both operands are read through ``kernel.cols``: ``[m_e, feature]`` for MXFP8 and + half precision (an NT launch), or the materialized ``[feature, m_e]`` transposes for + tensor-wise FP8 (a TN launch). ``dW`` is the packed ``[G*N_out, K_in]`` output. + + The kernel's contraction K (the per-expert token count) is baked in at compile + time, so experts are grouped into **K-buckets** by their own padded token count + (``bucketize_wgrad_tokens``): each expert rounds up to its own BLOCK_M multiple + (floor 512), and experts sharing a rounded K share one launch. A single large + expert therefore no longer forces every small expert to contract over its token + count (which wasted work and overflowed the int32 launch-arg limit under skew). + Balanced token counts collapse to a single bucket. + + Per bucket, the experts' operands are read in place when none needs padding, + otherwise zero-padded on the token axis to the bucket's ``k_pad`` and stacked + (bucket-local, so the kernel's ``expert*K`` addressing uses the bucket-local + index). Output goes straight to ``dW`` when the bucket is every expert in order, + else to a bucket-local buffer copied into each expert's global rows. The + hand-scheduled 4-wave MFMA pipeline is untouched -- K stays compile-time per launch. """ - from .half_prec_gemm import doGemm as _half_prec_doGemm, _BLOCK_M, _BLOCK_N - + dy_datas, dy_scales = zip(*(kernel.cols(t) for t in dy_tensors)) + x_datas, x_scales = zip(*(kernel.cols(t) for t in x_tensors)) + axis = kernel.wgrad_token_axis + n_out = dy_datas[0].size(1 - axis) + k_in = x_datas[0].size(1 - axis) num_groups = len(m_splits) device = dW.device - # BF16/FP16 share the kernel; take the dtype from the operands. - input_dtype = dy_cols[0].dtype - if input_dtype not in (torch.bfloat16, torch.float16): - raise FlyDSLUnsupportedError( - f"FlyDSL half-precision grouped wgrad requires bf16/fp16 operands, got {input_dtype}" - ) - dW = _validate_or_allocate_output( dW, shape=torch.Size((num_groups * n_out, k_in)), - dtype=output_dtype, + dtype=dW.dtype, device=device, - backend_name="BF16 grouped wgrad NT", + backend_name=f"{kernel.name} grouped wgrad", ) - num_pid_m = n_out // _BLOCK_M - num_pid_n = k_in // _BLOCK_N - buckets = gd.bucketize_wgrad_tokens(m_splits, _BLOCK_M, max_buckets=_wgrad_max_buckets()) - buckets = gd.drop_empty_experts(buckets, m_splits, dW, n_out) + # One (pid_m, pid_n, expert) descriptor per tile of each expert's [N_out, K_in] + # weight gradient; the kernel offsets the output row by expert*N_out. + num_pid_m = n_out // kernel.block_m + num_pid_n = k_in // kernel.block_n + buckets = grouped_utils.bucketize_wgrad_tokens( + m_splits, kernel.block_m, max_buckets=_wgrad_max_buckets() + ) + buckets = grouped_utils.drop_empty_experts(buckets, m_splits, dW, n_out) for k_pad, ids in buckets: bsize = len(ids) - dy_b = [dy_cols[e] for e in ids] - x_b = [x_cols[e] for e in ids] + dy_b = [dy_datas[e] for e in ids] + x_b = [x_datas[e] for e in ids] dy_data, dy_ptrs, x_data, x_ptrs = _wgrad_operands( dy_b, x_b, - all(t.size(0) == k_pad for t in dy_b + x_b), - lambda ts: gd.uniform_pad_concat(ts, k_pad), + all(t.size(axis) == k_pad for t in dy_b + x_b), + lambda ts: _pad_tokens(ts, k_pad, axis), device, ) + dy_data = _as_kernel_dtype(dy_data, dy_tensors[0]) + x_data = _as_kernel_dtype(x_data, x_tensors[0]) + # Scales are read in place when the operands are; otherwise padded alongside them. + scale_k_pad = None if dy_ptrs is not None else k_pad + dy_scale = kernel.scale_arg([dy_scales[e] for e in ids], device, scale_k_pad) + x_scale = kernel.scale_arg([x_scales[e] for e in ids], device, scale_k_pad) # When the bucket holds all experts in order (the balanced case), bucket-local - # and global expert indices match, so the kernel writes dW directly. Otherwise - # it writes a bucket-local buffer that is copied into dW below. + # and global expert indices match, so the kernel writes dW directly. direct = ids == list(range(num_groups)) - dW_b = dW if direct else torch.empty(bsize * n_out, k_in, dtype=output_dtype, device=device) - - work_queue, num_tiles = gd.build_wgrad_work_queue(bsize, num_pid_m, num_pid_n, device) - sched = _select_grouped_sched(num_tiles, device) - work_counter = gd.work_counter(sched, device) + dW_b = dW if direct else torch.empty(bsize * n_out, k_in, dtype=dW.dtype, device=device) - _half_prec_doGemm( + work_queue, num_tiles = grouped_utils.build_wgrad_work_queue( + bsize, num_pid_m, num_pid_n, device + ) + sched = grouped_utils.select_grouped_sched(num_tiles, device) + if axis == 0: + a_shape, b_shape = (bsize * k_pad, n_out), (bsize * k_pad, k_in) + else: # transposed operands stack the bucket's experts on the feature axis + a_shape, b_shape = (bsize * n_out, k_pad), (bsize * k_in, k_pad) + kernel.launch( dy_data, + dy_scale, x_data, + x_scale, dW_b, layout="NT", m=n_out, n=k_in, k=k_pad, - input_dtype=input_dtype, - label="half-precision grouped wgrad", group=False, num_groups=bsize, work_queue=work_queue, - work_counter=work_counter, + work_counter=grouped_utils.work_counter(sched, device), sched=sched, wgrad=True, - a_data_ptrs=dy_ptrs, num_work=num_tiles, + a_data_ptrs=dy_ptrs, b_data_ptrs=x_ptrs, - a_shape=(bsize * k_pad, n_out) if dy_ptrs is not None else None, - b_shape=(bsize * k_pad, k_in) if x_ptrs is not None else None, + a_shape=a_shape if dy_ptrs is not None else None, + b_shape=b_shape if x_ptrs is not None else None, ) if not direct: - # Scatter each bucket-local expert block to its global dW row. + # Scatter each bucket-local expert block to its global dW rows. for j, e in enumerate(ids): dW[e * n_out : (e + 1) * n_out].copy_(dW_b[j * n_out : (j + 1) * n_out]) - return dW - -def _fp8_scale_arg(scale): - # int64 = per-expert scale_inv pointer array, read in place by the kernel; otherwise - # a length-G fp32 vector of the per-expert scales. - if scale.dtype == torch.int64: - return scale - return scale.to(torch.float32).contiguous() +def _packed_wgrad_view(wgrad_list, n_out, k_in): + """Return a ``[G*n_out, k_in]`` view over ``wgrad_list`` if it is one packed run. - -def _run_fp8_grouped( - a_data, - a_scale_vec, # [num_groups] fp32 per-expert inverse scales, or int64 scale_inv pointers - b_data, - b_scale_vec, # [num_groups] fp32 per-expert inverse scales, or int64 scale_inv pointers - D, - *, - m_splits, - m_total, - n, - k, - output_dtype: torch.dtype, - bias=None, - b_data_ptrs=None, - a_data_ptrs=None, -): - """Dispatch a packed tensor-wise FP8 grouped GEMM (TN) through the persistent kernel. - - Tensor-wise FP8 is TN-only; grouped dgrad/wgrad reach this same TN kernel via - host-materialized transpose backings (see the backward helpers). Experts pack on - B's output-N axis ([G*N, K]); each expert carries one scalar inverse scale, so the - scales arrive as length-``num_groups`` vectors the kernel indexes by ``expert``, or - as int64 arrays of per-expert ``scale_inv`` pointers it reads in place. - ``m_splits`` is the host-side per-expert token list; the descriptor queue is built - + shape-cached from it (no GPU->CPU sync). ``a_data_ptrs`` (optional) reads A per - expert in place: ``a_data`` is then one representative expert payload and each - entry is that expert's base shifted by its packed row start (see - ``grouped_gemm_utils.stacked_operand``). + The non-fused grouped-wgrad destinations are adjacent ``[n_out, k_in]`` slices of + one contiguous ``[G, n_out, k_in]`` buffer, so the kernel can write its packed + output directly into that storage, eliding the per-expert copy-out. Return ``None`` + (the caller allocates a scratch buffer and scatters) whenever the destinations are + not a single contiguous, row-contiguous, correctly-shaped run. """ - from .fp8_gemm import doGemm as _fp8_doGemm - - num_groups = len(m_splits) - - a_fp8_dtype = getattr(a_data, "_fp8_dtype", None) - b_fp8_dtype = getattr(b_data, "_fp8_dtype", None) - if a_data.dtype == torch.uint8: - a_data = reinterpret_as_fp8_tensor(a_data, a_fp8_dtype) - if b_data.dtype == torch.uint8: - b_data = reinterpret_as_fp8_tensor(b_data, b_fp8_dtype) - - D = _validate_or_allocate_output( - D, - shape=torch.Size((m_total, n)), - dtype=output_dtype, - device=a_data.device, - backend_name="FP8 grouped TN", - ) + first = wgrad_list[0] + elem = first.element_size() + expected_ptr = first.data_ptr() + for w in wgrad_list: + if ( + w.dtype != first.dtype + or tuple(w.shape) != (n_out, k_in) + or not w.is_contiguous() + or w.data_ptr() != expected_ptr + ): + return None + expected_ptr += w.numel() * elem + # One contiguous run anchored at first.data_ptr(); span all rows over its storage. + return torch.as_strided(first, (len(wgrad_list) * n_out, k_in), (k_in, 1)) - from .fp8_gemm import _BLOCK_M, _BLOCK_N - num_pid_n = n // _BLOCK_N - work_queue, num_tiles = gd.build_fwd_dgrad_work_queue( - m_splits, num_pid_n, _BLOCK_M, a_data.device - ) - sched = _select_grouped_sched(num_tiles, a_data.device) - work_counter = gd.work_counter(sched, a_data.device) +def _grouped_gemm_flydsl(A, B, out, *, layout, m_splits, single_output): + """Run ``general_grouped_gemm``'s plain GEMM on FlyDSL, or raise if unsupported. - epilogue, bias_arg = _resolve_epilogue(bias, n) - _fp8_doGemm( - a_data, - b_data, - D.view(m_total, n), - _fp8_scale_arg(a_scale_vec), - _fp8_scale_arg(b_scale_vec), - epilogue=epilogue, - bias=bias_arg, - group=True, - num_groups=num_groups, - work_queue=work_queue, - work_counter=work_counter, - b_data_ptrs=b_data_ptrs, - a_data_ptrs=a_data_ptrs, - num_work=num_tiles, - sched=sched, - m=m_total, - n=n, - k=k, - a_shape=(m_total, a_data.size(-1)) if a_data_ptrs is not None else None, + TE's grouped layouts map onto the three phases: TN forward ``out = B_e @ A_e^T`` + (A = weights, B = activations), NN dgrad ``out = B_e @ A_e`` (A = weights, + B = grad_output), NT wgrad ``out[e] = B_e^T @ A_e`` (A = activations, + B = grad_output). Raises ``FlyDSLUnsupportedError`` for any operand or config the + kernels cannot serve. + """ + if layout not in ("TN", "NN", "NT") or m_splits is None: + raise FlyDSLUnsupportedError(f"FlyDSL grouped GEMM does not support layout {layout}") + kernel = _grouped_kernel_for(list(A) + list(B)) + if layout == "NT": + if single_output: + raise FlyDSLUnsupportedError("FlyDSL grouped wgrad writes per-expert outputs") + k_in = kernel.cols(A[0])[0].size(1 - kernel.wgrad_token_axis) + n_out = out[0].size(0) + # Write straight into the packed wgrad storage when ``out`` is one contiguous + # run; otherwise launch into scratch and scatter. + dW = _packed_wgrad_view(out, n_out, k_in) + scatter = dW is None + if scatter: + dW = torch.empty(len(out) * n_out, k_in, dtype=out[0].dtype, device=out[0].device) + _run_grouped_wgrad(kernel, B, A, dW, m_splits=m_splits) + if scatter: + for g, w in enumerate(out): + w.copy_(dW[g * n_out : (g + 1) * n_out]) + return + if not single_output: + raise FlyDSLUnsupportedError("FlyDSL grouped fwd/dgrad writes one packed output") + _run_grouped(kernel, B, A, out[0], m_splits=m_splits, layout=layout) + + +def _warn_fallback_enabled() -> bool: + return os.environ.get("NVTE_FLYDSL_GEMM_WARN_FALLBACK", "0").lower() not in ( + "", + "0", + "false", + "no", + "off", ) - return D -def _run_fp8_grouped_wgrad( - dyt_cols, # list[G] of per-expert dY^T (host-transposed) [N_out, m_e] - dy_scale_vec, # [num_groups] fp32 inverse scales for dY, or list[G] of per-expert scale_inv - xt_cols, # list[G] of per-expert X^T (host-transposed) [K_in, m_e] - x_scale_vec, # [num_groups] fp32 inverse scales for X, or list[G] of per-expert scale_inv - dW, # output [G*N_out, K_in] - *, - m_splits, - n_out, - k_in, - dy_fp8_dtype, - x_fp8_dtype, - output_dtype, +def general_grouped_gemm_flydsl( + A, + B, + out, + quantization_params, + out_dtype, + layout="TN", + m_splits=None, + gelu=False, + grad=False, + accumulate=False, + bias=None, + use_bias=False, + use_split_accumulator=False, + D_dtype=None, + single_output=False, + **kwargs, ): - """Grouped tensor-wise FP8 wgrad: dW_e = dY_e^T @ X_e, via K-bucketed TN launches. - - Tensor-wise FP8 is TN-NORMAL-read: the contraction axis (tokens m) is the - trailing dim, so -- unlike the MXFP8/BF16 NT path -- operands arrive already - transposed (the FP8 ``_transpose`` backing is a materialized row-major - transpose). Per expert: A = dY^T [N_out, m_e], B = X^T [K_in, m_e]. - - K-bucketing (see :func:`_run_bf16_grouped_wgrad`): experts are grouped by their - own padded token count instead of a single global max. Each expert's tokens are - zero-padded on the TRAILING (contraction) axis up to the bucket's ``k_pad``, and - the bucket's experts stack on the feature axis to A [bsize*N_out, k_pad] / - B [bsize*K_in, k_pad]. Per-expert scalar scales are sliced to the bucket. Output - goes to a bucket-local buffer then scatters to global dW rows (direct-write when - the bucket is all experts in order). The 4-wave MFMA pipeline is untouched. + """Drop-in replacement for ``general_grouped_gemm`` using the FlyDSL grouped kernels. + + Serves the plain grouped GEMM over BF16/FP16, tensor-wise FP8, and MXFP8 operands + (gfx950) in all three layouts. Anything else -- bias, GELU, accumulation, an + explicit output dtype, another quantized storage, or a shape the kernel tiling + cannot cover -- falls back to ``general_grouped_gemm``, mirroring the dense FlyDSL + fallback in ``cpp_extensions.gemm``. ``NVTE_FLYDSL_GEMM_WARN_FALLBACK=1`` warns on + fallback. Returns ``general_grouped_gemm``'s ``(out, bias, gelu_input)`` triple. """ - from .fp8_gemm import doGemm as _fp8_doGemm, _BLOCK_M, _BLOCK_N - - num_groups = len(m_splits) - device = dW.device - - dW = _validate_or_allocate_output( - dW, - shape=torch.Size((num_groups * n_out, k_in)), - dtype=output_dtype, - device=device, - backend_name="FP8 grouped wgrad TN", - ) - - def _pad_cols(t, cols): - # Zero-pad the trailing (token/contraction) axis of a [feature, m_e] view. - if t.size(1) == cols: - return t.contiguous() - out = t.new_zeros((t.size(0), cols)) - out[:, : t.size(1)].copy_(t) - return out - - num_pid_m = n_out // _BLOCK_M - num_pid_n = k_in // _BLOCK_N - # Per-expert scale_inv lists are read in place through per-bucket pointer arrays. - scale_lists = isinstance(dy_scale_vec, (list, tuple)) - if not scale_lists: - dy_scale_vec = dy_scale_vec.to(torch.float32).contiguous() - x_scale_vec = x_scale_vec.to(torch.float32).contiguous() - buckets = gd.bucketize_wgrad_tokens(m_splits, _BLOCK_M, max_buckets=_wgrad_max_buckets()) - - buckets = gd.drop_empty_experts(buckets, m_splits, dW, n_out) + from ...cpp_extensions.gemm import _empty_tensor, general_grouped_gemm - for k_pad, ids in buckets: - bsize = len(ids) - dy_b = [dyt_cols[e] for e in ids] - x_b = [xt_cols[e] for e in ids] - # FP8 operands are [feature, m_e] transposes, so tokens pad on the trailing axis. - dy_data, dy_ptrs, x_data, x_ptrs = _wgrad_operands( - dy_b, - x_b, - all(t.size(1) == k_pad for t in dy_b + x_b), - lambda ts: torch.cat([_pad_cols(t, k_pad) for t in ts], dim=0), - device, + try: + if gelu or accumulate or use_bias or D_dtype is not None: + raise FlyDSLUnsupportedError( + "FlyDSL grouped GEMM does not support gelu, accumulate, bias, or D_dtype" + ) + _grouped_gemm_flydsl( + A, B, out, layout=layout, m_splits=m_splits, single_output=single_output ) - if dy_data.dtype == torch.uint8: - dy_data = reinterpret_as_fp8_tensor(dy_data, dy_fp8_dtype) - if x_data.dtype == torch.uint8: - x_data = reinterpret_as_fp8_tensor(x_data, x_fp8_dtype) - - if scale_lists: - dy_scale_b = gd.expert_ptrs([dy_scale_vec[e] for e in ids], device) - x_scale_b = gd.expert_ptrs([x_scale_vec[e] for e in ids], device) - else: - idx = torch.tensor(ids, device=device) - dy_scale_b = dy_scale_vec[idx] - x_scale_b = x_scale_vec[idx] - - direct = ids == list(range(num_groups)) - dW_b = dW if direct else torch.empty(bsize * n_out, k_in, dtype=output_dtype, device=device) - - work_queue, num_tiles = gd.build_wgrad_work_queue(bsize, num_pid_m, num_pid_n, device) - sched = _select_grouped_sched(num_tiles, device) - work_counter = gd.work_counter(sched, device) - - _fp8_doGemm( - dy_data, - x_data, - dW_b, - dy_scale_b, - x_scale_b, - group=False, - num_groups=bsize, - work_queue=work_queue, - work_counter=work_counter, - sched=sched, - wgrad=True, - m=n_out, - n=k_in, - k=k_pad, - a_data_ptrs=dy_ptrs, - num_work=num_tiles, - b_data_ptrs=x_ptrs, - # TN over transposes: the bucket's experts stack on the feature axis. - a_shape=(bsize * n_out, k_pad) if dy_ptrs is not None else None, - b_shape=(bsize * k_in, k_pad) if x_ptrs is not None else None, + except FlyDSLUnsupportedError as exc: + if _warn_fallback_enabled(): + warnings.warn( + "[FLYDSL WARNING]: FlyDSL grouped GEMM does not support this configuration " + f"(layout={layout}); falling back to the default backend. Reason: {exc}", + UserWarning, + stacklevel=2, + ) + return general_grouped_gemm( + A, + B, + out, + quantization_params, + out_dtype, + layout=layout, + m_splits=m_splits, + gelu=gelu, + grad=grad, + accumulate=accumulate, + bias=bias, + use_bias=use_bias, + use_split_accumulator=use_split_accumulator, + D_dtype=D_dtype, + single_output=single_output, ) - - if not direct: - for j, e in enumerate(ids): - dW[e * n_out : (e + 1) * n_out].copy_(dW_b[j * n_out : (j + 1) * n_out]) - - return dW + empty_tensors = [_empty_tensor()] * len(A) + return out, empty_tensors, empty_tensors def _select_fp8_storage_for_layout(A, transa, B, transb): diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py index b078fe7fa5..11bc8694e5 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/grouped_gemm_utils.py @@ -36,6 +36,7 @@ from flydsl.expr.typing import T from .exceptions import FlyDSLUnsupportedError +from .gemm_common_utils import get_num_cus, placeholder_arg # --------------------------------------------------------------------------- @@ -165,12 +166,7 @@ def select_grouped_sched(num_tiles: int, device) -> str: raise ValueError( f"NVTE_FLYDSL_GROUPED_SCHED must be auto|static|worksteal, got {override!r}" ) - return "static" if num_tiles <= _GROUPED_SCHED_ROUNDS * _num_cus(device) else "worksteal" - - -@functools.lru_cache(maxsize=None) -def _num_cus(device): - return torch.cuda.get_device_properties(device).multi_processor_count + return "static" if num_tiles <= _GROUPED_SCHED_ROUNDS * get_num_cus(device) else "worksteal" # Per-call metadata caches. Keys include the current stream: a cached device tensor is @@ -271,11 +267,6 @@ def build_wgrad_work_queue(num_groups, num_pid_m, num_pid_n, device): ) -@functools.lru_cache(maxsize=None) -def _unused_counter(device): - return torch.zeros(1, dtype=torch.int32, device=device) - - def work_counter(sched, device): """The atomic tile counter a grouped launch needs for schedule ``sched``. @@ -284,7 +275,7 @@ def work_counter(sched, device): """ if sched == "worksteal": return torch.zeros(1, dtype=torch.int32, device=device) - return _unused_counter(device) + return placeholder_arg(device, torch.int32) class _SpanCAI: diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py index 922f58d85a..6742324e18 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py @@ -7,11 +7,11 @@ FP16 and BF16 share one source-level kernel generator: the algorithm is identical and only the ``v_mfma_f32_16x16x32_{f16,bf16}`` opcode differs. That opcode is selected at compile time via ``mfma_suffix``, so each (dtype, K, -output, layout) combination still compiles to its own cached binary with no +output, layout) combination still compiles to its own cached kernel with no runtime dtype branch. All supported layouts share the same generator while compiling to separate -cached binaries: +cached kernels: TN: A [M,K] normal read, B [N,K] normal read NN: A [M,K] normal read, B [K,N] transpose read @@ -37,9 +37,15 @@ from flydsl.expr.typing import Vector as Vec # Transformer Engine-local FlyDSL utilities. -from . import grouped_gemm_utils as gd +from . import grouped_gemm_utils as grouped_utils from .exceptions import FlyDSLUnsupportedError -from .gemm_common_utils import get_num_cus, require_block_tiling, require_launch_size +from .gemm_common_utils import ( + compiled_launch, + get_num_cus, + placeholder_arg, + require_block_tiling, + require_launch_size, +) from .fp16_gemm_utils import ( G2SLoader, S2RLoader, @@ -418,7 +424,7 @@ def kernel_gemm( # ``grouped`` gates the shared descriptor queue / work-steal driver (group # fwd/dgrad or the wgrad specialization); const_expr folds it all out of the - # dense binary. + # dense compiled kernel. grouped = group or wgrad # Grouped drain bound: the exact number of descriptors the host wrote to # ``work_queue`` (``num_work``), so no slot past the queue end is ever read. @@ -446,13 +452,13 @@ def kernel_gemm( # LDS broadcast scratch (one int32): thread 0 writes the atomically # claimed slot, a barrier broadcasts it to the whole workgroup. The # ``if tx==0`` claim control flow must stay inline (AST-rewritten here). - claim_iter, claim_view = gd.make_claim_broadcast(lds.claim_slot) + claim_iter, claim_view = grouped_utils.make_claim_broadcast(lds.claim_slot) def _claim_next_slot(): if tx_i32 == fx.Int32(0): old = llvm.AtomicRMWOp( llvm.AtomicBinOp.add, - gd.global_i32_ptr(work_counter, fx.Int32(0)), + grouped_utils.global_i32_ptr(work_counter, fx.Int32(0)), fx.as_ir_value(fx.Int32(1)), llvm.AtomicOrdering.monotonic, syncscope="agent", @@ -463,15 +469,16 @@ def _claim_next_slot(): return fx.Vector(claim_view.load())[0] def _desc_field(slot, field): - return gd.read_desc_field(wq_div, desc_ld_atom, slot, field) + return grouped_utils.read_desc_field(wq_div, desc_ld_atom, slot, field) def process_tile(pid_m, pid_n, expert): bx_m = pid_m * BLOCK_M by_n = pid_n * BLOCK_N # The flattened/XCD-swizzled block coordinates are i32, while global - # address arithmetic below is expressed in MLIR index type. Convert - # once here and use these index-typed tile bases for every address. + # address arithmetic below is 64-bit (fx.Int64, which FlyDSL prefers over the + # deprecated fx.Index). Convert once here and use these i64 tile bases for + # every address, so large operands cannot overflow i32 offsets. bx_m_idx = fx.Int64(bx_m) by_n_idx = fx.Int64(by_n) @@ -506,12 +513,14 @@ def process_tile(pid_m, pid_n, expert): # Issue both per-expert base-pointer loads before either is made uniform, so a # tile waits on one memory round trip rather than two. if const_expr(act_ptrs): - a_base_raw = gd.read_ptr_at(a_data_ptrs, expert) + a_base_raw = grouped_utils.read_ptr_at(a_data_ptrs, expert) if const_expr(group or act_ptrs): - b_base_raw = gd.read_ptr_at(b_data_ptrs, expert) + b_base_raw = grouped_utils.read_ptr_at(b_data_ptrs, expert) if const_expr(act_ptrs): - a_base_i64 = gd.make_uniform_i64(a_base_raw) - a_div_tile = gd.rebase_buffer_view(a_base_i64, fx.Uint8.ir_type, a_ptr_layout) + a_base_i64 = grouped_utils.make_uniform_i64(a_base_raw) + a_div_tile = grouped_utils.rebase_buffer_view( + a_base_i64, fx.Uint8.ir_type, a_ptr_layout + ) else: a_div_tile = a_div a_g2s = G2SLoader( @@ -527,13 +536,17 @@ def process_tile(pid_m, pid_n, expert): # wgrad bases are pre-shifted (act_ptrs), so they keep the packed column # and contraction shift. Dense keeps the shared packed view. if const_expr(group): - b_base_i64 = gd.make_uniform_i64(b_base_raw) - b_div_tile = gd.rebase_buffer_view(b_base_i64, fx.Uint8.ir_type, b_ptr_layout) + b_base_i64 = grouped_utils.make_uniform_i64(b_base_raw) + b_div_tile = grouped_utils.rebase_buffer_view( + b_base_i64, fx.Uint8.ir_type, b_ptr_layout + ) b_data_col = by_n_idx b_k_off = fx.Int64(0) elif const_expr(act_ptrs): - b_base_i64 = gd.make_uniform_i64(b_base_raw) - b_div_tile = gd.rebase_buffer_view(b_base_i64, fx.Uint8.ir_type, b_ptr_layout) + b_base_i64 = grouped_utils.make_uniform_i64(b_base_raw) + b_div_tile = grouped_utils.rebase_buffer_view( + b_base_i64, fx.Uint8.ir_type, b_ptr_layout + ) b_data_col = by_n_b_idx b_k_off = b_expert_k_off else: @@ -1439,51 +1452,6 @@ def _cached_launch( ) -@functools.lru_cache(maxsize=None) -def _dense_placeholder(device: torch.device, dtype: torch.dtype, numel: int = 1): - """One cached 1-element buffer per (device, dtype) for dense kernel-signature - placeholders. - - The dense half-precision binary requires the grouped/epilogue tensor arguments - (Bias/Aux/work_queue/work_counter/b_data_ptrs) to be present but never - dereferences them. Re-allocating them via ``torch.zeros`` on every call is pure - host overhead that dominates wall time at small M; since the contents are never - read, a single shared read-only buffer per (device, dtype) is safe to reuse. - """ - return torch.zeros(numel, dtype=dtype, device=device) - - -# Per-key cache of flyc.compile'd launchers. flydsl's JitFunction.__call__ rebuilds -# a full cache key (sig.bind + globals snapshot + arg-signature resolve) on every -# launch -- ~190us/call, which dominates wall time at small M where the kernel is -# short. flyc.compile returns a CompiledFunction whose __call__ skips all of that -# (ctypes pointer/scalar update + invoke, ~5us), baking constexpr values at compile -# time so one CompiledFunction per launch_key serves every pointer/scalar/stream -# variation. M/N are runtime scalar args (not baked), so the key excludes them -- -# identical to the MXFP8/FP8 paths. -# NOTE: the fast path intentionally bypasses flydsl's globals-drift safety check; -# this module never mutates its captured globals between launches, so it is safe. -_COMPILED_LAUNCH_CACHE = {} - - -def _compiled_launch(key, jit_fn, args): - """Return (and cache) a flyc.compile'd fast launcher for ``jit_fn``. - - ``key`` uniquely identifies the compiled specialization (the same tuple the - JitFunction builder's lru_cache uses). ``args`` is the full positional arg list - for the first call; used once to compile, then the cached CompiledFunction is - reused positionally for all later calls with the same ABI. - """ - # A compiled launcher is built for the current device; key on it so a launch on - # another device compiles its own instead of reusing this one. - key = (key, torch.cuda.current_device()) - cf = _COMPILED_LAUNCH_CACHE.get(key) - if cf is None: - cf = flyc.compile(jit_fn, *args) - _COMPILED_LAUNCH_CACHE[key] = cf - return cf - - def _half_prec_matmul( a: torch.Tensor, b: torch.Tensor, @@ -1748,7 +1716,7 @@ def doGemm( raise ValueError(f"{label} epilogue {epilogue} does not accept an aux tensor") if stream is None: - stream = gd.current_stream_handle(A.device) + stream = grouped_utils.current_stream_handle(A.device) # Preserve the original validated byte-addressed G2L path. These are # metadata-only dtype/flatten views of the already-contiguous row-major @@ -1762,14 +1730,14 @@ def doGemm( if needs_bias: Bias_arg = bias.contiguous().view(-1) else: - Bias_arg = _dense_placeholder(A.device, torch.float32) + Bias_arg = placeholder_arg(A.device, torch.float32) if needs_aux: Aux_arg = aux.view(-1) else: - Aux_arg = _dense_placeholder(A.device, C.dtype) + Aux_arg = placeholder_arg(A.device, C.dtype) # Grouped runtime tensors are required by the kernel signature; the dense - # binary never dereferences them but still needs placeholder buffers. + # compiled kernel never dereferences them but still needs placeholder buffers. if group or wgrad: # num_work is the kernel's drain bound: it must be the queue's descriptor count. if work_queue.numel() != 3 * int(num_work): @@ -1780,8 +1748,8 @@ def doGemm( work_queue_arg = work_queue.contiguous().view(-1) work_counter_arg = work_counter.contiguous().view(-1) else: - work_queue_arg = _dense_placeholder(A.device, torch.int32) - work_counter_arg = _dense_placeholder(A.device, torch.int32) + work_queue_arg = placeholder_arg(A.device, torch.int32) + work_counter_arg = placeholder_arg(A.device, torch.int32) # Per-expert weight base pointers (grouped TN/NN; wgrad's pre-shifted X with # act_ptrs): an int64[G] device array read in-kernel so the operand needs no @@ -1789,7 +1757,7 @@ def doGemm( if b_data_ptrs is not None: b_data_ptrs_arg = b_data_ptrs.contiguous().view(-1) else: - b_data_ptrs_arg = _dense_placeholder(A.device, torch.int64) + b_data_ptrs_arg = placeholder_arg(A.device, torch.int64) # act_ptrs also selects how the kernel reads B: outside grouped fwd/dgrad, B is # read through b_data_ptrs only when act_ptrs is set. So activation pointers are # grouped-only, and wgrad takes the dY and X pointer arrays together or not at all. @@ -1802,7 +1770,7 @@ def doGemm( if act_ptrs: a_data_ptrs_arg = a_data_ptrs.contiguous().view(-1) else: - a_data_ptrs_arg = _dense_placeholder(A.device, torch.int64) + a_data_ptrs_arg = placeholder_arg(A.device, torch.int64) launch_key = ( "half_prec_gemm", @@ -1836,7 +1804,7 @@ def doGemm( # flyc.compile fast path: skips flydsl's per-call cache-key rebuild (sig.bind + # globals snapshot + arg resolve), the dominant host cost at small M. M/N are # runtime scalars, so one compiled launcher per launch_key serves all shapes. - _compiled_launch( + compiled_launch( launch_key, _cached_launch( K_runtime, diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index c41b7b6aa2..f2ac5b3d19 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -34,9 +34,15 @@ from flydsl.expr.typing import Vector as Vec # Transformer Engine-local FlyDSL utilities. -from . import grouped_gemm_utils as gd +from . import grouped_gemm_utils as grouped_utils from .exceptions import FlyDSLUnsupportedError -from .gemm_common_utils import get_num_cus, require_block_tiling, require_launch_size +from .gemm_common_utils import ( + compiled_launch, + get_num_cus, + placeholder_arg, + require_block_tiling, + require_launch_size, +) from .fp8_gemm_utils import ( G2SLoader, S2RLoader, @@ -60,46 +66,6 @@ SCALE_GROUP_SIZE = 32 -@functools.lru_cache(maxsize=None) -def _dummy_arg(dtype: torch.dtype, device: torch.device): - """One-element placeholder for an unused uniform-signature kernel arg. - - Cached per (dtype, device) and shared across calls: the kernel never - dereferences these buffers, so a single zero tensor avoids a per-launch - allocate+zero of a fresh one. - """ - return torch.zeros(1, dtype=dtype, device=device) - - -# Per-key cache of flyc.compile'd launchers. flydsl's JitFunction.__call__ rebuilds -# a full cache key (sig.bind + globals snapshot + arg-signature resolve) on every -# launch; flyc.compile returns a CompiledFunction whose __call__ skips that, baking -# constexpr values at compile time so one CompiledFunction per (shape-class, dtype, -# layout) key serves every pointer/scalar/stream variation. -# NOTE: the fast path intentionally bypasses flydsl's globals-drift safety check; -# these kernel modules never mutate their captured globals between launches, so it -# is safe here. Keyed identically to the underlying @lru_cache launcher builders. -_COMPILED_LAUNCH_CACHE = {} - - -def _compiled_launch(key, jit_fn, args): - """Return (and cache) a flyc.compile'd fast launcher for ``jit_fn``. - - ``key`` uniquely identifies the compiled specialization (same tuple the - JitFunction's lru_cache uses). ``args`` is the full positional arg list for the - first call; it is used once to compile, then the cached CompiledFunction is - reused positionally for all later calls with the same ABI. - """ - # A compiled launcher is built for the current device; key on it so a launch on - # another device compiles its own instead of reusing this one. - key = (key, torch.cuda.current_device()) - cf = _COMPILED_LAUNCH_CACHE.get(key) - if cf is None: - cf = flyc.compile(jit_fn, *args) - _COMPILED_LAUNCH_CACHE[key] = cf - return cf - - def _debug_enabled() -> bool: value = os.getenv("DEBUG_FLYDSL_MXFP8_GEMM", "") return value.lower() not in ("", "0", "false", "no", "off") @@ -146,7 +112,7 @@ def _compile_mx32_scale_pack_kernel( ) # Select source addressing before FlyDSL captures the kernel. The emitted - # rowwise and columnwise binaries contain no runtime orientation branch. + # rowwise and columnwise compiled kernels contain no runtime orientation branch. if source_colwise: def _source_offset(source_k32, source_row): @@ -213,9 +179,11 @@ def kernel_pack_mx32_scales(src_ptrs: fx.Tensor, expert_table: fx.Tensor, dst: f # K128 block, so one expert: make its base wave-uniform and rebase a buffer # view, keeping the buffer loads of the packed path. block = k128 if source_colwise else tile - expert = fx.Int32(gd.load_i32_at(expert_table, block)) - base = gd.make_uniform_i64(gd.read_ptr_at(src_ptrs, expert)) - src_div = gd.rebase_buffer_view(base, fx.Uint8.ir_type, fx.make_layout(dim * qk, 1)) + expert = fx.Int32(grouped_utils.load_i32_at(expert_table, block)) + base = grouped_utils.make_uniform_i64(grouped_utils.read_ptr_at(src_ptrs, expert)) + src_div = grouped_utils.rebase_buffer_view( + base, fx.Uint8.ir_type, fx.make_layout(dim * qk, 1) + ) _store_word(dst, src_div, coords) @flyc.jit @@ -270,7 +238,7 @@ def _cached_mx32_scale_pack_launch( stride1: int, per_expert: bool = False, ): - """Cache orientation-and-stride-specialized fused pack binaries.""" + """Cache orientation-and-stride-specialized fused pack compiled kernels.""" return _compile_mx32_scale_pack_kernel(dim, qk, source_colwise, stride0, stride1, per_expert) @@ -314,12 +282,12 @@ def pack_mx32_scales_for_hk( device=scales_u8.device, ) if stream is None: - stream = gd.current_stream_handle(scales_u8.device) + stream = grouped_utils.current_stream_handle(scales_u8.device) stride0, stride1 = (int(x) for x in scales_u8.stride()) pack_key = ("mx32_scale_pack", dim, qk, bool(source_colwise), stride0, stride1) pack_args = (scales_u8, packed, stream) - _compiled_launch( + compiled_launch( pack_key, _cached_mx32_scale_pack_launch(dim, qk, bool(source_colwise), stride0, stride1), pack_args, @@ -341,7 +309,9 @@ def pack_mx32_scales_for_hk_experts( back to packing a real concat when the tensors cannot be addressed in place. """ # stacked_operand requires one common trailing shape, so rank 2 holds for all. - stacked = gd.stacked_operand(scales, scales[0].device) if scales[0].dim() == 2 else None + stacked = ( + grouped_utils.stacked_operand(scales, scales[0].device) if scales[0].dim() == 2 else None + ) if stacked is None: return pack_mx32_scales_for_hk( torch.cat(scales, dim=0), source_colwise=source_colwise, stream=stream @@ -361,16 +331,16 @@ def pack_mx32_scales_for_hk_experts( dim, qk = total, inner # Experts stack on dim; each thread's bytes share one 64-row tile. block = 64 - expert_table = gd.block_expert_table(extents, block, representative.device) + expert_table = grouped_utils.block_expert_table(extents, block, representative.device) packed = torch.empty((qk // 4, dim), dtype=torch.int32, device=representative.device) if stream is None: - stream = gd.current_stream_handle(representative.device) + stream = grouped_utils.current_stream_handle(representative.device) # The shifted pointers address the contiguous concatenation, so its strides apply. stride0, stride1 = inner, 1 pack_key = ("mx32_scale_pack_experts", dim, qk, bool(source_colwise), stride0, stride1) pack_args = (src_ptrs, expert_table, packed, stream) - _compiled_launch( + compiled_launch( pack_key, _cached_mx32_scale_pack_launch( dim, qk, bool(source_colwise), stride0, stride1, per_expert=True @@ -710,7 +680,7 @@ def kernel_gemm( # (pid_m, pid_n) output tile; the body below runs once. # # Grouped path (group is True, packed TN, const_expr-folded so the dense - # binary is untouched): the launch is a fixed num-CU grid of persistent + # compiled kernel is untouched): the launch is a fixed num-CU grid of persistent # workgroups. Each workgroup claims a queue slot by atomically fetch-adding # the shared counter (``work_counter[0]``), reads the precomputed # ``(pid_m, pid_n, expert)`` descriptor at that slot from ``work_queue``, and @@ -744,7 +714,7 @@ def kernel_gemm( # a 1-element view. Thread 0 writes via ptr_store; every thread reads via # the view after a workgroup barrier. (Descriptor reads reuse the int32 # ``scale_ld_atom`` as their copy atom.) - claim_iter, claim_view = gd.make_claim_broadcast(lds.claim_slot) + claim_iter, claim_view = grouped_utils.make_claim_broadcast(lds.claim_slot) def _claim_next_slot(): # One atomic per workgroup: thread 0 fetch-adds the shared counter @@ -754,7 +724,7 @@ def _claim_next_slot(): if tx_i32 == fx.Int32(0): old = llvm.AtomicRMWOp( llvm.AtomicBinOp.add, - gd.global_i32_ptr(work_counter, fx.Int32(0)), + grouped_utils.global_i32_ptr(work_counter, fx.Int32(0)), fx.as_ir_value(fx.Int32(1)), llvm.AtomicOrdering.monotonic, syncscope="agent", @@ -765,15 +735,16 @@ def _claim_next_slot(): return fx.Vector(claim_view.load())[0] def _desc_field(slot, field): - return gd.read_desc_field(wq_div, scale_ld_atom, slot, field) + return grouped_utils.read_desc_field(wq_div, scale_ld_atom, slot, field) def process_tile(pid_m, pid_n, expert): bx_m = pid_m * BLOCK_M by_n = pid_n * BLOCK_N # The flattened/XCD-swizzled block coordinates are i32, while global - # address arithmetic below is expressed in MLIR index type. Convert - # once here and use these index-typed tile bases for every address. + # address arithmetic below is 64-bit (fx.Int64, which FlyDSL prefers over the + # deprecated fx.Index). Convert once here and use these i64 tile bases for + # every address, so large operands cannot overflow i32 offsets. bx_m_idx = fx.Int64(bx_m) by_n_idx = fx.Int64(by_n) @@ -878,12 +849,12 @@ def process_tile(pid_m, pid_n, expert): # Issue both per-expert base-pointer loads before either is made uniform, so a # tile waits on one memory round trip rather than two. if const_expr(act_ptrs): - a_base_raw = gd.read_ptr_at(a_data_ptrs, expert) + a_base_raw = grouped_utils.read_ptr_at(a_data_ptrs, expert) if const_expr(group or act_ptrs): - b_base_raw = gd.read_ptr_at(b_data_ptrs, expert) + b_base_raw = grouped_utils.read_ptr_at(b_data_ptrs, expert) if const_expr(act_ptrs): - a_base_i64 = gd.make_uniform_i64(a_base_raw) - a_div_tile = gd.rebase_buffer_view(a_base_i64, a_f8_ir_t, a_ptr_layout) + a_base_i64 = grouped_utils.make_uniform_i64(a_base_raw) + a_div_tile = grouped_utils.rebase_buffer_view(a_base_i64, a_f8_ir_t, a_ptr_layout) else: a_div_tile = a_div a_g2s = G2SLoader( @@ -899,13 +870,13 @@ def process_tile(pid_m, pid_n, expert): # TN and NN. wgrad bases are pre-shifted (act_ptrs), so they keep the # packed column and contraction shift. Dense keeps the shared ``b_div``. if const_expr(group): - b_base_i64 = gd.make_uniform_i64(b_base_raw) - b_div_tile = gd.rebase_buffer_view(b_base_i64, b_f8_ir_t, b_ptr_layout) + b_base_i64 = grouped_utils.make_uniform_i64(b_base_raw) + b_div_tile = grouped_utils.rebase_buffer_view(b_base_i64, b_f8_ir_t, b_ptr_layout) b_data_col = by_n_idx b_k_off = fx.Int64(0) elif const_expr(act_ptrs): - b_base_i64 = gd.make_uniform_i64(b_base_raw) - b_div_tile = gd.rebase_buffer_view(b_base_i64, b_f8_ir_t, b_ptr_layout) + b_base_i64 = grouped_utils.make_uniform_i64(b_base_raw) + b_div_tile = grouped_utils.rebase_buffer_view(b_base_i64, b_f8_ir_t, b_ptr_layout) b_data_col = by_n_b_idx b_k_off = b_expert_k_off else: @@ -2093,7 +2064,7 @@ def do_gemm( raise ValueError(f"MXFP8 epilogue {epilogue} does not accept an aux tensor") if stream is None: - stream = gd.current_stream_handle(A.device) + stream = grouped_utils.current_stream_handle(A.device) # A non-contiguous output cannot receive the kernel's writes: the flat # descriptor below would target a throwaway ``.contiguous()`` copy. Reject @@ -2109,19 +2080,19 @@ def do_gemm( C_arg = C.view(-1) # DEFAULT keeps the kernel signature uniform with dummy 1-element buffers. These # placeholders are never dereferenced by the kernel, so a cached shared buffer - # (see _dummy_arg) is used rather than a fresh allocate+zero every call. + # (see ``placeholder_arg``) is used rather than a fresh allocate+zero every call. if needs_bias: Bias_arg = bias.contiguous().view(-1) else: - Bias_arg = _dummy_arg(torch.float32, A.device) + Bias_arg = placeholder_arg(A.device, torch.float32) if needs_aux: Aux_arg = aux.view(-1) else: - Aux_arg = _dummy_arg(C.dtype, A.device) + Aux_arg = placeholder_arg(A.device, C.dtype) # Grouped runtime tensors are required by the kernel signature; the dense - # binary never dereferences them but still needs placeholder buffers. wgrad - # uses the same persistent descriptor queue as the grouped fwd/dgrad path. + # compiled kernel never dereferences them but still needs placeholder + # buffers. wgrad uses the same persistent descriptor queue as the grouped fwd/dgrad path. if group or wgrad: # num_work is the kernel's drain bound: it must be the queue's descriptor count. if work_queue.numel() != 3 * int(num_work): @@ -2132,8 +2103,8 @@ def do_gemm( work_queue_arg = work_queue.contiguous().view(-1) work_counter_arg = work_counter.contiguous().view(-1) else: - work_queue_arg = _dummy_arg(torch.int32, A.device) - work_counter_arg = _dummy_arg(torch.int32, A.device) + work_queue_arg = placeholder_arg(A.device, torch.int32) + work_counter_arg = placeholder_arg(A.device, torch.int32) # Per-expert B base pointers: the weights in grouped fwd/dgrad, or wgrad's # pre-shifted X with act_ptrs. Dense passes an unused placeholder so the @@ -2141,7 +2112,7 @@ def do_gemm( if b_data_ptrs is not None: b_data_ptrs_arg = b_data_ptrs.contiguous().view(-1) else: - b_data_ptrs_arg = _dummy_arg(torch.int64, A.device) + b_data_ptrs_arg = placeholder_arg(A.device, torch.int64) # act_ptrs also selects how the kernel reads B: outside grouped fwd/dgrad, B is # read through b_data_ptrs only when act_ptrs is set. So activation pointers are # grouped-only, and wgrad takes the dY and X pointer arrays together or not at all. @@ -2154,7 +2125,7 @@ def do_gemm( if act_ptrs: a_data_ptrs_arg = a_data_ptrs.contiguous().view(-1) else: - a_data_ptrs_arg = _dummy_arg(torch.int64, A.device) + a_data_ptrs_arg = placeholder_arg(A.device, torch.int64) launch_key = ( "mxfp8_gemm", @@ -2187,7 +2158,7 @@ def do_gemm( int(num_work), stream, ) - _compiled_launch( + compiled_launch( launch_key, _cached_launch( K_runtime, A.dtype, B.dtype, C.dtype, layout, epilogue, group, sched, wgrad, act_ptrs @@ -2209,7 +2180,7 @@ def _cached_launch( wgrad: bool = False, act_ptrs: bool = False, ): - """Cache independent TN/NN/NT binaries with no runtime layout argument.""" + """Cache independent TN/NN/NT compiled kernels with no runtime layout argument.""" return _compile_kernel( K, a_fp8_dtype, @@ -2342,7 +2313,7 @@ def mxfp8_matmul( sched: str = "worksteal", wgrad: bool = False, ): - """Normalize scale orientation and launch a compile-time layout binary. + """Normalize scale orientation and launch a kernel compiled for a fixed layout. Wrapper-visible contracts: diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 7195b12863..da5f180d27 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -35,7 +35,7 @@ ) from ._common import can_reconstruct_wgrad_input_from_original, WeightGradStore from . import _split_quantization -from ..quantization import FP8GlobalStateManager, QuantizerRole +from ..quantization import FP8GlobalStateManager, QuantizerRole, is_flydsl_grouped_gemm_enabled from ..utils import ( divide, cast_if_needed, @@ -90,71 +90,6 @@ from transformer_engine.pytorch.triton_kernels.grouped_gemm import general_grouped_gemm_triton import os - -def _flydsl_rows(t): - """``t`` as 2-D rows; GroupedLinear's per-expert tensors already are, so skip the view.""" - return t if t.dim() == 2 else t.reshape(-1, t.shape[-1]) - - -def _flydsl_fp8_transposes(tensors): - """The materialized FP8 transposes of ``tensors`` for the FlyDSL TN-only kernel. - - Raises ``FlyDSLUnsupportedError`` (so the caller falls back to the default grouped - GEMM) when a transpose is missing or marked stale by ``_transpose_invalid``. - """ - from ..flydsl_kernels.gemm import FlyDSLUnsupportedError - - for t in tensors: - if t._transpose is None or t._transpose_invalid: - raise FlyDSLUnsupportedError("FP8 operand has no valid transpose") - return [t._transpose for t in tensors] - - -def _flydsl_row_operand(datas, m_splits): - """Per-expert ``[m_e, K]`` row payloads as one FlyDSL grouped-kernel operand. - - Returns ``(a_data, a_data_ptrs)``: each expert's rows are read in place through - base pointers shifted to their packed row (``grouped_gemm_utils.stacked_operand``), - so the kernel sees the ``[M_total, K]`` layout without a concat. Falls back to - ``(concat_or_view_rows(datas), None)`` when the payloads are not row-major - ``[m_e, K]`` blocks. - """ - from ..flydsl_kernels.gemm.grouped_gemm_utils import concat_or_view_rows, stacked_operand - - if all(d.dim() == 2 and d.size(0) == int(m) for d, m in zip(datas, m_splits)): - stacked = stacked_operand(datas, datas[0].device) - if stacked is not None: - return stacked - return concat_or_view_rows(datas), None - - -def _packed_wgrad_view(wgrad_list, num_gemms, n_out, k_in): - """Return a [num_gemms*n_out, k_in] view over ``wgrad_list`` if it is one packed run. - - The non-fused grouped-wgrad destinations are ``[wgrad_packed[g] for g in range]`` - -- adjacent ``[n_out, k_in]`` slices of a single contiguous ``[num_gemms, n_out, - k_in]`` buffer. When that layout holds, the grouped wgrad kernel can write its - ``[num_gemms*n_out, k_in]`` output directly into that storage, eliding the - per-expert copy-out. Return ``None`` (caller allocates a scratch buffer and - scatters) whenever the destinations are not a single contiguous, row-contiguous, - correctly-shaped run -- e.g. fused main_grad accumulation into separate params. - """ - first = wgrad_list[0] - elem = first.element_size() - expected_ptr = first.data_ptr() - for w in wgrad_list: - if ( - w.dtype != first.dtype - or tuple(w.shape) != (n_out, k_in) - or not w.is_contiguous() - or w.data_ptr() != expected_ptr - ): - return None - expected_ptr += w.numel() * elem - # One contiguous run anchored at first.data_ptr(); span all rows over its storage. - return torch.as_strided(first, (num_gemms * n_out, k_in), (k_in, 1)) - - __all__ = ["GroupedLinear", "is_module_grouped_tensor_path_supported"] @@ -815,8 +750,11 @@ def _is_blockwise_fp8_triton_grouped_gemm_supported( ) @staticmethod - def _flydsl_grouped_common_ok( + def _is_flydsl_grouped_gemm_supported( *, + fp8, + recipe, + activation_dtype, use_bias, use_grouped_gemm_triton, backward_override, @@ -828,17 +766,20 @@ def _flydsl_grouped_common_ok( in_features, out_features, ) -> bool: - """Dtype-independent gate shared by every FlyDSL grouped GEMM path. + """ROCm FlyDSL grouped GEMM: the opt-in and precision (see + ``is_flydsl_grouped_gemm_enabled``) plus this module's per-call constraints. - HIP extension, gfx950, the ``NVTE_USE_GROUPED_GEMM_FLYDSL=1`` opt-in, - 128-aligned N/K for the tile geometry, and the unsupported-feature - exclusions. The per-dtype gates add only the quantization condition - (BF16/FP16: none; tensor-wise FP8: current/delayed scaling; MXFP8). + It does not quantize: ``general_grouped_gemm_flydsl`` consumes TE's quantized + operands and still falls back per GEMM on anything its kernels cannot serve. """ + if fp8: + enabled = recipe is not None and is_flydsl_grouped_gemm_enabled(recipe) + else: + # One half-precision kernel covers BF16 and FP16. + half_precision = activation_dtype in (torch.bfloat16, torch.float16) + enabled = half_precision and is_flydsl_grouped_gemm_enabled() return ( - IS_HIP_EXTENSION - and os.getenv("NVTE_USE_GROUPED_GEMM_FLYDSL", "0") == "1" - and get_device_compute_capability() == (9, 5) + enabled and in_features % 128 == 0 and out_features % 128 == 0 and not ( @@ -853,513 +794,6 @@ def _flydsl_grouped_common_ok( ) ) - @staticmethod - def _flydsl_grouped_kind( - *, - fp8, - recipe, - use_bias, - use_grouped_gemm_triton, - backward_override, - cpu_offloading, - save_original_input, - debug, - unpad_output, - actual_m_splits, - in_features, - out_features, - ): - """Return which FlyDSL grouped path applies: 'mxfp8' | 'fp8' | 'bf16' | None.""" - if not _GroupedLinear._flydsl_grouped_common_ok( - use_bias=use_bias, - use_grouped_gemm_triton=use_grouped_gemm_triton, - backward_override=backward_override, - cpu_offloading=cpu_offloading, - save_original_input=save_original_input, - debug=debug, - unpad_output=unpad_output, - actual_m_splits=actual_m_splits, - in_features=in_features, - out_features=out_features, - ): - return None - if not fp8: - # Unquantized grouped GEMM. The 'bf16' kind covers BF16 and FP16 (one - # half-precision kernel); the caller drops it for other activation dtypes. - return "bf16" - if recipe is None: - return None - if recipe.mxfp8(): - return "mxfp8" - if recipe.float8_current_scaling() or recipe.delayed(): - return "fp8" - return None - - @staticmethod - def _fwd_flydsl_mxfp8_grouped( - *, - inputmats, - weights, - out, - m_splits, - m_splits_tensor, - ): - """Dispatch the FlyDSL MXFP8 grouped GEMM over already-quantized operands. - - Called inline at the forward GEMM dispatch. It does not quantize -- it - consumes TE's quantized operands (``inputmats`` the per-group MXFP8 - activations from ``tex.split_quantize``, ``weights`` the per-expert MXFP8 - weights): activations, weights, and their scales are read per expert in place - through per-expert base pointers (no concat), then one TN launch is issued via - :func:`_run_mxfp8_grouped`, writing into ``out``. The descriptor queue is - built host-side from ``m_splits``. - """ - from ..flydsl_kernels.gemm import _run_mxfp8_grouped - from ..flydsl_kernels.gemm.grouped_gemm_utils import expert_ptrs - - # TN operands: activations are read in place from each expert's rows (no - # [M_total, K] concat). Each MXFP8 operand carries its rowwise fp8 payload and - # E8M0 scale; the scales are passed per expert and packed in place. - a_data, a_data_ptrs = _flydsl_row_operand([im._rowwise_data for im in inputmats], m_splits) - a_scale = [im._rowwise_scale_inv for im in inputmats] - # Weights are independently-allocated per-expert slabs. The grouped-TN - # kernel reads each expert's weight from a per-expert base pointer, so the - # host passes an int64[G] address array instead of concatenating the G - # weights into one packed buffer. b_data is a per-expert representative used - # only for dtype/shape. - b_data_ptrs = expert_ptrs([w._rowwise_data for w in weights], a_data.device) - b_data = weights[0]._rowwise_data - b_scale = [w._rowwise_scale_inv for w in weights] - - # Preserve the native MXFP8 dtype so the wrapper reinterprets the uint8 - # payload to the correct fp8 view at the kernel boundary. - a_data._fp8_dtype = inputmats[0]._fp8_dtype - b_data._fp8_dtype = weights[0]._fp8_dtype - - n = weights[0]._rowwise_data.size(0) - k = weights[0]._rowwise_data.size(-1) - m_total = sum(int(m) for m in m_splits) - - _run_mxfp8_grouped( - a_data, - a_scale, - b_data, - b_scale, - out, - m_splits=m_splits, - m_total=m_total, - n=n, - k=k, - output_dtype=out.dtype, - b_data_ptrs=b_data_ptrs, - a_data_ptrs=a_data_ptrs, - ) - - @staticmethod - def _dgrad_flydsl_mxfp8_grouped( - *, - grad_output, - weights, - dgrad, - m_splits, - m_splits_tensor, - ): - """FlyDSL grouped dgrad: dX = dY @ W through the packed NN kernel. - - Rides off backward's existing quantization -- ``grad_output`` are the - per-group MXFP8 activations (rowwise) and ``weights`` the per-expert MXFP8 - weights with columnwise payloads ensured by the caller. dY is the NN kernel-A - ``[M_total, N_out]`` and W (columnwise) the contraction-packed kernel-B - ``[G*N_out, K_in]``, both read per expert in place through per-expert base - pointers; one grouped NN launch writes into ``dgrad`` ``[M_total, K_in]``. No - quantization here. - """ - from ..flydsl_kernels.gemm import _run_mxfp8_grouped - from ..flydsl_kernels.gemm.grouped_gemm_utils import expert_ptrs - - # NN kernel-A = dY, rowwise (contraction over N_out), read in place from each - # expert's rows. NN kernel-B = W, columnwise (transpose-read, contraction on - # the N_out axis): independent per-expert [N_out, K_in] columnwise slabs - # addressed via per-expert base pointers. Neither is concatenated. - a_data, a_data_ptrs = _flydsl_row_operand([g._rowwise_data for g in grad_output], m_splits) - a_scale = [g._rowwise_scale_inv for g in grad_output] - b_data_ptrs = expert_ptrs([w._columnwise_data for w in weights], a_data.device) - b_data = weights[0]._columnwise_data - b_scale = [w._columnwise_scale_inv for w in weights] - - a_data._fp8_dtype = grad_output[0]._fp8_dtype - b_data._fp8_dtype = weights[0]._fp8_dtype - - n_out = grad_output[0]._rowwise_data.size(-1) # per-expert contraction (N_out) - k_in = dgrad.size(-1) # output width (K_in) - m_total = sum(int(m) for m in m_splits) - - _run_mxfp8_grouped( - a_data, - a_scale, - b_data, - b_scale, - dgrad, - m_splits=m_splits, - m_total=m_total, - n=k_in, # NN output width is K_in - k=n_out, # NN per-expert contraction is N_out - output_dtype=dgrad.dtype, - layout="NN", - b_data_ptrs=b_data_ptrs, - a_data_ptrs=a_data_ptrs, - ) - - @staticmethod - def _wgrad_flydsl_mxfp8_grouped( - *, - grad_output, - inputmats, - wgrad_list, - m_splits, - ): - """FlyDSL grouped wgrad: dW = dY^T @ X per expert, through the packed NT kernel. - - Both operands transpose-read and contract over the token axis, so both use - their COLUMNWISE MXFP8 payloads (dY from ``grad_output``, X from - ``inputmats``). The wrapper pads each expert's tokens to a uniform M_pad and - writes into the packed ``[G*N_out, K_in]`` gradient; per-expert results land - in ``wgrad_list[g]`` (a view of that packed buffer). No quantization here. - """ - from ..flydsl_kernels.gemm import _run_mxfp8_grouped_wgrad - - num_gemms = len(inputmats) - dy_cols = [g._columnwise_data for g in grad_output] - dy_scales = [g._columnwise_scale_inv for g in grad_output] - x_cols = [x._columnwise_data for x in inputmats] - x_scales = [x._columnwise_scale_inv for x in inputmats] - - n_out = grad_output[0]._columnwise_data.size(-1) # per-expert output feature N - k_in = inputmats[0]._columnwise_data.size(-1) # per-expert input feature K - - # Write straight into the packed wgrad storage when wgrad_list is one contiguous - # run; otherwise launch into scratch and scatter. - dW = _packed_wgrad_view(wgrad_list, num_gemms, n_out, k_in) - scatter = dW is None - if scatter: - dW = torch.empty( - num_gemms * n_out, k_in, dtype=wgrad_list[0].dtype, device=wgrad_list[0].device - ) - - _run_mxfp8_grouped_wgrad( - dy_cols, - dy_scales, - x_cols, - x_scales, - dW, - m_splits=m_splits, - n_out=n_out, - k_in=k_in, - dy_fp8_dtype=grad_output[0]._fp8_dtype, - x_fp8_dtype=inputmats[0]._fp8_dtype, - output_dtype=wgrad_list[0].dtype, - ) - - if scatter: - for g in range(num_gemms): - wgrad_list[g].copy_(dW[g * n_out : (g + 1) * n_out]) - - @staticmethod - def _fwd_flydsl_bf16_grouped(*, inputmats, weights, out, m_splits, m_splits_tensor): - """FlyDSL BF16 grouped forward (TN): Y_e = X_e @ W_e^T, one packed launch. - - Unquantized analog of :meth:`_fwd_flydsl_mxfp8_grouped`: activations are the - per-group BF16 splits and weights the per-expert BF16 tensors, both read in - place through per-expert base pointers. No scales. - """ - from ..flydsl_kernels.gemm import _run_bf16_grouped - from ..flydsl_kernels.gemm.grouped_gemm_utils import expert_ptrs - - # Activations are read in place from each expert's rows (no [M_total, K] concat). - a_data, a_data_ptrs = _flydsl_row_operand([_flydsl_rows(im) for im in inputmats], m_splits) - # Weights are independent per-expert [N_out, K] tensors; the grouped-TN kernel - # reads each from a per-expert base pointer, so pass an int64[G] address array - # instead of concatenating the G weights. b_data is a per-expert representative - # (shape only; the kernel addresses weights via the pointers). - b_data_ptrs = expert_ptrs(list(weights), a_data.device) - b_data = weights[0] - n = weights[0].size(0) - k = weights[0].size(-1) - m_total = sum(int(m) for m in m_splits) - - _run_bf16_grouped( - a_data, - b_data, - out, - m_splits=m_splits, - m_total=m_total, - n=n, - k=k, - output_dtype=out.dtype, - b_data_ptrs=b_data_ptrs, - a_data_ptrs=a_data_ptrs, - ) - - @staticmethod - def _dgrad_flydsl_bf16_grouped(*, grad_output, weights, dgrad, m_splits, m_splits_tensor): - """FlyDSL BF16 grouped dgrad (NN): dX_e = dY_e @ W_e, one packed launch.""" - from ..flydsl_kernels.gemm import _run_bf16_grouped - from ..flydsl_kernels.gemm.grouped_gemm_utils import expert_ptrs - - # dY is read in place from each expert's rows (no [M_total, N_out] concat). - a_data, a_data_ptrs = _flydsl_row_operand([_flydsl_rows(g) for g in grad_output], m_splits) - # Weights are independent per-expert [N_out, K_in] tensors; the grouped-NN - # kernel reads each from a per-expert base pointer (no concat). b_data is a - # per-expert representative. - b_data_ptrs = expert_ptrs(list(weights), a_data.device) - b_data = weights[0] - n_out = grad_output[0].shape[-1] # per-expert contraction (N_out) - k_in = dgrad.size(-1) # output width (K_in) - m_total = sum(int(m) for m in m_splits) - - _run_bf16_grouped( - a_data, - b_data, - dgrad, - m_splits=m_splits, - m_total=m_total, - n=k_in, - k=n_out, - output_dtype=dgrad.dtype, - layout="NN", - b_data_ptrs=b_data_ptrs, - a_data_ptrs=a_data_ptrs, - ) - - @staticmethod - def _wgrad_flydsl_bf16_grouped(*, grad_output, inputmats, wgrad_list, m_splits): - """FlyDSL BF16 grouped wgrad (NT): dW_e = dY_e^T @ X_e, one packed launch.""" - from ..flydsl_kernels.gemm import _run_bf16_grouped_wgrad - - num_gemms = len(inputmats) - dy_cols = [_flydsl_rows(g) for g in grad_output] - x_cols = [_flydsl_rows(x) for x in inputmats] - n_out = dy_cols[0].shape[-1] - k_in = x_cols[0].shape[-1] - - # Write straight into the packed wgrad storage when wgrad_list is one contiguous - # run; otherwise launch into scratch and scatter. - dW = _packed_wgrad_view(wgrad_list, num_gemms, n_out, k_in) - scatter = dW is None - if scatter: - dW = torch.empty( - num_gemms * n_out, k_in, dtype=wgrad_list[0].dtype, device=wgrad_list[0].device - ) - _run_bf16_grouped_wgrad( - dy_cols, - x_cols, - dW, - m_splits=m_splits, - n_out=n_out, - k_in=k_in, - output_dtype=wgrad_list[0].dtype, - ) - if scatter: - for g in range(num_gemms): - wgrad_list[g].copy_(dW[g * n_out : (g + 1) * n_out]) - - @staticmethod - def _fwd_flydsl_fp8_grouped(*, inputmats, weights, out, m_splits, m_splits_tensor): - """FlyDSL tensor-wise FP8 grouped forward (TN): Y_e = X_e @ W_e^T. - - Experts pack weights on the output-N axis ([G*N, K]); each operand carries - one scalar inverse scale, read in place through per-expert pointers. - """ - from ..flydsl_kernels.gemm import _run_fp8_grouped - from ..flydsl_kernels.gemm.grouped_gemm_utils import expert_ptrs - - # Activations are read in place from each expert's payload (no [M_total, K] concat). - a_data, a_data_ptrs = _flydsl_row_operand([im._data for im in inputmats], m_splits) - # Weights are independent per-expert [N, K] FP8 tensors; the grouped-TN kernel - # reads each from a per-expert base pointer, so pass an int64[G] address array - # instead of concatenating the G weights. b_data is a per-expert representative. - b_data_ptrs = expert_ptrs([w._data for w in weights], a_data.device) - b_data = weights[0]._data - a_data._fp8_dtype = inputmats[0]._fp8_dtype - b_data._fp8_dtype = weights[0]._fp8_dtype - # Per-expert scale_inv pointers, read in place by the kernel (no gather). - a_scale = expert_ptrs([im._scale_inv for im in inputmats], a_data.device) - b_scale = expert_ptrs([w._scale_inv for w in weights], a_data.device) - - n = weights[0]._data.size(0) - k = weights[0]._data.size(-1) - m_total = sum(int(m) for m in m_splits) - - _run_fp8_grouped( - a_data, - a_scale, - b_data, - b_scale, - out, - m_splits=m_splits, - m_total=m_total, - n=n, - k=k, - output_dtype=out.dtype, - b_data_ptrs=b_data_ptrs, - a_data_ptrs=a_data_ptrs, - ) - - @staticmethod - def _dgrad_flydsl_fp8_grouped(*, grad_output, weights, dgrad, m_splits, m_splits_tensor): - """FlyDSL tensor-wise FP8 grouped dgrad: dX_e = dY_e @ W_e, via the TN kernel. - - FP8 is TN-only, so dgrad reaches the same TN kernel using materialized - transpose backings: A = dY rowwise ``_data`` [M_total, N_out], B = W^T - columnwise ``_transpose`` [G*N_out... ]. The TN kernel computes A @ B^T with - B packed on the output axis = K_in, i.e. per expert dY @ (W^T)^T = dY @ W. - """ - from ..flydsl_kernels.gemm import _run_fp8_grouped - from ..flydsl_kernels.gemm.grouped_gemm_utils import expert_ptrs - - # A = dY (rowwise fp8 data); B = W columnwise transpose, the materialized - # [K_in, N_out] packed on K_in (the TN GEMM output axis). Both are independent - # per-expert slabs, addressed via per-expert base pointers (no concat). - a_data, a_data_ptrs = _flydsl_row_operand([g._data for g in grad_output], m_splits) - w_transposes = _flydsl_fp8_transposes(weights) - b_data_ptrs = expert_ptrs(w_transposes, a_data.device) - b_data = w_transposes[0] - a_data._fp8_dtype = grad_output[0]._fp8_dtype - b_data._fp8_dtype = weights[0]._fp8_dtype - # Per-expert scale_inv pointers, read in place by the kernel (no gather). - a_scale = expert_ptrs([g._scale_inv for g in grad_output], a_data.device) - b_scale = expert_ptrs([w._scale_inv for w in weights], a_data.device) - - n_out = grad_output[0]._data.size(-1) # per-expert contraction (N_out) - k_in = dgrad.size(-1) # output width (K_in) - m_total = sum(int(m) for m in m_splits) - - _run_fp8_grouped( - a_data, - a_scale, - b_data, - b_scale, - dgrad, - m_splits=m_splits, - m_total=m_total, - n=k_in, - k=n_out, - output_dtype=dgrad.dtype, - b_data_ptrs=b_data_ptrs, - a_data_ptrs=a_data_ptrs, - ) - - @staticmethod - def _wgrad_flydsl_fp8_grouped(*, grad_output, inputmats, wgrad_list, m_splits): - """FlyDSL tensor-wise FP8 grouped wgrad: dW_e = dY_e^T @ X_e, via the TN kernel. - - Both operands reach the TN kernel already transposed: A = dY^T (the fp8 - ``_transpose`` backing, [N_out, m_e]), B = X^T ([K_in, m_e]). The kernel - computes A @ B^T = dY^T @ X = dW. - """ - from ..flydsl_kernels.gemm import _run_fp8_grouped_wgrad - - num_gemms = len(inputmats) - dyt_cols = _flydsl_fp8_transposes(grad_output) - xt_cols = _flydsl_fp8_transposes(inputmats) - # Per-expert scale_inv tensors; the wgrad runner reads them in place by pointer. - dy_scale = [g._scale_inv for g in grad_output] - x_scale = [x._scale_inv for x in inputmats] - - # Operands may carry only the columnwise (_transpose) backing. The - # transpose is [feature, m_e]: dY^T is [N_out, m_e], X^T is [K_in, m_e]. - n_out = dyt_cols[0].size(0) # per-expert output feature N - k_in = xt_cols[0].size(0) # per-expert input feature K - - # wgrad_list[g] are normally adjacent [n_out, k_in] slices of one packed - # [num_gemms, n_out, k_in] buffer (the non-fused ``wgrad_packed`` path). When - # that holds, point the kernel's [num_gemms*n_out, k_in] output straight at - # that storage so it writes each dW_e in place -- eliding the per-expert - # copy-out scatter. Fall back to a fresh buffer + scatter only when the - # destinations are not one contiguous run (e.g. fused main_grad accumulation). - dW = _packed_wgrad_view(wgrad_list, num_gemms, n_out, k_in) - scatter = dW is None - if scatter: - dW = torch.empty( - num_gemms * n_out, k_in, dtype=wgrad_list[0].dtype, device=wgrad_list[0].device - ) - _run_fp8_grouped_wgrad( - dyt_cols, - dy_scale, - xt_cols, - x_scale, - dW, - m_splits=m_splits, - n_out=n_out, - k_in=k_in, - dy_fp8_dtype=grad_output[0]._fp8_dtype, - x_fp8_dtype=inputmats[0]._fp8_dtype, - output_dtype=wgrad_list[0].dtype, - ) - if scatter: - for g in range(num_gemms): - wgrad_list[g].copy_(dW[g * n_out : (g + 1) * n_out]) - - @classmethod - def _flydsl_grouped_dispatch(cls, phase, kind): - """Return the FlyDSL grouped helper for ``phase`` and dtype ``kind``. - - ``phase`` is 'fwd' | 'dgrad' | 'wgrad' and ``kind`` is the value from - :meth:`_flydsl_grouped_kind` ('mxfp8' | 'fp8' | 'bf16'). All helpers of a - phase share a keyword signature, so the caller passes one kwarg set. - """ - table = { - "fwd": { - "mxfp8": cls._fwd_flydsl_mxfp8_grouped, - "fp8": cls._fwd_flydsl_fp8_grouped, - "bf16": cls._fwd_flydsl_bf16_grouped, - }, - "dgrad": { - "mxfp8": cls._dgrad_flydsl_mxfp8_grouped, - "fp8": cls._dgrad_flydsl_fp8_grouped, - "bf16": cls._dgrad_flydsl_bf16_grouped, - }, - "wgrad": { - "mxfp8": cls._wgrad_flydsl_mxfp8_grouped, - "fp8": cls._wgrad_flydsl_fp8_grouped, - "bf16": cls._wgrad_flydsl_bf16_grouped, - }, - } - return table[phase][kind] - - @classmethod - def _try_flydsl_grouped(cls, phase, kind, **kwargs): - """Run the FlyDSL grouped helper for ``phase``/``kind``; report if it ran. - - Returns True when FlyDSL handled the GEMM. A ``FlyDSLUnsupportedError`` - (e.g. a per-expert dimension the kernel tiling cannot cover) means the - config is unsupported, so the helper returns False and the caller runs the - default grouped GEMM instead -- mirroring the dense GEMM fallback in - ``cpp_extensions.gemm``. ``NVTE_FLYDSL_GEMM_WARN_FALLBACK`` warns on fallback. - """ - from ..flydsl_kernels.gemm import FlyDSLUnsupportedError - - try: - cls._flydsl_grouped_dispatch(phase, kind)(**kwargs) - return True - except FlyDSLUnsupportedError as exc: - warn_fallback = os.environ.get("NVTE_FLYDSL_GEMM_WARN_FALLBACK", "0").lower() not in ( - "", - "0", - "false", - "no", - "off", - ) - if warn_fallback: - warnings.warn( - "[FLYDSL WARNING]: FlyDSL grouped GEMM does not support this " - f"configuration (kind={kind}, phase={phase}); falling back to the " - f"default backend. Reason: {exc}", - UserWarning, - stacklevel=2, - ) - return False - @staticmethod def _forward_blockwise_fp8_triton( ctx, @@ -1772,13 +1206,12 @@ def forward( pow2_grad=blockwise_recipe.fp8_quant_bwd_grad.power_2_scale, ) - # FlyDSL grouped GEMM (ROCm gfx950) opt-in, covering MXFP8, tensor-wise FP8, - # and BF16/FP16. It does not quantize -- it consumes TE's existing quantized - # operands and is dispatched inline at the forward GEMM (see below). Returns - # the dtype kind ('mxfp8' | 'fp8' | 'bf16') or None when unsupported. - flydsl_grouped_kind = _GroupedLinear._flydsl_grouped_kind( + # FlyDSL grouped GEMM (ROCm gfx950) opt-in: a drop-in for general_grouped_gemm + # over TE's already-quantized operands. + use_grouped_gemm_flydsl = _GroupedLinear._is_flydsl_grouped_gemm_supported( fp8=fp8, recipe=blockwise_recipe, + activation_dtype=activation_dtype, use_bias=use_bias, use_grouped_gemm_triton=use_grouped_gemm_triton, backward_override=backward_override, @@ -1790,14 +1223,6 @@ def forward( in_features=weights[0].size(-1), out_features=weights[0].size(0), ) - # The half-precision kernel handles bf16 and fp16. - # Drop the kind for any other (e.g. fp32) activation dtype. - if flydsl_grouped_kind == "bf16" and activation_dtype not in ( - torch.bfloat16, - torch.float16, - ): - flydsl_grouped_kind = None - use_grouped_gemm_flydsl = flydsl_grouped_kind is not None backward_needs_input = is_grad_enabled and weight_requires_grad if backward_override is None and save_original_input and backward_needs_input: @@ -1947,19 +1372,12 @@ def forward( # tensors (like scales), but bulk allocation shares storage across all tensors, # so if scales can't be offloaded, nothing in the group can be offloaded. fused_padding_kwargs = {} - if ( - actual_m_splits is not None - and IS_HIP_EXTENSION - and inp_view.shape[0] == sum(actual_m_splits) - ): + if actual_m_splits is not None and IS_HIP_EXTENSION \ + and inp_view.shape[0] == sum(actual_m_splits): fused_padding_kwargs["valid_split_sections"] = actual_m_splits inputmats = tex.split_quantize( - inp_view, - m_splits, - input_quantizers, - disable_bulk_allocation=cpu_offloading, - **fused_padding_kwargs, - ) + inp_view, m_splits, input_quantizers, + disable_bulk_allocation=cpu_offloading, **fused_padding_kwargs) elif debug: inputmats = DebugQuantizer.multi_tensor_quantize( inp_view, input_quantizers, m_splits, activation_dtype @@ -2027,35 +1445,30 @@ def forward( use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator # Perform GEMM - ran_flydsl_fwd = use_grouped_gemm_flydsl and _GroupedLinear._try_flydsl_grouped( - "fwd", - flydsl_grouped_kind, - inputmats=inputmats, - weights=weights_fp8, - out=out, + if use_grouped_gemm_triton: + general_grouped_gemm_func = general_grouped_gemm_triton + kwargs = {"m_splits_tensor": m_splits_tensor} + elif use_grouped_gemm_flydsl: + from ..flydsl_kernels.gemm import general_grouped_gemm_flydsl + + general_grouped_gemm_func = general_grouped_gemm_flydsl + kwargs = {} + else: + general_grouped_gemm_func = general_grouped_gemm + kwargs = {} + general_grouped_gemm_func( + weights_fp8, + inputmats, + [out], + output_quantizers, + activation_dtype, + single_output=True, m_splits=m_splits, - m_splits_tensor=m_splits_tensor, + bias=biases, + use_bias=use_bias, + use_split_accumulator=use_split_accumulator, + **kwargs, ) - if not ran_flydsl_fwd: - if use_grouped_gemm_triton: - general_grouped_gemm_func = general_grouped_gemm_triton - kwargs = {"m_splits_tensor": m_splits_tensor} - else: - general_grouped_gemm_func = general_grouped_gemm - kwargs = {} - general_grouped_gemm_func( - weights_fp8, - inputmats, - [out], - output_quantizers, - activation_dtype, - single_output=True, - m_splits=m_splits, - bias=biases, - use_bias=use_bias, - use_split_accumulator=use_split_accumulator, - **kwargs, - ) output_unpadded = False if ( @@ -2083,10 +1496,6 @@ def forward( if is_grad_enabled: ctx.grouped_tensor_supported = False - # FlyDSL grouped forward routes dgrad through the grouped NN path; the - # shared backward branches on this kind (wgrad still uses the default - # where a dtype/feature is unsupported). - ctx.flydsl_grouped_kind = flydsl_grouped_kind ctx.weight_quantizers = weight_quantizers ctx.weights_shape_1 = weights[0].shape[1] @@ -2181,6 +1590,7 @@ def forward( ctx.save_original_input = save_original_input ctx.input_quantizers = input_quantizers ctx.use_grouped_gemm_triton = use_grouped_gemm_triton + ctx.use_grouped_gemm_flydsl = use_grouped_gemm_flydsl ctx.num_input_tensors = len(inputmats) ctx.dgrad_out = dgrad_out @@ -2646,56 +2056,39 @@ def backward( rowwise_usage=quantizer.rowwise_usage, columnwise_usage=quantizer.columnwise_usage, ) - _flydsl_kind = getattr(ctx, "flydsl_grouped_kind", None) - # dgrad dX = dY @ W through the FlyDSL grouped NN kernel, riding off - # the already-quantized grad_output (rowwise) and weights (columnwise, - # ensured above). Falls back to the default backend when FlyDSL does - # not support this config. - ran_flydsl_dgrad = _flydsl_kind is not None and _GroupedLinear._try_flydsl_grouped( - "dgrad", - _flydsl_kind, - grad_output=grad_output, - weights=weights_for_dgrad, - dgrad=dgrad, + if ctx.use_grouped_gemm_triton: + general_grouped_gemm_func = general_grouped_gemm_triton + kwargs = {"m_splits_tensor": ctx.m_splits_tensor} + elif ctx.use_grouped_gemm_flydsl: + from ..flydsl_kernels.gemm import general_grouped_gemm_flydsl + + general_grouped_gemm_func = general_grouped_gemm_flydsl + kwargs = {} + else: + general_grouped_gemm_func = general_grouped_gemm + kwargs = {} + general_grouped_gemm_func( + weights_for_dgrad, + grad_output, + [dgrad], + ctx.grad_input_quantizers, + ctx.activation_dtype, + single_output=True, + layout="NN", m_splits=ctx.m_splits, - m_splits_tensor=ctx.m_splits_tensor, + grad=True, + use_split_accumulator=dgrad_gemm_use_split_accumulator, + **kwargs, ) - if not ran_flydsl_dgrad: - if ctx.use_grouped_gemm_triton: - general_grouped_gemm_func = general_grouped_gemm_triton - kwargs = {"m_splits_tensor": ctx.m_splits_tensor} - else: - general_grouped_gemm_func = general_grouped_gemm - kwargs = {} - general_grouped_gemm_func( - weights_for_dgrad, - grad_output, - [dgrad], - ctx.grad_input_quantizers, - ctx.activation_dtype, - single_output=True, - layout="NN", - m_splits=ctx.m_splits, - grad=True, - use_split_accumulator=dgrad_gemm_use_split_accumulator, - **kwargs, - ) - if ( - ctx.actual_m_splits is not None - and ctx.actual_m_splits != ctx.m_splits - and not ctx.output_unpadded - ): + if ctx.actual_m_splits is not None and ctx.actual_m_splits != ctx.m_splits \ + and not ctx.output_unpadded: dgrad_unpadded = torch.empty( (sum(ctx.actual_m_splits), dgrad.shape[-1]), - dtype=dgrad.dtype, - device=dgrad.device, + dtype=dgrad.dtype, device=dgrad.device, ) tex.fused_multi_row_unpadding( - dgrad, - dgrad_unpadded, - ctx.m_splits, - ctx.actual_m_splits, + dgrad, dgrad_unpadded, ctx.m_splits, ctx.actual_m_splits, ) dgrad = dgrad_unpadded @@ -2797,67 +2190,54 @@ def backward( inputmats_dequant.append(cast_if_needed(inputmat, ctx.activation_dtype)) inputmats = inputmats_dequant - # FlyDSL grouped wgrad (NT) opt-in, plain path only; fused wgrad - # accumulation, bias grad, and delayed wgrad fall through to the - # default grouped GEMM below. - _flydsl_kind = getattr(ctx, "flydsl_grouped_kind", None) - use_flydsl_wgrad = ( - _flydsl_kind is not None + if ctx.use_grouped_gemm_triton: + general_grouped_gemm_func = general_grouped_gemm_triton + kwargs = {"m_splits_tensor": ctx.m_splits_tensor} + elif ( + ctx.use_grouped_gemm_flydsl and not ctx.fuse_wgrad_accumulation - and not ctx.use_bias # bias-grad-in-wgrad-GEMM not supported yet and not (ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute()) - ) - ran_flydsl_wgrad = use_flydsl_wgrad and _GroupedLinear._try_flydsl_grouped( - "wgrad", - _flydsl_kind, - grad_output=grad_output, - inputmats=inputmats, - wgrad_list=wgrad_list, + ): + # FlyDSL wgrad writes plain per-expert gradients; fused main_grad + # accumulation and delayed wgrad stay on the default grouped GEMM. + from ..flydsl_kernels.gemm import general_grouped_gemm_flydsl + + general_grouped_gemm_func = general_grouped_gemm_flydsl + kwargs = {} + else: + general_grouped_gemm_func = general_grouped_gemm + kwargs = {} + grouped_gemm_wgrad = functools.partial( + general_grouped_gemm_func, + quantization_params=ctx.grad_weight_quantizers, + out_dtype=ctx.activation_dtype, + layout="NT", + grad=True, m_splits=ctx.m_splits, + use_bias=ctx.use_bias if grad_biases[0] is None else None, + bias=biases, + use_split_accumulator=wgrad_gemm_use_split_accumulator, + accumulate=( + accumulate_wgrad_into_param_main_grad + if not is_dist_weight + and not getattr(ctx, "origin_weights_overwrite_main_grad", False) + else False + ), + **kwargs, ) - if ran_flydsl_wgrad: - clear_tensor_data(*inputmats) - grouped_gemm_wgrad = None + # WGRAD + if ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute(): + ctx.wgrad_store.put([inputmats, grad_output, wgrad_list], grouped_gemm_wgrad) else: - if ctx.use_grouped_gemm_triton: - general_grouped_gemm_func = general_grouped_gemm_triton - kwargs = {"m_splits_tensor": ctx.m_splits_tensor} - else: - general_grouped_gemm_func = general_grouped_gemm - kwargs = {} - grouped_gemm_wgrad = functools.partial( - general_grouped_gemm_func, - quantization_params=ctx.grad_weight_quantizers, - out_dtype=ctx.activation_dtype, - layout="NT", - grad=True, - m_splits=ctx.m_splits, - use_bias=ctx.use_bias if grad_biases[0] is None else None, - bias=biases, - use_split_accumulator=wgrad_gemm_use_split_accumulator, - accumulate=( - accumulate_wgrad_into_param_main_grad - if not is_dist_weight - and not getattr(ctx, "origin_weights_overwrite_main_grad", False) - else False - ), - **kwargs, - ) - # WGRAD - if ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute(): - ctx.wgrad_store.put( - [inputmats, grad_output, wgrad_list], grouped_gemm_wgrad - ) - else: - _, grad_biases_, _ = grouped_gemm_wgrad(inputmats, grad_output, wgrad_list) + _, grad_biases_, _ = grouped_gemm_wgrad(inputmats, grad_output, wgrad_list) - for i in range(ctx.num_gemms): - if grad_biases[i] is None: - grad_biases[i] = grad_biases_[i] - del grad_biases_ + for i in range(ctx.num_gemms): + if grad_biases[i] is None: + grad_biases[i] = grad_biases_[i] + del grad_biases_ - # Deallocate input tensor - clear_tensor_data(*inputmats) + # Deallocate input tensor + clear_tensor_data(*inputmats) def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): if ctx.weights_requires_grad: diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 009b771d9b..68a8085ad1 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -331,6 +331,22 @@ def get_default_recipe() -> Recipe: return get_default_fp8_recipe() +def is_flydsl_grouped_gemm_enabled(recipe: Optional[Recipe] = None) -> bool: + """Whether the FlyDSL grouped GEMM (ROCm gfx950 opt-in) is enabled for ``recipe``. + + Covers the opt-in and the precisions the kernels implement: MXFP8 and tensor-wise + FP8 (current/delayed scaling) when ``recipe`` is given, else unquantized BF16/FP16. + GroupedLinear adds its per-call checks on top of this. + """ + if not ( + IS_HIP_EXTENSION + and os.getenv("NVTE_USE_GROUPED_GEMM_FLYDSL", "0") == "1" + and get_device_compute_capability() == (9, 5) + ): + return False + return recipe is None or recipe.mxfp8() or recipe.float8_current_scaling() or recipe.delayed() + + def get_align_size_for_quantization(recipe: Recipe) -> int: """Get the alignment used to pad grouped quantized operations. @@ -346,12 +362,7 @@ def get_align_size_for_quantization(recipe: Recipe) -> int: # The FlyDSL grouped GEMM (gfx950 opt-in) tiles every expert with BLOCK_M=256, so # per-expert token counts must be 256-aligned for the recipes it handles (MXFP8 and # tensor-wise FP8); other recipes fall back and keep their own alignment. - if ( - IS_HIP_EXTENSION - and get_device_compute_capability() == (9, 5) - and os.environ.get("NVTE_USE_GROUPED_GEMM_FLYDSL", "0") == "1" - and (recipe.mxfp8() or recipe.float8_current_scaling() or recipe.delayed()) - ): + if is_flydsl_grouped_gemm_enabled(recipe): return 256 if recipe.mxfp8(): # HipKittens grouped GEMM requires 256-aligned expert dimensions.