diff --git a/benchmarks/microbenchmarks/benchmark_grouped_gemm.py b/benchmarks/microbenchmarks/benchmark_grouped_gemm.py index a576568a0..249b64ae6 100755 --- a/benchmarks/microbenchmarks/benchmark_grouped_gemm.py +++ b/benchmarks/microbenchmarks/benchmark_grouped_gemm.py @@ -79,6 +79,13 @@ def generate_deepseekv3_test_cases(): ) +def generate_qwen3_235b_test_cases(): + # Qwen3-235B-A22B: 128 routed experts, top-8, moe_intermediate_size 1536, hidden 4096. + return _generate_moe_test_cases( + "Qwen3-235B", n_routed_experts=128, moe_intermediate_size=1536, hidden_size=4096 + ) + + def generate_deepseekv2_test_cases(): return _generate_moe_test_cases( "DSV2", n_routed_experts=160, moe_intermediate_size=1536, hidden_size=5120 diff --git a/benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py b/benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py new file mode 100644 index 000000000..b26c109c8 --- /dev/null +++ b/benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py @@ -0,0 +1,545 @@ +#!/usr/bin/env python +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +"""Compare permute-free gather-GEMM vs permute + grouped GEMM backends on ROCm. + +Backends: + - permute_free: TE permute-free (FlyDSL) route-list gather-in-GEMM (align + GEMM), no permute + - hipblaslt: TE multistream hipBLASLt grouped GEMM after moe_permute + - ck: TE CK grouped GEMM after moe_permute (NVTE_USE_CUTLASS_GROUPED_GEMM=1) + - triton: AITER Triton grouped GEMM after moe_permute + +Run from this directory:: + + python benchmark_perm_free_grouped_gemm.py + python benchmark_perm_free_grouped_gemm.py --quick --csv +""" + +from __future__ import annotations + +import os +from typing import Callable, Dict, List, Tuple + +import torch +from torch.utils.cpp_extension import IS_HIP_EXTENSION + +from benchmark_grouped_gemm import ( + generate_deepseekv2_lite_test_cases, + generate_deepseekv2_test_cases, + generate_deepseekv3_test_cases, + generate_grok_v2_test_cases, + generate_qwen3_235b_test_cases, +) +from utils import compute_tflops, make_metric_record, make_parser, run_benchmarks, time_func + +DEFAULT_TOPK = 8 +BACKENDS = ("permute_free", "hipblaslt", "ck", "triton") +# ``permute_free_act`` is an opt-in, GateUP-only backend: the permute-free FC1 gather-GEMM +# with the gated SiLU activation (``silu(gate) * up``) fused into the epilogue. It is not in +# the default sweep (it only applies to gate+up shapes); request it via ``--backends``. +KNOWN_BACKENDS = BACKENDS + ("permute_free_act",) + +# Training phases benchmarked (and reported) separately in --train mode. +PHASES = ("fwd", "dgrad", "wgrad") + +# Short backend labels for the per-phase --train metric names. +_BACKEND_SHORT = { + "permute_free": "PermuteFree", + "permute_free_act": "PermuteFreeAct", + "hipblaslt": "hipBLASLt", + "ck": "CK", + "triton": "Triton", +} + + +def _require_rocm(): + if not IS_HIP_EXTENSION or not torch.cuda.is_available(): + raise RuntimeError("This benchmark requires ROCm (HIP extension) and a CUDA device.") + + +def _make_routing(num_tokens: int, num_experts: int, topk: int, device: str, seed: int): + from transformer_engine.pytorch.moe import MoERoutingMetadata + + gen = torch.Generator(device=device) + gen.manual_seed(seed) + logits = torch.randn(num_tokens, num_experts, device=device, generator=gen) + probs = torch.softmax(logits, dim=-1) + topk_weights, topk_ids = torch.topk(probs, k=topk, dim=-1) + topk_ids = topk_ids.to(torch.int32) + topk_weights = topk_weights.to(torch.float32) + # MoERoutingMetadata now takes a boolean routing_map as its primary input; build it + # from the sampled topk_ids so the permute-free and permute backends see identical + # routing. Passing ``topk`` tightens the block-padded over-allocation (em_max). topk_ids/ + # topk_weights are attached for the permute (moe_permute) backends. + routing_map = torch.zeros(num_tokens, num_experts, dtype=torch.bool, device=device) + routing_map.scatter_(1, topk_ids.to(torch.long), True) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts, topk=topk) + routing.topk_ids = topk_ids + routing.topk_weights = topk_weights + return routing + + +def _m_splits_from_topk(topk_ids: torch.Tensor, num_experts: int) -> List[int]: + counts = torch.bincount(topk_ids.reshape(-1).long(), minlength=num_experts) + return [int(v) for v in counts.tolist()] + + +def _permute_hidden( + hidden: torch.Tensor, topk_ids: torch.Tensor, num_tokens: int, topk: int +) -> torch.Tensor: + from transformer_engine.pytorch import moe_permute + + num_out_tokens = num_tokens * topk + permuted, _ = moe_permute(hidden, topk_ids, num_out_tokens, map_type="index") + return permuted + + +def _grouped_gemm_hip( + permuted: torch.Tensor, + weights: torch.Tensor, + m_splits: List[int], + *, + use_ck: bool, +) -> torch.Tensor: + from transformer_engine.pytorch.cpp_extensions import general_grouped_gemm + + prev = os.environ.get("NVTE_USE_CUTLASS_GROUPED_GEMM") + os.environ["NVTE_USE_CUTLASS_GROUPED_GEMM"] = "1" if use_ck else "0" + try: + b = len(m_splits) + n = int(weights.shape[1]) + sum_m = sum(m_splits) + xs = list(torch.split(permuted.view(sum_m, -1), m_splits)) + weight_list = [weights[i] for i in range(b)] + out = torch.empty((sum_m, n), device=permuted.device, dtype=permuted.dtype) + general_grouped_gemm( + A=weight_list, + B=xs, + out=[out], + quantization_params=[None] * b, + out_dtype=permuted.dtype, + single_output=True, + m_splits=m_splits, + use_bias=False, + bias=None, + layout="TN", + ) + return out + finally: + if prev is None: + os.environ.pop("NVTE_USE_CUTLASS_GROUPED_GEMM", None) + else: + os.environ["NVTE_USE_CUTLASS_GROUPED_GEMM"] = prev + + +def _grouped_gemm_triton( + permuted: torch.Tensor, + weights: torch.Tensor, + m_splits: List[int], +) -> torch.Tensor: + from transformer_engine.pytorch.triton_kernels.grouped_gemm import general_grouped_gemm_triton + + b = len(m_splits) + n = int(weights.shape[1]) + sum_m = sum(m_splits) + xs = list(torch.split(permuted.view(sum_m, -1), m_splits)) + weight_list = [weights[i] for i in range(b)] + out = torch.empty((sum_m, n), device=permuted.device, dtype=permuted.dtype) + general_grouped_gemm_triton( + A=weight_list, + B=xs, + out=[out], + quantization_params=[None] * b, + out_dtype=permuted.dtype, + single_output=True, + m_splits=m_splits, + use_bias=False, + bias=None, + layout="TN", + ) + return out + + +def _permute_free_gemm( + hidden: torch.Tensor, + weights: torch.Tensor, + routing, +) -> torch.Tensor: + from transformer_engine.pytorch.moe import ( + permute_free_grouped_gemm_bf16, + ) + + return permute_free_grouped_gemm_bf16(hidden, weights, routing) + + +def _permute_free_act_gemm( + hidden: torch.Tensor, + weights: torch.Tensor, + routing, +) -> torch.Tensor: + """Permute-free FC1 gather-GEMM + standalone SiLU gated activation (no in-kernel fusion). + + ``weights`` is the gate+up projection ``[E, 2F, K]``; returns the F-wide activated buffer. + """ + from transformer_engine.pytorch.moe import ( + permute_free_gated_act_fwd, + permute_free_grouped_gemm_bf16, + ) + + preact = permute_free_grouped_gemm_bf16(hidden, weights, routing) + return permute_free_gated_act_fwd(preact, routing, activation="silu") + + +def _build_backend_fns( + hidden: torch.Tensor, + weights: torch.Tensor, + routing, + m_splits: List[int], + num_tokens: int, + topk: int, + is_gated: bool = False, +) -> Dict[str, Callable[[], torch.Tensor]]: + fns: Dict[str, Callable[[], torch.Tensor]] = {} + + def hipblaslt_fn(): + permuted = _permute_hidden(hidden, routing.topk_ids, num_tokens, topk) + return _grouped_gemm_hip(permuted, weights, m_splits, use_ck=False) + + def ck_fn(): + permuted = _permute_hidden(hidden, routing.topk_ids, num_tokens, topk) + return _grouped_gemm_hip(permuted, weights, m_splits, use_ck=True) + + def triton_fn(): + permuted = _permute_hidden(hidden, routing.topk_ids, num_tokens, topk) + return _grouped_gemm_triton(permuted, weights, m_splits) + + fns["hipblaslt"] = hipblaslt_fn + fns["ck"] = ck_fn + fns["triton"] = triton_fn + fns["permute_free"] = lambda: _permute_free_gemm(hidden, weights, routing) + # GateUP-only: permute-free FC1 GEMM + standalone SiLU gated activation. + if is_gated: + fns["permute_free_act"] = lambda: _permute_free_act_gemm(hidden, weights, routing) + return fns + + +def _build_train_phase_fns( + hidden: torch.Tensor, + weights: torch.Tensor, + routing, + m_splits: List[int], + num_tokens: int, + topk: int, + dtype: torch.dtype, +) -> Dict[str, Dict[str, Callable[[], torch.Tensor]]]: + """Per-phase (fwd / dgrad / wgrad) closures for each backend. + + Returns ``{backend: {phase: fn}}`` so each phase can be timed independently. + + - permute_free: gather fwd + gather dgrad + fused wgrad (FlyDSL / Triton), matching the + ``_GroupedLinear`` permute-free autograd path. Every kernel gathers from token-major + tensors, so no standalone permute/unpermute is charged. + - hipblaslt/ck: the traditional permuted grouped GEMMs. The token permute is charged to + ``fwd`` and the dgrad unpermute to ``dgrad`` (both produce a token-major-equivalent + result); ``wgrad`` reuses the operands already permuted in fwd, so it is GEMM-only -- + mirroring a real training iteration where the permute is paid once. (Triton GG is + excluded from --train: its grouped backend does not expose the NN/NT grad layouts.) + """ + from transformer_engine.pytorch import moe_permute, moe_unpermute + from transformer_engine.pytorch.cpp_extensions import general_grouped_gemm + from transformer_engine.pytorch.moe import ( + permute_free_grouped_gemm_bf16, + permute_free_grouped_gemm_bf16_dgrad, + permute_free_grouped_gemm_bf16_wgrad, + ) + + device = hidden.device + b = len(m_splits) + n = int(weights.shape[1]) + k = int(weights.shape[2]) + total_m = sum(m_splits) + weight_list = [weights[i] for i in range(b)] + weights_shape = (b, n, k) + + # Run the permute-free forward once (outside timing) to finalize the block-padded align on + # ``routing`` (the fwd enforces a v3 block-size floor, which fixes em_max) and to learn the + # padded slot extent. The dgrad/wgrad now consume the block-padded ``[em_max, out_features]`` + # slot gradient -- i.e. the gradient of this fwd output -- not a dense route-ordered + # ``[num_routes]`` one. + em_max = int(permute_free_grouped_gemm_bf16(hidden, weights, routing).shape[0]) + + # Upstream gradients. ``grad_out_pf`` is the block-padded slot grad for the permute-free + # dgrad/wgrad; the traditional path keeps its expert-sorted ``[total_m, n]`` grad. + # Allocated once, outside the timed region. + grad_out_pf = torch.randn(em_max, n, dtype=dtype, device=device) + grad_out_perm = torch.randn(total_m, n, dtype=dtype, device=device) + grad_splits = list(torch.split(grad_out_perm, m_splits)) + + # Pre-permuted activations reused by the traditional dgrad/wgrad phases. The permute + # itself is (re)charged inside the traditional fwd phase below. + permuted_hidden, row_id_map = moe_permute( + hidden, routing.topk_ids, total_m, map_type="index" + ) + xs_perm = list(torch.split(permuted_hidden.view(total_m, -1), m_splits)) + + # --- permute-free phases (gather-in-GEMM; token-major in/out) --- + def pf_fwd(): + return permute_free_grouped_gemm_bf16(hidden, weights, routing) + + def pf_dgrad(): + return permute_free_grouped_gemm_bf16_dgrad(grad_out_pf, weights, routing) + + def pf_wgrad(): + return permute_free_grouped_gemm_bf16_wgrad( + hidden, grad_out_pf, weights_shape, routing + ) + + # --- traditional (permute + grouped GEMM) phases --- + def _with_ck_env(use_ck: bool, fn: Callable[[], torch.Tensor]): + def run(): + prev = os.environ.get("NVTE_USE_CUTLASS_GROUPED_GEMM") + os.environ["NVTE_USE_CUTLASS_GROUPED_GEMM"] = "1" if use_ck else "0" + try: + return fn() + finally: + if prev is None: + os.environ.pop("NVTE_USE_CUTLASS_GROUPED_GEMM", None) + else: + os.environ["NVTE_USE_CUTLASS_GROUPED_GEMM"] = prev + + return run + + def _trad_fwd(): + permuted, _ = moe_permute(hidden, routing.topk_ids, total_m, map_type="index") + xs = list(torch.split(permuted.view(total_m, -1), m_splits)) + out = torch.empty((total_m, n), device=device, dtype=dtype) + general_grouped_gemm( + A=weight_list, B=xs, out=[out], quantization_params=[None] * b, + out_dtype=dtype, single_output=True, m_splits=m_splits, + use_bias=False, bias=None, layout="TN", + ) + return out + + def _trad_dgrad(): + dx_buf = torch.empty((total_m, k), device=device, dtype=dtype) + dxs = list(torch.split(dx_buf, m_splits)) + general_grouped_gemm( + A=weight_list, B=grad_splits, out=dxs, quantization_params=[None] * b, + out_dtype=dtype, single_output=False, m_splits=m_splits, grad=False, + use_bias=False, bias=None, layout="NN", + ) + return moe_unpermute(dx_buf, row_id_map, map_type="index") + + def _trad_wgrad(): + dw = torch.empty((b, n, k), device=device, dtype=dtype) + dws = [dw[i] for i in range(b)] + general_grouped_gemm( + A=xs_perm, B=grad_splits, out=dws, quantization_params=[None] * b, + out_dtype=dtype, single_output=False, m_splits=m_splits, grad=False, + use_bias=False, bias=None, layout="NT", + ) + return dw + + def _trad_phase_fns(use_ck: bool) -> Dict[str, Callable[[], torch.Tensor]]: + return { + "fwd": _with_ck_env(use_ck, _trad_fwd), + "dgrad": _with_ck_env(use_ck, _trad_dgrad), + "wgrad": _with_ck_env(use_ck, _trad_wgrad), + } + + return { + "permute_free": {"fwd": pf_fwd, "dgrad": pf_dgrad, "wgrad": pf_wgrad}, + "hipblaslt": _trad_phase_fns(use_ck=False), + "ck": _trad_phase_fns(use_ck=True), + } + + +def bench_moe_grouped_gemm_backends( + Case: str, + B: int, + M: int, + N: int, + K: int, + dtype: torch.dtype, + topk: int = DEFAULT_TOPK, + seed: int = 0, + backends: Tuple[str, ...] = BACKENDS, + train: bool = False, +): + _require_rocm() + device = "cuda" + num_tokens = M + num_experts = B + + effective_topk = min(topk, num_experts) + if effective_topk < 1: + raise ValueError(f"Need at least one expert (B={num_experts}).") + + hidden = torch.randn(num_tokens, K, dtype=dtype, device=device) + weights = torch.randn(num_experts, N, K, dtype=dtype, device=device) + routing = _make_routing(num_tokens, num_experts, effective_topk, device, seed) + m_splits = _m_splits_from_topk(routing.topk_ids, num_experts) + + total_m = num_tokens * effective_topk + + if train: + # Per-phase analysis: time fwd / dgrad / wgrad independently for each backend. + phase_fns = _build_train_phase_fns( + hidden, weights, routing, m_splits, num_tokens, effective_topk, dtype + ) + # Each of fwd / dgrad / wgrad performs a ~2*total_m*N*K FLOP GEMM. + phase_flops = 2 * total_m * N * K + + # Warmup every selected phase (also builds the permute-free align buffers so the + # dgrad/wgrad timings reflect cache reuse, as in a real iteration). + for name in backends: + if name not in phase_fns: + continue + for phase in PHASES: + phase_fns[name][phase]() + torch.cuda.synchronize() + + # Emit records phase-major so the printout groups fwd, then dgrad, then wgrad. + records = [] + for phase in PHASES: + for name in backends: + if name not in phase_fns: + continue + ms, measurement = time_func(phase_fns[name][phase]) + records.append( + make_metric_record( + f"{phase} \u00b7 {_BACKEND_SHORT[name]}", + ms, + "TFLOPS", + compute_tflops(phase_flops, ms), + measurement=measurement, + ) + ) + return records + + # GateUP shapes carry a gate+up projection (out_features == 2F, even), which is the only + # case where the fused gated-activation permute-free backend applies. + is_gated = Case.endswith("GateUP") + backend_fns = _build_backend_fns( + hidden, weights, routing, m_splits, num_tokens, effective_topk, is_gated=is_gated + ) + fwd_flops = 2 * total_m * N * K + + # Permute-free backends emit a padded/route-major (or F-wide, for the activation fusion) + # buffer, so their output shape differs from the permuted [total_m, N] result. + _permute_free_names = ("permute_free", "permute_free_act") + + # Warmup + correctness spot-check (permute-free vs hipblaslt layout differs; skip allclose) + for name in backends: + if name not in backend_fns: + continue + out = backend_fns[name]() + torch.cuda.synchronize() + if name not in _permute_free_names: + assert out.shape == (total_m, N), f"{name}: bad shape {out.shape}" + + records = [] + for name in backends: + if name not in backend_fns: + continue + label = { + "permute_free": "PermuteFree Gather-GEMM", + "permute_free_act": "PermuteFree Gather-GEMM + SiLU", + "hipblaslt": "Permute+hipBLASLt Grouped GEMM", + "ck": "Permute+CK Grouped GEMM", + "triton": "Permute+AITER Triton Grouped GEMM", + }[name] + ms, measurement = time_func(backend_fns[name]) + records.append( + make_metric_record( + label, + ms, + "TFLOPS", + compute_tflops(fwd_flops, ms), + measurement=measurement, + ) + ) + return records + + +def _filter_test_cases(test_cases, quick: bool): + if not quick: + return test_cases + allowed_m = {512, 1024} + filtered = [c for c in test_cases if c["M"] in allowed_m and c["Case"].endswith("-Down")] + return filtered[:4] if filtered else test_cases[:2] + + +def main(): + base = make_parser(description="Benchmark MoE grouped GEMM backends on ROCm.") + base.add_argument("--quick", action="store_true", help="Run a small subset of shapes.") + base.add_argument( + "--train", + action="store_true", + help="Benchmark a full training iteration (fwd + dgrad + wgrad) instead of fwd only.", + ) + base.add_argument("--topk", type=int, default=DEFAULT_TOPK, help="MoE top-k routing.") + base.add_argument( + "--backends", + type=str, + default=",".join(BACKENDS), + help=f"Comma-separated backends (default: {','.join(BACKENDS)}).", + ) + base.add_argument( + "--case-prefix", + type=str, + default=None, + help="Only run cases whose Case name starts with this prefix (e.g. DSV3).", + ) + base.add_argument( + "--include-dsv3-gateup", + action="store_true", + help="Also benchmark DSV3-GateUP (skipped by default; known to run on gfx950).", + ) + args = base.parse_args() + + _require_rocm() + + test_cases = ( + generate_deepseekv2_lite_test_cases() + + generate_deepseekv2_test_cases() + + generate_deepseekv3_test_cases(include_gateup=args.include_dsv3_gateup) + + generate_grok_v2_test_cases() + + generate_qwen3_235b_test_cases() + ) + test_cases = _filter_test_cases(test_cases, args.quick) + if args.case_prefix: + prefix = args.case_prefix + test_cases = [ + c + for c in test_cases + if c["Case"].startswith(prefix) + and not (prefix == "DSV2-" and c["Case"].startswith("DSV2-Lite")) + ] + for case in test_cases: + case["topk"] = min(args.topk, case["B"]) + case["seed"] = 42 + + selected_backends = tuple(b.strip() for b in args.backends.split(",") if b.strip()) + for b in selected_backends: + if b not in KNOWN_BACKENDS: + raise ValueError(f"Unknown backend {b!r}, expected one of {KNOWN_BACKENDS}") + + def bench_fn(**case): + return bench_moe_grouped_gemm_backends( + **case, backends=selected_backends, train=args.train + ) + + run_benchmarks( + test_cases=test_cases, + bench_fn=bench_fn, + param_columns=["Case", "B", "M", "N", "K", "dtype", "topk"], + args=args, + ) + + +if __name__ == "__main__": + main() diff --git a/setup.py b/setup.py index 385c9eef9..c12303180 100644 --- a/setup.py +++ b/setup.py @@ -195,6 +195,14 @@ def setup_requirements() -> Tuple[List[str], List[str]]: ] test_reqs: List[str] = ["pytest>=8.2.1"] + # Optional FlyDSL dependency for ROCm PyTorch builds. + if ( + rocm_build() + and "pytorch" in frameworks + and bool(int(os.getenv("NVTE_USE_FLYDSL", "0"))) + ): + install_reqs.extend(["flydsl>=0.2.4,<0.3"]) + # Framework-specific requirements if not bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))): if "pytorch" in frameworks: diff --git a/tests/pytorch/flydsl_kernels/test_gemm.py b/tests/pytorch/flydsl_kernels/test_gemm.py new file mode 100644 index 000000000..48510880f --- /dev/null +++ b/tests/pytorch/flydsl_kernels/test_gemm.py @@ -0,0 +1,551 @@ +# Copyright (c) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +# +# License for AMD contributions = MIT. See LICENSE for more information + +"""User-facing FlyDSL GEMM tests -- ``general_gemm()`` under ``NVTE_USE_FLYDSL=1``. + +Exercises the same public entry point used by TE ``Linear`` / +``LayerNormLinear``. Coverage mirrors the Triton user-facing GEMM tests for the +currently supported FlyDSL surface: + +- fp32 / fp16 / bf16 regular tensors +- same-format and mixed-format tensor-wise FP8 +- same-format and mixed-format MXFP8 +- TN / NN / NT layouts +- batched multidimensional FP8 flattening + +Fused BIAS and BGRADB epilogues are intentionally not included yet because the +FlyDSL GEMM path does not currently support them. + +Each test compares the FlyDSL path against two independent references: + +1. ``torch.matmul`` on dequantized inputs, independent of hipBLASLt behavior. +2. The native C++ ``tex.generic_gemm`` backend through the same + ``general_gemm`` public surface. + +FlyDSL kernels currently require tile-aligned launch dimensions, so the test +shapes are aligned to the 256x256x128 kernel contract rather than reusing the +odd-sized Triton edge-mask cases. +""" + +import os + +import pytest +import torch + +from transformer_engine.pytorch import Float8Tensor +from transformer_engine.pytorch.cpp_extensions.gemm import general_gemm +from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer +from transformer_engine.pytorch.tensor.mxfp8_tensor import ( + MXFP8Quantizer, + MXFP8Tensor, +) +import transformer_engine_torch as tex + + +# --- Feature detection -------------------------------------------------------- + +major, minor = torch.cuda.get_device_capability() + +# The current FlyDSL MXFP8 implementation uses the gfx950 fp8-scaled MFMA. +has_mxfp8_support = major == 9 and minor >= 5 + +requires_mxfp8_support = pytest.mark.skipif( + not has_mxfp8_support, + reason="FlyDSL MXFP8 requires gfx950+ fp8-scaled MFMA support", +) + + +# --- Test parameters ---------------------------------------------------------- + +# The current FlyDSL kernels have no M/N edge masks and specialize K in K128 +# tiles. Keep all dimensions aligned to exercise the supported production path. +FLYDSL_SHAPES = [ + (512, 512, 512), + (512, 1024, 512), + (1024, 512, 1024), +] + +MXFP8_SHAPES = [ + (512, 512, 512), + (512, 1024, 512), +] + +LAYOUTS = ["TN", "NN", "NT"] + +FP8_FORMAT_COMBOS = [ + (tex.DType.kFloat8E4M3, tex.DType.kFloat8E4M3), + (tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2), + (tex.DType.kFloat8E5M2, tex.DType.kFloat8E4M3), + (tex.DType.kFloat8E5M2, tex.DType.kFloat8E5M2), +] + +FP8_FORMAT_IDS = [ + "e4m3_e4m3", + "e4m3_e5m2", + "e5m2_e4m3", + "e5m2_e5m2", +] + +REGULAR_DTYPES = [torch.float32, torch.float16, torch.bfloat16] + + +# --- Fixtures ----------------------------------------------------------------- + +@pytest.fixture(autouse=True) +def cleanup_env(): + """Save and restore FlyDSL-related environment variables between tests.""" + old_flydsl = os.environ.get("NVTE_USE_FLYDSL") + old_mxfp8 = os.environ.get("NVTE_ROCM_ENABLE_MXFP8") + + yield + + if old_flydsl is None: + os.environ.pop("NVTE_USE_FLYDSL", None) + else: + os.environ["NVTE_USE_FLYDSL"] = old_flydsl + + if old_mxfp8 is None: + os.environ.pop("NVTE_ROCM_ENABLE_MXFP8", None) + else: + os.environ["NVTE_ROCM_ENABLE_MXFP8"] = old_mxfp8 + + +# --- Helpers ------------------------------------------------------------------ + +def get_shapes(layout, M, K, N): + """Return the A/B storage shapes used by TE's public GEMM tests.""" + if layout == "TN": + return (M, K), (N, K) + if layout == "NN": + return (M, K), (K, M) + if layout == "NT": + return (M, K), (M, K) + raise ValueError(f"Unsupported layout: {layout}") + + +def compute_pytorch_reference(A_ref, B_ref, layout): + """Compute the equivalent public-layout GEMM with ``torch.matmul``.""" + if layout == "TN": + return torch.matmul(B_ref, A_ref.T) + if layout == "NN": + return torch.matmul(B_ref, A_ref) + if layout == "NT": + return torch.matmul(B_ref.T, A_ref) + raise ValueError(f"Unsupported layout: {layout}") + + +def create_fp8_tensors(M, K, N, layout, fp8_dtype_a, fp8_dtype_b): + """Create independently typed Float8Tensor inputs and references.""" + A_shape, B_shape = get_shapes(layout, M, K, N) + A_f32 = torch.randn(A_shape, dtype=torch.float32, device="cuda") * 0.5 + B_f32 = torch.randn(B_shape, dtype=torch.float32, device="cuda") * 0.5 + + A_fp8 = Float8Quantizer( + scale=torch.full((1,), 1.0, dtype=torch.float32, device="cuda"), + amax=torch.empty((1,), dtype=torch.float32, device="cuda"), + fp8_dtype=fp8_dtype_a, + )(A_f32) + B_fp8 = Float8Quantizer( + scale=torch.full((1,), 1.0, dtype=torch.float32, device="cuda"), + amax=torch.empty((1,), dtype=torch.float32, device="cuda"), + fp8_dtype=fp8_dtype_b, + )(B_f32) + + return A_fp8, B_fp8, A_fp8.dequantize(), B_fp8.dequantize() + + +def _make_mxfp8_quantizer(fp8_dtype): + """Create one independently typed MXFP8 quantizer with both orientations.""" + quantizer = MXFP8Quantizer(fp8_dtype=fp8_dtype) + quantizer.set_usage(rowwise=True, columnwise=True) + return quantizer + + +def create_mxfp8_tensors( + M, + K, + N, + layout, + fp8_dtype_a, + fp8_dtype_b, +): + """Create independently typed MXFP8Tensor inputs and references.""" + A_shape, B_shape = get_shapes(layout, M, K, N) + A_f32 = torch.randn(A_shape, dtype=torch.float32, device="cuda") * 0.5 + B_f32 = torch.randn(B_shape, dtype=torch.float32, device="cuda") * 0.5 + + A_mxfp8 = _make_mxfp8_quantizer(fp8_dtype_a)(A_f32) + B_mxfp8 = _make_mxfp8_quantizer(fp8_dtype_b)(B_f32) + + return ( + A_mxfp8, + B_mxfp8, + A_mxfp8.dequantize(), + B_mxfp8.dequantize(), + ) + + +def call_gemm(A, B, layout, out_dtype, use_flydsl=True): + """Call ``general_gemm`` through either FlyDSL or the native C++ path.""" + os.environ["NVTE_USE_FLYDSL"] = "1" if use_flydsl else "0" + + output, bias_grad, gelu_input, extra_output = general_gemm( + A=A, + B=B, + out_dtype=out_dtype, + layout=layout, + bias=None, + quantization_params=None, + gelu=False, + grad=False, + accumulate=False, + ) + + assert bias_grad is None + assert gelu_input is None + assert extra_output is None + return output + + +def assert_gemm_close(actual, expected, *, atol, rtol): + """Compare through FP32 so output narrowing does not hide diagnostics.""" + torch.testing.assert_close( + actual.float(), + expected.float(), + atol=atol, + rtol=rtol, + equal_nan=False, + ) + + +# ============================================================================== +# Approach 1: FlyDSL vs PyTorch torch.matmul reference +# ============================================================================== + +@pytest.mark.parametrize("M, K, N", FLYDSL_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +@pytest.mark.parametrize( + "dtype", + REGULAR_DTYPES, + ids=["fp32", "fp16", "bf16"], +) +def test_flydsl_vs_pytorch_regular(M, K, N, layout, dtype): + """Test regular FlyDSL GEMM against an FP32 PyTorch reference.""" + torch.manual_seed(42) + + A_shape, B_shape = get_shapes(layout, M, K, N) + A = torch.randn(A_shape, dtype=dtype, device="cuda") * 0.5 + B = torch.randn(B_shape, dtype=dtype, device="cuda") * 0.5 + + output = call_gemm( + A, + B, + layout, + out_dtype=dtype, + use_flydsl=True, + ) + expected = compute_pytorch_reference(A.float(), B.float(), layout) + + assert_gemm_close(output, expected, atol=1e-3, rtol=1e-2) + + +@pytest.mark.parametrize("M, K, N", FLYDSL_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +@pytest.mark.parametrize( + "fp8_format", + FP8_FORMAT_COMBOS, + ids=FP8_FORMAT_IDS, +) +def test_flydsl_vs_pytorch_fp8(M, K, N, layout, fp8_format): + """Test same-format and mixed-format tensor-wise FP8 FlyDSL GEMMs.""" + torch.manual_seed(42) + + fp8_dtype_a, fp8_dtype_b = fp8_format + A_fp8, B_fp8, A_deq, B_deq = create_fp8_tensors( + M, + K, + N, + layout, + fp8_dtype_a, + fp8_dtype_b, + ) + + output = call_gemm( + A_fp8, + B_fp8, + layout, + out_dtype=torch.float32, + use_flydsl=True, + ) + expected = compute_pytorch_reference( + A_deq.float(), + B_deq.float(), + layout, + ) + + assert_gemm_close(output, expected, atol=5e-3, rtol=1e-2) + + +@requires_mxfp8_support +@pytest.mark.parametrize("M, K, N", MXFP8_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +@pytest.mark.parametrize( + "fp8_format", + FP8_FORMAT_COMBOS, + ids=FP8_FORMAT_IDS, +) +def test_flydsl_vs_pytorch_mxfp8(M, K, N, layout, fp8_format): + """Test same-format and mixed-format MXFP8 FlyDSL GEMMs.""" + os.environ["NVTE_ROCM_ENABLE_MXFP8"] = "1" + torch.manual_seed(42) + + fp8_dtype_a, fp8_dtype_b = fp8_format + A_mxfp8, B_mxfp8, A_deq, B_deq = create_mxfp8_tensors( + M, + K, + N, + layout, + fp8_dtype_a, + fp8_dtype_b, + ) + + output = call_gemm( + A_mxfp8, + B_mxfp8, + layout, + out_dtype=torch.float32, + use_flydsl=True, + ) + expected = compute_pytorch_reference( + A_deq.float(), + B_deq.float(), + layout, + ) + + assert_gemm_close(output, expected, atol=5e-3, rtol=1e-2) + + +# ============================================================================== +# Approach 2: FlyDSL vs native C++ ``generic_gemm`` reference +# ============================================================================== + +@pytest.mark.parametrize("M, K, N", FLYDSL_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +@pytest.mark.parametrize( + "dtype", + REGULAR_DTYPES, + ids=["fp32", "fp16", "bf16"], +) +def test_flydsl_vs_cpp_regular(M, K, N, layout, dtype): + """Test regular FlyDSL GEMM against the native C++ backend.""" + torch.manual_seed(42) + + A_shape, B_shape = get_shapes(layout, M, K, N) + A = torch.randn(A_shape, dtype=dtype, device="cuda") * 0.5 + B = torch.randn(B_shape, dtype=dtype, device="cuda") * 0.5 + + flydsl_out = call_gemm( + A, + B, + layout, + out_dtype=dtype, + use_flydsl=True, + ) + cpp_out = call_gemm( + A, + B, + layout, + out_dtype=dtype, + use_flydsl=False, + ) + + assert_gemm_close(flydsl_out, cpp_out, atol=1e-3, rtol=1e-2) + + +@pytest.mark.parametrize("M, K, N", FLYDSL_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +@pytest.mark.parametrize( + "fp8_format", + FP8_FORMAT_COMBOS, + ids=FP8_FORMAT_IDS, +) +def test_flydsl_vs_cpp_fp8(M, K, N, layout, fp8_format): + """Test same-format and mixed-format FP8 against native C++.""" + torch.manual_seed(42) + + fp8_dtype_a, fp8_dtype_b = fp8_format + A_fp8, B_fp8, _, _ = create_fp8_tensors( + M, + K, + N, + layout, + fp8_dtype_a, + fp8_dtype_b, + ) + + flydsl_out = call_gemm( + A_fp8, + B_fp8, + layout, + out_dtype=torch.float32, + use_flydsl=True, + ) + cpp_out = call_gemm( + A_fp8, + B_fp8, + layout, + out_dtype=torch.float32, + use_flydsl=False, + ) + + assert_gemm_close(flydsl_out, cpp_out, atol=5e-3, rtol=1e-2) + + +@requires_mxfp8_support +@pytest.mark.parametrize("M, K, N", MXFP8_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +@pytest.mark.parametrize( + "fp8_format", + FP8_FORMAT_COMBOS, + ids=FP8_FORMAT_IDS, +) +def test_flydsl_vs_cpp_mxfp8(M, K, N, layout, fp8_format): + """Test same-format and mixed-format MXFP8 against native C++.""" + os.environ["NVTE_ROCM_ENABLE_MXFP8"] = "1" + torch.manual_seed(42) + + fp8_dtype_a, fp8_dtype_b = fp8_format + A_mxfp8, B_mxfp8, _, _ = create_mxfp8_tensors( + M, + K, + N, + layout, + fp8_dtype_a, + fp8_dtype_b, + ) + + flydsl_out = call_gemm( + A_mxfp8, + B_mxfp8, + layout, + out_dtype=torch.float32, + use_flydsl=True, + ) + cpp_out = call_gemm( + A_mxfp8, + B_mxfp8, + layout, + out_dtype=torch.float32, + use_flydsl=False, + ) + + assert_gemm_close(flydsl_out, cpp_out, atol=5e-3, rtol=1e-2) + + +# ============================================================================== +# Batched multidimensional FP8 coverage +# ============================================================================== + +@pytest.mark.parametrize( + "batch_size, M, K, N", + [ + (2, 256, 512, 256), + (4, 256, 512, 256), + ], +) +@pytest.mark.parametrize( + "fp8_format", + FP8_FORMAT_COMBOS, + ids=FP8_FORMAT_IDS, +) +def test_flydsl_vs_pytorch_fp8_multidim( + batch_size, + M, + K, + N, + fp8_format, +): + """Exercise flatten-leading-dim semantics for multidimensional FP8.""" + torch.manual_seed(42) + + fp8_dtype_a, fp8_dtype_b = fp8_format + + # TN layout: the wrapper flattens all leading dimensions into rows. + A_f32 = ( + torch.randn( + batch_size, + M, + K, + dtype=torch.float32, + device="cuda", + ) + * 0.5 + ) + B_f32 = ( + torch.randn( + batch_size, + N, + K, + dtype=torch.float32, + device="cuda", + ) + * 0.5 + ) + + A_fp8 = Float8Quantizer( + scale=torch.full((1,), 1.0, dtype=torch.float32, device="cuda"), + amax=torch.empty((1,), dtype=torch.float32, device="cuda"), + fp8_dtype=fp8_dtype_a, + )(A_f32) + B_fp8 = Float8Quantizer( + scale=torch.full((1,), 1.0, dtype=torch.float32, device="cuda"), + amax=torch.empty((1,), dtype=torch.float32, device="cuda"), + fp8_dtype=fp8_dtype_b, + )(B_f32) + + output = call_gemm( + A_fp8, + B_fp8, + layout="TN", + out_dtype=torch.float32, + use_flydsl=True, + ) + + A_flat = A_fp8.dequantize().reshape(-1, K) + B_flat = B_fp8.dequantize().reshape(-1, K) + expected = torch.matmul(B_flat, A_flat.T) + + assert_gemm_close(output, expected, atol=5e-3, rtol=1e-2) + + +if __name__ == "__main__": + # Quick smoke tests using one case from each supported input family. + os.environ["NVTE_USE_FLYDSL"] = "1" + os.environ["NVTE_ROCM_ENABLE_MXFP8"] = "1" + + test_flydsl_vs_pytorch_regular( + 256, + 512, + 256, + "TN", + torch.float16, + ) + test_flydsl_vs_pytorch_fp8( + 256, + 512, + 256, + "TN", + (tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2), + ) + + if has_mxfp8_support: + test_flydsl_vs_pytorch_mxfp8( + 256, + 512, + 256, + "TN", + (tex.DType.kFloat8E5M2, tex.DType.kFloat8E4M3), + ) + + print("All FlyDSL GEMM smoke tests passed!") diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 3a4405533..2bc5356e8 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -1331,6 +1331,150 @@ def test_linear_accuracy(dtype, bs, model, return_bias, bias): assert_allclose(te_output, torch_output, tolerance, rtol[dtype]) +@pytest.mark.parametrize("dtype", param_types) +@pytest.mark.parametrize("bs", batch_sizes) +@pytest.mark.parametrize("model", ["small", "126m"]) +@pytest.mark.parametrize( + "fp8_recipe", + [ + None, + recipe.Float8CurrentScaling(), + recipe.DelayedScaling(), + recipe.MXFP8BlockScaling(), + ], +) +def test_linear_accuracy_flydsl( + dtype, + bs, + model, + fp8_recipe, +): + """Compare FlyDSL and native TE Linear forward, dgrad, and wgrad.""" + + if not IS_HIP_EXTENSION: + pytest.skip("FlyDSL GEMM is only supported on HIP.") + + fp8 = fp8_recipe is not None + config = model_configs[model] + + if isinstance(fp8_recipe, recipe.MXFP8BlockScaling): + if not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + elif fp8 and not fp8_available: + pytest.skip(reason_for_no_fp8) + + if config.max_seqlen_q % 16 != 0 and fp8: + pytest.skip("FP8 requires sequence length to be divisible by 16.") + + # Validate the GEMM backend, not quantized parameter storage. + # FlyDSL GEMM does not currently support bias. + with quantized_model_init(enabled=False, recipe=fp8_recipe): + linear_ref = Linear( + config.hidden_size, + 4 * config.hidden_size, + bias=False, + params_dtype=dtype, + device="cuda", + ).eval() + + linear_flydsl = Linear( + config.hidden_size, + 4 * config.hidden_size, + bias=False, + params_dtype=dtype, + device="cuda", + ).eval() + + with torch.no_grad(): + linear_flydsl.weight.copy_(linear_ref.weight) + + input_shape = ( + config.max_seqlen_q, + bs, + config.hidden_size, + ) + + inp_ref = torch.randn( + input_shape, + dtype=dtype, + device="cuda", + requires_grad=True, + ) + inp_flydsl = inp_ref.detach().clone().requires_grad_(True) + + try: + # Native TE backend. + os.environ.pop("NVTE_USE_FLYDSL", None) + os.environ.pop("NVTE_FLYDSL_GEMM_WARN_FALLBACK", None) + + reset_rng_states() + FP8GlobalStateManager.reset() + + with autocast(enabled=fp8, recipe=fp8_recipe): + out_ref = linear_ref(inp_ref) + + out_ref.sum().backward() + torch.cuda.synchronize() + + # FlyDSL backend. + os.environ["NVTE_USE_FLYDSL"] = "1" + os.environ["NVTE_FLYDSL_GEMM_WARN_FALLBACK"] = "1" + + reset_rng_states() + FP8GlobalStateManager.reset() + + with autocast(enabled=fp8, recipe=fp8_recipe): + out_flydsl = linear_flydsl(inp_flydsl) + + out_flydsl.sum().backward() + torch.cuda.synchronize() + + finally: + os.environ.pop("NVTE_USE_FLYDSL", None) + os.environ.pop("NVTE_FLYDSL_GEMM_WARN_FALLBACK", None) + FP8GlobalStateManager.reset() + + tols = dtype_tols(dtype) + atol = tols["atol"] + rtol = tols["rtol"] + + if fp8: + atol = max(atol, 1e-2) + rtol = max(rtol, 1e-2) + + torch.testing.assert_close( + out_flydsl, + out_ref, + atol=atol, + rtol=rtol, + ) + + torch.testing.assert_close( + inp_flydsl.grad, + inp_ref.grad, + atol=atol, + rtol=rtol, + ) + + # Wgrad is an NT GEMM with a reduction over the flattened + # sequence/batch dimension. FlyDSL and the native TE backend may use + # different FP32 accumulation orders, so allow the small expected + # non-associative rounding difference. + wgrad_atol = atol + wgrad_rtol = rtol + + if dtype == torch.float32 and not fp8: + wgrad_atol = max(wgrad_atol, 1e-4) + wgrad_rtol = max(wgrad_rtol, 1e-4) + + torch.testing.assert_close( + linear_flydsl.weight.grad, + linear_ref.weight.grad, + atol=wgrad_atol, + rtol=wgrad_rtol, + ) + + @pytest.mark.parametrize("dtype", param_types) @pytest.mark.parametrize("bs", batch_sizes) @pytest.mark.parametrize("model", ["small"]) diff --git a/tests/pytorch/test_perm_free_grouped_linear.py b/tests/pytorch/test_perm_free_grouped_linear.py new file mode 100644 index 000000000..7bfa0a679 --- /dev/null +++ b/tests/pytorch/test_perm_free_grouped_linear.py @@ -0,0 +1,1160 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for the route-list permute-free bf16 MoE gather-GEMM kernels.""" + +import pytest +import torch +from torch.utils.cpp_extension import IS_HIP_EXTENSION + +from transformer_engine.pytorch import GroupedLinear +from transformer_engine.pytorch.moe import ( + MoERoutingMetadata, + PermuteFreeMetadata, + is_permute_free_grouped_gemm_enabled, + permute_free_gated_act_bwd, + permute_free_grouped_gemm_backward, + permute_free_grouped_gemm_bf16, + permute_free_grouped_gemm_bf16_dgrad, + permute_free_grouped_gemm_bf16_fc2, + permute_free_grouped_gemm_bf16_fc2_dgrad, + permute_free_grouped_gemm_bf16_fc2_wgrad, + permute_free_grouped_gemm_bf16_wgrad, + permute_free_grouped_gemm_forward, + prepare_moe_align, +) +from transformer_engine.pytorch.moe.permute_free_grouped_gemm import ( + _FLYDSL_FWD_BLOCK_M, + _expert_per_route, +) + +pytestmark = pytest.mark.skipif( + not (IS_HIP_EXTENSION and torch.cuda.is_available()), + reason="Permute-free grouped GEMM tests require ROCm and CUDA device.", +) + + +@pytest.fixture(autouse=True) +def _cuda_sync_after_test(): + yield + if torch.cuda.is_available(): + torch.cuda.synchronize() + + +def _rel_l2(actual: torch.Tensor, ref: torch.Tensor) -> float: + return ((actual.float() - ref.float()).norm() / ref.float().norm().clamp_min(1e-8)).item() + + +def _random_routing_map(num_recv_tokens, num_experts, max_hits, device, seed): + """Boolean [num_recv_tokens, num_experts] with 1..max_hits local experts per token.""" + gen = torch.Generator(device=device).manual_seed(seed) + routing_map = torch.zeros(num_recv_tokens, num_experts, dtype=torch.bool, device=device) + for t in range(num_recv_tokens): + k = int(torch.randint(1, max_hits + 1, (1,), device=device, generator=gen).item()) + experts = torch.randperm(num_experts, device=device, generator=gen)[:k] + routing_map[t, experts] = True + # Guarantee every expert owns at least one route so per-expert refs are exercised. + for e in range(num_experts): + if not routing_map[:, e].any(): + routing_map[int(e) % num_recv_tokens, e] = True + return routing_map + + +def _dense_route_order(routing_map): + """Expert-sorted (token, expert) dense route list, matching ``moe_align_route_list``.""" + tok, exp = routing_map.nonzero(as_tuple=True) + order = torch.argsort(exp, stable=True) + return tok[order].to(torch.int64), exp[order].to(torch.int64) + +# --------------------------------------------------------------------------- +# Route-list gather-GEMM: fwd / dgrad / wgrad +# --------------------------------------------------------------------------- +def test_route_list_fwd(): + torch.manual_seed(7) + num_recv_tokens, in_features, out_features = 128, 128, 256 + num_experts, max_hits = 8, 3 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=1) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + + hidden = torch.randn(num_recv_tokens, in_features, device="cuda", dtype=torch.bfloat16) + weights = torch.randn( + num_experts, out_features, in_features, device="cuda", dtype=torch.bfloat16 + ) + + out_full = permute_free_grouped_gemm_bf16(hidden, weights, routing) + + route_to_token, route_expert = _dense_route_order(routing_map) + num_routes = route_to_token.numel() + # Block-padded canonical layout: the output is [em_max, out] in expert-sorted, block-padded + # slot order. Padded slot s holds hidden[sorted_slot_ids[s]] @ W[expert(s)]^T for the valid + # slots (sorted_slot_ids[s] < num_recv_tokens); padding slots carry inert dead values and + # are dropped only at the final padded->token gather-combine. + em_max = int(routing.sorted_slot_ids.shape[0]) + assert out_full.shape == (em_max, out_features) + slot_token = routing.sorted_slot_ids + slot_expert = _expert_per_route(routing, em_max) + valid = slot_token < num_recv_tokens + tok = slot_token.to(torch.int64).clamp_(0, num_recv_tokens - 1) + ref = torch.einsum( + "rk,rnk->rn", + hidden[tok].float(), + weights[slot_expert.to(torch.int64)].float(), + ) + assert int(valid.sum().item()) == num_routes + assert _rel_l2(out_full[valid], ref[valid]) < 2e-2 + + +def test_route_list_dgrad(): + torch.manual_seed(11) + num_recv_tokens, in_features, out_features = 128, 96, 128 + num_experts, max_hits = 8, 4 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=2) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + num_routes = int(routing_map.sum().item()) + # Block-padded canonical: the incoming route grad lives in the same [em_max] block-padded + # slot order as the FC1 forward output (the FC2 dgrad hands it back in this layout). Prepare + # the align at the forward block size so the standalone dgrad reuses the same slot layout. + routing = prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(routing.sorted_slot_ids.shape[0]) + slot_token = routing.sorted_slot_ids + slot_expert = _expert_per_route(routing, em_max) + valid = slot_token < num_recv_tokens + + weights = torch.randn( + num_experts, out_features, in_features, device="cuda", dtype=torch.bfloat16 + ) + grad = torch.zeros(em_max, out_features, device="cuda", dtype=torch.bfloat16) + grad[valid] = torch.randn( + int(valid.sum().item()), out_features, device="cuda", dtype=torch.bfloat16 + ) + + dA = permute_free_grouped_gemm_bf16_dgrad(grad, weights, routing) + assert dA.shape == (num_recv_tokens, in_features) + + # dA[t] = sum_{padded slots s with token==t, valid} grad[s] @ W[expert(s)] + per_slot = torch.einsum( + "rn,rnk->rk", grad.float(), weights[slot_expert.to(torch.int64)].float() + ) + per_slot = per_slot * valid[:, None].float() + tok = slot_token.to(torch.int64).clamp_(0, num_recv_tokens - 1) + ref = torch.zeros(num_recv_tokens, in_features, device="cuda", dtype=torch.float32) + ref.index_add_(0, tok, per_slot) + assert int(valid.sum().item()) == num_routes + assert _rel_l2(dA, ref) < 2e-2 + + +def test_route_list_wgrad(): + torch.manual_seed(17) + num_recv_tokens, in_features, out_features = 128, 96, 128 + num_experts, max_hits = 8, 4 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=3) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + num_routes = int(routing_map.sum().item()) + # Block-padded canonical: the incoming route grad lives in the [em_max] block-padded slot + # order (same layout the forward writes / FC2 dgrad hands back). Prepare the forward align so + # the wgrad reads grad at block_start[e]*block_size_m + within-rank. + routing = prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(routing.sorted_slot_ids.shape[0]) + slot_token = routing.sorted_slot_ids + slot_expert = _expert_per_route(routing, em_max) + valid = slot_token < num_recv_tokens + + hidden = torch.randn(num_recv_tokens, in_features, device="cuda", dtype=torch.bfloat16) + weights_shape = (num_experts, out_features, in_features) + grad = torch.zeros(em_max, out_features, device="cuda", dtype=torch.bfloat16) + grad[valid] = torch.randn( + int(valid.sum().item()), out_features, device="cuda", dtype=torch.bfloat16 + ) + + dW = permute_free_grouped_gemm_bf16_wgrad(hidden, grad, weights_shape, routing) + assert dW.shape == weights_shape + + # dW[e] = sum_{valid slots s with expert==e} outer(grad[s], hidden[slot_token[s]]) + tok = slot_token.to(torch.int64).clamp_(0, num_recv_tokens - 1) + per_slot = torch.einsum("rn,rk->rnk", grad.float(), hidden[tok].float()) + per_slot = per_slot * valid[:, None, None].float() + ref = torch.zeros(num_experts, out_features, in_features, device="cuda", dtype=torch.float32) + # Past-end padding slots carry expert id -1; clamp for the scatter (their rows are zeroed). + ref.index_add_(0, slot_expert.to(torch.int64).clamp_min(0), per_slot) + assert int(valid.sum().item()) == num_routes + assert _rel_l2(dW, ref) < 2e-2 + + +def test_route_list_fc2_wgrad(): + """FC2 wgrad via operand-swap + transpose; ``out``/``accumulate`` fold in bf16 or fp32.""" + torch.manual_seed(19) + num_recv_tokens, in_features, out_features = 128, 96, 128 # F=in, H=out (W2 is [E, H, F]) + num_experts, max_hits = 8, 4 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=4) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + num_routes = int(routing_map.sum().item()) + # Block-padded canonical: fc2_input (the FC1 output) lives in the [em_max] block-padded slot + # order; grad_output stays token-space and is gathered internally. + routing = prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(routing.sorted_slot_ids.shape[0]) + slot_token = routing.sorted_slot_ids + slot_expert = _expert_per_route(routing, em_max) + valid = slot_token < num_recv_tokens + + fc2_input = torch.zeros(em_max, in_features, device="cuda", dtype=torch.bfloat16) + fc2_input[valid] = torch.randn( + int(valid.sum().item()), in_features, device="cuda", dtype=torch.bfloat16 + ) + grad_output = torch.randn(num_recv_tokens, out_features, device="cuda", dtype=torch.bfloat16) + weights_shape = (num_experts, out_features, in_features) # [E, H, F] + + # dW2[e] = sum_{valid slots s with expert==e} outer(grad_output[slot_token[s]], fc2_input[s]) + tok = slot_token.to(torch.int64).clamp_(0, num_recv_tokens - 1) + per_slot = torch.einsum("rh,rf->rhf", grad_output[tok].float(), fc2_input.float()) + per_slot = per_slot * valid[:, None, None].float() + ref = torch.zeros(num_experts, out_features, in_features, device="cuda", dtype=torch.float32) + # Past-end padding slots carry expert id -1; clamp for the scatter (their rows are zeroed). + ref.index_add_(0, slot_expert.to(torch.int64).clamp_min(0), per_slot) + + # Fresh output (no transpose needed downstream). + dW2 = permute_free_grouped_gemm_bf16_fc2_wgrad(fc2_input, grad_output, weights_shape, routing) + assert dW2.shape == weights_shape + assert _rel_l2(dW2, ref) < 2e-2 + + # Direct fp32 accumulate into an existing buffer: overwrite then add == 2x. + out = torch.zeros(weights_shape, device="cuda", dtype=torch.float32) + permute_free_grouped_gemm_bf16_fc2_wgrad( + fc2_input, grad_output, weights_shape, routing, out=out, accumulate=False + ) + assert _rel_l2(out, ref) < 2e-2 + permute_free_grouped_gemm_bf16_fc2_wgrad( + fc2_input, grad_output, weights_shape, routing, out=out, accumulate=True + ) + assert _rel_l2(out, 2.0 * ref) < 2e-2 + + +def test_route_list_fc2_wgrad_recompute_from_preact(): + """FC2 wgrad recompute-from-preact: rebuild act = act(gate)*up*prob from the saved 2F + pre-activation into a transient buffer, feed the unchanged wgrad, and match a stored-act + reference (the backward then checkpoints only the 2F preact, never the F-wide act).""" + from transformer_engine.pytorch.moe import prepare_moe_align + + torch.manual_seed(29) + num_recv_tokens, in_features, out_features = 128, 96, 128 # F=in, H=out (W2 is [E, H, F]) + num_experts, max_hits = 8, 4 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=6) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + + # Block-padded canonical: preact / act / grad all live in the [em_max] block-padded slot + # order (padding slots have token sentinel >= num_recv_tokens and are dropped by the kernel). + em_max = int(routing.sorted_slot_ids.shape[0]) + num_routes = int(routing_map.sum().item()) + slot_token = routing.sorted_slot_ids + slot_expert = _expert_per_route(routing, em_max) + valid = slot_token < num_recv_tokens + tok = slot_token.to(torch.int64).clamp_(0, num_recv_tokens - 1) + exp = slot_expert.to(torch.int64).clamp_min(0) + + preact = torch.zeros(em_max, 2 * in_features, device="cuda", dtype=torch.bfloat16) + preact[valid] = torch.randn( + int(valid.sum().item()), 2 * in_features, device="cuda", dtype=torch.bfloat16 + ) + grad_output = torch.randn(num_recv_tokens, out_features, device="cuda", dtype=torch.bfloat16) + probs = torch.rand(num_recv_tokens, num_experts, device="cuda", dtype=torch.float32) + weights_shape = (num_experts, out_features, in_features) # [E, H, F] + + # Reference act in the block-padded slot layout using the kernel's own slot arrays. + g = preact[:, :in_features].float() + u = preact[:, in_features:].float() + p = probs[tok, exp] + act_ref = torch.nn.functional.silu(g) * u * p[:, None] + act_ref = act_ref * valid[:, None].float() + act_stored = act_ref.to(torch.bfloat16) + + # Stored-act baseline: feed the materialized activation directly. + dW2_stored = permute_free_grouped_gemm_bf16_fc2_wgrad( + act_stored, grad_output, weights_shape, routing + ) + # Recompute path: feed preact + probs, rebuild act in-flight. + dW2_rc = permute_free_grouped_gemm_bf16_fc2_wgrad( + None, grad_output, weights_shape, routing, + preact=preact, dispatched_probs=probs, activation="silu", + ) + + # fp32 reference dW2[e] = sum_{valid slots s: exp==e} outer(grad_output[tok_s], act_ref[s]). + ref = torch.zeros(num_experts, out_features, in_features, device="cuda", dtype=torch.float32) + per_route = torch.einsum("rh,rf->rhf", grad_output[tok].float(), act_ref) + ref.index_add_(0, exp, per_route) + + assert _rel_l2(dW2_stored, ref) < 2e-2 + assert _rel_l2(dW2_rc, ref) < 2e-2 + # Recompute should also track the stored-act path closely (same GEMM, act rebuilt). + assert _rel_l2(dW2_rc, dW2_stored) < 1e-2 + + +def test_fc2_backward_dispatch_recompute_matches_stored(): + """FC2 backward dispatch: wgrad from the saved 2F preact (+ fc2_activation) matches the + legacy stored-F activation path.""" + from transformer_engine.pytorch.moe import ( + permute_free_grouped_gemm_backward, + prepare_moe_align, + ) + + torch.manual_seed(31) + num_recv_tokens, in_features, out_features = 128, 96, 128 # F=in, H=out (W2 is [E, H, F]) + num_experts, max_hits = 8, 4 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=8) + routing = PermuteFreeMetadata( + routing_map=routing_map, num_experts=num_experts, route_space=True, activation="silu" + ) + prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + + # Block-padded canonical: preact / act live in the [em_max] block-padded slot order. + em_max = int(routing.sorted_slot_ids.shape[0]) + slot_token = routing.sorted_slot_ids + slot_expert = _expert_per_route(routing, em_max) + valid = slot_token < num_recv_tokens + tok = slot_token.to(torch.int64).clamp_(0, num_recv_tokens - 1) + exp = slot_expert.to(torch.int64).clamp_min(0) + + preact = torch.zeros(em_max, 2 * in_features, device="cuda", dtype=torch.bfloat16) + preact[valid] = torch.randn( + int(valid.sum().item()), 2 * in_features, device="cuda", dtype=torch.bfloat16 + ) + grad_output = torch.randn(num_recv_tokens, out_features, device="cuda", dtype=torch.bfloat16) + probs = torch.rand(num_recv_tokens, num_experts, device="cuda", dtype=torch.float32) + weights = [ + torch.randn(out_features, in_features, device="cuda", dtype=torch.bfloat16) + for _ in range(num_experts) + ] + + g = preact[:, :in_features].float() + u = preact[:, in_features:].float() + p = probs[tok, exp] + act_stored = (torch.nn.functional.silu(g) * u * p[:, None] * valid[:, None].float()).to( + torch.bfloat16 + ) + + stored = permute_free_grouped_gemm_backward( + grad_output, + routing=routing, + weights=weights, + num_gemms=num_experts, + hidden_states=act_stored, + requires_wgrad=True, + ) + recompute = permute_free_grouped_gemm_backward( + grad_output, + routing=routing, + weights=weights, + num_gemms=num_experts, + hidden_states=preact, + requires_wgrad=True, + dispatched_probs=probs, + fc2_activation="silu", + ) + assert recompute.wgrad_stacked.shape == (num_experts, out_features, in_features) + assert _rel_l2(recompute.wgrad_stacked, stored.wgrad_stacked) < 1e-2 + + +def test_fc2_backward_dispatch_fused_dgrad_wgrad_matches_split(): + """FC2 backward with both grads + a gated activation: the act-bwd re-emits the F-wide + fc2_input for the wgrad (one fused pass over the 2F preact). Both dgrad and wgrad must match + running dgrad-only and wgrad-only (separate recompute) independently.""" + from transformer_engine.pytorch.moe import ( + permute_free_grouped_gemm_backward, + prepare_moe_align, + ) + + torch.manual_seed(43) + num_recv_tokens, in_features, out_features = 128, 96, 128 # F=in, H=out (W2 is [E, H, F]) + num_experts, max_hits = 8, 4 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=12) + routing = PermuteFreeMetadata( + routing_map=routing_map, num_experts=num_experts, route_space=True, activation="silu" + ) + prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + + em_max = int(routing.sorted_slot_ids.shape[0]) + valid = routing.sorted_slot_ids < num_recv_tokens + + preact = torch.zeros(em_max, 2 * in_features, device="cuda", dtype=torch.bfloat16) + preact[valid] = torch.randn( + int(valid.sum().item()), 2 * in_features, device="cuda", dtype=torch.bfloat16 + ) + grad_output = torch.randn(num_recv_tokens, out_features, device="cuda", dtype=torch.bfloat16) + probs = torch.rand( + num_recv_tokens, num_experts, device="cuda", dtype=torch.float32, requires_grad=True + ) + weights = [ + torch.randn(out_features, in_features, device="cuda", dtype=torch.bfloat16) + for _ in range(num_experts) + ] + + common = dict( + routing=routing, weights=weights, num_gemms=num_experts, + hidden_states=preact, fc2_activation="silu", dispatched_probs=probs, + ) + # Fused: dgrad + wgrad in one call -> the wgrad reuses the act-bwd's fc2_input. + fused = permute_free_grouped_gemm_backward( + grad_output, requires_dgrad=True, requires_wgrad=True, **common + ) + # Split references: dgrad-only (act-bwd) and wgrad-only (own recompute-from-preact). + dref = permute_free_grouped_gemm_backward( + grad_output, requires_dgrad=True, requires_wgrad=False, **common + ) + wref = permute_free_grouped_gemm_backward( + grad_output, requires_dgrad=False, requires_wgrad=True, **common + ) + + assert fused.dgrad.shape == (em_max, 2 * in_features) + assert fused.wgrad_stacked.shape == (num_experts, out_features, in_features) + # dgrad is bit-identical (same act-bwd, EMIT_ACT only adds a store). + assert torch.equal(fused.dgrad, dref.dgrad) + assert torch.equal(fused.grad_probs, dref.grad_probs) + # wgrad from the fused (re-emitted) act matches the standalone recompute to bf16 rounding. + assert _rel_l2(fused.wgrad_stacked, wref.wgrad_stacked) < 1e-2 + + +def test_fc1_fc2_gated_pipeline(monkeypatch): + """End-to-end FC1 raw 2F -> FC2 standalone gated activation: forward outputs and backward grads + match a PyTorch reference through the permute-free GroupedLinear modules.""" + import dataclasses + + from transformer_engine.pytorch.moe import prepare_moe_align + + monkeypatch.setenv("NVTE_PERMUTE_FREE_GROUPED_GEMM", "1") + torch.manual_seed(37) + # FFN width is a multiple of the FlyDSL block_k (64) for kernel alignment. + num_recv_tokens, hidden, ffn = 128, 128, 128 + num_experts, max_hits = 8, 3 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=9) + fc1_meta = PermuteFreeMetadata( + routing_map=routing_map, num_experts=num_experts, activation="silu" + ) + # Prepare at the FlyDSL-compatible forward block_m (>= 128, multiple of 128). FC1 and FC2 share the + # same block-padded slot layout, so the derived fc2_meta inherits this align. + prepare_moe_align(fc1_meta, _FLYDSL_FWD_BLOCK_M) + fc2_meta = dataclasses.replace(fc1_meta, route_space=True) + num_routes = int(routing_map.sum().item()) + m_splits = [num_recv_tokens // num_experts] * num_experts + m_splits[-1] += num_recv_tokens - sum(m_splits) + + inp = torch.randn(num_recv_tokens, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True) + probs = torch.rand(num_recv_tokens, num_experts, device="cuda", dtype=torch.float32, requires_grad=True) + + fc1 = GroupedLinear( + num_experts, hidden, 2 * ffn, bias=False, params_dtype=torch.bfloat16, device="cuda" + ) + fc2 = GroupedLinear( + num_experts, ffn, hidden, bias=False, params_dtype=torch.bfloat16, device="cuda" + ) + torch.manual_seed(38) + with torch.no_grad(): + for i in range(num_experts): + getattr(fc1, f"weight{i}").normal_(0, 0.05) + getattr(fc2, f"weight{i}").normal_(0, 0.05) + + preact = fc1(inp, m_splits, permute_free_metadata=fc1_meta) + assert preact.shape[1] == 2 * ffn + preact.retain_grad() # non-leaf: keep its grad so we can sanity-check the FC2->FC1 edge + out = fc2(preact, m_splits, permute_free_metadata=fc2_meta, dispatched_probs=probs) + + # Reference: FC1 raw 2F [gate|up] -> silu(gate)*up*prob -> FC2, summed over each token's + # experts (the fused FC2 prologue must reproduce this token-space output). Computed on the + # CPU so no GPU (re)allocation happens between the forward and the backward -- the FlyDSL + # FC2 dgrad autotune is sensitive to heap layout shifts (a latent OOB read faults only when + # a large device alloc/free reshuffles the pool between fwd and bwd). + with torch.no_grad(): + w1 = torch.stack([getattr(fc1, f"weight{i}").cpu() for i in range(num_experts)]).float() + w2 = torch.stack([getattr(fc2, f"weight{i}").cpu() for i in range(num_experts)]).float() + pre_ref = torch.einsum("th,enh->ten", inp.detach().cpu().float(), w1) # [T, E, 2F] + gate, up = pre_ref[..., :ffn], pre_ref[..., ffn:] + act = torch.nn.functional.silu(gate) * up * probs.detach().cpu().float()[..., None] + y = torch.einsum("tef,ehf->teh", act, w2) # [T, E, H] + ref = (y * routing_map.cpu().float()[..., None]).sum(dim=1) # [T, H] + rel = (out.detach().cpu().float() - ref).norm() / ref.norm().clamp_min(1e-8) + assert rel < 3e-2, rel.item() + + loss = out.float().sum() + loss.backward() + + assert inp.grad is not None + assert probs.grad is not None + for i in range(num_experts): + assert getattr(fc1, f"weight{i}").grad is not None + assert getattr(fc2, f"weight{i}").grad is not None + assert getattr(fc2, f"weight{i}").grad.abs().sum() > 0 + # Sanity: dense route head got a non-zero FC1 output grad (FC2 act-bwd -> FC1). + assert preact.grad[:num_routes].abs().sum() > 0 + + # Numerical check on probs.grad (locks the gated-act backward's per-slot (token, expert) + # mapping): with loss = out.sum(), dL/dprob[t, e] = routing_map[t, e] * . A wrong expert-per-slot map (e.g. dense route index misread as a + # block-padded slot) silently corrupts this even though probs.grad stays non-zero. + with torch.no_grad(): + act_noprob = torch.nn.functional.silu(gate) * up # [T, E, F] (CPU float) + w2sum = w2.sum(dim=1) # [E, F] = sum over output features + dprob_ref = routing_map.cpu().float() * torch.einsum("tef,ef->te", act_noprob, w2sum) + rel_p = (probs.grad.detach().cpu() - dprob_ref).norm() / dprob_ref.norm().clamp_min(1e-8) + assert rel_p < 3e-2, rel_p.item() + + +def test_route_list_fwd_dgrad_wgrad_consistency(): + """fwd + dgrad + wgrad against a single autograd reference in dense route-ordered space.""" + torch.manual_seed(23) + # The gather/dgrad tiles need both contraction dims (in for fwd, out for dgrad) to be a + # multiple of the FlyDSL block_k (64) and >= 128 (K_ITERS >= 2), so use 128-wide features. + num_recv_tokens, in_features, out_features = 96, 128, 128 + num_experts, max_hits = 6, 3 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=5) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + route_to_token, route_expert = _dense_route_order(routing_map) + num_routes = route_to_token.numel() + # Prepare the fwd/dgrad align (block-padded slot layout) up front so we can map the dense + # per-route grad into padded-slot order for the block-padded dgrad. + routing = prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(routing.sorted_slot_ids.shape[0]) + valid = routing.sorted_slot_ids < num_recv_tokens + assert int(valid.sum().item()) == num_routes + + hidden = torch.randn(num_recv_tokens, in_features, device="cuda", dtype=torch.bfloat16) + weights = torch.randn( + num_experts, out_features, in_features, device="cuda", dtype=torch.bfloat16 + ) + + # Forward output is block-padded [em_max]; the valid padded slots are expert-ascending, the + # same order as the dense route list, so out_full[valid] lines up with the dense ref. + out_full = permute_free_grouped_gemm_bf16(hidden, weights, routing) + out = out_full[valid] + # One per-route gradient in the block-padded [em_max] slot layout that both the dgrad and the + # wgrad now route-read (padded slot = block_start[e]*block_size_m + within-rank). The valid + # slots are expert-ascending, matching the dense autograd reference order. ``grad`` keeps a + # dense route-ordered [num_routes] copy only to seed the dense reference backward. + grad = torch.randn(num_routes, out_features, device="cuda", dtype=torch.bfloat16) + grad_bp = torch.zeros(em_max, out_features, device="cuda", dtype=torch.bfloat16) + grad_bp[valid] = grad + dA = permute_free_grouped_gemm_bf16_dgrad(grad_bp, weights, routing) + dW = permute_free_grouped_gemm_bf16_wgrad(hidden, grad_bp, weights.shape, routing) + + # Autograd reference in dense route-ordered space. + ref_hidden = hidden.float().clone().requires_grad_(True) + ref_w = weights.float().clone().requires_grad_(True) + ref_out = torch.einsum( + "rk,rnk->rn", ref_hidden[route_to_token], ref_w[route_expert] + ) + assert _rel_l2(out, ref_out) < 2e-2 + ref_out.backward(grad.float()) + + assert _rel_l2(dA, ref_hidden.grad) < 2e-2 + assert _rel_l2(dW, ref_w.grad) < 2e-2 + + +def test_route_list_fc2_fwd(): + """FC2 forward (F-wide input): per-route GEMM + token gather-combine.""" + torch.manual_seed(51) + num_recv_tokens, in_features, out_features = 128, 128, 128 + num_experts, max_hits = 8, 3 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=21) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + routing = prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(routing.sorted_slot_ids.shape[0]) + slot_token = routing.sorted_slot_ids + slot_expert = _expert_per_route(routing, em_max) + valid = slot_token < num_recv_tokens + tok = slot_token.to(torch.int64).clamp_(0, num_recv_tokens - 1) + exp = slot_expert.to(torch.int64).clamp_min(0) + + fc2_input = torch.zeros(em_max, in_features, device="cuda", dtype=torch.bfloat16) + fc2_input[valid] = torch.randn( + int(valid.sum().item()), in_features, device="cuda", dtype=torch.bfloat16 + ) + weights = torch.randn( + num_experts, out_features, in_features, device="cuda", dtype=torch.bfloat16 + ) + + out = permute_free_grouped_gemm_bf16_fc2(fc2_input, weights, routing) + assert out.shape == (num_recv_tokens, out_features) + + per_slot = torch.einsum("rf,rhf->rh", fc2_input.float(), weights[exp].float()) + per_slot = per_slot * valid[:, None].float() + ref = torch.zeros(num_recv_tokens, out_features, device="cuda", dtype=torch.float32) + ref.index_add_(0, tok, per_slot) + assert _rel_l2(out, ref) < 2e-2 + + +def test_route_list_fc2_fwd_standalone_act(): + """FC2 forward dispatch: standalone gated-act recompute + plain GEMM (no in-kernel fusion).""" + torch.manual_seed(53) + num_recv_tokens, in_features, out_features = 128, 128, 128 + num_experts, max_hits = 8, 3 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=23) + routing = PermuteFreeMetadata( + routing_map=routing_map, num_experts=num_experts, route_space=True, activation="silu" + ) + prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(routing.sorted_slot_ids.shape[0]) + slot_token = routing.sorted_slot_ids + slot_expert = _expert_per_route(routing, em_max) + valid = slot_token < num_recv_tokens + tok = slot_token.to(torch.int64).clamp_(0, num_recv_tokens - 1) + exp = slot_expert.to(torch.int64).clamp_min(0) + + preact = torch.zeros(em_max, 2 * in_features, device="cuda", dtype=torch.bfloat16) + preact[valid] = torch.randn( + int(valid.sum().item()), 2 * in_features, device="cuda", dtype=torch.bfloat16 + ) + probs = torch.rand(num_recv_tokens, num_experts, device="cuda", dtype=torch.float32) + weights = torch.randn( + num_experts, out_features, in_features, device="cuda", dtype=torch.bfloat16 + ) + + out = permute_free_grouped_gemm_forward( + preact, weights, routing, activation="silu", dispatched_probs=probs + ) + assert out.shape == (num_recv_tokens, out_features) + + g = preact[:, :in_features].float() + u = preact[:, in_features:].float() + act = torch.nn.functional.silu(g) * u * probs[tok, exp][:, None] * valid[:, None].float() + per_slot = torch.einsum("rf,rhf->rh", act, weights[exp].float()) * valid[:, None].float() + ref = torch.zeros(num_recv_tokens, out_features, device="cuda", dtype=torch.float32) + ref.index_add_(0, tok, per_slot) + assert _rel_l2(out, ref) < 2e-2 + + +def test_route_list_fc2_dgrad(): + """FC2 dgrad: token-space grad gathered per route into block-padded ``[em_max, F]``.""" + torch.manual_seed(57) + num_recv_tokens, in_features, out_features = 128, 128, 128 + num_experts, max_hits = 8, 4 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=25) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + routing = prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(routing.sorted_slot_ids.shape[0]) + slot_token = routing.sorted_slot_ids + slot_expert = _expert_per_route(routing, em_max) + valid = slot_token < num_recv_tokens + tok = slot_token.to(torch.int64).clamp_(0, num_recv_tokens - 1) + exp = slot_expert.to(torch.int64).clamp_min(0) + + grad_output = torch.randn(num_recv_tokens, out_features, device="cuda", dtype=torch.bfloat16) + weights = torch.randn( + num_experts, out_features, in_features, device="cuda", dtype=torch.bfloat16 + ) + + dgrad = permute_free_grouped_gemm_bf16_fc2_dgrad(grad_output, weights, routing) + assert dgrad.shape == (em_max, in_features) + + per_slot = torch.einsum("rh,rhf->rf", grad_output[tok].float(), weights[exp].float()) + assert _rel_l2(dgrad[valid], per_slot[valid]) < 2e-2 + + +def test_route_list_gated_act_bwd(): + """Standalone gated-activation backward: dpre + dprob vs autograd reference.""" + torch.manual_seed(59) + num_recv_tokens, in_features = 128, 128 + num_experts, max_hits = 8, 4 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=27) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + routing = prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(routing.sorted_slot_ids.shape[0]) + slot_token = routing.sorted_slot_ids + slot_expert = _expert_per_route(routing, em_max) + valid = slot_token < num_recv_tokens + tok = slot_token.to(torch.int64).clamp_(0, num_recv_tokens - 1) + exp = slot_expert.to(torch.int64).clamp_min(0) + + gate = torch.randn(em_max, in_features, device="cuda", dtype=torch.float32) + up = torch.randn(em_max, in_features, device="cuda", dtype=torch.float32) + gate = (gate * valid[:, None].float()).requires_grad_(True) + up = (up * valid[:, None].float()).requires_grad_(True) + probs = torch.rand( + num_recv_tokens, num_experts, device="cuda", dtype=torch.float32, requires_grad=True + ) + grad_out = torch.zeros(em_max, in_features, device="cuda", dtype=torch.bfloat16) + grad_out[valid] = torch.randn( + int(valid.sum().item()), in_features, device="cuda", dtype=torch.bfloat16 + ) + + pr = probs[tok, exp] + act = torch.nn.functional.silu(gate) * up * pr[:, None] * valid[:, None].float() + (act * grad_out.float()).sum().backward() + dpre_ref = torch.cat([gate.grad, up.grad], dim=1) + + preact = torch.cat([gate.detach(), up.detach()], dim=1).to(torch.bfloat16) + dpre, dprob = permute_free_gated_act_bwd( + grad_out, preact, routing, activation="silu", dispatched_probs=probs.detach() + ) + assert dpre.shape == (em_max, 2 * in_features) + assert dprob.shape == (num_recv_tokens, num_experts) + assert _rel_l2(dpre[valid], dpre_ref[valid]) < 2e-2 + assert _rel_l2(dprob, probs.grad) < 2e-2 + + +def test_route_list_wgrad_out_accumulate(): + """FC1 wgrad in-kernel fold into bf16/fp32 ``out`` (overwrite + accumulate).""" + torch.manual_seed(61) + num_recv_tokens, in_features, out_features = 128, 128, 128 + num_experts, max_hits = 8, 4 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=29) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + routing = prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(routing.sorted_slot_ids.shape[0]) + slot_token = routing.sorted_slot_ids + slot_expert = _expert_per_route(routing, em_max) + valid = slot_token < num_recv_tokens + tok = slot_token.to(torch.int64).clamp_(0, num_recv_tokens - 1) + + hidden = torch.randn(num_recv_tokens, in_features, device="cuda", dtype=torch.bfloat16) + weights_shape = (num_experts, out_features, in_features) + grad = torch.zeros(em_max, out_features, device="cuda", dtype=torch.bfloat16) + grad[valid] = torch.randn( + int(valid.sum().item()), out_features, device="cuda", dtype=torch.bfloat16 + ) + + per_slot = torch.einsum("rn,rk->rnk", grad.float(), hidden[tok].float()) + per_slot = per_slot * valid[:, None, None].float() + ref = torch.zeros(num_experts, out_features, in_features, device="cuda", dtype=torch.float32) + ref.index_add_(0, slot_expert.to(torch.int64).clamp_min(0), per_slot) + + for dtype in (torch.bfloat16, torch.float32): + out = torch.zeros(weights_shape, device="cuda", dtype=dtype) + ret = permute_free_grouped_gemm_bf16_wgrad( + hidden, grad, weights_shape, routing, out=out, accumulate=False + ) + assert ret is out + assert _rel_l2(out, ref) < 2e-2 + permute_free_grouped_gemm_bf16_wgrad( + hidden, grad, weights_shape, routing, out=out, accumulate=True + ) + assert _rel_l2(out, 2.0 * ref) < 2e-2 + + +def test_permute_free_forward_dispatch(): + """``permute_free_grouped_gemm_forward`` routes FC1 plain GEMM and FC2 standalone-act paths.""" + torch.manual_seed(63) + num_recv_tokens, hidden_dim, ffn = 128, 128, 128 + num_experts, max_hits = 8, 3 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=31) + + fc1_routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + fc1_routing = prepare_moe_align(fc1_routing, _FLYDSL_FWD_BLOCK_M) + hidden = torch.randn(num_recv_tokens, hidden_dim, device="cuda", dtype=torch.bfloat16) + w1 = torch.randn(num_experts, 2 * ffn, hidden_dim, device="cuda", dtype=torch.bfloat16) + res_fc1 = permute_free_grouped_gemm_forward(hidden, w1, fc1_routing) + direct = permute_free_grouped_gemm_bf16(hidden, w1, fc1_routing) + valid_fc1 = fc1_routing.sorted_slot_ids < num_recv_tokens + assert _rel_l2(res_fc1[valid_fc1], direct[valid_fc1]) < 1e-3 + + fc2_routing = PermuteFreeMetadata( + routing_map=routing_map, num_experts=num_experts, route_space=True, activation="silu" + ) + prepare_moe_align(fc2_routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(fc2_routing.sorted_slot_ids.shape[0]) + slot_token = fc2_routing.sorted_slot_ids + slot_expert = _expert_per_route(fc2_routing, em_max) + valid = slot_token < num_recv_tokens + tok = slot_token.to(torch.int64).clamp_(0, num_recv_tokens - 1) + exp = slot_expert.to(torch.int64).clamp_min(0) + + preact = torch.zeros(em_max, 2 * ffn, device="cuda", dtype=torch.bfloat16) + preact[valid] = torch.randn( + int(valid.sum().item()), 2 * ffn, device="cuda", dtype=torch.bfloat16 + ) + probs = torch.rand(num_recv_tokens, num_experts, device="cuda", dtype=torch.float32) + w2 = torch.randn(num_experts, hidden_dim, ffn, device="cuda", dtype=torch.bfloat16) + out = permute_free_grouped_gemm_forward( + preact, w2, fc2_routing, activation="silu", dispatched_probs=probs + ) + assert out.shape == (num_recv_tokens, hidden_dim) + + g = preact[:, :ffn].float() + u = preact[:, ffn:].float() + act = torch.nn.functional.silu(g) * u * probs[tok, exp][:, None] * valid[:, None].float() + per_slot = torch.einsum("rf,rhf->rh", act, w2[exp].float()) * valid[:, None].float() + ref = torch.zeros(num_recv_tokens, hidden_dim, device="cuda", dtype=torch.float32) + ref.index_add_(0, tok, per_slot) + assert _rel_l2(out, ref) < 2e-2 + + +def test_is_permute_free_grouped_gemm_enabled(monkeypatch): + """Env gate: True only on ROCm with ``NVTE_PERMUTE_FREE_GROUPED_GEMM=1``.""" + monkeypatch.setenv("NVTE_PERMUTE_FREE_GROUPED_GEMM", "1") + assert is_permute_free_grouped_gemm_enabled() is True + monkeypatch.setenv("NVTE_PERMUTE_FREE_GROUPED_GEMM", "0") + assert is_permute_free_grouped_gemm_enabled() is False + monkeypatch.delenv("NVTE_PERMUTE_FREE_GROUPED_GEMM", raising=False) + assert is_permute_free_grouped_gemm_enabled() is False + + +def test_slot_expert_ids_cache(): + """``prepare_moe_align`` caches per-slot expert ids; ``_expert_per_route`` reuses the cache.""" + torch.manual_seed(67) + num_recv_tokens, num_experts, max_hits = 128, 8, 3 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=33) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + + # Not built until the align runs. + assert routing.slot_expert_ids is None + prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(routing.sorted_slot_ids.shape[0]) + + # Populated at align time with the right shape/dtype. + assert routing.slot_expert_ids is not None + assert routing.slot_expert_ids.shape == (em_max,) + assert routing.slot_expert_ids.dtype == torch.int32 + + # Matches a manual block-broadcast of ``expert_ids`` (expert_ids[slot // block_m]). + block_m = int(routing.block_size_m) + pos = torch.arange(em_max, device="cuda", dtype=torch.int64) + block_idx = (pos // block_m).clamp(max=routing.expert_ids.numel() - 1) + manual = routing.expert_ids[block_idx].to(torch.int32) + assert torch.equal(routing.slot_expert_ids, manual) + + # ``_expert_per_route`` returns the *cached* tensor (no redundant recompute). + assert _expert_per_route(routing, em_max) is routing.slot_expert_ids + + # Valid slots carry the expert-sorted dense route experts. + _, route_expert = _dense_route_order(routing_map) + valid = routing.sorted_slot_ids < num_recv_tokens + assert torch.equal(routing.slot_expert_ids[valid].to(torch.int64), route_expert) + + # Rebuild-on-early-return: drop the cache but keep the align, then re-prepare with the same + # block_m (hits the cached-align early return) -- it must repopulate rather than stay None. + routing.slot_expert_ids = None + prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + assert routing.slot_expert_ids is not None + assert torch.equal(routing.slot_expert_ids, manual) + + +def test_permute_free_backward_fc1_dispatch(): + """``permute_free_grouped_gemm_backward`` FC1 path (``route_space=False``): dgrad + wgrad match + the standalone kernels, and ``wgrad_out`` folds in place (``wgrad_applied``).""" + torch.manual_seed(71) + num_recv_tokens, in_features, out_features = 128, 128, 128 + num_experts, max_hits = 8, 3 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=35) + routing = PermuteFreeMetadata(routing_map=routing_map, num_experts=num_experts) # FC1 + routing = prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(routing.sorted_slot_ids.shape[0]) + valid = routing.sorted_slot_ids < num_recv_tokens + + hidden = torch.randn(num_recv_tokens, in_features, device="cuda", dtype=torch.bfloat16) + weights = [ + torch.randn(out_features, in_features, device="cuda", dtype=torch.bfloat16) + for _ in range(num_experts) + ] + W = torch.stack(weights, dim=0) + grad = torch.zeros(em_max, out_features, device="cuda", dtype=torch.bfloat16) + grad[valid] = torch.randn( + int(valid.sum().item()), out_features, device="cuda", dtype=torch.bfloat16 + ) + + res = permute_free_grouped_gemm_backward( + grad, routing=routing, weights=weights, num_gemms=num_experts, + hidden_states=hidden, requires_dgrad=True, requires_wgrad=True, + ) + dgrad_ref = permute_free_grouped_gemm_bf16_dgrad(grad, W, routing) + wgrad_ref = permute_free_grouped_gemm_bf16_wgrad( + hidden, grad, (num_experts, out_features, in_features), routing + ) + assert res.dgrad.shape == (num_recv_tokens, in_features) + assert res.wgrad_stacked.shape == (num_experts, out_features, in_features) + assert res.wgrad_applied is False + assert res.grad_probs is None + assert _rel_l2(res.dgrad, dgrad_ref) < 1e-3 + assert _rel_l2(res.wgrad_stacked, wgrad_ref) < 1e-3 + + # wgrad_out fold (fp32 sink, overwrite): wgrad_applied True and the buffer is returned. + wgrad_out = torch.zeros( + num_experts, out_features, in_features, device="cuda", dtype=torch.float32 + ) + res2 = permute_free_grouped_gemm_backward( + grad, routing=routing, weights=weights, num_gemms=num_experts, + hidden_states=hidden, requires_wgrad=True, + wgrad_out=wgrad_out, wgrad_accumulate=False, + ) + assert res2.wgrad_applied is True + assert res2.wgrad_stacked is wgrad_out + assert res2.dgrad is None + assert _rel_l2(wgrad_out, wgrad_ref) < 5e-3 + + +def test_permute_free_backward_fc2_dgrad_probs(): + """FC2 backward dispatch dgrad path: dgrad + grad_probs match the standalone act-bwd, and + grad_probs is suppressed when ``dispatched_probs`` does not require grad.""" + torch.manual_seed(73) + num_recv_tokens, in_features, out_features = 128, 128, 128 # F=in, H=out + num_experts, max_hits = 8, 4 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=37) + routing = PermuteFreeMetadata( + routing_map=routing_map, num_experts=num_experts, route_space=True, activation="silu" + ) + prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(routing.sorted_slot_ids.shape[0]) + valid = routing.sorted_slot_ids < num_recv_tokens + + preact = torch.zeros(em_max, 2 * in_features, device="cuda", dtype=torch.bfloat16) + preact[valid] = torch.randn( + int(valid.sum().item()), 2 * in_features, device="cuda", dtype=torch.bfloat16 + ) + grad_output = torch.randn(num_recv_tokens, out_features, device="cuda", dtype=torch.bfloat16) + weights = [ + torch.randn(out_features, in_features, device="cuda", dtype=torch.bfloat16) + for _ in range(num_experts) + ] + W = torch.stack(weights, dim=0) + + probs = torch.rand( + num_recv_tokens, num_experts, device="cuda", dtype=torch.float32, requires_grad=True + ) + res = permute_free_grouped_gemm_backward( + grad_output, routing=routing, weights=weights, num_gemms=num_experts, + hidden_states=preact, requires_dgrad=True, + fc2_activation="silu", dispatched_probs=probs, + ) + # Reference: fc2_dgrad -> standalone gated act-bwd. + dgrad_f = permute_free_grouped_gemm_bf16_fc2_dgrad(grad_output, W, routing) + dpre_ref, dprob_ref = permute_free_gated_act_bwd( + dgrad_f, preact, routing, activation="silu", dispatched_probs=probs.detach() + ) + assert res.dgrad.shape == (em_max, 2 * in_features) + assert res.grad_probs is not None + assert res.grad_probs.shape == (num_recv_tokens, num_experts) + assert _rel_l2(res.dgrad[valid], dpre_ref[valid]) < 2e-2 + assert _rel_l2(res.grad_probs, dprob_ref) < 2e-2 + + # No prob grad requested -> grad_probs suppressed. + res_nop = permute_free_grouped_gemm_backward( + grad_output, routing=routing, weights=weights, num_gemms=num_experts, + hidden_states=preact, requires_dgrad=True, + fc2_activation="silu", dispatched_probs=probs.detach(), + ) + assert res_nop.grad_probs is None + assert _rel_l2(res_nop.dgrad[valid], dpre_ref[valid]) < 2e-2 + + +def test_route_list_gated_act_bwd_gelu(): + """Gelu (tanh) variant of the standalone gated-activation backward (gelu path).""" + torch.manual_seed(79) + num_recv_tokens, in_features = 128, 128 + num_experts, max_hits = 8, 3 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=39) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + routing = prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(routing.sorted_slot_ids.shape[0]) + slot_token = routing.sorted_slot_ids + slot_expert = _expert_per_route(routing, em_max) + valid = slot_token < num_recv_tokens + tok = slot_token.to(torch.int64).clamp_(0, num_recv_tokens - 1) + exp = slot_expert.to(torch.int64).clamp_min(0) + + gate = torch.randn(em_max, in_features, device="cuda", dtype=torch.float32) + up = torch.randn(em_max, in_features, device="cuda", dtype=torch.float32) + gate = (gate * valid[:, None].float()).requires_grad_(True) + up = (up * valid[:, None].float()).requires_grad_(True) + probs = torch.rand( + num_recv_tokens, num_experts, device="cuda", dtype=torch.float32, requires_grad=True + ) + grad_out = torch.zeros(em_max, in_features, device="cuda", dtype=torch.bfloat16) + grad_out[valid] = torch.randn( + int(valid.sum().item()), in_features, device="cuda", dtype=torch.bfloat16 + ) + + pr = probs[tok, exp] + act = ( + torch.nn.functional.gelu(gate, approximate="tanh") + * up + * pr[:, None] + * valid[:, None].float() + ) + (act * grad_out.float()).sum().backward() + dpre_ref = torch.cat([gate.grad, up.grad], dim=1) + + preact = torch.cat([gate.detach(), up.detach()], dim=1).to(torch.bfloat16) + dpre, dprob = permute_free_gated_act_bwd( + grad_out, preact, routing, activation="gelu", dispatched_probs=probs.detach() + ) + assert dpre.shape == (em_max, 2 * in_features) + assert _rel_l2(dpre[valid], dpre_ref[valid]) < 3e-2 + assert _rel_l2(dprob, probs.grad) < 3e-2 + + +# --------------------------------------------------------------------------- +# Module-level weight-gradient accumulation: grouped weight (-> main_grad) vs. +# separate per-expert weights (-> autograd .grad). +# --------------------------------------------------------------------------- +def _make_grouped_linear(num_gemms, in_features, out_features, *, grouped, fuse_wgrad): + mod = GroupedLinear( + num_gemms, + in_features, + out_features, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + fuse_wgrad_accumulation=fuse_wgrad, + single_grouped_weight=grouped, + ) + return mod + + +def test_grouped_weight_main_grad_matches_ungrouped_grad(monkeypatch): + """Permute-free wgrad lands in the grouped param's ``main_grad`` and matches the + ungrouped autograd ``.grad`` expert-for-expert (same kernel dW, different sink).""" + monkeypatch.setenv("NVTE_PERMUTE_FREE_GROUPED_GEMM", "1") + # single_grouped_weight requires this gate, else the module falls back to per-expert params. + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + torch.manual_seed(41) + + # The fwd/dgrad tiles need both contraction dims (in for fwd, out for dgrad) >= 128 and a + # multiple of block_k (64); use 128-wide features. + num_recv_tokens, in_features, out_features = 128, 128, 128 + num_experts, max_hits = 8, 3 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=13) + num_routes = int(routing_map.sum().item()) + m_splits = [num_recv_tokens // num_experts] * num_experts + + # inp requires grad (as in real training, where it comes from a prior layer): for the + # grouped path the detached per-expert views do not drive autograd, so the activation is + # what makes the Function output require grad and triggers the backward. + inp = torch.randn( + num_recv_tokens, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) + # Shared initial weights for both modules (so the kernel dW is identical). + W = torch.randn(num_experts, out_features, in_features, device="cuda", dtype=torch.bfloat16) + # Shared upstream gradient over the dense route-ordered range. + g = torch.randn(num_routes, out_features, device="cuda", dtype=torch.bfloat16) + + # --- ungrouped: separate per-expert params, wgrad via autograd .grad --- + mod_a = _make_grouped_linear( + num_experts, in_features, out_features, grouped=False, fuse_wgrad=False + ) + with torch.no_grad(): + for i in range(num_experts): + getattr(mod_a, f"weight{i}").copy_(W[i]) + routing_a = PermuteFreeMetadata(routing_map=routing_map, num_experts=num_experts) + out_a = mod_a(inp, m_splits, permute_free_metadata=routing_a) + # Block-padded canonical: the module output is [em_max, out] in padded slot order; the valid + # (expert-ascending) slots line up with the dense upstream grad g. Slicing [:num_routes] + # would cut across the padding gaps, so gather the valid slots instead. + valid_a = routing_a.sorted_slot_ids < num_recv_tokens + (out_a[valid_a].float() * g.float()).sum().backward() + grad_a = torch.stack( + [getattr(mod_a, f"weight{i}").grad.float() for i in range(num_experts)], dim=0 + ) + + # --- grouped: single grouped param, wgrad accumulated into main_grad --- + mod_b = _make_grouped_linear( + num_experts, in_features, out_features, grouped=True, fuse_wgrad=True + ) + with torch.no_grad(): + for i, view in enumerate(mod_b._get_weight_tensors()): + view.copy_(W[i]) + mod_b.weight.main_grad = torch.zeros( + num_experts, out_features, in_features, device="cuda", dtype=torch.float32 + ) + mod_b.weight.grad_added_to_main_grad = False + routing_b = PermuteFreeMetadata(routing_map=routing_map, num_experts=num_experts) + out_b = mod_b(inp, m_splits, permute_free_metadata=routing_b) + valid_b = routing_b.sorted_slot_ids < num_recv_tokens + (out_b[valid_b].float() * g.float()).sum().backward() + + # The grouped param collects grad in main_grad, not in .grad. + assert mod_b.weight.grad is None + assert getattr(mod_b.weight, "grad_added_to_main_grad", False) is True + main_grad = mod_b.weight.main_grad.view(num_experts, out_features, in_features).float() + + # Both paths fold the same bf16 kernel dW into their sink (grouped -> fp32 main_grad, + # ungrouped -> bf16 .grad), so they match to within bf16 rounding (~1e-3 relative). + assert _rel_l2(main_grad, grad_a) < 5e-3 + # And every expert actually received a non-zero gradient (guards against frozen experts). + for e in range(num_experts): + assert main_grad[e].abs().sum() > 0, f"expert {e} has a zero main_grad (frozen)." + + +def test_grouped_weight_nonfused_grad_matches_ungrouped(monkeypatch): + """single_grouped_weight without fuse_wgrad_accumulation routes the [E, out, in] wgrad into + the grouped param's autograd ``.grad`` (the GEMM only sees detached per-expert views), and it + must match the ungrouped per-expert ``.grad`` expert-for-expert and accumulate across + backwards.""" + monkeypatch.setenv("NVTE_PERMUTE_FREE_GROUPED_GEMM", "1") + # single_grouped_weight requires this gate, else the module falls back to per-expert params. + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + torch.manual_seed(43) + + # 128-wide features to satisfy the fwd/dgrad tile (K>=128, multiple of block_k=64). + num_recv_tokens, in_features, out_features = 128, 128, 128 + num_experts, max_hits = 8, 3 + m_splits = [num_recv_tokens // num_experts] * num_experts + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=17) + num_routes = int(routing_map.sum().item()) + # inp drives autograd for the grouped path (detached views carry no grad edge). + inp = torch.randn( + num_recv_tokens, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) + W = torch.randn(num_experts, out_features, in_features, device="cuda", dtype=torch.bfloat16) + g = torch.randn(num_routes, out_features, device="cuda", dtype=torch.bfloat16) + + # --- ungrouped reference: separate per-expert params, wgrad via autograd .grad --- + mod_a = _make_grouped_linear( + num_experts, in_features, out_features, grouped=False, fuse_wgrad=False + ) + with torch.no_grad(): + for i in range(num_experts): + getattr(mod_a, f"weight{i}").copy_(W[i]) + routing_a = PermuteFreeMetadata(routing_map=routing_map, num_experts=num_experts) + out_a = mod_a(inp, m_splits, permute_free_metadata=routing_a) + # Block-padded canonical: gather the valid (expert-ascending) slots rather than slicing + # [:num_routes], which would cut across the padding gaps. + valid_a = routing_a.sorted_slot_ids < num_recv_tokens + (out_a[valid_a].float() * g.float()).sum().backward() + grad_a = torch.stack( + [getattr(mod_a, f"weight{i}").grad.float() for i in range(num_experts)], dim=0 + ) + + # --- grouped, no fusion: wgrad accumulates into the grouped param's .grad --- + mod_b = _make_grouped_linear( + num_experts, in_features, out_features, grouped=True, fuse_wgrad=False + ) + with torch.no_grad(): + for i, view in enumerate(mod_b._get_weight_tensors()): + view.copy_(W[i]) + routing_b = PermuteFreeMetadata(routing_map=routing_map, num_experts=num_experts) + out_b = mod_b(inp, m_splits, permute_free_metadata=routing_b) + valid_b = routing_b.sorted_slot_ids < num_recv_tokens + (out_b[valid_b].float() * g.float()).sum().backward() + + # Grad lands in the grouped param's .grad (not main_grad), shape [E, out, in]. + assert getattr(mod_b.weight, "grad_added_to_main_grad", False) is False + assert mod_b.weight.grad is not None + assert tuple(mod_b.weight.grad.shape) == (num_experts, out_features, in_features) + grad_b = mod_b.weight.grad.view(num_experts, out_features, in_features).float() + assert _rel_l2(grad_b, grad_a) < 1e-3 + for e in range(num_experts): + assert grad_b[e].abs().sum() > 0, f"expert {e} has a zero grad (frozen)." + + # A second backward without zero_grad accumulates in place (grad-accumulation semantics). + out_b2 = mod_b(inp, m_splits, permute_free_metadata=routing_b) + (out_b2[valid_b].float() * g.float()).sum().backward() + grad_b2 = mod_b.weight.grad.view(num_experts, out_features, in_features).float() + assert _rel_l2(grad_b2, 2.0 * grad_a) < 1e-3 diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 6a1e84dfe..0e5b03523 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -9,6 +9,7 @@ from typing import Iterable, Optional, Tuple, Union, List import os import functools +import warnings import torch from torch.utils.cpp_extension import IS_HIP_EXTENSION import transformer_engine_torch as tex @@ -493,7 +494,38 @@ def general_gemm( } if not _is_nvfp4_row_scaled_tensor(A) and not _is_nvfp4_row_scaled_tensor(B): - out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*args, **kwargs) + use_gemm_flydsl = (IS_HIP_EXTENSION + and get_device_compute_capability() == (9, 5) + and bool(int(os.environ.get("NVTE_USE_FLYDSL", "0")))) + if use_gemm_flydsl: + # Lazy import keeps FlyDSL off the normal Transformer Engine import path. + from ..flydsl_kernels.gemm import ( + FlyDSLUnsupportedError, + te_generic_gemm_flydsl, + ) + + try: + out, bias_grad, gelu_input, extra_output = te_generic_gemm_flydsl( + *args, + **kwargs, + ) + 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 GEMM does not support this configuration; " + f"falling back to the default backend. Reason: {exc}", + UserWarning, + stacklevel=2, + ) + + out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*args, **kwargs) + else: + out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*args, **kwargs) else: if _is_nvfp4_row_scaled_tensor(A): raise NotImplementedError("Row-scaled NVFP4 GEMM does not support row-scaled A.") diff --git a/transformer_engine/pytorch/flydsl_kernels/__init__.py b/transformer_engine/pytorch/flydsl_kernels/__init__.py new file mode 100644 index 000000000..c64b988c6 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/__init__.py @@ -0,0 +1,3 @@ +from . import gemm + +__all__ = ["gemm"] diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py b/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py new file mode 100644 index 000000000..4eae105ef --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""FlyDSL GEMM kernels (dense, non-grouped) for BF16/FP16/FP32/FP8/MXFP8.""" + +from .exceptions import FlyDSLUnsupportedError +from .gemm_wrappers import te_generic_gemm_flydsl + +__all__ = [ + "FlyDSLUnsupportedError", + "te_generic_gemm_flydsl", +] diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py new file mode 100644 index 000000000..7764803a4 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py @@ -0,0 +1,1622 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""FlyDSL BF16 TN/NN/NT 4-wave GEMM kernel for Transformer Engine. + +All supported layouts share one source-level kernel generator while compiling +to separate cached binaries: + + TN: A [M,K] normal read, B [N,K] normal read + NN: A [M,K] normal read, B [K,N] transpose read + NT: A [K,M] transpose read, B [K,N] transpose read + +The layout is a Python-only cache key. Global addressing and LDS fragment +reconstruction are selected while building each specialization, so no runtime +layout branch is emitted in the GEMM kernel. + +This module imports ``flydsl`` at import time and must therefore be imported +lazily only after FlyDSL availability has been confirmed. +""" + +import functools + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.compiler.ast_rewriter import ASTRewriter +from flydsl._mlir.dialects import llvm +from flydsl.expr import arith, buffer_ops, const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec + +# Transformer Engine-local FlyDSL utilities. +from .exceptions import FlyDSLUnsupportedError +from .fp16_gemm_utils import ( + G2SLoader, + S2RLoader, + compute_global_bf16_transpose_swizzle, + compute_global_swizzle, + make_bf16_byte_buffer_tensor, + pack_i32x4_i32x8, + swizzle_128, +) +from .fp8_gemm_utils import wait_barrier + + +_BLOCK_M = 256 +_BLOCK_N = 256 +_BLOCK_K = 64 + +# Public metadata consumed by wrappers. +BLOCK_M = _BLOCK_M +BLOCK_N = _BLOCK_N +BLOCK_K = _BLOCK_K + +NUM_THREADS = 256 +WARP_SIZE = 64 +NUM_WAVES = NUM_THREADS // WARP_SIZE + +SUBTILE_M = 64 +SUBTILE_N = 64 + +MFMA_M = 16 +MFMA_N = 16 + +SUBTILES_PER_WAVE = 4 +MFMA_M_PER_SUBTILE = SUBTILE_M // MFMA_M +MFMA_N_PER_SUBTILE = SUBTILE_N // MFMA_N +ACCS_PER_WAVE = SUBTILES_PER_WAVE * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + +ELEM_BYTES = 2 +VEC_BYTES = 16 + +LDS_ELEMS_A = BLOCK_M * BLOCK_K +LDS_ELEMS_B = BLOCK_N * BLOCK_K +LDS_BYTES_A = LDS_ELEMS_A * ELEM_BYTES +LDS_BYTES_B = LDS_ELEMS_B * ELEM_BYTES + +LOAD_PASSES_A = LDS_BYTES_A // (NUM_THREADS * VEC_BYTES) +LOAD_PASSES_B = LDS_BYTES_B // (NUM_THREADS * VEC_BYTES) +LOAD_PASSES_A_SUBTILE = LOAD_PASSES_A // 2 +LOAD_PASSES_B_SUBTILE = LOAD_PASSES_B // 2 +PASSES_PER_A_MI = LOAD_PASSES_A_SUBTILE // MFMA_M_PER_SUBTILE + +LDS_SYM_A0 = "bf16_pp_smem_a0" +LDS_SYM_A1 = "bf16_pp_smem_a1" +LDS_SYM_B0 = "bf16_pp_smem_b0" +LDS_SYM_B1 = "bf16_pp_smem_b1" +LDS_ALIAS_DOMAIN = '#llvm.alias_scope_domain' +SCOPE_IDS = ("a0", "a1", "b0", "b1") + +assert BLOCK_K == 64 +# DO NOT CHANGE THE FOLLOWING LINE. +assert NUM_THREADS == 256 +assert LOAD_PASSES_A * NUM_THREADS * VEC_BYTES == LDS_BYTES_A +assert LOAD_PASSES_B * NUM_THREADS * VEC_BYTES == LDS_BYTES_B +assert LOAD_PASSES_A % 2 == 0 +assert LOAD_PASSES_B % 2 == 0 + + + +def swizzle_xor16(row, col_in_bytes): + """XOR swizzle for the LDS K-byte coordinate.""" + chunk = col_in_bytes // fx.Index(VEC_BYTES) + byte_in_chunk = col_in_bytes % fx.Index(VEC_BYTES) + row_bits = (row % fx.Index(16)) // fx.Index(2) + swz_chunk = chunk ^ row_bits + return swz_chunk * fx.Index(VEC_BYTES) + byte_in_chunk + + +def _encode_waitcnt(vmcnt=63, lgkmcnt=15): + """Encode the CDNA4/gfx950 ``S_WAITCNT`` SIMM16 operand. + + ``rocdl.s_waitcnt`` accepts the raw 16-bit immediate operand of the + 32-bit ``S_WAITCNT`` ISA instruction. On CDNA4, that SIMM16 field is: + + SIMM16[3:0] = vmcnt[3:0] + SIMM16[6:4] = expcnt[2:0] + SIMM16[11:8] = lgkmcnt[3:0] + SIMM16[15:14] = vmcnt[5:4] + + ``vmcnt`` is therefore one six-bit counter split across two noncontiguous + fields; bits [5:4] are placed in SIMM16[15:14], while bits [3:0] remain + in SIMM16[3:0]. + + A wait-counter field set to its maximum representable value is effectively + unconstrained: the instruction does not wait on that counter. This helper + always encodes ``expcnt=7`` and defaults to ``vmcnt=63`` and ``lgkmcnt=15``, + so callers specify only the counters on which they intend to wait. + + For example, ``_encode_waitcnt(lgkmcnt=0)`` returns ``0xC07F``, which the + assembler renders as ``s_waitcnt lgkmcnt(0)``. + See: https://llvm.org/docs/AMDGPU/gfx9_waitcnt.html + """ + if not 0 <= vmcnt <= 63: + raise ValueError(f"vmcnt must be in [0, 63], got {vmcnt}") + if not 0 <= lgkmcnt <= 15: + raise ValueError(f"lgkmcnt must be in [0, 15], got {lgkmcnt}") + + return ( + (7 << 4) # expcnt=7 -> SIMM16[6:4] (unconstrained) + | (vmcnt & 0x0F) # vmcnt[3:0] -> SIMM16[3:0] + | ((lgkmcnt & 0x0F) << 8) # lgkmcnt[3:0] -> SIMM16[11:8] + | ((vmcnt & 0x30) << 10) # vmcnt[5:4] -> SIMM16[15:14] + ) + + +# Keep the documented gfx950 encoding invariant executable and import-time cheap. +assert _encode_waitcnt(lgkmcnt=0) == 0xC07F + + +def _barrier(vmcnt=63, lgkmcnt=15): + if vmcnt != 63 or lgkmcnt != 15: + rocdl.s_waitcnt(_encode_waitcnt(vmcnt=vmcnt, lgkmcnt=lgkmcnt)) + rocdl.s_barrier() + +def _min(a, b): + return arith.select(a < b, a, b) + + +def _divmod(a, b): + return a // b, a % b + + +def _xcd_swizzle(num_pid_m, num_pid_n): + NUM_XCDS = 8 + WGM = 4 + NUM_CUS = 32 * NUM_XCDS + SWIZZLE_THRESHOLD = 4 * NUM_CUS + + wgid = fx.block_idx.x + num_wg = num_pid_m * num_pid_n + + # Simple row-major path. + simple_m, simple_n = _divmod(wgid, num_pid_n) + + # XCD-remapped grouped-M path. + intra_xcd, xcd = _divmod(wgid, NUM_XCDS) + wgid_remap = xcd * (num_wg // NUM_XCDS) + intra_xcd + num_wgid_in_group = WGM * num_pid_n + group_id, intra_group = _divmod(wgid_remap, num_wgid_in_group) + first_pid_m = group_id * WGM + group_size_m = _min(num_pid_m - first_pid_m, WGM) + pid_n, intra_group_m = _divmod(intra_group, group_size_m) + pid_m = first_pid_m + intra_group_m + + use_simple = (num_wg < SWIZZLE_THRESHOLD) | (num_wg % NUM_XCDS != 0) + return ( + arith.select(use_simple, simple_m, pid_m), + arith.select(use_simple, simple_n, pid_n), + ) + + +def _compile_kernel( + K: int, + output_dtype: torch.dtype, + layout: str, + use_xcd_remap: bool = True, +): + """Build one compile-time-specialized TN, NN, or NT BF16 kernel. + + ``K`` must contain at least four K64 tiles. Runtime M/N are expected to + be exact multiples of ``BLOCK_M``/``BLOCK_N``; the kernel has no edge masks. + """ + if layout not in ("TN", "NN", "NT"): + raise ValueError(f"Unsupported BF16 kernel layout: {layout}") + + a_transpose_read = layout == "NT" + b_transpose_read = layout in ("NN", "NT") + + BLOCK_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K + NUM_THREADS = 256 + WARP_SIZE = 64 + + SUBTILE_M = 64 + SUBTILE_N = 64 + + MFMA_M = 16 + MFMA_N = 16 + + SUBTILES_PER_WAVE = 4 + MFMA_M_PER_SUBTILE = SUBTILE_M // MFMA_M + MFMA_N_PER_SUBTILE = SUBTILE_N // MFMA_N + ACCS_PER_WAVE = SUBTILES_PER_WAVE * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + + ELEM_BYTES = 2 + VEC_BYTES = 16 + + if output_dtype == torch.float16: + output_element_bytes = 2 + output_fx_dtype = fx.Float16 + elif output_dtype == torch.bfloat16: + output_element_bytes = 2 + output_fx_dtype = fx.BFloat16 + elif output_dtype == torch.float32: + output_element_bytes = 4 + output_fx_dtype = fx.Float32 + else: + raise TypeError( + "FlyDSL BF16 GEMM output dtype must be torch.float16, " + f"torch.bfloat16, or torch.float32, got {output_dtype}" + ) + + LDS_ELEMS_A = BLOCK_M * BLOCK_K + LDS_ELEMS_B = BLOCK_N * BLOCK_K + LDS_BYTES_A = LDS_ELEMS_A * ELEM_BYTES + LDS_BYTES_B = LDS_ELEMS_B * ELEM_BYTES + + LOAD_PASSES_A = LDS_BYTES_A // (NUM_THREADS * VEC_BYTES) + LOAD_PASSES_B = LDS_BYTES_B // (NUM_THREADS * VEC_BYTES) + LOAD_PASSES_A_SUBTILE = LOAD_PASSES_A // 2 + LOAD_PASSES_B_SUBTILE = LOAD_PASSES_B // 2 + + assert K % BLOCK_K == 0, f"K must be a multiple of {BLOCK_K}, got {K}" + NUM_K_TILES = K // BLOCK_K + assert NUM_K_TILES >= 4, f"K={K} gives {NUM_K_TILES} K64 tiles; the two-page pipeline needs at least 4" + + LDS_ELEMS_HALF = (BLOCK_M // 2) * BLOCK_K + LDS_BYTES_HALF = LDS_ELEMS_HALF * ELEM_BYTES + LOAD_PASSES_HALF = LDS_BYTES_HALF // (NUM_THREADS * VEC_BYTES) + assert LOAD_PASSES_HALF == LOAD_PASSES_A_SUBTILE == LOAD_PASSES_B_SUBTILE + + # Resolve layout-specific addressing and fragment reads before capture. + Q0_SCHED_DSRD = 4 if a_transpose_read else 2 + PREFETCH_SCHED_DSRD = 8 if a_transpose_read else 4 + + if a_transpose_read: + 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) + ) + + def _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, + ): + del load_frag_half_at_byte_base, lane_mod_16 + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + local_m_tile = ( + subtile_m_idx * fx.Index(SUBTILE_M) + + fx.Index(mi * MFMA_M) + - fx.Index(sm * (BLOCK_M // 2)) + ) + return load_transposed_frag_half(lds_a[sm], local_m_tile, half) + else: + def _a_leading_dim_bytes(c_m): + del c_m + return K * ELEM_BYTES + + 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(K * ELEM_BYTES) + + k_base * fx.Index(ELEM_BYTES) + ) + + def _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, + ): + 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)) + return load_frag_half_at_byte_base( + lds_a[sm], + half_row * fx.Index(BLOCK_K * ELEM_BYTES), + half, + ) + + if b_transpose_read: + 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) + ) + + def _load_b_ni( + load_transposed_frag, + load_normal_b_frag, + lds_b, + sn, + ni, + reg_subtile_n_idx0, + lane_mod_16, + ): + del load_normal_b_frag, lane_mod_16 + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + local_n_tile = ( + subtile_n_idx * fx.Index(SUBTILE_N) + + fx.Index(ni * MFMA_N) + - fx.Index(sn * (BLOCK_N // 2)) + ) + return load_transposed_frag(lds_b[sn], local_n_tile) + else: + def _b_leading_dim_bytes(c_n): + del c_n + return K * ELEM_BYTES + + 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(K * ELEM_BYTES) + + k_base * fx.Index(ELEM_BYTES) + ) + + def _load_b_ni( + load_transposed_frag, + load_normal_b_frag, + lds_b, + sn, + ni, + reg_subtile_n_idx0, + 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 + ) + return load_normal_b_frag(lds_b, b_row_addr, sn) + + # Resolve global staging maps before FlyDSL captures ``kernel_gemm``. + # BF16 uses K64, so each transpose-read half-page is two independent + # [K64, X64] slices with 128-byte physical rows. + if a_transpose_read: + def _a_global_offsets(lane, wave_id, c_m): + return compute_global_bf16_transpose_swizzle( + lane, + wave_id, + _a_leading_dim_bytes(c_m), + LOAD_PASSES_HALF, + ) + else: + def _a_global_offsets(lane, wave_id, c_m): + del c_m + return compute_global_swizzle( + lane, + wave_id, + K * ELEM_BYTES, + LOAD_PASSES_HALF, + preshuffled=False, + ) + + if b_transpose_read: + def _b_global_offsets(lane, wave_id, c_n): + return compute_global_bf16_transpose_swizzle( + lane, + wave_id, + _b_leading_dim_bytes(c_n), + LOAD_PASSES_HALF, + ) + else: + def _b_global_offsets(lane, wave_id, c_n): + del c_n + return compute_global_swizzle( + lane, + wave_id, + K * ELEM_BYTES, + LOAD_PASSES_HALF, + preshuffled=False, + ) + + @fx.struct + class SharedStorage: + # Preserve the passing TN byte-staging contract exactly. A BF16 K64 + # half-page is 128 rows x 128 bytes = 16 KiB. + a0_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a0_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a1_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a1_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b0_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + 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] + + @flyc.kernel(known_block_size=[NUM_THREADS, 1, 1]) + def kernel_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + ): + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + lds_a0 = (lds.a0_0, lds.a0_1) + lds_a1 = (lds.a1_0, lds.a1_1) + lds_b0 = (lds.b0_0, lds.b0_1) + lds_b1 = (lds.b1_0, lds.b1_1) + + # A/B arrive as contiguous uint8 byte views of the original + # row-major BF16 tensors. This preserves the validated 16-byte + # BufferCopyLDS128b path and byte-based address arithmetic. + gA = make_bf16_byte_buffer_tensor(A) + gB = make_bf16_byte_buffer_tensor(B) + 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") + + 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 BF16 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) + + # C can exceed the signed-i32 element/byte offset range for large M*N. + # Bias the buffer descriptor base once per CTA using an index/i64 GEP, + # then store with only tile-local i32 offsets. This keeps the hot store + # instruction form unchanged while avoiding i32 wrap in buffer_store(). + c_n_idx_for_base = fx.Index(c_n) + c_tile_base_elems = bx_m_idx * c_n_idx_for_base + by_n_idx + c_tile_base_bytes = c_tile_base_elems * fx.Index(output_element_bytes) + c_rsrc = buffer_ops.create_buffer_resource( + C, + max_size=True, + base_byte_offset=c_tile_base_bytes, + ) + + 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, + ) + + 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(): + # 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): + # BF16 uses v_mfma_f32_16x16x32_bf16, not the MXFP8 K128 + # instruction. A 128-X half-page is therefore two independent + # swizzled [K64, X64] BF16 slices. One ds_read_b64_tr_b16 returns + # four BF16 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_bf16( + lds_page, + base, + other, + immediate_offset=immediate_offset, + ) + + 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 _acc_idx(subtile_id, mi, ni): + return subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + mi * MFMA_N_PER_SUBTILE + ni + + def _bf16_k32_frag(full_frag, k32): + # A/B 16x64 BF16 wave fragments are i32x8. Each K32 MFMA + # consumes one contiguous i32x4 slice (eight BF16 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_bf16_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_bf16 " + f"a[{acc_pin}:{acc_pin + 3}], " + f"$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}]" + ), + ( + f"v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}}," + f"~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}" + ), + has_side_effects=True, + ) + + def pinned_mfma(acc_idx, a_frag, b_frag): + """Accumulate one logical 16x16x64 BF16 product into pinned AGPRs.""" + for k32 in range_constexpr(2): + _pinned_bf16_mfma_once( + acc_idx, + _bf16_k32_frag(a_frag, k32), + _bf16_k32_frag(b_frag, k32), + ) + + 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 = _bf16_k32_frag(a_frags[mi], k32) + for nj in range_constexpr(2): + _pinned_bf16_mfma_once( + _acc_idx(subtile_id, mi, n_base + nj), + a_k32, + _bf16_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 = _bf16_k32_frag(a_frags[mi], k32) + for ni in range_constexpr(4): + _pinned_bf16_mfma_once( + _acc_idx(subtile_id, mi, ni), + a_k32, + _bf16_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 + + 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 + for ii in range_constexpr(4): + row = row_base + fx.Index(ii) + c_idx = row * fx.Index(c_n) + col + value = Vec(acc)[ii] + if const_expr(output_dtype != torch.float32): + value = value.to(output_fx_dtype) + buffer_ops.buffer_store(value, c_rsrc, 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), + ) + + 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_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. 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) + + 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_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]. + + 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_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) + + 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_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]. + + 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. + + 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) + + 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_a_subtile_pass(k_refill, 1, 0, refill_a) + mfma_2n_4mi_k32(1, 2, 0, a00, a01, a02, a03, b12, b13) + + 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, 1, refill_a) + # Keep this refill slot compute-free. + + 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, 2, refill_a) + mfma_2n_4mi_k32(1, 2, 1, a00, a01, a02, a03, b12, b13) + + 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, 3, refill_a) + # Keep this refill slot compute-free. + 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, lgkmcnt=0) + rocdl.sched_barrier(0) + + 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_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + # K32 slice 0 already covers K[0:32]. + + 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_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + # K32 slice 1 already covers K[32:64]. + + 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_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + # K32 slice 0 already covers K[0:32]. + + 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_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + # K32 slice 1 already covers K[32:64]. + + 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) + + return next_a0_regs, next_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) + + 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) + + 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) + + 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) + + 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) + + rocdl.sched_barrier(0) + _barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, 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_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, 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_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, 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_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, 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) + + 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) + + return next_a0_regs, next_b0_regs + + def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): + _barrier(vmcnt=0, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + # 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) + + 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) + + # 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) + + 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 + 3 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(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, + ) + + # 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) + 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) + + + @flyc.jit + def launch_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + # The integration only dispatches aligned shapes; no partial-tile masking exists. + grid_x = (c_m // BLOCK_M) * (c_n // BLOCK_N) + kernel_gemm( + A, + B, + C, + c_m, + c_n, + 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) + + return launch_gemm + +@functools.lru_cache(maxsize=None) +def _cached_launch( + K: int, + output_dtype: torch.dtype, + layout: str, + use_xcd_remap: bool = True, +): + return _compile_kernel( + K, + output_dtype, + layout, + use_xcd_remap=use_xcd_remap, + ) + + +def bf16_matmul( + a: torch.Tensor, + b: torch.Tensor, + c: torch.Tensor, + *, + layout: str, + m: int, + n: int, + k: int, + stream=None, +): + """Launch the wrapper-selected BF16 TN/NN/NT specialization.""" + if layout not in ("TN", "NN", "NT"): + raise ValueError(f"Unsupported BF16 layout: {layout}") + if a.ndim != 2 or b.ndim != 2: + raise ValueError( + f"FlyDSL BF16 expects rank-2 operands, got A{tuple(a.shape)} " + f"and B{tuple(b.shape)}" + ) + if a.dtype != torch.bfloat16 or b.dtype != torch.bfloat16: + raise TypeError( + "FlyDSL BF16 GEMM expects torch.bfloat16 operands, " + f"got A={a.dtype}, B={b.dtype}" + ) + if not a.is_contiguous() or not b.is_contiguous(): + raise FlyDSLUnsupportedError( + f"FlyDSL BF16 {layout} requires original contiguous row-major " + f"operands, got A stride={tuple(a.stride())}, " + f"B stride={tuple(b.stride())}" + ) + + m = int(m) + n = int(n) + k = int(k) + + expected_shapes = { + "TN": ((m, k), (n, k)), + "NN": ((m, k), (k, n)), + "NT": ((k, m), (k, n)), + } + expected_a, expected_b = expected_shapes[layout] + if tuple(a.shape) != expected_a or tuple(b.shape) != expected_b: + raise ValueError( + f"FlyDSL BF16 {layout} physical operands do not match contract: " + f"A{tuple(a.shape)} expected {expected_a}; " + f"B{tuple(b.shape)} expected {expected_b}" + ) + + 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( + "FlyDSL BF16 output must be float16, bfloat16, or float32, " + f"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 " + f"{a.device}, {b.device}, and {c.device}" + ) + if not c.is_contiguous(): + raise ValueError("FlyDSL BF16 GEMM requires contiguous output storage") + + doGemm( + a, + b, + c, + layout=layout, + m=m, + n=n, + k=k, + stream=stream, + ) + +def doGemm( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + *, + layout: str, + m: int, + n: int, + k: int, + stream=None, + use_xcd_remap: bool = True, +): + """Launch one cached K/output/layout-specialized BF16 core. + + A and B are passed unchanged from ``gemm_wrappers.py``. Their pointers + reference the original rowwise allocations: + + TN: A backing [M,K], B backing [N,K] + NN: A backing [M,K], B backing [K,N] + 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. + """ + if layout not in ("TN", "NN", "NT"): + raise ValueError(f"Unsupported BF16 layout: {layout}") + + M_runtime = int(m) + N_runtime = int(n) + K_runtime = int(k) + + if A.dtype != torch.bfloat16 or B.dtype != torch.bfloat16: + raise TypeError( + f"BF16 {layout} requires BF16 inputs, got {A.dtype} and {B.dtype}" + ) + if C.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise TypeError(f"Unsupported BF16 output dtype: {C.dtype}") + + if M_runtime % _BLOCK_M != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL BF16 GEMM requires M to be a multiple of {_BLOCK_M}, " + f"got M={M_runtime}" + ) + if N_runtime % _BLOCK_N != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL BF16 GEMM requires N to be a multiple of {_BLOCK_N}, " + f"got N={N_runtime}" + ) + if K_runtime % _BLOCK_K != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL BF16 GEMM requires K to be a multiple of {_BLOCK_K}, " + f"got K={K_runtime}" + ) + + num_k_tiles = K_runtime // _BLOCK_K + if num_k_tiles < 4: + raise FlyDSLUnsupportedError( + f"FlyDSL BF16 GEMM requires at least 4 K{_BLOCK_K} tiles, " + f"got K={K_runtime} ({num_k_tiles} tiles)" + ) + + if tuple(C.shape) != (M_runtime, N_runtime): + raise ValueError( + f"C shape {tuple(C.shape)} != expected {(M_runtime, N_runtime)}" + ) + + if stream is None: + stream = torch.cuda.current_stream() + + launch = _cached_launch( + K_runtime, + C.dtype, + layout, + bool(use_xcd_remap), + ) + # 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. + A_arg = A.view(torch.uint8).view(-1) + B_arg = B.view(torch.uint8).view(-1) + C_arg = C.view(-1) + + launch( + A_arg, + B_arg, + C_arg, + M_runtime, + N_runtime, + stream=stream, + ) + + +@ASTRewriter.transform +def dense_mma_pipeline_bf16( + lds, + a_g2s, + b_g2s, + a_s2r, + b_s2r, + mfma, + store_c, + A0_gl_offset, + A1_gl_offset, + B0_gl_offset, + B1_gl_offset, + a_k_step, + b_k_step, + block_m, + block_n, + wave_m, + wave_n, + K, + BLOCK_M, + BLOCK_N, + nt_vmcnt, + a_g2s_hi=None, +): + """Shared 4-quadrant pipelined MMA loop + store epilogue for the fixed-K bf16 tile (NT/NN/TN). + + ``a_g2s_hi`` (optional): a second A global->LDS loader used ONLY for the upper LDS + half-tile (rows [LDS_BLOCK_M, BLOCK_M)). Defaults to ``a_g2s`` (contiguous dense/grouped + GEMM, where both halves share one swizzle). A gathering GEMM passes a distinct loader whose + per-lane offsets are redirected through the gather index for the upper rows. + """ + a_g2s_hi = a_g2s if a_g2s_hi is None else a_g2s_hi + K_ITERS = K // BLOCK_K + assert K_ITERS >= 2, f"K_ITERS={K_ITERS} too small; need K >= {2 * BLOCK_K}" + N_TILES_A = BLOCK_M // 128 + N_TILES_B = BLOCK_N // 256 + N_ACCUMS = N_TILES_A * N_TILES_B + LDS_BLOCK_M = BLOCK_M // 2 + LDS_BLOCK_N = BLOCK_N // 2 + N_LDS_STEPS_A = LDS_BLOCK_M // 64 + N_LDS_STEPS_B = LDS_BLOCK_N // 64 + + a_cur0 = lds.A_lds_cur_0 + a_cur1 = lds.A_lds_cur_1 + a_next0 = lds.A_lds_next_0 + a_next1 = lds.A_lds_next_1 + b_cur0 = lds.B_lds_cur_0 + b_cur1 = lds.B_lds_cur_1 + b_next0 = lds.B_lds_next_0 + b_next1 = lds.B_lds_next_1 + + c00_frag = [mfma.zero_value] * N_ACCUMS + c01_frag = [mfma.zero_value] * N_ACCUMS + c10_frag = [mfma.zero_value] * N_ACCUMS + c11_frag = [mfma.zero_value] * N_ACCUMS + + b_g2s.load(b_cur0, B0_gl_offset + 0 * b_k_step) + a_g2s.load(a_cur0, A0_gl_offset + 0 * a_k_step) + b_g2s.load(b_cur1, B1_gl_offset + 0 * b_k_step) + a_g2s_hi.load(a_cur1, A1_gl_offset + 0 * a_k_step) + + if wave_m == 1: + rocdl.s_barrier() + wait_barrier(N_LDS_STEPS_A + N_LDS_STEPS_B) + + b_g2s.load(b_next0, B0_gl_offset + 1 * b_k_step) + a_g2s.load(a_next0, A0_gl_offset + 1 * a_k_step) + b_g2s.load(b_next1, B1_gl_offset + 1 * b_k_step) + + wait_barrier(N_LDS_STEPS_A + 2 * N_LDS_STEPS_B) + + for k in range_constexpr(K_ITERS - 2): + b0_frag = b_s2r.load(b_cur0) + a0_frag = a_s2r.load(a_cur0) + a_g2s_hi.load(a_next1, A1_gl_offset + (k + 1) * a_k_step) + rocdl.s_barrier() + + rocdl.s_setprio(1) + c00_frag = mfma.call(a0_frag, b0_frag, c00_frag) + rocdl.s_setprio(0) + rocdl.s_barrier() + + b1_frag = b_s2r.load(b_cur1) + b_g2s.load(b_cur0, B0_gl_offset + (k + 2) * b_k_step) + rocdl.s_barrier() + + rocdl.s_setprio(1) + c01_frag = mfma.call(a0_frag, b1_frag, c01_frag) + rocdl.s_setprio(0) + rocdl.s_barrier() + + a1_frag = a_s2r.load(a_cur1) + a_g2s.load(a_cur0, A0_gl_offset + (k + 2) * a_k_step) + rocdl.s_barrier() + + rocdl.s_setprio(1) + c10_frag = mfma.call(a1_frag, b0_frag, c10_frag) + rocdl.s_setprio(0) + rocdl.s_barrier() + + b_g2s.load(b_cur1, B1_gl_offset + (k + 2) * b_k_step) + wait_barrier(2 * N_LDS_STEPS_A + N_LDS_STEPS_B) + + rocdl.s_setprio(1) + c11_frag = mfma.call(a1_frag, b1_frag, c11_frag) + rocdl.s_setprio(0) + rocdl.s_barrier() + + if const_expr(nt_vmcnt >= 0): + llvm.inline_asm( + res=None, + operands_=[], + asm_string=f"s_waitcnt vmcnt({nt_vmcnt})", + constraints="", + has_side_effects=True, + ) + a_cur0, a_next0 = a_next0, a_cur0 + a_cur1, a_next1 = a_next1, a_cur1 + b_cur0, b_next0 = b_next0, b_cur0 + b_cur1, b_next1 = b_next1, b_cur1 + + k = K_ITERS - 2 + b0_frag = b_s2r.load(b_cur0) + a0_frag = a_s2r.load(a_cur0) + rocdl.s_barrier() + rocdl.s_setprio(1) + c00_frag = mfma.call(a0_frag, b0_frag, c00_frag) + rocdl.s_setprio(0) + rocdl.s_barrier() + + b1_frag = b_s2r.load(b_cur1) + rocdl.s_barrier() + rocdl.s_setprio(1) + c01_frag = mfma.call(a0_frag, b1_frag, c01_frag) + rocdl.s_setprio(0) + rocdl.s_barrier() + + a1_frag = a_s2r.load(a_cur1) + rocdl.s_barrier() + rocdl.s_setprio(1) + c10_frag = mfma.call(a1_frag, b0_frag, c10_frag) + rocdl.s_setprio(0) + rocdl.s_barrier() + + b0_frag = b_s2r.load(b_next0) + a_g2s_hi.load(a_next1, A1_gl_offset + (k + 1) * a_k_step) + rocdl.s_barrier() + rocdl.s_setprio(1) + c11_frag = mfma.call(a1_frag, b1_frag, c11_frag) + rocdl.s_setprio(0) + rocdl.s_barrier() + + a_cur0, a_next0 = a_next0, a_cur0 + a_cur1, a_next1 = a_next1, a_cur1 + b_cur0, b_next0 = b_next0, b_cur0 + b_cur1, b_next1 = b_next1, b_cur1 + + a0_frag = a_s2r.load(a_cur0) + wait_barrier(0) + rocdl.s_setprio(1) + c00_frag = mfma.call(a0_frag, b0_frag, c00_frag) + rocdl.s_setprio(0) + rocdl.s_barrier() + + b1_frag = b_s2r.load(b_cur1) + rocdl.s_barrier() + rocdl.s_setprio(1) + c01_frag = mfma.call(a0_frag, b1_frag, c01_frag) + rocdl.s_setprio(0) + rocdl.s_barrier() + + a1_frag = a_s2r.load(a_cur1) + rocdl.s_barrier() + rocdl.s_setprio(1) + c10_frag = mfma.call(a1_frag, b0_frag, c10_frag) + c11_frag = mfma.call(a1_frag, b1_frag, c11_frag) + rocdl.s_setprio(0) + rocdl.s_barrier() + + wave_n_offset = wave_n * (N_TILES_B * 32) + wave_m_offset = wave_m * (N_TILES_A * 32) + base_row = block_m * BLOCK_M + wave_m_offset + base_col = block_n * BLOCK_N + wave_n_offset + store_c.store(c00_frag, base_row + 0, base_col + 0) + store_c.store(c01_frag, base_row + 0, base_col + LDS_BLOCK_N) + store_c.store(c10_frag, base_row + LDS_BLOCK_M, base_col + 0) + store_c.store(c11_frag, base_row + LDS_BLOCK_M, base_col + LDS_BLOCK_N) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/exceptions.py b/transformer_engine/pytorch/flydsl_kernels/gemm/exceptions.py new file mode 100644 index 000000000..b7fc19a23 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/exceptions.py @@ -0,0 +1,6 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +class FlyDSLUnsupportedError(RuntimeError): + """The GEMM request is valid but unsupported by the available FlyDSL kernels.""" diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py new file mode 100644 index 000000000..ace8ecda4 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py @@ -0,0 +1,1430 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""FlyDSL FP16 TN/NN/NT 4-wave GEMM kernel for Transformer Engine. + +All supported layouts share one source-level kernel generator while compiling +to separate cached binaries: + + TN: A [M,K] normal read, B [N,K] normal read + NN: A [M,K] normal read, B [K,N] transpose read + NT: A [K,M] transpose read, B [K,N] transpose read + +The layout is a Python-only cache key. Global addressing and LDS fragment +reconstruction are selected while building each specialization, so no runtime +layout branch is emitted in the GEMM kernel. + +This module imports ``flydsl`` at import time and must therefore be imported +lazily only after FlyDSL availability has been confirmed. +""" + +import functools + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir.dialects import llvm +from flydsl.expr import arith, buffer_ops, const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec + +# Transformer Engine-local FlyDSL utilities. +from .exceptions import FlyDSLUnsupportedError +from .fp16_gemm_utils import ( + G2SLoader, + S2RLoader, + compute_global_fp16_transpose_swizzle, + compute_global_swizzle, + make_fp16_byte_buffer_tensor, + pack_i32x4_i32x8, + swizzle_128, +) + + +_BLOCK_M = 256 +_BLOCK_N = 256 +_BLOCK_K = 64 + +# Public metadata consumed by wrappers. +BLOCK_M = _BLOCK_M +BLOCK_N = _BLOCK_N +BLOCK_K = _BLOCK_K + +NUM_THREADS = 256 +WARP_SIZE = 64 +NUM_WAVES = NUM_THREADS // WARP_SIZE + +SUBTILE_M = 64 +SUBTILE_N = 64 + +MFMA_M = 16 +MFMA_N = 16 + +SUBTILES_PER_WAVE = 4 +MFMA_M_PER_SUBTILE = SUBTILE_M // MFMA_M +MFMA_N_PER_SUBTILE = SUBTILE_N // MFMA_N +ACCS_PER_WAVE = SUBTILES_PER_WAVE * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + +ELEM_BYTES = 2 +VEC_BYTES = 16 + +LDS_ELEMS_A = BLOCK_M * BLOCK_K +LDS_ELEMS_B = BLOCK_N * BLOCK_K +LDS_BYTES_A = LDS_ELEMS_A * ELEM_BYTES +LDS_BYTES_B = LDS_ELEMS_B * ELEM_BYTES + +LOAD_PASSES_A = LDS_BYTES_A // (NUM_THREADS * VEC_BYTES) +LOAD_PASSES_B = LDS_BYTES_B // (NUM_THREADS * VEC_BYTES) +LOAD_PASSES_A_SUBTILE = LOAD_PASSES_A // 2 +LOAD_PASSES_B_SUBTILE = LOAD_PASSES_B // 2 +PASSES_PER_A_MI = LOAD_PASSES_A_SUBTILE // MFMA_M_PER_SUBTILE + +LDS_SYM_A0 = "fp16_pp_smem_a0" +LDS_SYM_A1 = "fp16_pp_smem_a1" +LDS_SYM_B0 = "fp16_pp_smem_b0" +LDS_SYM_B1 = "fp16_pp_smem_b1" +LDS_ALIAS_DOMAIN = '#llvm.alias_scope_domain' +SCOPE_IDS = ("a0", "a1", "b0", "b1") + +assert BLOCK_K == 64 +# DO NOT CHANGE THE FOLLOWING LINE. +assert NUM_THREADS == 256 +assert LOAD_PASSES_A * NUM_THREADS * VEC_BYTES == LDS_BYTES_A +assert LOAD_PASSES_B * NUM_THREADS * VEC_BYTES == LDS_BYTES_B +assert LOAD_PASSES_A % 2 == 0 +assert LOAD_PASSES_B % 2 == 0 + + + +def swizzle_xor16(row, col_in_bytes): + """XOR swizzle for the LDS K-byte coordinate.""" + chunk = col_in_bytes // fx.Index(VEC_BYTES) + byte_in_chunk = col_in_bytes % fx.Index(VEC_BYTES) + row_bits = (row % fx.Index(16)) // fx.Index(2) + swz_chunk = chunk ^ row_bits + return swz_chunk * fx.Index(VEC_BYTES) + byte_in_chunk + + +def _encode_waitcnt(vmcnt=63, lgkmcnt=15): + """Encode the CDNA4/gfx950 ``S_WAITCNT`` SIMM16 operand. + + ``rocdl.s_waitcnt`` accepts the raw 16-bit immediate operand of the + 32-bit ``S_WAITCNT`` ISA instruction. On CDNA4, that SIMM16 field is: + + SIMM16[3:0] = vmcnt[3:0] + SIMM16[6:4] = expcnt[2:0] + SIMM16[11:8] = lgkmcnt[3:0] + SIMM16[15:14] = vmcnt[5:4] + + ``vmcnt`` is therefore one six-bit counter split across two noncontiguous + fields; bits [5:4] are placed in SIMM16[15:14], while bits [3:0] remain + in SIMM16[3:0]. + + A wait-counter field set to its maximum representable value is effectively + unconstrained: the instruction does not wait on that counter. This helper + always encodes ``expcnt=7`` and defaults to ``vmcnt=63`` and ``lgkmcnt=15``, + so callers specify only the counters on which they intend to wait. + + For example, ``_encode_waitcnt(lgkmcnt=0)`` returns ``0xC07F``, which the + assembler renders as ``s_waitcnt lgkmcnt(0)``. + See: https://llvm.org/docs/AMDGPU/gfx9_waitcnt.html + """ + if not 0 <= vmcnt <= 63: + raise ValueError(f"vmcnt must be in [0, 63], got {vmcnt}") + if not 0 <= lgkmcnt <= 15: + raise ValueError(f"lgkmcnt must be in [0, 15], got {lgkmcnt}") + + return ( + (7 << 4) # expcnt=7 -> SIMM16[6:4] (unconstrained) + | (vmcnt & 0x0F) # vmcnt[3:0] -> SIMM16[3:0] + | ((lgkmcnt & 0x0F) << 8) # lgkmcnt[3:0] -> SIMM16[11:8] + | ((vmcnt & 0x30) << 10) # vmcnt[5:4] -> SIMM16[15:14] + ) + + +# Keep the documented gfx950 encoding invariant executable and import-time cheap. +assert _encode_waitcnt(lgkmcnt=0) == 0xC07F + + +def _barrier(vmcnt=63, lgkmcnt=15): + if vmcnt != 63 or lgkmcnt != 15: + rocdl.s_waitcnt(_encode_waitcnt(vmcnt=vmcnt, lgkmcnt=lgkmcnt)) + rocdl.s_barrier() + +def _min(a, b): + return arith.select(a < b, a, b) + + +def _divmod(a, b): + return a // b, a % b + + +def _xcd_swizzle(num_pid_m, num_pid_n): + NUM_XCDS = 8 + WGM = 4 + NUM_CUS = 32 * NUM_XCDS + SWIZZLE_THRESHOLD = 4 * NUM_CUS + + wgid = fx.block_idx.x + num_wg = num_pid_m * num_pid_n + + # Simple row-major path. + simple_m, simple_n = _divmod(wgid, num_pid_n) + + # XCD-remapped grouped-M path. + intra_xcd, xcd = _divmod(wgid, NUM_XCDS) + wgid_remap = xcd * (num_wg // NUM_XCDS) + intra_xcd + num_wgid_in_group = WGM * num_pid_n + group_id, intra_group = _divmod(wgid_remap, num_wgid_in_group) + first_pid_m = group_id * WGM + group_size_m = _min(num_pid_m - first_pid_m, WGM) + pid_n, intra_group_m = _divmod(intra_group, group_size_m) + pid_m = first_pid_m + intra_group_m + + use_simple = (num_wg < SWIZZLE_THRESHOLD) | (num_wg % NUM_XCDS != 0) + return ( + arith.select(use_simple, simple_m, pid_m), + arith.select(use_simple, simple_n, pid_n), + ) + + +def _compile_kernel( + K: int, + output_dtype: torch.dtype, + layout: str, + use_xcd_remap: bool = True, +): + """Build one compile-time-specialized TN, NN, or NT FP16 kernel. + + ``K`` must contain at least four K64 tiles. Runtime M/N are expected to + be exact multiples of ``BLOCK_M``/``BLOCK_N``; the kernel has no edge masks. + """ + if layout not in ("TN", "NN", "NT"): + raise ValueError(f"Unsupported FP16 kernel layout: {layout}") + + a_transpose_read = layout == "NT" + b_transpose_read = layout in ("NN", "NT") + + BLOCK_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K + NUM_THREADS = 256 + WARP_SIZE = 64 + + SUBTILE_M = 64 + SUBTILE_N = 64 + + MFMA_M = 16 + MFMA_N = 16 + + SUBTILES_PER_WAVE = 4 + MFMA_M_PER_SUBTILE = SUBTILE_M // MFMA_M + MFMA_N_PER_SUBTILE = SUBTILE_N // MFMA_N + ACCS_PER_WAVE = SUBTILES_PER_WAVE * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + + ELEM_BYTES = 2 + VEC_BYTES = 16 + + if output_dtype == torch.float16: + output_element_bytes = 2 + output_fx_dtype = fx.Float16 + elif output_dtype == torch.bfloat16: + output_element_bytes = 2 + output_fx_dtype = fx.BFloat16 + elif output_dtype == torch.float32: + output_element_bytes = 4 + output_fx_dtype = fx.Float32 + else: + raise TypeError( + "FlyDSL FP16 GEMM output dtype must be torch.float16, " + f"torch.bfloat16, or torch.float32, got {output_dtype}" + ) + + LDS_ELEMS_A = BLOCK_M * BLOCK_K + LDS_ELEMS_B = BLOCK_N * BLOCK_K + LDS_BYTES_A = LDS_ELEMS_A * ELEM_BYTES + LDS_BYTES_B = LDS_ELEMS_B * ELEM_BYTES + + LOAD_PASSES_A = LDS_BYTES_A // (NUM_THREADS * VEC_BYTES) + LOAD_PASSES_B = LDS_BYTES_B // (NUM_THREADS * VEC_BYTES) + LOAD_PASSES_A_SUBTILE = LOAD_PASSES_A // 2 + LOAD_PASSES_B_SUBTILE = LOAD_PASSES_B // 2 + + assert K % BLOCK_K == 0, f"K must be a multiple of {BLOCK_K}, got {K}" + NUM_K_TILES = K // BLOCK_K + assert NUM_K_TILES >= 4, f"K={K} gives {NUM_K_TILES} K64 tiles; the two-page pipeline needs at least 4" + + LDS_ELEMS_HALF = (BLOCK_M // 2) * BLOCK_K + LDS_BYTES_HALF = LDS_ELEMS_HALF * ELEM_BYTES + LOAD_PASSES_HALF = LDS_BYTES_HALF // (NUM_THREADS * VEC_BYTES) + assert LOAD_PASSES_HALF == LOAD_PASSES_A_SUBTILE == LOAD_PASSES_B_SUBTILE + + # Resolve layout-specific addressing and fragment reads before capture. + Q0_SCHED_DSRD = 4 if a_transpose_read else 2 + PREFETCH_SCHED_DSRD = 8 if a_transpose_read else 4 + + if a_transpose_read: + 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) + ) + + def _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, + ): + del load_frag_half_at_byte_base, lane_mod_16 + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + local_m_tile = ( + subtile_m_idx * fx.Index(SUBTILE_M) + + fx.Index(mi * MFMA_M) + - fx.Index(sm * (BLOCK_M // 2)) + ) + return load_transposed_frag_half(lds_a[sm], local_m_tile, half) + else: + def _a_leading_dim_bytes(c_m): + del c_m + return K * ELEM_BYTES + + 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(K * ELEM_BYTES) + + k_base * fx.Index(ELEM_BYTES) + ) + + def _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, + ): + 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)) + return load_frag_half_at_byte_base( + lds_a[sm], + half_row * fx.Index(BLOCK_K * ELEM_BYTES), + half, + ) + + if b_transpose_read: + 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) + ) + + def _load_b_ni( + load_transposed_frag, + load_normal_b_frag, + lds_b, + sn, + ni, + reg_subtile_n_idx0, + lane_mod_16, + ): + del load_normal_b_frag, lane_mod_16 + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + local_n_tile = ( + subtile_n_idx * fx.Index(SUBTILE_N) + + fx.Index(ni * MFMA_N) + - fx.Index(sn * (BLOCK_N // 2)) + ) + return load_transposed_frag(lds_b[sn], local_n_tile) + else: + def _b_leading_dim_bytes(c_n): + del c_n + return K * ELEM_BYTES + + 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(K * ELEM_BYTES) + + k_base * fx.Index(ELEM_BYTES) + ) + + def _load_b_ni( + load_transposed_frag, + load_normal_b_frag, + lds_b, + sn, + ni, + reg_subtile_n_idx0, + 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 + ) + return load_normal_b_frag(lds_b, b_row_addr, sn) + + # Resolve global staging maps before FlyDSL captures ``kernel_gemm``. + # FP16 uses K64, so each transpose-read half-page is two independent + # [K64, X64] slices with 128-byte physical rows. + if a_transpose_read: + def _a_global_offsets(lane, wave_id, c_m): + return compute_global_fp16_transpose_swizzle( + lane, + wave_id, + _a_leading_dim_bytes(c_m), + LOAD_PASSES_HALF, + ) + else: + def _a_global_offsets(lane, wave_id, c_m): + del c_m + return compute_global_swizzle( + lane, + wave_id, + K * ELEM_BYTES, + LOAD_PASSES_HALF, + preshuffled=False, + ) + + if b_transpose_read: + def _b_global_offsets(lane, wave_id, c_n): + return compute_global_fp16_transpose_swizzle( + lane, + wave_id, + _b_leading_dim_bytes(c_n), + LOAD_PASSES_HALF, + ) + else: + def _b_global_offsets(lane, wave_id, c_n): + del c_n + return compute_global_swizzle( + lane, + wave_id, + K * ELEM_BYTES, + LOAD_PASSES_HALF, + preshuffled=False, + ) + + @fx.struct + class SharedStorage: + # Preserve the passing TN byte-staging contract exactly. A FP16 K64 + # half-page is 128 rows x 128 bytes = 16 KiB. + a0_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a0_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a1_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a1_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b0_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + 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] + + @flyc.kernel(known_block_size=[NUM_THREADS, 1, 1]) + def kernel_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + ): + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + lds_a0 = (lds.a0_0, lds.a0_1) + lds_a1 = (lds.a1_0, lds.a1_1) + lds_b0 = (lds.b0_0, lds.b0_1) + lds_b1 = (lds.b1_0, lds.b1_1) + + # A/B arrive as contiguous uint8 byte views of the original + # row-major FP16 tensors. This preserves the validated 16-byte + # BufferCopyLDS128b path and byte-based address arithmetic. + gA = make_fp16_byte_buffer_tensor(A) + gB = make_fp16_byte_buffer_tensor(B) + 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") + + 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 BF16 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) + + # C can exceed the signed-i32 element/byte offset range for large M*N. + # Bias the buffer descriptor base once per CTA using an index/i64 GEP, + # then store with only tile-local i32 offsets. This keeps the hot store + # instruction form unchanged while avoiding i32 wrap in buffer_store(). + c_n_idx_for_base = fx.Index(c_n) + c_tile_base_elems = bx_m_idx * c_n_idx_for_base + by_n_idx + c_tile_base_bytes = c_tile_base_elems * fx.Index(output_element_bytes) + c_rsrc = buffer_ops.create_buffer_resource( + C, + max_size=True, + base_byte_offset=c_tile_base_bytes, + ) + + 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, + ) + + 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(): + # 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): + # FP16 uses v_mfma_f32_16x16x32_f16, not the MXFP8 K128 + # instruction. A 128-X half-page is therefore two independent + # swizzled [K64, X64] FP16 slices. One ds_read_b64_tr_b16 returns + # four FP16 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_fp16( + lds_page, + base, + other, + immediate_offset=immediate_offset, + ) + + 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 _acc_idx(subtile_id, mi, ni): + return subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + mi * MFMA_N_PER_SUBTILE + ni + + def _fp16_k32_frag(full_frag, k32): + # A/B 16x64 FP16 wave fragments are i32x8. Each K32 MFMA + # consumes one contiguous i32x4 slice (eight FP16 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_fp16_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_f16 " + f"a[{acc_pin}:{acc_pin + 3}], " + f"$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}]" + ), + ( + f"v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}}," + f"~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}" + ), + has_side_effects=True, + ) + + def pinned_mfma(acc_idx, a_frag, b_frag): + """Accumulate one logical 16x16x64 FP16 product into pinned AGPRs.""" + for k32 in range_constexpr(2): + _pinned_fp16_mfma_once( + acc_idx, + _fp16_k32_frag(a_frag, k32), + _fp16_k32_frag(b_frag, k32), + ) + + 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 = _fp16_k32_frag(a_frags[mi], k32) + for nj in range_constexpr(2): + _pinned_fp16_mfma_once( + _acc_idx(subtile_id, mi, n_base + nj), + a_k32, + _fp16_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 = _fp16_k32_frag(a_frags[mi], k32) + for ni in range_constexpr(4): + _pinned_fp16_mfma_once( + _acc_idx(subtile_id, mi, ni), + a_k32, + _fp16_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 + + 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 + for ii in range_constexpr(4): + row = row_base + fx.Index(ii) + c_idx = row * fx.Index(c_n) + col + value = Vec(acc)[ii] + if const_expr(output_dtype != torch.float32): + value = value.to(output_fx_dtype) + buffer_ops.buffer_store(value, c_rsrc, 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), + ) + + 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_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. 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) + + 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_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]. + + 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_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) + + 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_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]. + + 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. + + 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) + + 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_a_subtile_pass(k_refill, 1, 0, refill_a) + mfma_2n_4mi_k32(1, 2, 0, a00, a01, a02, a03, b12, b13) + + 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, 1, refill_a) + # Keep this refill slot compute-free. + + 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, 2, refill_a) + mfma_2n_4mi_k32(1, 2, 1, a00, a01, a02, a03, b12, b13) + + 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, 3, refill_a) + # Keep this refill slot compute-free. + 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, lgkmcnt=0) + rocdl.sched_barrier(0) + + 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_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + # K32 slice 0 already covers K[0:32]. + + 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_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + # K32 slice 1 already covers K[32:64]. + + 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_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + # K32 slice 0 already covers K[0:32]. + + 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_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + # K32 slice 1 already covers K[32:64]. + + 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) + + return next_a0_regs, next_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) + + 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) + + 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) + + 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) + + 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) + + rocdl.sched_barrier(0) + _barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, 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_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, 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_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, 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_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, 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) + + 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) + + return next_a0_regs, next_b0_regs + + def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): + _barrier(vmcnt=0, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + # 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) + + 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) + + # 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) + + 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 + 3 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(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, + ) + + # 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) + 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) + + + @flyc.jit + def launch_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + # The integration only dispatches aligned shapes; no partial-tile masking exists. + grid_x = (c_m // BLOCK_M) * (c_n // BLOCK_N) + kernel_gemm( + A, + B, + C, + c_m, + c_n, + 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) + + return launch_gemm + +@functools.lru_cache(maxsize=None) +def _cached_launch( + K: int, + output_dtype: torch.dtype, + layout: str, + use_xcd_remap: bool = True, +): + return _compile_kernel( + K, + output_dtype, + layout, + use_xcd_remap=use_xcd_remap, + ) + + +def fp16_matmul( + a: torch.Tensor, + b: torch.Tensor, + c: torch.Tensor, + *, + layout: str, + m: int, + n: int, + k: int, + stream=None, +): + """Launch the wrapper-selected BF16 TN/NN/NT specialization.""" + if layout not in ("TN", "NN", "NT"): + raise ValueError(f"Unsupported FP16 layout: {layout}") + if a.ndim != 2 or b.ndim != 2: + raise ValueError( + f"FlyDSL BF16 expects rank-2 operands, got A{tuple(a.shape)} " + f"and B{tuple(b.shape)}" + ) + if a.dtype != torch.float16 or b.dtype != torch.float16: + raise TypeError( + "FlyDSL FP16 GEMM expects torch.float16 operands, " + f"got A={a.dtype}, B={b.dtype}" + ) + if not a.is_contiguous() or not b.is_contiguous(): + raise FlyDSLUnsupportedError( + f"FlyDSL BF16 {layout} requires original contiguous row-major " + f"operands, got A stride={tuple(a.stride())}, " + f"B stride={tuple(b.stride())}" + ) + + m = int(m) + n = int(n) + k = int(k) + + expected_shapes = { + "TN": ((m, k), (n, k)), + "NN": ((m, k), (k, n)), + "NT": ((k, m), (k, n)), + } + expected_a, expected_b = expected_shapes[layout] + if tuple(a.shape) != expected_a or tuple(b.shape) != expected_b: + raise ValueError( + f"FlyDSL BF16 {layout} physical operands do not match contract: " + f"A{tuple(a.shape)} expected {expected_a}; " + f"B{tuple(b.shape)} expected {expected_b}" + ) + + 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( + "FlyDSL FP16 output must be float16, bfloat16, or float32, " + f"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 " + f"{a.device}, {b.device}, and {c.device}" + ) + if not c.is_contiguous(): + raise ValueError("FlyDSL FP16 GEMM requires contiguous output storage") + + doGemm( + a, + b, + c, + layout=layout, + m=m, + n=n, + k=k, + stream=stream, + ) + +def doGemm( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + *, + layout: str, + m: int, + n: int, + k: int, + stream=None, + use_xcd_remap: bool = True, +): + """Launch one cached K/output/layout-specialized FP16 core. + + A and B are passed unchanged from ``gemm_wrappers.py``. Their pointers + reference the original rowwise allocations: + + TN: A backing [M,K], B backing [N,K] + NN: A backing [M,K], B backing [K,N] + 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. + """ + if layout not in ("TN", "NN", "NT"): + raise ValueError(f"Unsupported FP16 layout: {layout}") + + M_runtime = int(m) + N_runtime = int(n) + K_runtime = int(k) + + if A.dtype != torch.float16 or B.dtype != torch.float16: + raise TypeError( + f"BF16 {layout} requires BF16 inputs, got {A.dtype} and {B.dtype}" + ) + if C.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise TypeError(f"Unsupported FP16 output dtype: {C.dtype}") + + if M_runtime % _BLOCK_M != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP16 GEMM requires M to be a multiple of {_BLOCK_M}, " + f"got M={M_runtime}" + ) + if N_runtime % _BLOCK_N != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP16 GEMM requires N to be a multiple of {_BLOCK_N}, " + f"got N={N_runtime}" + ) + if K_runtime % _BLOCK_K != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP16 GEMM requires K to be a multiple of {_BLOCK_K}, " + f"got K={K_runtime}" + ) + + num_k_tiles = K_runtime // _BLOCK_K + if num_k_tiles < 4: + raise FlyDSLUnsupportedError( + f"FlyDSL FP16 GEMM requires at least 4 K{_BLOCK_K} tiles, " + f"got K={K_runtime} ({num_k_tiles} tiles)" + ) + + if tuple(C.shape) != (M_runtime, N_runtime): + raise ValueError( + f"C shape {tuple(C.shape)} != expected {(M_runtime, N_runtime)}" + ) + + if stream is None: + stream = torch.cuda.current_stream() + + launch = _cached_launch( + K_runtime, + C.dtype, + layout, + bool(use_xcd_remap), + ) + # 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. + A_arg = A.view(torch.uint8).view(-1) + B_arg = B.view(torch.uint8).view(-1) + C_arg = C.view(-1) + + launch( + A_arg, + B_arg, + C_arg, + M_runtime, + N_runtime, + stream=stream, + ) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm_utils.py new file mode 100644 index 000000000..99a7d5200 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm_utils.py @@ -0,0 +1,243 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors +"""Byte-staging helpers for the BF16 four-wave GEMM.""" + +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm as _llvm, vector +from flydsl.expr import const_expr, range_constexpr +from flydsl.expr.typing import Vector as Vec +from flydsl.expr.utils.arith import _to_raw as as_mlir_value + + +def cdiv(numer: int, denom: int) -> int: + return (numer + denom - 1) // denom + + +ceildiv = cdiv + + +def divmod(a, b): + return (a // b, a % b) + + +def swizzle_128(row, col_in_bytes): + """HK 128-byte row XOR swizzle; ``col_in_bytes`` is a byte coordinate.""" + offset = row * 128 + col_in_bytes + swizzle = ((offset % (16 * 128)) >> 8) << 4 + swizzled_offset = offset ^ swizzle + return swizzled_offset // 128, swizzled_offset % 128 + + +def make_bf16_buffer_tensor(arg_bf16): + """Create a BF16 BufferDesc directly from the wrapper-provided tensor.""" + return fx.rocdl.make_buffer_tensor(arg_bf16, max_size=False) + + +# Backward-compatible name used by fp16_gemm.py. +# Keep the exact existing behavior; this is only a symbol alias. +make_bf16_byte_buffer_tensor = make_bf16_buffer_tensor + +make_fp16_byte_buffer_tensor = make_bf16_byte_buffer_tensor + + +def compute_global_swizzle( + lane_id, + wave_id, + row_stride_bytes, + n_rounds, + preshuffled=False, +): + offsets = [] + n_waves = fx.block_dim.x // 64 + for round in range_constexpr(n_rounds): + if const_expr(preshuffled): + raise AssertionError( + "BF16 first-pass port does not support preshuffled operands" + ) + row = lane_id // 8 + wave_id * 8 + round * (n_waves * 8) + col_bytes = (lane_id % 8) * 16 + r, c = swizzle_128(row, col_bytes) + offsets.append(r * row_stride_bytes + c) + return offsets + + +def compute_global_bf16_transpose_swizzle( + lane_id, + wave_id, + leading_dim_bytes, + n_rounds, +): + """Offsets for a K-major BF16 source staged for ``ds_read_b64_tr_b16``. + + One 128-row output half-page is represented in LDS as two independent + swizzled ``[K64, X64]`` slices. Each slice is 64 rows by 128 bytes, so the + complete half-page remains 16 KiB and preserves the existing four-pass + 16-byte/thread DMA cadence. + + The returned offsets are relative to the source tile base: + ``source[k, x_base]`` for a contiguous K-major BF16 matrix. + """ + offsets = [] + n_waves = fx.block_dim.x // 64 + for round in range_constexpr(n_rounds): + linear_row = lane_id // 8 + wave_id * 8 + round * (n_waves * 8) + col_bytes = (lane_id % 8) * 16 + + slice_idx = linear_row // 64 + physical_k = linear_row % 64 + + # XOR swizzle is self-inverse for this layout. Map the physical LDS + # chunk back to its logical K/X-byte source coordinate. + logical_k, logical_x_bytes = swizzle_128(physical_k, col_bytes) + offsets.append( + logical_k * leading_dim_bytes + + slice_idx * 64 * 2 + + logical_x_bytes + ) + return offsets + + +compute_global_fp16_transpose_swizzle = compute_global_bf16_transpose_swizzle + +class G2SLoader: + """Issue native 16-byte BF16 BufferDesc-to-BF16 LDS copies.""" + + def __init__(self, gl_src, gl_offsets, n_load_steps, lds_dtype, wave_id): + self.g2lds_atom = fx.make_copy_atom(fx.rocdl.BufferCopyLDS128b(), 128) + self.LdsPtr_t = fx.PointerType.get(lds_dtype, 2, 512) + self.gl_src = gl_src + self.gl_offsets = gl_offsets + self.n_load_steps = n_load_steps + self.wave_id = wave_id + self.n_waves = fx.block_dim.x // 64 + + def _lds_dst_at(self, lds_dst, step): + step_off = self.wave_id * 1024 + step * (self.n_waves * 1024) + base_i32 = fx.Int32(fx.ptrtoint(lds_dst.ptr)) + lds_ptr = fx.inttoptr( + self.LdsPtr_t, + base_i32 + fx.Int32(step_off), + ) + return fx.make_view(lds_ptr, fx.make_layout(1, 1)) + + def load(self, lds_dst, byte_offset): + for step in range_constexpr(self.n_load_steps): + src = fx.slice( + self.gl_src, + (None, fx.Int32(self.gl_offsets[step])), + ) + fx.copy( + self.g2lds_atom, + src, + self._lds_dst_at(lds_dst, step), + soffset=fx.Int32(byte_offset), + ) + + def load_one(self, lds_dst, byte_offset, step): + src = fx.slice( + self.gl_src, + (None, fx.Int32(self.gl_offsets[step])), + ) + fx.copy( + self.g2lds_atom, + src, + self._lds_dst_at(lds_dst, step), + soffset=fx.Int32(byte_offset), + ) + + +def pack_i32x4_i32x8(lo, hi): + return lo.shuffle(hi, list(range(8))) + + +class S2RLoader: + """LDS readers used to assemble BF16 K64 fragments.""" + + def __init__(self, wave_idx, n_tiles): + self.lane_id = fx.thread_idx.x % 64 + self.wave_idx = wave_idx + self.n_tiles = n_tiles + + def _vec_load_16bytes(self, lds_src, offset): + ptr_off = fx.add_offset(lds_src.ptr, fx.make_int_tuple(offset)) + i8_iter = fx.recast_iter(fx.Uint8, ptr_off) + return fx.make_view(i8_iter, fx.make_layout(16, 1)).load() + + def load_one(self, lds_src, lds_offset): + return self._vec_load_16bytes( + lds_src, + lds_offset, + ).bitcast(fx.Int32) + + def _ds_read_b64_tr_b16( + self, + lds_src, + byte_offset, + immediate_offset=0, + ): + """Issue one gfx950 ``ds_read_b64_tr_b16`` and return i32x2.""" + if immediate_offset == 0: + asm = "ds_read_b64_tr_b16 $0, $1 offset:0\n" + elif immediate_offset == 0x1000: + asm = "ds_read_b64_tr_b16 $0, $1 offset:4096\n" + else: + raise ValueError( + "ds_read_b64_tr_b16 supports immediate offsets 0 and 0x1000, " + f"got {immediate_offset:#x}" + ) + + base_i32 = fx.Int32(fx.ptrtoint(lds_src.ptr)) + addr_i32 = base_i32 + fx.Int32(byte_offset) + raw_type = ir.VectorType.get( + [2], + ir.IntegerType.get_signless(32), + ) + raw = _llvm.inline_asm( + raw_type, + [as_mlir_value(addr_i32)], + asm, + "=v,v,~{memory}", + has_side_effects=True, + ) + return Vec( + vector.BitCastOp(raw_type, raw).result, + (2,), + fx.Int32, + ) + + def load_one_transpose_bf16( + self, + lds_src, + first_byte_offset, + second_byte_offset, + immediate_offset=0, + ): + """Return one i32x4 K32 BF16 fragment from two transpose reads.""" + lo = self._ds_read_b64_tr_b16( + lds_src, + first_byte_offset, + immediate_offset, + ) + hi = self._ds_read_b64_tr_b16( + lds_src, + second_byte_offset, + immediate_offset, + ) + return lo.shuffle(hi, [0, 1, 2, 3]) + + + def load_one_transpose_fp16( + self, + lds_src, + first_byte_offset, + second_byte_offset, + immediate_offset=0, + ): + """Return one i32x4 K32 FP16 fragment from two transpose reads.""" + return self.load_one_transpose_bf16( + lds_src, + first_byte_offset, + second_byte_offset, + immediate_offset, + ) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py new file mode 100644 index 000000000..6cbd61102 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py @@ -0,0 +1,1091 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""FlyDSL FP32 4-wave GEMM kernel for Transformer Engine. + +The kernel specializes on K at compile time because the K32 loop is fully +hand-unrolled. M/N are runtime launch dimensions. The private optimized core +consumes A and B as FP32 tensors shaped [M, K] and [N, K], and writes FP32 C +shaped [M, N]. The public ``fp32_matmul`` entry point accepts Transformer +Engine's TN contract and performs the required private adaptation. + +This module imports ``flydsl`` at import time and must therefore be imported +lazily only after FlyDSL availability has been confirmed. +""" + +import functools + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir.dialects import llvm +from flydsl.expr import arith, buffer_ops, const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec + +# Transformer Engine-local FlyDSL utilities. +from .exceptions import FlyDSLUnsupportedError +from .fp16_gemm_utils import ( + G2SLoader, + S2RLoader, + compute_global_swizzle, + make_bf16_byte_buffer_tensor as make_fp32_byte_buffer_tensor, + pack_i32x4_i32x8, + swizzle_128, +) + + +_BLOCK_M = 256 +_BLOCK_N = 256 +_BLOCK_K = 32 + +BLOCK_M = _BLOCK_M +BLOCK_N = _BLOCK_N +BLOCK_K = _BLOCK_K + +NUM_THREADS = 256 +WARP_SIZE = 64 +NUM_WAVES = NUM_THREADS // WARP_SIZE + +SUBTILE_M = 64 +SUBTILE_N = 64 + +MFMA_M = 16 +MFMA_N = 16 + +SUBTILES_PER_WAVE = 4 +MFMA_M_PER_SUBTILE = SUBTILE_M // MFMA_M +MFMA_N_PER_SUBTILE = SUBTILE_N // MFMA_N +ACCS_PER_WAVE = SUBTILES_PER_WAVE * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + +ELEM_BYTES = 4 +VEC_BYTES = 16 + +LDS_ELEMS_A = BLOCK_M * BLOCK_K +LDS_ELEMS_B = BLOCK_N * BLOCK_K +LDS_BYTES_A = LDS_ELEMS_A * ELEM_BYTES +LDS_BYTES_B = LDS_ELEMS_B * ELEM_BYTES + +LOAD_PASSES_A = LDS_BYTES_A // (NUM_THREADS * VEC_BYTES) +LOAD_PASSES_B = LDS_BYTES_B // (NUM_THREADS * VEC_BYTES) +LOAD_PASSES_A_SUBTILE = LOAD_PASSES_A // 2 +LOAD_PASSES_B_SUBTILE = LOAD_PASSES_B // 2 +PASSES_PER_A_MI = LOAD_PASSES_A_SUBTILE // MFMA_M_PER_SUBTILE + +LDS_SYM_A0 = "fp32_pp_smem_a0" +LDS_SYM_A1 = "fp32_pp_smem_a1" +LDS_SYM_B0 = "fp32_pp_smem_b0" +LDS_SYM_B1 = "fp32_pp_smem_b1" +LDS_ALIAS_DOMAIN = '#llvm.alias_scope_domain' +SCOPE_IDS = ("a0", "a1", "b0", "b1") + +assert BLOCK_K == 32 +# DO NOT CHANGE THE FOLLOWING LINE. +assert NUM_THREADS == 256 +assert LOAD_PASSES_A * NUM_THREADS * VEC_BYTES == LDS_BYTES_A +assert LOAD_PASSES_B * NUM_THREADS * VEC_BYTES == LDS_BYTES_B +assert LOAD_PASSES_A % 2 == 0 +assert LOAD_PASSES_B % 2 == 0 + + +def make_fp32_inputs(M, N, K, device="cuda"): + """Generate FP32 A[M,K] and B[N,K] inputs.""" + A = (torch.randn(M, K, device=device) * 0.5).to(torch.float32) + B = (torch.randn(N, K, device=device) * 0.5).to(torch.float32) + return A, B + + +def swizzle_xor16(row, col_in_bytes): + """XOR swizzle for the LDS K-byte coordinate.""" + chunk = col_in_bytes // fx.Index(VEC_BYTES) + byte_in_chunk = col_in_bytes % fx.Index(VEC_BYTES) + row_bits = (row % fx.Index(16)) // fx.Index(2) + swz_chunk = chunk ^ row_bits + return swz_chunk * fx.Index(VEC_BYTES) + byte_in_chunk + + +def _encode_waitcnt(vmcnt=63, lgkmcnt=15): + """Encode the CDNA4/gfx950 ``S_WAITCNT`` SIMM16 operand. + + ``rocdl.s_waitcnt`` accepts the raw 16-bit immediate operand of the + 32-bit ``S_WAITCNT`` ISA instruction. On CDNA4, that SIMM16 field is: + + SIMM16[3:0] = vmcnt[3:0] + SIMM16[6:4] = expcnt[2:0] + SIMM16[11:8] = lgkmcnt[3:0] + SIMM16[15:14] = vmcnt[5:4] + + ``vmcnt`` is therefore one six-bit counter split across two noncontiguous + fields; bits [5:4] are placed in SIMM16[15:14], while bits [3:0] remain + in SIMM16[3:0]. + + A wait-counter field set to its maximum representable value is effectively + unconstrained: the instruction does not wait on that counter. This helper + always encodes ``expcnt=7`` and defaults to ``vmcnt=63`` and ``lgkmcnt=15``, + so callers specify only the counters on which they intend to wait. + + For example, ``_encode_waitcnt(lgkmcnt=0)`` returns ``0xC07F``, which the + assembler renders as ``s_waitcnt lgkmcnt(0)``. + See: https://llvm.org/docs/AMDGPU/gfx9_waitcnt.html + """ + if not 0 <= vmcnt <= 63: + raise ValueError(f"vmcnt must be in [0, 63], got {vmcnt}") + if not 0 <= lgkmcnt <= 15: + raise ValueError(f"lgkmcnt must be in [0, 15], got {lgkmcnt}") + + return ( + (7 << 4) # expcnt=7 -> SIMM16[6:4] (unconstrained) + | (vmcnt & 0x0F) # vmcnt[3:0] -> SIMM16[3:0] + | ((lgkmcnt & 0x0F) << 8) # lgkmcnt[3:0] -> SIMM16[11:8] + | ((vmcnt & 0x30) << 10) # vmcnt[5:4] -> SIMM16[15:14] + ) + + +# Keep the documented gfx950 encoding invariant executable and import-time cheap. +assert _encode_waitcnt(lgkmcnt=0) == 0xC07F + + +def _barrier(vmcnt=63, lgkmcnt=15): + if vmcnt != 63 or lgkmcnt != 15: + rocdl.s_waitcnt(_encode_waitcnt(vmcnt=vmcnt, lgkmcnt=lgkmcnt)) + rocdl.s_barrier() + +def _min(a, b): + return arith.select(a < b, a, b) + + +def _divmod(a, b): + return a // b, a % b + + +def _xcd_swizzle(num_pid_m, num_pid_n): + NUM_XCDS = 8 + WGM = 4 + NUM_CUS = 32 * NUM_XCDS + SWIZZLE_THRESHOLD = 4 * NUM_CUS + + wgid = fx.block_idx.x + num_wg = num_pid_m * num_pid_n + + # Simple row-major path. + simple_m, simple_n = _divmod(wgid, num_pid_n) + + # XCD-remapped grouped-M path. + intra_xcd, xcd = _divmod(wgid, NUM_XCDS) + wgid_remap = xcd * (num_wg // NUM_XCDS) + intra_xcd + num_wgid_in_group = WGM * num_pid_n + group_id, intra_group = _divmod(wgid_remap, num_wgid_in_group) + first_pid_m = group_id * WGM + group_size_m = _min(num_pid_m - first_pid_m, WGM) + pid_n, intra_group_m = _divmod(intra_group, group_size_m) + pid_m = first_pid_m + intra_group_m + + use_simple = (num_wg < SWIZZLE_THRESHOLD) | (num_wg % NUM_XCDS != 0) + return ( + arith.select(use_simple, simple_m, pid_m), + arith.select(use_simple, simple_n, pid_n), + ) + + +def _compile_kernel(K: int, use_xcd_remap: bool = True): + """Build the specialized 4-wave kernel for compile-time ``K``. + + ``K`` must contain at least four K32 tiles. Runtime M/N are expected to + be exact multiples of ``BLOCK_M``/``BLOCK_N``; the kernel has no edge masks. + """ + BLOCK_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K + NUM_THREADS = 256 + WARP_SIZE = 64 + + SUBTILE_M = 64 + SUBTILE_N = 64 + + MFMA_M = 16 + MFMA_N = 16 + + SUBTILES_PER_WAVE = 4 + MFMA_M_PER_SUBTILE = SUBTILE_M // MFMA_M + MFMA_N_PER_SUBTILE = SUBTILE_N // MFMA_N + ACCS_PER_WAVE = SUBTILES_PER_WAVE * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + + ELEM_BYTES = 4 + VEC_BYTES = 16 + + LDS_ELEMS_A = BLOCK_M * BLOCK_K + LDS_ELEMS_B = BLOCK_N * BLOCK_K + LDS_BYTES_A = LDS_ELEMS_A * ELEM_BYTES + LDS_BYTES_B = LDS_ELEMS_B * ELEM_BYTES + + LOAD_PASSES_A = LDS_BYTES_A // (NUM_THREADS * VEC_BYTES) + LOAD_PASSES_B = LDS_BYTES_B // (NUM_THREADS * VEC_BYTES) + LOAD_PASSES_A_SUBTILE = LOAD_PASSES_A // 2 + LOAD_PASSES_B_SUBTILE = LOAD_PASSES_B // 2 + + assert K % BLOCK_K == 0, f"K must be a multiple of {BLOCK_K}, got {K}" + NUM_K_TILES = K // BLOCK_K + assert NUM_K_TILES >= 4, f"K={K} gives {NUM_K_TILES} K32 tiles; the two-page pipeline needs at least 4" + + LDS_ELEMS_HALF = (BLOCK_M // 2) * BLOCK_K + LDS_BYTES_HALF = LDS_ELEMS_HALF * ELEM_BYTES + LOAD_PASSES_HALF = LDS_BYTES_HALF // (NUM_THREADS * VEC_BYTES) + assert LOAD_PASSES_HALF == LOAD_PASSES_A_SUBTILE == LOAD_PASSES_B_SUBTILE + + @fx.struct + class SharedStorage: + # Each logical 256x64 FP32 page is two independent 128x64 half-pages. + # Store LDS as bytes so BufferCopyLDS128b sees i8 on both source and + # destination. Each half-page remains exactly 16 KiB. + a0_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a0_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a1_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a1_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b0_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + 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] + + @flyc.kernel(known_block_size=[NUM_THREADS, 1, 1]) + def kernel_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + ): + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + lds_a0 = (lds.a0_0, lds.a0_1) + lds_a1 = (lds.a1_0, lds.a1_1) + lds_b0 = (lds.b0_0, lds.b0_1) + lds_b1 = (lds.b1_0, lds.b1_1) + + # A/B arrive as contiguous uint8 byte views. Keeping staging byte-addressed + # preserves the original 16-byte G2L instruction cadence and vmcnt values. + gA = make_fp32_byte_buffer_tensor(A) + gB = make_fp32_byte_buffer_tensor(B) + 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") + + 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 * ELEM_BYTES, LOAD_PASSES_HALF, preshuffled=False) + gl_off_b = compute_global_swizzle(lane, wave_id, K * ELEM_BYTES, LOAD_PASSES_HALF, preshuffled=False) + 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) + + # C can exceed the signed-i32 element/byte offset range for large M*N. + # Bias the buffer descriptor base once per CTA using an index/i64 GEP, + # then store with only tile-local i32 offsets. This keeps the hot store + # instruction form unchanged while avoiding i32 wrap in buffer_store(). + c_n_idx_for_base = fx.Index(c_n) + c_tile_base_elems = bx_m_idx * c_n_idx_for_base + by_n_idx + c_tile_base_bytes = c_tile_base_elems * fx.Index(4) # C is FP32. + c_rsrc = buffer_ops.create_buffer_resource( + C, + max_size=True, + base_byte_offset=c_tile_base_bytes, + ) + + 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, + ) + + 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(): + # Eight refill VMEM operations overlap four independent 8-MFMA + # groups (two K16 halves x two N-halves). + for _ in range_constexpr(4): + rocdl.sched_vmem(2) + rocdl.sched_mfma(32) + 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(2) + rocdl.sched_mfma(32) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q_prefetch_4n(): + # Eight two-read prefetch groups overlap four complete-quadrant + # 16-MFMA groups (two K16 halves for each of Q2 and Q3). + for _ in range_constexpr(4): + rocdl.sched_dsrd(4) + rocdl.sched_mfma(64) + 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 = (bx_m_idx + fx.Index(subtile * (BLOCK_M // 2))) * fx.Index(K * ELEM_BYTES) + k_base * fx.Index(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(K * ELEM_BYTES) + k_base * fx.Index(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): + 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 _fp32_k4_operand(full_frag, k32_half, k4): + # A/B 16x32 FP32 wave fragments are i32x8: one FP32 value per + # VGPR and eight K4 MFMA steps per logical K32 tile. Keep the + # existing two-half schedule by grouping four K4 steps per half. + return Vec(full_frag)[k32_half * 4 + k4] + + def _pinned_fp32_mfma_once(acc_idx, a_k4, b_k4): + acc_pin = PIN_ACC_BASE + acc_idx * 4 + llvm.InlineAsmOp( + None, + [arith._to_raw(a_k4), arith._to_raw(b_k4)], + ( + f"v_mfma_f32_16x16x4_f32 " + f"a[{acc_pin}:{acc_pin + 3}], " + f"$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}]" + ), + ( + f"v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}}," + f"~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}" + ), + has_side_effects=True, + ) + + def pinned_mfma(acc_idx, a_frag, b_frag): + """Accumulate one logical 16x16x32 FP32 product into pinned AGPRs.""" + for k32_half in range_constexpr(2): + for k4 in range_constexpr(4): + _pinned_fp32_mfma_once( + acc_idx, + _fp32_k4_operand(a_frag, k32_half, k4), + _fp32_k4_operand(b_frag, k32_half, k4), + ) + + def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag): + # The final logical K32 update is eight in-place K4 FP32 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_half, a0, a1, a2, a3, b0, b1): + """Issue four K4 steps (one K16 half) for a 4x2 accumulator slab.""" + a_frags = (a0, a1, a2, a3) + b_frags = (b0, b1) + for k4 in range_constexpr(4): + for mi in range_constexpr(4): + a_k4 = _fp32_k4_operand(a_frags[mi], k32_half, k4) + for nj in range_constexpr(2): + _pinned_fp32_mfma_once( + _acc_idx(subtile_id, mi, n_base + nj), + a_k4, + _fp32_k4_operand(b_frags[nj], k32_half, k4), + ) + + def mfma_4n_4mi_k32(subtile_id, k32_half, a0, a1, a2, a3, b0, b1, b2, b3): + """Issue four K4 steps (one K16 half) for a complete 4x4 quadrant.""" + a_frags = (a0, a1, a2, a3) + b_frags = (b0, b1, b2, b3) + for k4 in range_constexpr(4): + for mi in range_constexpr(4): + a_k4 = _fp32_k4_operand(a_frags[mi], k32_half, k4) + for ni in range_constexpr(4): + _pinned_fp32_mfma_once( + _acc_idx(subtile_id, mi, ni), + a_k4, + _fp32_k4_operand(b_frags[ni], k32_half, k4), + ) + + 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 + for ii in range_constexpr(4): + row = row_base + fx.Index(ii) + c_idx = row * fx.Index(c_n) + col + buffer_ops.buffer_store(Vec(acc)[ii], c_rsrc, 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 * 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): + 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. Compute is K16-half-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) + + 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_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:16]. + + 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_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) + + 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_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[16:32]. + + 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. + + 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) + + 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_a_subtile_pass(k_refill, 1, 0, refill_a) + mfma_2n_4mi_k32(1, 2, 0, a00, a01, a02, a03, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 1, refill_b) + # K32 slice 0 already covers K[0:16]. + + stage_a_subtile_pass(k_refill, 1, 1, refill_a) + # Keep this refill slot compute-free. + + 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, 2, refill_a) + mfma_2n_4mi_k32(1, 2, 1, a00, a01, a02, a03, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 3, refill_b) + # K32 slice 1 already covers K[16:32]. + + stage_a_subtile_pass(k_refill, 1, 3, refill_a) + # Keep this refill slot compute-free. + 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, lgkmcnt=0) + rocdl.sched_barrier(0) + + 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_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + # K32 slice 0 already covers K[0:16]. + + 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_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + # K32 slice 1 already covers K[16: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_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + # K32 slice 0 already covers K[0:16]. + + 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_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + # K32 slice 1 already covers K[16:32]. + + 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) + + return next_a0_regs, next_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) + + 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) + + 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) + + 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) + + 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) + + rocdl.sched_barrier(0) + _barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, 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_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, 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_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, 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_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, 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) + + 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) + + return next_a0_regs, next_b0_regs + + def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): + _barrier(vmcnt=0, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + # 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) + + 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) + + # 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) + + 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 + 3 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) + + b0_regs = load_b_subtile_regs(lds_b0, 0) + + # Main HK loop: exactly one logical K32 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, + a0_regs, + b0_regs, + ) + hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs) + 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) + + + @flyc.jit + def launch_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + # The integration only dispatches aligned shapes; no partial-tile masking exists. + grid_x = (c_m // BLOCK_M) * (c_n // BLOCK_N) + kernel_gemm( + A, + B, + C, + c_m, + c_n, + 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) + + return launch_gemm + +@functools.lru_cache(maxsize=None) +def _cached_launch(K: int, use_xcd_remap: bool = True): + return _compile_kernel(K, use_xcd_remap=use_xcd_remap) + + + +def fp32_matmul( + a: torch.Tensor, + b: torch.Tensor, + c: torch.Tensor, + stream=None, +): + """TE-facing TN FP32 GEMM adapter. + + Public/backend contract: + a: [M, K] FP32 + b: [K, N] FP32 + c: [M, N] FP32 output + + The optimized core streams both operands with K contiguous and therefore + privately consumes B as [N, K]. In the normal TE TN path, ``b`` is a + transpose view of contiguous rowwise weight storage, so ``b.T`` is already + contiguous and does not require a physical transpose. + """ + if a.ndim != 2 or b.ndim != 2: + raise ValueError( + f"FlyDSL FP32 TN expects rank-2 operands, got A{tuple(a.shape)} " + f"and B{tuple(b.shape)}" + ) + + m, k = a.shape + kb, n = 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 TypeError( + 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 " + f"{a.device}, {b.device}, and {c.device}" + ) + if not c.is_contiguous(): + raise ValueError("FlyDSL FP32 GEMM requires contiguous output storage") + + b_hk = b.transpose(0, 1).contiguous() + doGemm(a, b_hk, c, stream=stream) + + +def doGemm( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + stream=None, + use_xcd_remap: bool = True, +): + """Launch the private K-specialized FP32 core. + + A and B are shaped [M, K] and [N, K]; C is shaped [M, N]. M and N + remain runtime values, while K selects the cached compile-time specialization. + """ + 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}" + assert A.dtype == torch.float32 and B.dtype == torch.float32 + assert C.dtype == torch.float32 + if M_runtime % _BLOCK_M != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP32 GEMM requires M to be a multiple of {_BLOCK_M}, " + f"got M={M_runtime}" + ) + + if N_runtime % _BLOCK_N != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP32 GEMM requires N to be a multiple of {_BLOCK_N}, " + f"got N={N_runtime}" + ) + + if K_runtime % _BLOCK_K != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP32 GEMM requires K to be a multiple of {_BLOCK_K}, " + f"got K={K_runtime}" + ) + + num_k_tiles = K_runtime // _BLOCK_K + if num_k_tiles < 4: + raise FlyDSLUnsupportedError( + f"FlyDSL FP32 GEMM requires at least 4 K{_BLOCK_K} tiles, " + f"got K={K_runtime} ({num_k_tiles} tiles)" + ) + assert C.shape == (M_runtime, N_runtime) + if stream is None: + stream = torch.cuda.current_stream() + + A_arg = A.contiguous().view(torch.uint8).view(-1) + B_arg = B.contiguous().view(torch.uint8).view(-1) + C_arg = C.view(-1) + launch = _cached_launch(int(K_runtime), bool(use_xcd_remap)) + launch(A_arg, B_arg, C_arg, M_runtime, N_runtime, stream=stream) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py new file mode 100644 index 000000000..45c7ad66e --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -0,0 +1,1178 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""FlyDSL tensor-wise FP8 4-wave GEMM kernel for Transformer Engine. + +The kernel specializes on K at compile time because the K128 loop is fully +hand-unrolled. M/N are runtime launch dimensions. The public entry point and private optimized core consume independently typed +FP8 E4M3 or E5M2 A/B tensors shaped [M, K] and [N, K], one FP32 inverse scale +per operand, and write float16, bfloat16, or float32 C shaped [M, N]. Operand +normalization is performed by the Transformer Engine wrapper. + +This module imports ``flydsl`` at import time and must therefore be imported +lazily only after FlyDSL availability has been confirmed. +""" + +import functools + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir.dialects import llvm +from flydsl.expr import arith, buffer_ops, const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec + +from .exceptions import FlyDSLUnsupportedError + +# Transformer Engine-local FlyDSL utilities. +from .fp8_gemm_utils import ( + G2SLoader, + S2RLoader, + compute_global_swizzle, + make_fp8_buffer_tensor, + pack_i32x4_i32x8, + swizzle_128, +) + + + +_BLOCK_M = 256 +_BLOCK_N = 256 +_BLOCK_K = 128 + +BLOCK_M = _BLOCK_M +BLOCK_N = _BLOCK_N +BLOCK_K = _BLOCK_K + +NUM_THREADS = 256 +WARP_SIZE = 64 +NUM_WAVES = NUM_THREADS // WARP_SIZE + +SUBTILE_M = 64 +SUBTILE_N = 64 + +MFMA_M = 16 +MFMA_N = 16 + +SUBTILES_PER_WAVE = 4 +MFMA_M_PER_SUBTILE = SUBTILE_M // MFMA_M +MFMA_N_PER_SUBTILE = SUBTILE_N // MFMA_N +ACCS_PER_WAVE = SUBTILES_PER_WAVE * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + +ELEM_BYTES = 1 +VEC_BYTES = 16 + +LDS_ELEMS_A = BLOCK_M * BLOCK_K +LDS_ELEMS_B = BLOCK_N * BLOCK_K +LDS_BYTES_A = LDS_ELEMS_A * ELEM_BYTES +LDS_BYTES_B = LDS_ELEMS_B * ELEM_BYTES + +LOAD_PASSES_A = LDS_BYTES_A // (NUM_THREADS * VEC_BYTES) +LOAD_PASSES_B = LDS_BYTES_B // (NUM_THREADS * VEC_BYTES) +LOAD_PASSES_A_SUBTILE = LOAD_PASSES_A // 2 +LOAD_PASSES_B_SUBTILE = LOAD_PASSES_B // 2 +PASSES_PER_A_MI = LOAD_PASSES_A_SUBTILE // MFMA_M_PER_SUBTILE + +LDS_SYM_A0 = "fp8_pp_smem_a0" +LDS_SYM_A1 = "fp8_pp_smem_a1" +LDS_SYM_B0 = "fp8_pp_smem_b0" +LDS_SYM_B1 = "fp8_pp_smem_b1" +LDS_ALIAS_DOMAIN = '#llvm.alias_scope_domain' +SCOPE_IDS = ("a0", "a1", "b0", "b1") + +assert BLOCK_K == 128 +# DO NOT CHANGE THE FOLLOWING LINE. +assert NUM_THREADS == 256 +assert LOAD_PASSES_A * NUM_THREADS * VEC_BYTES == LDS_BYTES_A +assert LOAD_PASSES_B * NUM_THREADS * VEC_BYTES == LDS_BYTES_B +assert LOAD_PASSES_A % 2 == 0 +assert LOAD_PASSES_B % 2 == 0 + + +def swizzle_xor16(row, col_in_bytes): + """XOR swizzle for the LDS K-byte coordinate.""" + chunk = col_in_bytes // fx.Index(VEC_BYTES) + byte_in_chunk = col_in_bytes % fx.Index(VEC_BYTES) + row_bits = (row % fx.Index(16)) // fx.Index(2) + swz_chunk = chunk ^ row_bits + return swz_chunk * fx.Index(VEC_BYTES) + byte_in_chunk + + +def _encode_waitcnt(vmcnt=63, lgkmcnt=15): + """Encode the CDNA4/gfx950 ``S_WAITCNT`` SIMM16 operand. + + ``rocdl.s_waitcnt`` accepts the raw 16-bit immediate operand of the + 32-bit ``S_WAITCNT`` ISA instruction. On CDNA4, that SIMM16 field is: + + SIMM16[3:0] = vmcnt[3:0] + SIMM16[6:4] = expcnt[2:0] + SIMM16[11:8] = lgkmcnt[3:0] + SIMM16[15:14] = vmcnt[5:4] + + ``vmcnt`` is therefore one six-bit counter split across two noncontiguous + fields; bits [5:4] are placed in SIMM16[15:14], while bits [3:0] remain + in SIMM16[3:0]. + + A wait-counter field set to its maximum representable value is effectively + unconstrained: the instruction does not wait on that counter. This helper + always encodes ``expcnt=7`` and defaults to ``vmcnt=63`` and ``lgkmcnt=15``, + so callers specify only the counters on which they intend to wait. + + For example, ``_encode_waitcnt(lgkmcnt=0)`` returns ``0xC07F``, which the + assembler renders as ``s_waitcnt lgkmcnt(0)``. + See: https://llvm.org/docs/AMDGPU/gfx9_waitcnt.html + """ + if not 0 <= vmcnt <= 63: + raise ValueError(f"vmcnt must be in [0, 63], got {vmcnt}") + if not 0 <= lgkmcnt <= 15: + raise ValueError(f"lgkmcnt must be in [0, 15], got {lgkmcnt}") + + return ( + (7 << 4) # expcnt=7 -> SIMM16[6:4] (unconstrained) + | (vmcnt & 0x0F) # vmcnt[3:0] -> SIMM16[3:0] + | ((lgkmcnt & 0x0F) << 8) # lgkmcnt[3:0] -> SIMM16[11:8] + | ((vmcnt & 0x30) << 10) # vmcnt[5:4] -> SIMM16[15:14] + ) + + + +# Keep the documented gfx950 encoding invariant executable and import-time cheap. +assert _encode_waitcnt(lgkmcnt=0) == 0xC07F + +def _barrier(vmcnt=63, lgkmcnt=15): + if vmcnt != 63 or lgkmcnt != 15: + rocdl.s_waitcnt(_encode_waitcnt(vmcnt=vmcnt, lgkmcnt=lgkmcnt)) + rocdl.s_barrier() + +def _min(a, b): + return arith.select(a < b, a, b) + + +def _divmod(a, b): + return a // b, a % b + + +def _xcd_swizzle(num_pid_m, num_pid_n): + NUM_XCDS = 8 + WGM = 4 + NUM_CUS = 32 * NUM_XCDS + SWIZZLE_THRESHOLD = 4 * NUM_CUS + + wgid = fx.block_idx.x + num_wg = num_pid_m * num_pid_n + + # Simple row-major path. + simple_m, simple_n = _divmod(wgid, num_pid_n) + + # XCD-remapped grouped-M path. + intra_xcd, xcd = _divmod(wgid, NUM_XCDS) + wgid_remap = xcd * (num_wg // NUM_XCDS) + intra_xcd + num_wgid_in_group = WGM * num_pid_n + group_id, intra_group = _divmod(wgid_remap, num_wgid_in_group) + first_pid_m = group_id * WGM + group_size_m = _min(num_pid_m - first_pid_m, WGM) + pid_n, intra_group_m = _divmod(intra_group, group_size_m) + pid_m = first_pid_m + intra_group_m + + use_simple = (num_wg < SWIZZLE_THRESHOLD) | (num_wg % NUM_XCDS != 0) + return ( + arith.select(use_simple, simple_m, pid_m), + arith.select(use_simple, simple_n, pid_n), + ) + + +def _compile_kernel( + K: int, + a_fp8_dtype: torch.dtype, + b_fp8_dtype: torch.dtype, + output_dtype: torch.dtype, + use_xcd_remap: bool = True, +): + """Build the specialized kernel for compile-time K, A/B FP8 types, and output dtype. + + ``K`` must contain at least four K128 tiles. Runtime M/N are expected to + be exact multiples of ``BLOCK_M``/``BLOCK_N``; the kernel has no edge masks. + """ + BLOCK_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K + + fp8_input_types = { + torch.float8_e4m3fn: (fx.Float8E4M3FN, 0), + torch.float8_e5m2: (fx.Float8E5M2, 1), + } + try: + a_fx_dtype, a_matrix_format = fp8_input_types[a_fp8_dtype] + b_fx_dtype, b_matrix_format = fp8_input_types[b_fp8_dtype] + except KeyError as exc: + raise TypeError( + "FlyDSL FP8 input dtype must be torch.float8_e4m3fn or " + f"torch.float8_e5m2, got A={a_fp8_dtype}, B={b_fp8_dtype}" + ) from exc + + if output_dtype == torch.float16: + output_element_bytes = 2 + output_fx_dtype = fx.Float16 + elif output_dtype == torch.bfloat16: + output_element_bytes = 2 + output_fx_dtype = fx.BFloat16 + elif output_dtype == torch.float32: + output_element_bytes = 4 + output_fx_dtype = fx.Float32 + else: + raise TypeError( + "FlyDSL FP8 supports only float16, bfloat16, and float32 " + f"outputs, got {output_dtype}" + ) + NUM_THREADS = 256 + WARP_SIZE = 64 + + SUBTILE_M = 64 + SUBTILE_N = 64 + + MFMA_M = 16 + MFMA_N = 16 + + SUBTILES_PER_WAVE = 4 + MFMA_M_PER_SUBTILE = SUBTILE_M // MFMA_M + MFMA_N_PER_SUBTILE = SUBTILE_N // MFMA_N + ACCS_PER_WAVE = SUBTILES_PER_WAVE * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + + ELEM_BYTES = 1 + VEC_BYTES = 16 + + LDS_ELEMS_A = BLOCK_M * BLOCK_K + LDS_ELEMS_B = BLOCK_N * BLOCK_K + LDS_BYTES_A = LDS_ELEMS_A * ELEM_BYTES + LDS_BYTES_B = LDS_ELEMS_B * ELEM_BYTES + + LOAD_PASSES_A = LDS_BYTES_A // (NUM_THREADS * VEC_BYTES) + LOAD_PASSES_B = LDS_BYTES_B // (NUM_THREADS * VEC_BYTES) + LOAD_PASSES_A_SUBTILE = LOAD_PASSES_A // 2 + LOAD_PASSES_B_SUBTILE = LOAD_PASSES_B // 2 + + assert K % BLOCK_K == 0, f"K must be a multiple of {BLOCK_K}, got {K}" + NUM_K_TILES = K // BLOCK_K + assert NUM_K_TILES >= 4, f"K={K} gives {NUM_K_TILES} K128 tiles; the two-page pipeline needs at least 4" + + LDS_ELEMS_HALF = (BLOCK_M // 2) * BLOCK_K + LOAD_PASSES_HALF = LDS_ELEMS_HALF // (NUM_THREADS * VEC_BYTES) + assert LOAD_PASSES_HALF == LOAD_PASSES_A_SUBTILE == LOAD_PASSES_B_SUBTILE + + @fx.struct + class SharedStorage: + # Each logical 256x128 page is two independent 128x128 half-pages. + # The hot loop refills one 16-byte pass of one half-page at a time. + a0_0: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a0_1: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a1_0: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a1_1: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + b0_0: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + 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] + + @flyc.kernel(known_block_size=[NUM_THREADS, 1, 1]) + def kernel_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + A_scale_inv: fx.Tensor, + B_scale_inv: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + ): + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + lds_a0 = (lds.a0_0, lds.a0_1) + lds_a1 = (lds.a1_0, lds.a1_1) + lds_b0 = (lds.b0_0, lds.b0_1) + lds_b1 = (lds.b1_0, lds.b1_1) + + a_f8_ir_t = a_fx_dtype.ir_type + b_f8_ir_t = b_fx_dtype.ir_type + gA = make_fp8_buffer_tensor(A, a_f8_ir_t) + 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 = buffer_ops.create_buffer_resource(A_scale_inv, max_size=True) + b_scale_rsrc = buffer_ops.create_buffer_resource(B_scale_inv, max_size=True) + output_scale = ( + buffer_ops.buffer_load(a_scale_rsrc, fx.Index(0), vec_width=1, dtype=T.f32) + * buffer_ops.buffer_load(b_scale_rsrc, fx.Index(0), vec_width=1, dtype=T.f32) + ) + tx = gpu.thread_id("x") + + 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) + + # C can exceed the signed-i32 element/byte offset range for large M*N. + # Bias the buffer descriptor base once per CTA using an index/i64 GEP, + # then store with only tile-local i32 offsets. This keeps the hot store + # instruction form unchanged while avoiding i32 wrap in buffer_store(). + c_n_idx_for_base = fx.Index(c_n) + c_tile_base_elems = bx_m_idx * c_n_idx_for_base + by_n_idx + c_tile_base_bytes = c_tile_base_elems * fx.Index(output_element_bytes) + c_rsrc = buffer_ops.create_buffer_resource( + C, + max_size=True, + base_byte_offset=c_tile_base_bytes, + ) + + 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, + ) + + 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. + 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_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), + ], + ( + f"v_mfma_f32_16x16x128_f8f6f4 " + f"a[{acc_pin}:{acc_pin + 3}], " + f"$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}}}," + f"~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}" + ), + has_side_effects=True, + ) + + 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), + ], + ( + f"v_mfma_f32_16x16x128_f8f6f4 " + f"a[{dst_pin}:{dst_pin + 3}], " + f"$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}}}," + f"~{{a{dst_pin + 2}}},~{{a{dst_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 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 + for ii in range_constexpr(4): + row = row_base + fx.Index(ii) + c_idx = row * fx.Index(c_n) + col + value = Vec(acc)[ii] * output_scale + if output_dtype != torch.float32: + value = value.to(output_fx_dtype) + buffer_ops.buffer_store(value, c_rsrc, 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) + + 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) + + 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) + + 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) + + 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) + + 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) + + 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) + + 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) + + 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) + + rocdl.sched_barrier(0) + stage_b_subtile_pass(k_refill, 1, 0, refill_b) + mfma_2n(_acc_idx(1, 0, 0), a00, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 0, refill_a) + mfma_2n(_acc_idx(1, 0, 2), a00, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 1, refill_b) + mfma_2n(_acc_idx(1, 1, 0), a01, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 1, refill_a) + mfma_2n(_acc_idx(1, 1, 2), a01, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 2, refill_b) + mfma_2n(_acc_idx(1, 2, 0), a02, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 2, refill_a) + mfma_2n(_acc_idx(1, 2, 2), a02, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 3, refill_b) + mfma_2n(_acc_idx(1, 3, 0), a03, b10, b11) + + 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() + + # 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_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, 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_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, 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_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, 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_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, 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) + + 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) + + return next_a0_regs, next_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) + + 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) + + 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) + + 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) + + 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) + + rocdl.sched_barrier(0) + _barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, 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_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, 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_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, 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_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, 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) + + 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) + + return next_a0_regs, next_b0_regs + + def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): + _barrier(vmcnt=0, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + # 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) + + 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) + + # 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) + + 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 + 3 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) + + 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, + a0_regs, + b0_regs, + ) + hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs) + 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) + + + @flyc.jit + def launch_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + A_scale_inv: fx.Tensor, + B_scale_inv: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + # The integration only dispatches aligned shapes; no partial-tile masking exists. + grid_x = (c_m // BLOCK_M) * (c_n // BLOCK_N) + kernel_gemm( + A, + B, + C, + A_scale_inv, + B_scale_inv, + c_m, + c_n, + 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) + + return launch_gemm + +@functools.lru_cache(maxsize=None) +def _cached_launch( + K: int, + a_fp8_dtype: torch.dtype, + b_fp8_dtype: torch.dtype, + output_dtype: torch.dtype, + use_xcd_remap: bool = True, +): + return _compile_kernel( + K, + a_fp8_dtype, + b_fp8_dtype, + output_dtype, + use_xcd_remap=use_xcd_remap, + ) + + + +def fp8_matmul( + a: torch.Tensor, + a_scale_inv: torch.Tensor, + b: torch.Tensor, + b_scale_inv: torch.Tensor, + c: torch.Tensor, + stream=None, +): + """Launch TN tensor-wise FP8 GEMM using final kernel operand order. + + Contract: + a: [M, K] FP8 E4M3 or E5M2 activation payload + a_scale_inv: one-element FP32 inverse quantization scale + b: [N, K] FP8 E4M3 or E5M2 weight payload + b_scale_inv: one-element FP32 inverse quantization scale + c: [M, N] float16, bfloat16, or float32 output + + 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)} " + f"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( + "FlyDSL FP8 GEMM expects E4M3 or E5M2 payloads, " + f"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 TypeError( + "FlyDSL FP8 supports only float16, bfloat16, and float32 " + f"outputs, got {c.dtype}" + ) + if not c.is_contiguous(): + raise ValueError("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" + ) + + doGemm( + a, + b, + c, + a_scale_inv, + b_scale_inv, + stream=stream, + ) + +def doGemm( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + A_scale_inv: torch.Tensor, + B_scale_inv: torch.Tensor, + stream=None, + use_xcd_remap: bool = True, +): + """Launch tensor-wise FP8 GEMM with TE-style inverse input scales.""" + M_runtime, K_runtime = A.shape + N_runtime, Kb_runtime = B.shape + 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}" + assert C.dtype in ( + torch.float16, + torch.bfloat16, + torch.float32, + ), ( + "C dtype must be torch.float16, torch.bfloat16, or torch.float32, " + f"got {C.dtype}" + ) + assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" + if M_runtime % _BLOCK_M != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires M to be a multiple of {_BLOCK_M}, " + f"got M={M_runtime}" + ) + if N_runtime % _BLOCK_N != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires N to be a multiple of {_BLOCK_N}, " + f"got N={N_runtime}" + ) + if K_runtime % _BLOCK_K != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires K to be a multiple of {_BLOCK_K}, " + f"got K={K_runtime}" + ) + num_k_tiles = K_runtime // _BLOCK_K + if num_k_tiles < 4: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires at least 4 K{_BLOCK_K} tiles, " + f"got K={K_runtime} ({num_k_tiles} tiles)" + ) + 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})" + ) + if stream is None: + stream = torch.cuda.current_stream() + + A_arg = A.view(torch.uint8).contiguous().view(-1) + B_arg = B.view(torch.uint8).contiguous().view(-1) + C_arg = C.contiguous().view(-1) + A_scale_arg = A_scale_inv.contiguous().view(-1) + B_scale_arg = B_scale_inv.contiguous().view(-1) + + launch = _cached_launch( + int(K_runtime), + A.dtype, + B.dtype, + C.dtype, + bool(use_xcd_remap), + ) + launch( + A_arg, + B_arg, + C_arg, + A_scale_arg, + B_scale_arg, + M_runtime, + N_runtime, + stream=stream, + ) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_utils.py new file mode 100644 index 000000000..2c0e0534b --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_utils.py @@ -0,0 +1,417 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm as _llvm, vector +from flydsl._mlir.dialects.fly_rocdl import TargetAddressSpace +from flydsl.expr import arith, buffer_ops, const_expr, range_constexpr, rocdl +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec +from flydsl.expr.utils.arith import _to_raw as as_mlir_value + +# ceildiv is the canonical cdiv from the shared layer +def cdiv(numer: int, denom: int) -> int: + return (numer + denom - 1) // denom + + +ceildiv = cdiv + +def divmod(a, b): + """Integer divmod that works on DSL values (e.g. ``Int32``). + + The builtin ``divmod`` rejects DSL scalar types, so this uses the overloaded + ``//`` / ``%`` operators to emit the corresponding ops. + """ + return (a // b, a % b) + + +def preshuffle_b(b_t): + """Permute row-major ``B_T`` ``(N, K)`` for ``b_preshuffled=True``.""" + n, k = b_t.shape[-2:] + assert n % 16 == 0 and k % 64 == 0, f"need N%16==0 and K%64==0, got N={n} K={k}" + return b_t.reshape(n // 16, 16, k // 64, 4, 16).permute(0, 2, 3, 1, 4).contiguous() + + +def make_fp8_buffer_tensor(arg_i8, fp8_ir_t): + # max_size=False with no num_records_bytes: cosize(layout) becomes a + # runtime expression because TensorAdaptor defaults to layout-dynamic + # memref (post #554), so the descriptor adapts to the actual tensor + # extent and no longer bakes the first-call's shape into IR. + t_i8 = fx.rocdl.make_buffer_tensor(arg_i8, max_size=False) + iter_i8 = fx.get_iter(t_i8) + f8_buf_ptr_ty = fx.PointerType.get( + elem_ty=fp8_ir_t, + address_space=TargetAddressSpace.BufferDesc, + alignment=fx.PointerType(iter_i8.type).alignment, + ) + iter_f8 = fx.recast_iter(f8_buf_ptr_ty, iter_i8) + return fx.Tensor(fx.make_view(iter_f8, fx.get_layout(t_i8))) + + +def swizzle_128(row, col): + offset = row * 128 + col + swizzle = ((offset % (16 * 128)) >> 8) << 4 + swizzled_offset = offset ^ swizzle + return swizzled_offset // 128, swizzled_offset % 128 + + +def compute_global_swizzle(lane_id, wave_id, K, n_rounds, preshuffled): + offsets = [] + n_waves = fx.block_dim.x // 64 + for round in range_constexpr(n_rounds): + if const_expr(preshuffled): + row = lane_id % 8 + wave_id * 8 + round * (n_waves * 8) + col = (lane_id // 8) * 16 + offsets.append( + (row // 16) * (K * 16) + (row % 16) * 16 + (col // 64) * 1024 + ((col % 64) // 16) * 256 + (col % 16) + ) + else: + row = lane_id // 8 + wave_id * 8 + round * (n_waves * 8) + col = (lane_id % 8) * 16 + r, c = swizzle_128(row, col) + offsets.append(r * K + c) + return offsets + + +def compute_global_linear_128x128(lane_id, wave_id, leading_dim, n_rounds): + """Offsets for an unswizzled row-major 128x128 tile. + + This uses the same 16-byte/thread DMA decomposition as + ``compute_global_swizzle`` but does not XOR-permute the logical source + coordinates. It is used by the NN A path, whose LDS page is physically + [K128, M128] for the CDNA4 transpose-read instruction. + """ + offsets = [] + n_waves = fx.block_dim.x // 64 + for round in range_constexpr(n_rounds): + row = lane_id // 8 + wave_id * 8 + round * (n_waves * 8) + col = (lane_id % 8) * 16 + offsets.append(row * leading_dim + col) + return offsets + + +class G2SLoader: + def __init__(self, gl_src, gl_offsets, n_load_steps, lds_dtype, wave_id): + self.g2lds_atom = fx.make_copy_atom(fx.rocdl.BufferCopyLDS128b(), 128) + self.LdsPtr_t = fx.PointerType.get(lds_dtype, 2, 512) + self.gl_src = gl_src + self.gl_offsets = gl_offsets + self.n_load_steps = n_load_steps + self.wave_id = wave_id + self.n_waves = fx.block_dim.x // 64 + + def _lds_dst_at(self, lds_dst, step): + step_off = self.wave_id * 1024 + step * (self.n_waves * 1024) + base_i32 = fx.Int32(fx.ptrtoint(lds_dst.ptr)) + sum_i32 = base_i32 + fx.Int32(step_off) + lds_ptr = fx.inttoptr(self.LdsPtr_t, sum_i32) + return fx.make_view(lds_ptr, fx.make_layout(1, 1)) + + def load(self, lds_dst, k_offset): + for step in range_constexpr(self.n_load_steps): + src = fx.slice(self.gl_src, (None, fx.Int32(self.gl_offsets[step]))) + dst = self._lds_dst_at(lds_dst, step) + fx.copy(self.g2lds_atom, src, dst, soffset=fx.Int32(k_offset)) + + def load_one(self, lds_dst, k_offset, step): + src = fx.slice(self.gl_src, (None, fx.Int32(self.gl_offsets[step]))) + dst = self._lds_dst_at(lds_dst, step) + fx.copy(self.g2lds_atom, src, dst, soffset=fx.Int32(k_offset)) + + +class G2STransposeLoader: + """Stage a row-major 128x128 byte tile as swizzled physical [K, N]. + + The source is a row-major byte matrix ``[N, K]``. Each thread loads one + contiguous 16-byte K vector from global memory, then scatters those bytes + into the 128-byte XOR-swizzled LDS image consumed by + ``ds_read_b64_tr_b8``. + + One ``load_one`` call covers one of the four 4-KiB staging passes for a + 128x128 half-page. + """ + + def __init__(self, gl_src, leading_dim, wave_id): + self.gl_rsrc = buffer_ops.create_buffer_resource(gl_src, max_size=True) + self.leading_dim = fx.Int32(leading_dim) + self.wave_id = fx.Int32(wave_id) + self.lane_id = fx.thread_idx.x % 64 + self.n_waves = fx.block_dim.x // 64 + self.i8_lds_ptr_t = fx.PointerType.get( + elem_ty=ir.IntegerType.get_signless(8), + address_space=2, + alignment=1, + ) + + def _store_u8(self, lds_dst, byte_offset, value): + base_i32 = fx.Int32(fx.ptrtoint(lds_dst.ptr)) + addr_i32 = base_i32 + fx.Int32(byte_offset) + i8_ptr = fx.inttoptr(self.i8_lds_ptr_t, addr_i32) + view = fx.make_view(i8_ptr, fx.make_layout(1, 1)) + fx.memref_store_vec(Vec.filled(1, value, fx.Uint8), view) + + def load_one(self, lds_dst, global_n_base, k_base, step): + """Load one 16-byte/thread pass and transpose it into LDS. + + ``global_n_base`` is the first source N row of this 128-row half-page. + ``k_base`` is the first global K byte of the current K128 tile. + """ + row = ( + self.lane_id // fx.Int32(8) + + self.wave_id * fx.Int32(8) + + fx.Int32(step) * fx.Int32(self.n_waves * 8) + ) + col = (self.lane_id % fx.Int32(8)) * fx.Int32(16) + + global_byte = ( + (fx.Int32(global_n_base) + row) * self.leading_dim + + fx.Int32(k_base) + + col + ) + packed_i32x4 = buffer_ops.buffer_load( + self.gl_rsrc, + global_byte // fx.Int32(4), + vec_width=4, + dtype=T.i32, + ) + packed_u8x16 = Vec(packed_i32x4).bitcast(fx.Uint8) + + for byte_i in range_constexpr(16): + logical_k = col + fx.Int32(byte_i) + physical_k, physical_n = swizzle_128(logical_k, row) + lds_byte = physical_k * fx.Int32(128) + physical_n + self._store_u8(lds_dst, lds_byte, packed_u8x16[byte_i]) + + +def pack_i32x4_i32x8(lo, hi): + # Pack two i32x4 as one i32x8 + return lo.shuffle(hi, list(range(8))) + + +class S2RLoader: + def __init__(self, wave_idx, n_tiles): + self.lane_id = fx.thread_idx.x % 64 + self.wave_idx = wave_idx + self.n_tiles = n_tiles + + def _vec_load_16xf8(self, lds_src, offset): + off_tup = fx.make_int_tuple(offset) + ptr_off = fx.add_offset(lds_src.ptr, off_tup) + i8_iter = fx.recast_iter(fx.Uint8, ptr_off) + view = fx.make_view(i8_iter, fx.make_layout(16, 1)) + return view.load() + + def _vec_load_1xf8(self, lds_src, offset): + """Naive one-byte LDS load with direct dynamic byte addressing. + + Avoid ``make_int_tuple`` entirely because this FlyDSL build cannot + reliably infer tuple types from dynamic Index expressions. + """ + base_i32 = fx.Int32(fx.ptrtoint(lds_src.ptr)) + addr_i32 = base_i32 + fx.Int32(offset) + i8_lds_ptr_t = fx.PointerType.get( + elem_ty=ir.IntegerType.get_signless(8), + address_space=2, + alignment=1, + ) + i8_ptr = fx.inttoptr(i8_lds_ptr_t, addr_i32) + view = fx.make_view(i8_ptr, fx.make_layout(1, 1)) + return view.load() + + def load(self, lds_src, preshuffled=False): + frag = [] + for i in range_constexpr(self.n_tiles): + halves = [] + row = self.wave_idx * (self.n_tiles * 16) + i * 16 + self.lane_id % 16 + for step in range_constexpr(2): + col = (self.lane_id // 16) * 16 + step * 64 + if const_expr(preshuffled): + offset = (row // 8) * 1024 + (row % 8) * 16 + (col // 16) * 128 + else: + row_swz, col_swz = swizzle_128(row, col) + offset = row_swz * 128 + col_swz + v = self._vec_load_16xf8(lds_src, offset) + halves.append(v.bitcast(fx.Int32)) + frag.append(pack_i32x4_i32x8(halves[0], halves[1])) + return frag + + def load_one(self, lds_src, lds_offset): + v = self._vec_load_16xf8(lds_src, lds_offset) + return v.bitcast(fx.Int32) + + def _ds_read_b64_tr_b8(self, lds_src, byte_offset, immediate_offset=0): + """Issue one gfx950 ``ds_read_b64_tr_b8`` and return i32x2. + + ``immediate_offset`` is encoded in the DS instruction itself. The NN + K128 path uses 0 and 0x2000, where 0x2000 advances the logical K row by + 64 in a 128-byte-wide physical LDS image. + """ + if immediate_offset == 0: + asm = "ds_read_b64_tr_b8 $0, $1 offset:0\n" + elif immediate_offset == 0x2000: + asm = "ds_read_b64_tr_b8 $0, $1 offset:8192\n" + else: + raise ValueError( + "ds_read_b64_tr_b8 supports immediate offsets 0 and 0x2000, " + f"got {immediate_offset:#x}" + ) + + base_i32 = fx.Int32(fx.ptrtoint(lds_src.ptr)) + addr_i32 = base_i32 + fx.Int32(byte_offset) + raw_type = ir.VectorType.get([2], ir.IntegerType.get_signless(32)) + raw = _llvm.inline_asm( + raw_type, + [as_mlir_value(addr_i32)], + asm, + "=v,v,~{memory}", + has_side_effects=True, + ) + return Vec(vector.BitCastOp(raw_type, raw).result, (2,), fx.Int32) + + def load_one_transpose( + self, + lds_src, + first_byte_offset, + second_byte_offset, + immediate_offset=0, + ): + """Load one 16-byte portion of a K128 FP8 MFMA operand. + + Two transpose reads return four packed i32 values. Calling this once + with immediate 0 and once with immediate 0x2000 yields the two i32x4 + portions that concatenate into the production i32x8 MFMA fragment. + """ + lo = self._ds_read_b64_tr_b8( + lds_src, + first_byte_offset, + immediate_offset, + ) + hi = self._ds_read_b64_tr_b8( + lds_src, + second_byte_offset, + immediate_offset, + ) + return lo.shuffle(hi, [0, 1, 2, 3]) + + +class StoreC: + def __init__(self, A_scale, B_scale, C, c_rows, c_cols, c_idx_fn, n_tiles_a, n_tiles_b): + self.c_rows = c_rows + self.c_cols = c_cols + self.lane_id = fx.thread_idx.x % 64 + self.c_idx_fn = c_idx_fn + self.n_tiles_a = n_tiles_a + self.n_tiles_b = n_tiles_b + # Exact byte counts from compile-time shape (BF16 C output, FP32 scales). + # ``num_records_bytes`` is required when ``max_size=False`` -- see + # ``make_buffer_tensor`` docstring for the silent-OOB rationale. + c_nbytes = c_rows * c_cols * 2 # BFloat16 = 2 bytes + sa_nbytes = c_rows * 4 # Float32 row-wise scale + sb_nbytes = c_cols * 4 # Float32 col-wise scale + gC = fx.rocdl.make_buffer_tensor(C, max_size=False, num_records_bytes=c_nbytes) + gSA = fx.rocdl.make_buffer_tensor(A_scale, max_size=False, num_records_bytes=sa_nbytes) + gSB = fx.rocdl.make_buffer_tensor(B_scale, max_size=False, num_records_bytes=sb_nbytes) + self.c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) + self.sa_div = fx.logical_divide(gSA, fx.make_layout(1, 1)) + self.sb_div = fx.logical_divide(gSB, fx.make_layout(1, 1)) + + self.scale_atom_4 = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), fx.Float32) + self.scale_atom_1 = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) + self.out_atom_1 = fx.make_copy_atom(fx.rocdl.BufferCopy16b(), fx.BFloat16) + self.reg_f32_4 = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Float32) + self.reg_f32_1 = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) + self.reg_bf16_1 = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.BFloat16) + + def _load_scale_vec4(self, row): + fx.copy(self.scale_atom_4, fx.slice(self.sa_div, (None, fx.Int32(row))), self.reg_f32_4) + return Vec(fx.memref_load_vec(self.reg_f32_4)) + + def _load_scale_scalar(self, col): + fx.copy(self.scale_atom_1, fx.slice(self.sb_div, (None, fx.Int32(col))), self.reg_f32_1) + return Vec(fx.memref_load_vec(self.reg_f32_1))[0] + + def _store_bf16(self, value_bf16, c_index): + fx.memref_store_vec(Vec.filled(1, value_bf16, fx.BFloat16), self.reg_bf16_1) + fx.copy(self.out_atom_1, self.reg_bf16_1, fx.slice(self.c_div, (None, fx.Int32(c_index)))) + + def store(self, c_frag, base_row, base_col): + a_scales = [ + self._load_scale_vec4(base_row + i * 16 + (self.lane_id // 16) * 4) for i in range_constexpr(self.n_tiles_a) + ] + b_scales = [ + self._load_scale_scalar(base_col + i * 16 + self.lane_id % 16) for i in range_constexpr(self.n_tiles_b) + ] + for ti in range_constexpr(self.n_tiles_a): + row = base_row + ti * 16 + (self.lane_id // 16) * 4 + for tj in range_constexpr(self.n_tiles_b): + col = base_col + tj * 16 + self.lane_id % 16 + col_valid = col < self.c_cols + oob = fx.Int32(self.c_rows * self.c_cols) + vec_f32 = Vec(c_frag[self.c_idx_fn(ti, tj)]) + for i in range_constexpr(4): + scaled = (vec_f32[i] * (a_scales[ti][i] * b_scales[tj])).to(fx.BFloat16) + c_index = (row + i) * self.c_cols + col + self._store_bf16(scaled, arith.select(col_valid, c_index, oob)) + + +def wait_barrier(count): + _llvm.inline_asm( + res=None, + operands_=[], + asm_string=f"s_waitcnt vmcnt({count})\ns_barrier", + constraints="", + has_side_effects=True, + ) + + +class Mfma16x16x128: + def __init__(self, n_tiles_a, n_tiles_b): + self.atom = fx.make_mma_atom(fx.rocdl.cdna4.MFMA_Scale(16, 16, 128, fx.Float8E4M3FN)) + self.zero_value = Vec.filled(4, 0.0, fx.Float32) + self.n_tiles_a = n_tiles_a + self.n_tiles_b = n_tiles_b + + def idx(self, i, j): + return i * self.n_tiles_b + j + + def _make_operand_frag(self, value): + frag = fx.make_rmem_tensor(8, fx.Int32) + frag.store(Vec(value)) + return frag + + def _make_accum_frag(self, value): + frag = fx.make_rmem_tensor(4, fx.Float32) + frag.store(Vec(value)) + return frag + + def _do_mma(self, a, b, c): + a_frag = self._make_operand_frag(a) + b_frag = self._make_operand_frag(b) + c_frag = self._make_accum_frag(c) + fx.gemm(self.atom, c_frag, a_frag, b_frag, c_frag) + return c_frag.load().ir_value() + + def call(self, a, b, c, *, set_prio=True): + assert len(a) == self.n_tiles_a + assert len(b) == self.n_tiles_b + assert len(c) == self.n_tiles_a * self.n_tiles_b + + a_frags = [self._make_operand_frag(a[idx]) for idx in range_constexpr(self.n_tiles_a)] + b_frags = [self._make_operand_frag(b[idx]) for idx in range_constexpr(self.n_tiles_b)] + c_frags = [self._make_accum_frag(c[idx]) for idx in range_constexpr(self.n_tiles_a * self.n_tiles_b)] + if const_expr(set_prio): + rocdl.s_setprio(1) + for i in range_constexpr(self.n_tiles_a): + for j in range_constexpr(self.n_tiles_b): + cf = c_frags[self.idx(i, j)] + fx.gemm(self.atom, cf, a_frags[i], b_frags[j], cf) + if const_expr(set_prio): + rocdl.s_setprio(0) + rocdl.s_barrier() + return [c_frags[idx].load().ir_value() for idx in range_constexpr(self.n_tiles_a * self.n_tiles_b)] + + def call_one(self, a, b, c, i, j): + assert i < self.n_tiles_a and j < self.n_tiles_b + + return self._do_mma(a[i], b[j], c[self.idx(i, j)]) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py new file mode 100644 index 000000000..405ba1587 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -0,0 +1,1448 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""TE entry points for the FlyDSL GEMM backend.""" + +import os + +import torch +import transformer_engine_torch as tex + +from transformer_engine.pytorch.utils import get_device_compute_capability + +from .exceptions import FlyDSLUnsupportedError + +from .bf16_gemm import bf16_matmul +from .fp16_gemm import fp16_matmul +from .fp32_gemm import fp32_matmul +from .fp8_gemm import fp8_matmul +from .mxfp8_gemm import mxfp8_matmul + +# TODO: Some backend-independent GEMM wrapper utilities overlap with the +# Triton GEMM backend (PR#667), including operand classification, logical +# output-shape derivation, and quantized-storage inspection. Once both +# integrations stabilize, factor the genuinely common pieces into a shared +# GEMM wrapper utility module while preserving backend-specific layout and +# storage canonicalization. + +def _product(shape): + """Return the product of dimensions in ``shape``.""" + result = 1 + for dim in shape: + result *= dim + return result + + +def _get_gemm_output_shape(A, transa, B, transb) -> torch.Size: + """Compute TE's logical GEMM output shape. + + This matches TE's generic GEMM output-shape convention: the + physical GEMM is flattened to ``[M, N]``, while the returned tensor keeps + B's leading dimensions when ``transb`` is false. + """ + A_shape = A if isinstance(A, torch.Size) else A.shape + B_shape = B if isinstance(B, torch.Size) else B.shape + + if len(A_shape) < 2 or len(B_shape) < 2: + raise ValueError( + "FlyDSL GEMM expects both logical operands to have rank >= 2, " + f"got A={tuple(A_shape)} and B={tuple(B_shape)}" + ) + + A0 = _product(A_shape[:-1]) + A1 = A_shape[-1] + B1 = B_shape[-1] + + output_shape = [B1] if transb else list(B_shape[:-1]) + output_shape.append(A0 if transa else A1) + return torch.Size(output_shape) + + +def reinterpret_as_fp8_tensor( + a: torch.Tensor, + dtype: tex.DType, +) -> torch.Tensor: + """View TE's uint8 payload as the native torch FP8 dtype for this GPU.""" + capability = get_device_compute_capability() + + # gfx950 uses OCP FP8. gfx942 and earlier ROCm architectures use FNUZ. + use_ocp_fp8 = capability == (9, 5) + + if dtype == tex.DType.kFloat8E4M3: + torch_dtype = ( + torch.float8_e4m3fn + if use_ocp_fp8 + else torch.float8_e4m3fnuz + ) + elif dtype == tex.DType.kFloat8E5M2: + torch_dtype = ( + torch.float8_e5m2 + if use_ocp_fp8 + else torch.float8_e5m2fnuz + ) + else: + raise TypeError(f"Unsupported TE FP8 dtype: {dtype}") + + return a.view(torch_dtype) + +def _validate_common_epilogue( + *, + quantizer, + bias, + gelu, + grad, + accumulate, + alpha, + beta, +): + """Validate features not yet implemented by the FlyDSL GEMM backend.""" + if quantizer is not None: + raise NotImplementedError( + "FlyDSL GEMM output quantization is not implemented" + ) + + if float(alpha) != 1.0 or float(beta) != 0.0: + raise NotImplementedError( + "FlyDSL GEMM currently supports only alpha=1 and beta=0" + ) + + # TODO: Add accumulate option + if accumulate: + raise NotImplementedError( + "FlyDSL GEMM accumulation is not implemented" + ) + + # TODO: Add fused bias and BGRADB epilogues + if bias is not None and bias.numel() != 0: + raise NotImplementedError( + "FlyDSL GEMM bias is not implemented" + ) + + +def _classify_input(t): + """Classify a GEMM operand for the FlyDSL backend.""" + try: + from transformer_engine.pytorch.float8_tensor import Float8Tensor + from transformer_engine.pytorch.tensor.storage.float8_tensor_storage import ( + Float8TensorStorage, + ) + if isinstance(t, (Float8Tensor, Float8TensorStorage)): + return "fp8", t + except ImportError: + pass + + try: + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Tensor + from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import ( + MXFP8TensorStorage, + ) + if isinstance(t, (MXFP8Tensor, MXFP8TensorStorage)): + return "mxfp8", t + except ImportError: + pass + + try: + from transformer_engine.pytorch.quantized_tensor import ( + QuantizedTensorStorage, + ) + if isinstance(t, QuantizedTensorStorage): + raise ValueError( + f"The FlyDSL GEMM backend does not support " + f"{type(t).__name__}. Only Float8Tensor / " + f"Float8TensorStorage and MXFP8Tensor / " + f"MXFP8TensorStorage are implemented." + ) + except ImportError: + pass + + return "regular", None + + +def _reinterpret_fp8_payload(data, fp8_dtype, name): + """Reinterpret TE's uint8 payload using its ``tex.DType`` metadata.""" + if data is None: + raise FlyDSLUnsupportedError( + f"{name} does not contain the required FP8 payload" + ) + + if fp8_dtype not in ( + tex.DType.kFloat8E4M3, + tex.DType.kFloat8E5M2, + ): + raise TypeError( + f"{name} has unsupported TE FP8 dtype metadata: {fp8_dtype}" + ) + + # TE stores Float8Tensor payloads as uint8. Use TE's shared conversion + # helper so ROCm's correct native torch FP8 type is selected from tex.DType. + if data.dtype == torch.uint8: + return reinterpret_as_fp8_tensor(data, fp8_dtype) + + # A materialized payload may already have been reinterpreted. Accept it + # only when its TE metadata is one of the recognized FP8 enum values. + if data.element_size() == 1 and data.dtype.is_floating_point: + return data + + raise TypeError( + f"{name} FP8 storage must be uint8 or an already reinterpreted " + f"one-byte floating-point tensor, got {data.dtype}" + ) + + +def _valid_fp8_transpose(t): + """Return whether a TE Float8 operand has usable columnwise storage.""" + return ( + hasattr(t, "_transpose") + and t._transpose is not None + and not getattr(t, "_transpose_invalid", False) + ) + + +def _mxfp8_debug_enabled() -> bool: + value = os.getenv("DEBUG_FLYDSL_MXFP8_GEMM", "") + return value.lower() not in ("", "0", "false", "no", "off") + + +def _mxfp8_debug(message: str) -> None: + if _mxfp8_debug_enabled(): + print(f"[DEBUG_FLYDSL_MXFP8_GEMM] {message}") + + +def _fp8_debug_enabled() -> bool: + value = os.getenv("DEBUG_FLYDSL_FP8_GEMM", "") + return value.lower() not in ("", "0", "false", "no", "off") + + +def _fp8_debug(message: str) -> None: + if _fp8_debug_enabled(): + print(f"[DEBUG_FLYDSL_FP8_GEMM] {message}") + + +def _fp8_tensor_debug(name: str, tensor: torch.Tensor) -> None: + if not _fp8_debug_enabled(): + return + _fp8_debug( + f"{name}: shape={tuple(tensor.shape)}, stride={tuple(tensor.stride())}, " + f"dtype={tensor.dtype}, device={tensor.device}, " + f"contiguous={tensor.is_contiguous()}, data_ptr=0x{tensor.data_ptr():x}" + ) + + +def _fp8_scale_debug(name: str, scale: torch.Tensor) -> None: + if not _fp8_debug_enabled(): + return + value = scale.detach().float().reshape(-1).cpu().tolist() + _fp8_debug( + f"{name}: shape={tuple(scale.shape)}, dtype={scale.dtype}, " + f"device={scale.device}, data_ptr=0x{scale.data_ptr():x}, value={value}" + ) + + +def _canonicalize_blas_pair( + A_data: torch.Tensor, + transa: bool, + B_data: torch.Tensor, + transb: bool, +): + """Swap TE BLAS operand ownership without changing either tensor layout.""" + del transa, transb + return B_data, A_data + + +def _flatten_rowwise(t: torch.Tensor, name: str) -> torch.Tensor: + """Flatten all leading dimensions while preserving the final dimension.""" + if t.ndim < 2: + raise ValueError( + f"FlyDSL GEMM expects {name} to have rank >= 2, got {tuple(t.shape)}" + ) + return t.reshape(-1, t.shape[-1]) + + +def _flatten_columnwise(t: torch.Tensor, name: str) -> torch.Tensor: + """Flatten TE columnwise storage while preserving its leading dimension.""" + if t.ndim < 2: + raise ValueError( + f"FlyDSL GEMM expects {name} to have rank >= 2, got {tuple(t.shape)}" + ) + return t.reshape(t.shape[0], -1) + + +def _canonicalize_blas_operands( + A_data: torch.Tensor, + transa: bool, + B_data: torch.Tensor, + transb: bool, +): + """Convert TE's BLAS-shaped operands to FlyDSL row-major operands. + + TE's generic GEMM interface follows BLAS column-major interpretation. + FlyDSL kernels consume ordinary row-major matrices: + + a_flydsl: [M, K] + b_flydsl: [K, N] + + Operand ownership is swapped without creating tensor transpose views: + + a_flydsl = B + b_flydsl = A + """ + if transa and transb: + raise NotImplementedError( + "FlyDSL GEMM does not support transa=True, transb=True (TT)" + ) + + A_flat = _flatten_rowwise(A_data, "A") + B_flat = _flatten_rowwise(B_data, "B") + + a_flydsl, b_flydsl = _canonicalize_blas_pair( + A_flat, + transa, + B_flat, + transb, + ) + + m, k = a_flydsl.shape + kb, n = b_flydsl.shape + if kb != k: + layout = f"{'T' if transa else 'N'}{'T' if transb else 'N'}" + raise ValueError( + f"FlyDSL {layout} canonicalization produced incompatible operands: " + f"{tuple(a_flydsl.shape)} @ {tuple(b_flydsl.shape)}" + ) + + return a_flydsl, b_flydsl, m, n, k + + +def _validate_or_allocate_output( + D, + *, + shape, + dtype, + device, + backend_name, +): + if D is None: + return torch.empty(shape, dtype=dtype, device=device) + + if tuple(D.shape) != tuple(shape): + raise ValueError( + f"D shape {tuple(D.shape)} does not match expected {tuple(shape)}" + ) + if D.dtype != dtype: + raise TypeError( + f"FlyDSL {backend_name} requires {dtype} output, got {D.dtype}" + ) + if D.device != device: + raise ValueError( + f"D must be on {device}, got {D.device}" + ) + if not D.is_contiguous(): + raise ValueError( + f"FlyDSL {backend_name} requires contiguous output storage" + ) + return D + + +def _run_bf16_gemm( + A, + transa, + B, + transb, + D, + *, + output_dtype: torch.dtype, +): + """Dispatch BF16 using the original row-major operand allocations. + + No operand transpose view is created: + + TN: kernel A = TE B [M,K], kernel B = TE A [N,K] + NN: kernel A = TE B [M,K], kernel B = TE A [K,N] + NT: kernel A = TE B [K,M], kernel B = TE A [K,N] + """ + if not isinstance(A, torch.Tensor) or not isinstance(B, torch.Tensor): + raise TypeError("FlyDSL BF16 GEMM expects plain torch.Tensor operands") + if A.dtype != torch.bfloat16 or B.dtype != torch.bfloat16: + raise TypeError( + "FlyDSL BF16 GEMM requires torch.bfloat16 inputs, " + f"got A={A.dtype} and B={B.dtype}" + ) + if A.device != B.device: + raise ValueError( + f"A and B must be on the same device, got {A.device} and {B.device}" + ) + + dispatch = { + (True, False): "TN", + (False, False): "NN", + (False, True): "NT", + } + try: + layout = dispatch[(bool(transa), bool(transb))] + except KeyError as exc: + raise FlyDSLUnsupportedError( + "FlyDSL GEMM does not support transa=True, transb=True (TT)" + ) from exc + + output_shape = _get_gemm_output_shape(A, transa, B, transb) + + # Preserve the original row-major storage. This only collapses leading + # batch dimensions, matching the wrapper's existing regular-GEMM contract. + A_data = _flatten_rowwise(A, "A") + B_data = _flatten_rowwise(B, "B") + + # Kernel ownership is always swapped relative to TE's BLAS arguments. + a_flydsl = B_data + b_flydsl = A_data + + if layout == "TN": + m, k = a_flydsl.shape + n, kb = b_flydsl.shape + expected_a = (m, k) + expected_b = (n, k) + elif layout == "NN": + m, k = a_flydsl.shape + kb, n = b_flydsl.shape + expected_a = (m, k) + expected_b = (k, n) + else: + k, m = a_flydsl.shape + kb, n = b_flydsl.shape + expected_a = (k, m) + expected_b = (k, n) + + if kb != k: + raise FlyDSLUnsupportedError( + f"FlyDSL BF16 {layout} received incompatible row-major operands: " + f"a={tuple(a_flydsl.shape)}, b={tuple(b_flydsl.shape)}" + ) + if tuple(a_flydsl.shape) != expected_a or tuple(b_flydsl.shape) != expected_b: + raise FlyDSLUnsupportedError( + f"FlyDSL BF16 {layout} physical contract mismatch: " + f"a={tuple(a_flydsl.shape)} expected={expected_a}; " + f"b={tuple(b_flydsl.shape)} expected={expected_b}" + ) + + if _product(output_shape) != m * n: + raise RuntimeError( + f"FlyDSL BF16 logical output shape {tuple(output_shape)} " + f"does not match kernel output {(m, n)}" + ) + + D = _validate_or_allocate_output( + D, + shape=output_shape, + dtype=output_dtype, + device=A.device, + backend_name=f"BF16 {layout}", + ) + + bf16_matmul( + a_flydsl, + b_flydsl, + D.view(m, n), + layout=layout, + m=m, + n=n, + k=k, + ) + return D + + + +def _run_fp16_gemm( + A, + transa, + B, + transb, + D, + *, + output_dtype: torch.dtype, +): + """Dispatch FP16 using the original row-major operand allocations. + + No operand transpose view is created: + + TN: kernel A = TE B [M,K], kernel B = TE A [N,K] + NN: kernel A = TE B [M,K], kernel B = TE A [K,N] + NT: kernel A = TE B [K,M], kernel B = TE A [K,N] + """ + if not isinstance(A, torch.Tensor) or not isinstance(B, torch.Tensor): + raise TypeError("FlyDSL FP16 GEMM expects plain torch.Tensor operands") + if A.dtype != torch.float16 or B.dtype != torch.float16: + raise TypeError( + "FlyDSL FP16 GEMM requires torch.float16 inputs, " + f"got A={A.dtype} and B={B.dtype}" + ) + if A.device != B.device: + raise ValueError( + f"A and B must be on the same device, got {A.device} and {B.device}" + ) + + dispatch = { + (True, False): "TN", + (False, False): "NN", + (False, True): "NT", + } + try: + layout = dispatch[(bool(transa), bool(transb))] + except KeyError as exc: + raise FlyDSLUnsupportedError( + "FlyDSL GEMM does not support transa=True, transb=True (TT)" + ) from exc + + output_shape = _get_gemm_output_shape(A, transa, B, transb) + + # Preserve the original row-major storage. This only collapses leading + # batch dimensions, matching the wrapper's existing regular-GEMM contract. + A_data = _flatten_rowwise(A, "A") + B_data = _flatten_rowwise(B, "B") + + # Kernel ownership is always swapped relative to TE's BLAS arguments. + a_flydsl = B_data + b_flydsl = A_data + + if layout == "TN": + m, k = a_flydsl.shape + n, kb = b_flydsl.shape + expected_a = (m, k) + expected_b = (n, k) + elif layout == "NN": + m, k = a_flydsl.shape + kb, n = b_flydsl.shape + expected_a = (m, k) + expected_b = (k, n) + else: + k, m = a_flydsl.shape + kb, n = b_flydsl.shape + expected_a = (k, m) + expected_b = (k, n) + + if kb != k: + raise FlyDSLUnsupportedError( + f"FlyDSL FP16 {layout} received incompatible row-major operands: " + f"a={tuple(a_flydsl.shape)}, b={tuple(b_flydsl.shape)}" + ) + if tuple(a_flydsl.shape) != expected_a or tuple(b_flydsl.shape) != expected_b: + raise FlyDSLUnsupportedError( + f"FlyDSL FP16 {layout} physical contract mismatch: " + f"a={tuple(a_flydsl.shape)} expected={expected_a}; " + f"b={tuple(b_flydsl.shape)} expected={expected_b}" + ) + + if _product(output_shape) != m * n: + raise RuntimeError( + f"FlyDSL FP16 logical output shape {tuple(output_shape)} " + f"does not match kernel output {(m, n)}" + ) + + D = _validate_or_allocate_output( + D, + shape=output_shape, + dtype=output_dtype, + device=A.device, + backend_name=f"FP16 {layout}", + ) + + fp16_matmul( + a_flydsl, + b_flydsl, + D.view(m, n), + layout=layout, + m=m, + n=n, + k=k, + ) + return D + + +def _run_fp32_gemm( + A, + transa, + B, + transb, + D, +): + """Normalize FP32 TN/NN/NT inputs to the current kernel's TN interface. + + The existing FP32 entry point expects ordinary row-major GEMM operands: + + a_tn: [M, K] + b_tn: [K, N] + + TE provides BLAS-shaped operands, so ownership is swapped and only the + operands whose BLAS transpose flags require it are materialized: + + TN: a_tn = B + b_tn = A.T + + NN: a_tn = B + b_tn = A + + NT: a_tn = B.T + b_tn = A + + ``transpose(...).contiguous()`` is therefore used only for the FP32 + operands that are not already in the current TN kernel orientation. + BF16/FP16/FP8/MXFP8 dispatch is unchanged. + """ + if not isinstance(A, torch.Tensor) or not isinstance(B, torch.Tensor): + raise TypeError("FlyDSL FP32 GEMM expects plain torch.Tensor operands") + if A.dtype != torch.float32 or B.dtype != torch.float32: + raise TypeError( + "FlyDSL FP32 GEMM requires torch.float32 inputs, " + f"got A={A.dtype} and B={B.dtype}" + ) + if A.device != B.device: + raise ValueError( + f"A and B must be on the same device, got {A.device} and {B.device}" + ) + if bool(transa) and bool(transb): + raise FlyDSLUnsupportedError( + "FlyDSL GEMM does not support transa=True, transb=True (TT)" + ) + + output_shape = _get_gemm_output_shape(A, transa, B, transb) + + A_flat = _flatten_rowwise(A, "A") + B_flat = _flatten_rowwise(B, "B") + + # Standard BLAS-column-major -> row-major conversion: + # swap operands, then apply the original operand transpose flags. + # TODO: Optimize FP32 NN/NT execution. These layouts are currently + # materialized into the TN kernel contract with explicit transpose copies. + if bool(transb): + a_tn = B_flat.transpose(0, 1).contiguous() + else: + a_tn = B_flat + + if bool(transa): + b_tn = A_flat.transpose(0, 1).contiguous() + else: + b_tn = A_flat + + if not a_tn.is_contiguous(): + a_tn = a_tn.contiguous() + if not b_tn.is_contiguous(): + b_tn = b_tn.contiguous() + + if a_tn.ndim != 2 or b_tn.ndim != 2: + raise RuntimeError( + f"FlyDSL FP32 TN normalization produced rank mismatch: " + f"a={tuple(a_tn.shape)}, b={tuple(b_tn.shape)}" + ) + + m, k = a_tn.shape + kb, n = b_tn.shape + if kb != k: + layout = f"{'T' if transa else 'N'}{'T' if transb else 'N'}" + raise FlyDSLUnsupportedError( + f"FlyDSL FP32 {layout} could not normalize to TN: " + f"a_tn={tuple(a_tn.shape)}, b_tn={tuple(b_tn.shape)}" + ) + + if _product(output_shape) != m * n: + raise RuntimeError( + f"FlyDSL FP32 logical output shape {tuple(output_shape)} " + f"does not match normalized TN output {(m, n)}" + ) + + D = _validate_or_allocate_output( + D, + shape=output_shape, + dtype=torch.float32, + device=A.device, + backend_name="FP32 via TN core", + ) + + fp32_matmul( + a_tn, + b_tn, + D.view(m, n), + ) + return D + + + +def _get_fp8_rowwise_payload(t, name): + """Return TE's existing rowwise ``_data`` payload without copying.""" + data = getattr(t, "_data", None) + if data is None: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 requires existing {name} rowwise (_data) storage" + ) + return _reinterpret_fp8_payload( + data, + getattr(t, "_fp8_dtype", None), + f"{name}._data", + ) + + +def _get_fp8_columnwise_payload(t, name): + """Return TE's existing columnwise ``_transpose`` payload without copying.""" + if not _valid_fp8_transpose(t): + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 requires valid {name} columnwise (_transpose) storage" + ) + return _reinterpret_fp8_payload( + t._transpose, + getattr(t, "_fp8_dtype", None), + f"{name}._transpose", + ) + + +def _validate_fp8_kernel_operands( + kernel_a, + kernel_b, + *, + layout, + a_storage, + b_storage, +): + """Validate zero-copy physical operands before launching an FP8 kernel.""" + if kernel_a.ndim != 2 or kernel_b.ndim != 2: + raise ValueError( + f"FlyDSL FP8 {layout} expects rank-2 kernel operands, got " + f"{a_storage}={tuple(kernel_a.shape)} and " + f"{b_storage}={tuple(kernel_b.shape)}" + ) + if not kernel_a.is_contiguous() or not kernel_b.is_contiguous(): + raise ValueError( + f"FlyDSL FP8 {layout} requires contiguous {a_storage} and " + f"{b_storage}; refusing to materialize replacement operands" + ) + if kernel_a.device != kernel_b.device: + raise ValueError( + f"FlyDSL FP8 {layout} operands must be on the same device, got " + f"{kernel_a.device} and {kernel_b.device}" + ) + + +def _fp8_output_shape(D, m, n): + """Preserve TE's logical output shape when D is preallocated.""" + output_shape = D.shape if D is not None else torch.Size((m, n)) + if _product(output_shape) != m * n: + raise RuntimeError( + f"FlyDSL FP8 logical output shape {tuple(output_shape)} does not " + f"match flattened kernel shape {(m, n)}" + ) + return output_shape + + +def _select_mxfp8_data_and_scale( + t, + *, + will_transpose: bool, + name: str, +): + """Select the TE MXFP8 representation required by BLAS semantics.""" + if will_transpose: + data = getattr(t, "_columnwise_data", None) + scale = getattr(t, "_columnwise_scale_inv", None) + orientation = "columnwise" + else: + data = getattr(t, "_rowwise_data", None) + scale = getattr(t, "_rowwise_scale_inv", None) + orientation = "rowwise" + + _mxfp8_debug( + f"{name}: will_transpose={will_transpose}, " + f"selected={orientation}, data_present={data is not None}, " + f"scale_present={scale is not None}" + ) + + if data is None or scale is None: + raise RuntimeError( + f"{name} does not contain required {orientation} MXFP8 data and scales" + ) + + _mxfp8_debug( + f"{name} selected data shape={tuple(data.shape)}, " + f"dtype={data.dtype}, stride={tuple(data.stride())}; " + f"scale shape={tuple(scale.shape)}, dtype={scale.dtype}, " + f"stride={tuple(scale.stride())}" + ) + return data, scale + + +def _mxfp8_logical_shape(t, name: str) -> torch.Size: + """Return MXFP8 logical shape from a populated backing tensor. + + MXFP8TensorStorage is not a torch.Tensor and does not expose ``.shape``. + Rowwise and columnwise MXFP8 payloads retain the same logical row-major + shape, so either populated backing is sufficient for shape derivation. + """ + data = getattr(t, "_rowwise_data", None) + if data is None: + data = getattr(t, "_columnwise_data", None) + if data is None: + raise FlyDSLUnsupportedError( + f"{name} has neither rowwise nor columnwise MXFP8 data" + ) + return torch.Size(data.shape) + +def _flatten_mxfp8_scale( + t: torch.Tensor, + name: str, + *, + source_colwise: bool, +) -> torch.Tensor: + """Flatten a raw TE MXFP8 scale tensor without changing orientation. + + Rowwise source: + [..., K/32] -> [outer, K/32] + + Columnwise source: + [K/32, ...] -> [K/32, outer] + """ + if t.ndim < 2: + raise ValueError( + f"FlyDSL MXFP8 expects {name} scale rank >= 2, " + f"got {tuple(t.shape)}" + ) + + original_shape = tuple(t.shape) + if source_colwise: + t = t.reshape(t.shape[0], -1) + orientation = "columnwise" + else: + t = t.reshape(-1, t.shape[-1]) + orientation = "rowwise" + + _mxfp8_debug( + f"{name} {orientation} scale flatten: " + f"{original_shape} -> {tuple(t.shape)}, " + f"contiguous={t.is_contiguous()}" + ) + return t + + +def _run_mxfp8( + A, + transa, + B, + transb, + D, + *, + output_dtype: torch.dtype, +): + """Dispatch MXFP8 through exact TN/NN/NT physical contracts. + + TE owns BLAS-shaped operands. After the usual ownership swap, FlyDSL + kernels consume the selected backing directly: + + TN: a = B.rowwise [M, K] + b = A.rowwise [N, K] + + NN: a = B.rowwise [M, K] + b = A.columnwise [K, N] + + NT: a = B.columnwise [K, M] + b = A.columnwise [K, N] + + MXFP8 rowwise and columnwise payloads retain the same logical row-major + shape. Columnwise selection changes the quantization axis; the specialized + NN/NT kernels provide the required transpose-read semantics. + """ + a_fp8_dtype = getattr(A, "_fp8_dtype", None) + b_fp8_dtype = getattr(B, "_fp8_dtype", None) + supported_fp8_dtypes = ( + tex.DType.kFloat8E4M3, + tex.DType.kFloat8E5M2, + ) + if ( + a_fp8_dtype not in supported_fp8_dtypes + or b_fp8_dtype not in supported_fp8_dtypes + ): + raise FlyDSLUnsupportedError( + "FlyDSL MXFP8 supports E4M3 and E5M2 independently for A/B; " + f"got A={a_fp8_dtype} and B={b_fp8_dtype}" + ) + + layout = f"{'T' if transa else 'N'}{'T' if transb else 'N'}" + dispatch = { + (True, False): "TN", + (False, False): "NN", + (False, True): "NT", + } + try: + kernel_layout = dispatch[(bool(transa), bool(transb))] + except KeyError as exc: + raise FlyDSLUnsupportedError( + "FlyDSL GEMM does not support transa=True, transb=True (TT)" + ) from exc + + _mxfp8_debug( + f"entry: layout={layout}, common_kernel=" + f"{mxfp8_matmul.__module__}.{mxfp8_matmul.__name__}, " + f"A_type={type(A).__name__}, B_type={type(B).__name__}, " + f"D_provided={D is not None}" + ) + + # Resolve public shapes from actual payload tensors. Never access + # MXFP8TensorStorage.shape: the storage wrapper has no such attribute. + A_logical_shape = _mxfp8_logical_shape(A, "A") + B_logical_shape = _mxfp8_logical_shape(B, "B") + + # Match TE/C++ MXFP8 representation selection exactly: + # A: transa=True -> rowwise; transa=False -> columnwise + # B: transb=False -> rowwise; transb=True -> columnwise + A_source_colwise = not bool(transa) + B_source_colwise = bool(transb) + + A_data, A_scale = _select_mxfp8_data_and_scale( + A, + will_transpose=A_source_colwise, + name="A", + ) + B_data, B_scale = _select_mxfp8_data_and_scale( + B, + will_transpose=B_source_colwise, + name="B", + ) + + # Both MXFP8 payload orientations are stored row-major with the original + # logical shape. Flatten leading dimensions only; do not transpose or + # materialize selected columnwise payloads. + A_data = _flatten_rowwise(A_data, "A MXFP8 payload") + B_data = _flatten_rowwise(B_data, "B MXFP8 payload") + + 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) + + A_scale = _flatten_mxfp8_scale( + A_scale, + "A", + source_colwise=A_source_colwise, + ) + B_scale = _flatten_mxfp8_scale( + B_scale, + "B", + source_colwise=B_source_colwise, + ) + + # Kernel operand ownership is always swapped relative to TE: + # kernel a <- TE B + # kernel b <- TE A + if kernel_layout == "TN": + # Selected rowwise backings already match the TN normal-read contract: + # a [M,K], b [N,K] + a_flydsl = B_data + b_flydsl = A_data + a_scale = B_scale + b_scale = A_scale + + m, k = a_flydsl.shape + n, kb = b_flydsl.shape + expected_a_scale = (m, k // 32) + expected_b_scale = (n, k // 32) + + elif kernel_layout == "NN": + # A's columnwise MXFP8 payload is still row-major in its original + # shape, which is exactly the NN kernel's K-major [K,N] source. + a_flydsl = B_data + b_flydsl = A_data + a_scale = B_scale + b_scale = A_scale + + m, k = a_flydsl.shape + kb, n = b_flydsl.shape + expected_a_scale = (m, k // 32) + expected_b_scale = (k // 32, n) + + else: + # Both selected columnwise payloads directly satisfy the NT kernel's + # K-major contracts without tensor transposes or copies. + a_flydsl = B_data + b_flydsl = A_data + a_scale = B_scale + b_scale = A_scale + + k, m = a_flydsl.shape + kb, n = b_flydsl.shape + expected_a_scale = (k // 32, m) + expected_b_scale = (k // 32, n) + + if kb != k: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP8 {layout} selected incompatible payloads: " + f"a={tuple(a_flydsl.shape)} and b={tuple(b_flydsl.shape)}" + ) + + if a_flydsl.device != b_flydsl.device: + raise ValueError( + f"FlyDSL MXFP8 {layout} operands must be on the same device, got " + f"{a_flydsl.device} and {b_flydsl.device}" + ) + + if k % 32 != 0: + raise ValueError( + f"K={k} must be divisible by MXFP8 scale group size 32" + ) + + if tuple(a_scale.shape) != expected_a_scale: + raise ValueError( + f"FlyDSL MXFP8 {layout} a_scale shape " + f"{tuple(a_scale.shape)} != expected {expected_a_scale}" + ) + if tuple(b_scale.shape) != expected_b_scale: + raise ValueError( + f"FlyDSL MXFP8 {layout} b_scale shape " + f"{tuple(b_scale.shape)} != expected {expected_b_scale}" + ) + if a_scale.dtype != torch.uint8 or b_scale.dtype != torch.uint8: + raise TypeError("FlyDSL MXFP8 expects raw E8M0 scales as torch.uint8") + + # Derive the public result shape from payload shapes, not storage wrappers. + if D is not None: + output_shape = torch.Size(D.shape) + else: + output_shape = _get_gemm_output_shape( + A_logical_shape, + transa, + B_logical_shape, + transb, + ) + if _product(output_shape) != m * n: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP8 {layout} logical output shape " + f"{tuple(output_shape)} does not match kernel output {(m, n)}" + ) + + D = _validate_or_allocate_output( + D, + shape=output_shape, + dtype=output_dtype, + device=a_flydsl.device, + backend_name=f"MXFP8 {kernel_layout}", + ) + + _mxfp8_debug( + f"dispatch layout={layout}: " + f"a={tuple(a_flydsl.shape)}, stride={tuple(a_flydsl.stride())}; " + f"b={tuple(b_flydsl.shape)}, stride={tuple(b_flydsl.stride())}; " + f"a_scale={tuple(a_scale.shape)}; " + f"b_scale={tuple(b_scale.shape)}; " + f"M={m}, N={n}, K={k}" + ) + + mxfp8_matmul( + a_flydsl, + a_scale, + b_flydsl, + b_scale, + D.view(m, n), + layout=kernel_layout, + ) + return D + + +def _select_fp8_storage_for_layout(A, transa, B, transb): + """Select existing TE FP8 backings and normalize to one core contract. + + Tensor-wise FP8 columnwise storage is a materialized transpose, unlike + MXFP8 columnwise storage, which denotes a different quantization direction. + + After selecting the required TE backing and swapping BLAS operand ownership, + every supported layout produces the same kernel-visible operands: + + TN: B._data [M,K], A._data [N,K] + NN: B._data [M,K], A._transpose [N,K] + NT: B._transpose [M,K], A._transpose [N,K] + + No kernel-side transpose, transpose staging, or transpose-read is needed. + """ + layout = (bool(transa), bool(transb)) + + if layout == (True, False): # TN + A_payload = _get_fp8_rowwise_payload(A, "A") + A_storage = "A._data" + A_data = _flatten_rowwise(A_payload, A_storage) + + B_payload = _get_fp8_rowwise_payload(B, "B") + B_storage = "B._data" + B_data = _flatten_rowwise(B_payload, B_storage) + + elif layout == (False, False): # NN + A_payload = _get_fp8_columnwise_payload(A, "A") + A_storage = "A._transpose" + A_data = _flatten_columnwise(A_payload, A_storage) + + B_payload = _get_fp8_rowwise_payload(B, "B") + B_storage = "B._data" + B_data = _flatten_rowwise(B_payload, B_storage) + + elif layout == (False, True): # NT / dW + # Both selected transpose allocations are already materialized + # row-major [outer,K] backings for the normalized common core. + A_payload = _get_fp8_columnwise_payload(A, "A") + A_storage = "A._transpose" + A_data = _flatten_columnwise(A_payload, A_storage) + + B_payload = _get_fp8_columnwise_payload(B, "B") + B_storage = "B._transpose" + B_data = _flatten_columnwise(B_payload, B_storage) + + else: + raise FlyDSLUnsupportedError( + "FlyDSL GEMM does not support transa=True, transb=True (TT)" + ) + + return ( + A_data, + A_storage, + torch.Size(A_payload.shape), + B_data, + B_storage, + torch.Size(B_payload.shape), + ) + + +def _run_fp8( + A, + transa, + B, + transb, + D, + *, + output_dtype: torch.dtype, +): + """Normalize tensor-wise FP8 storage and invoke the common FP8 core.""" + supported_fp8_dtypes = ( + tex.DType.kFloat8E4M3, + tex.DType.kFloat8E5M2, + ) + a_fp8_dtype = getattr(A, "_fp8_dtype", None) + b_fp8_dtype = getattr(B, "_fp8_dtype", None) + if ( + a_fp8_dtype not in supported_fp8_dtypes + or b_fp8_dtype not in supported_fp8_dtypes + ): + raise FlyDSLUnsupportedError( + "FlyDSL FP8 supports E4M3 and E5M2 independently for A/B; " + f"got A={a_fp8_dtype} and B={b_fp8_dtype}" + ) + + A_scale_inv = getattr(A, "_scale_inv", None) + B_scale_inv = getattr(B, "_scale_inv", None) + for name, scale in ( + ("A._scale_inv", A_scale_inv), + ("B._scale_inv", B_scale_inv), + ): + if not isinstance(scale, torch.Tensor): + raise FlyDSLUnsupportedError(f"{name} is not populated") + if scale.dtype != torch.float32 or scale.numel() != 1: + raise FlyDSLUnsupportedError( + f"{name} must contain exactly one FP32 tensor-wise inverse " + f"scale, got dtype={scale.dtype}, shape={tuple(scale.shape)}" + ) + + layout = f"{'T' if transa else 'N'}{'T' if transb else 'N'}" + + ( + A_data, + A_storage, + A_payload_shape, + B_data, + B_storage, + B_payload_shape, + ) = _select_fp8_storage_for_layout( + A, + bool(transa), + B, + bool(transb), + ) + + _validate_fp8_kernel_operands( + A_data, + B_data, + layout=layout, + a_storage=A_storage, + b_storage=B_storage, + ) + + a_scale = B_scale_inv + b_scale = A_scale_inv + + # Storage selection is layout-specific; execution is not. Tensor-wise FP8 + # transpose backing is already a materialized row-major transpose, so all + # supported layouts normalize to the common [M,K] x [N,K] core contract. + matmul = fp8_matmul + kernel_layout = "common" + + a_flydsl = B_data + b_flydsl = A_data + + m, k = a_flydsl.shape + n, kb = b_flydsl.shape + + if not a_flydsl.is_contiguous() or not b_flydsl.is_contiguous(): + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 {layout} kernel contract requires contiguous final " + f"operands, got a={tuple(a_flydsl.shape)} " + f"stride={tuple(a_flydsl.stride())} and " + f"b={tuple(b_flydsl.shape)} stride={tuple(b_flydsl.stride())}" + ) + + if kb != k: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 {layout} selected incompatible physical backings: " + f"{B_storage}={tuple(B_data.shape)} and " + f"{A_storage}={tuple(A_data.shape)}; " + f"kernel operands are {tuple(a_flydsl.shape)} and " + f"{tuple(b_flydsl.shape)}" + ) + + if D is not None: + logical_output_shape = torch.Size(D.shape) + elif layout in ("TN", "NN"): + logical_output_shape = torch.Size((*B_payload_shape[:-1], n)) + else: + logical_output_shape = torch.Size((m, n)) + if _product(logical_output_shape) != m * n: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 {layout} logical output shape " + f"{tuple(logical_output_shape)} does not match kernel output " + f"shape {(m, n)}" + ) + + D = _validate_or_allocate_output( + D, + shape=logical_output_shape, + dtype=output_dtype, + device=a_flydsl.device, + backend_name=f"FP8 {layout} via {kernel_layout} core", + ) + + _fp8_debug( + f"dispatch entry: transa={bool(transa)}, transb={bool(transb)}, " + f"layout={layout}, normalized_core={matmul.__module__}.{matmul.__name__}" + ) + _fp8_debug(f"selected TE storage: A={A_storage}, B={B_storage}") + _fp8_tensor_debug(f"selected/{A_storage}", A_data) + _fp8_tensor_debug(f"selected/{B_storage}", B_data) + _fp8_tensor_debug("a_flydsl", a_flydsl) + _fp8_tensor_debug("b_flydsl", b_flydsl) + _fp8_scale_debug("a_scale", a_scale) + _fp8_scale_debug("b_scale", b_scale) + _fp8_debug(f"derived M={m}, N={n}, K={k}") + _fp8_tensor_debug("output/D", D) + + matmul( + a_flydsl, + a_scale, + b_flydsl, + b_scale, + D.view(m, n), + ) + return D + + +def te_generic_gemm_flydsl( + A, + transa, + B, + transb, + D, + quantizer, + output_dtype, + bias=None, + bias_type=None, + gelu=False, + gelu_in=None, + grad=False, + workspace=None, + workspaceSize=0, + accumulate=False, + use_split_accumulator=False, + comm_overlap=None, + comm_type=None, + extra_output=None, + bulk_overlap=False, + alpha=1.0, + beta=0.0, +): + """Run a supported FlyDSL GEMM through TE's generic GEMM interface. + + Supported layouts: + - TN: transa=True, transb=False + - NN: transa=False, transb=False + - NT: transa=False, transb=True + + TT is intentionally rejected. + + Supported dtypes: + - MXFP8 E4M3/E5M2 A/B combinations with FP16, BF16, or FP32 output + - tensor-wise E4M3/E5M2 FP8 A/B combinations with FP16, BF16, or FP32 output + - BF16 input with FP16, BF16, or FP32 output + - FP16 input with FP16, BF16, or FP32 output + - FP32 input with FP32 output + """ + del bias_type + del gelu_in + del workspace + del workspaceSize + del use_split_accumulator + del comm_overlap + del comm_type + del extra_output + del bulk_overlap + + if transa and transb: + raise NotImplementedError( + "FlyDSL GEMM does not support transa=True, transb=True (TT)" + ) + + _validate_common_epilogue( + quantizer=quantizer, + bias=bias, + gelu=gelu, + grad=grad, + accumulate=accumulate, + alpha=alpha, + beta=beta, + ) + + a_kind, _ = _classify_input(A) + b_kind, _ = _classify_input(B) + + if a_kind == "mxfp8" or b_kind == "mxfp8": + # Validate both are MXFP8 + if a_kind != b_kind: + raise ValueError( + "Mixed MXFP8 and non-MXFP8 FlyDSL GEMM inputs are not supported" + ) + + # Sanity: both operands must have at least one pre-quantized copy. + if getattr(A, '_rowwise_data', None) is None and getattr(A, '_columnwise_data', None) is None: + raise RuntimeError("MXFP8Tensor has neither rowwise nor columnwise data") + if getattr(B, '_rowwise_data', None) is None and getattr(B, '_columnwise_data', None) is None: + raise RuntimeError("MXFP8Tensor has neither rowwise nor columnwise data") + + mxfp8_output_dtypes = { + None: torch.float16, + tex.DType.kFloat16: torch.float16, + tex.DType.kBFloat16: torch.bfloat16, + tex.DType.kFloat32: torch.float32, + } + if output_dtype not in mxfp8_output_dtypes: + raise NotImplementedError( + "FlyDSL MXFP8 supports FP16, BF16, or FP32 output, " + f"got {output_dtype}" + ) + + D = _run_mxfp8( + A, + transa, + B, + transb, + D, + output_dtype=mxfp8_output_dtypes[output_dtype], + ) + return D, None, None, None + + if a_kind == "fp8" or b_kind == "fp8": + if a_kind != b_kind: + raise ValueError( + "Mixed regular FP8 and non-FP8 FlyDSL GEMM inputs are not supported" + ) + + fp8_output_dtypes = { + None: torch.float16, + tex.DType.kFloat16: torch.float16, + tex.DType.kBFloat16: torch.bfloat16, + tex.DType.kFloat32: torch.float32, + } + if output_dtype not in fp8_output_dtypes: + raise NotImplementedError( + "FlyDSL tensor-wise FP8 supports FP16, BF16, or FP32 output, " + f"got {output_dtype}" + ) + + D = _run_fp8( + A, + transa, + B, + transb, + D, + output_dtype=fp8_output_dtypes[output_dtype], + ) + return D, None, None, None + + if a_kind != "regular" or b_kind != "regular": + raise TypeError( + "Unsupported FlyDSL GEMM operand types: " + f"{type(A).__name__} and {type(B).__name__}" + ) + if not isinstance(A, torch.Tensor) or not isinstance(B, torch.Tensor): + raise TypeError( + "FlyDSL regular GEMM expects plain torch.Tensor operands" + ) + + if A.dtype == torch.bfloat16 and B.dtype == torch.bfloat16: + bf16_output_dtypes = { + None: torch.bfloat16, + tex.DType.kFloat16: torch.float16, + tex.DType.kBFloat16: torch.bfloat16, + tex.DType.kFloat32: torch.float32, + } + if output_dtype not in bf16_output_dtypes: + raise NotImplementedError( + "FlyDSL BF16 supports FP16, BF16, or FP32 output, " + f"got {output_dtype}" + ) + D = _run_bf16_gemm( + A, + transa, + B, + transb, + D, + output_dtype=bf16_output_dtypes[output_dtype], + ) + return D, None, None, None + + if A.dtype == torch.float16 and B.dtype == torch.float16: + fp16_output_dtypes = { + None: torch.float16, + tex.DType.kFloat16: torch.float16, + tex.DType.kBFloat16: torch.bfloat16, + tex.DType.kFloat32: torch.float32, + } + if output_dtype not in fp16_output_dtypes: + raise NotImplementedError( + "FlyDSL FP16 supports FP16, BF16, or FP32 output, " + f"got {output_dtype}" + ) + D = _run_fp16_gemm( + A, + transa, + B, + transb, + D, + output_dtype=fp16_output_dtypes[output_dtype], + ) + return D, None, None, None + + if A.dtype == torch.float32 and B.dtype == torch.float32: + if output_dtype not in (None, tex.DType.kFloat32): + raise NotImplementedError( + "FlyDSL FP32 currently supports only FP32 output, " + f"got {output_dtype}" + ) + D = _run_fp32_gemm( + A, + transa, + B, + transb, + D, + ) + return D, None, None, None + + raise NotImplementedError( + "FlyDSL GEMM currently supports only MXFP8, tensor-wise E4M3 FP8, " + "BF16, FP16, or FP32 inputs; " + f"got A={A.dtype} and B={B.dtype}" + ) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py new file mode 100644 index 000000000..76e82b5c5 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -0,0 +1,1830 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""FlyDSL MXFP8 TN/NN/NT 4-wave GEMM implementation. + +All supported MXFP8 layouts share one source-level kernel generator while +remaining separate compile-time specializations: + + TN: A [M,K] normal read, B [N,K] normal read + NN: A [M,K] normal read, B [K,N] transpose read + NT: A [K,M] transpose read, B [K,N] transpose read + +The layout is a Python-only cache key. It is never passed as a runtime kernel +argument. Global addressing, LDS fragment reads, scheduler directives, and +scale-source orientation are selected while building each specialized kernel, +so generated TN/NN/NT kernels contain no runtime layout branches. + +All operand payloads use direct ``BufferCopyLDS128b`` global-to-LDS staging. +Transpose variants differ only in K-major global addressing and +``ds_read_b64_tr_b8`` LDS-to-register fragment reconstruction. +""" + +import functools +import os + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir.dialects import llvm +from flydsl.expr import arith, buffer_ops, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec + +# Transformer Engine-local FlyDSL utilities. +from .exceptions import FlyDSLUnsupportedError +from .fp8_gemm_utils import ( + G2SLoader, + S2RLoader, + compute_global_swizzle, + make_fp8_buffer_tensor, + pack_i32x4_i32x8, + swizzle_128, +) + + +_BLOCK_M = 256 +_BLOCK_N = 256 +_BLOCK_K = 128 + +# Public metadata consumed by wrappers — keep. +BLOCK_M = _BLOCK_M +BLOCK_N = _BLOCK_N +BLOCK_K = _BLOCK_K +SCALE_GROUP_SIZE = 32 + + +def _debug_enabled() -> bool: + value = os.getenv("DEBUG_FLYDSL_MXFP8_GEMM", "") + return value.lower() not in ("", "0", "false", "no", "off") + + +def _debug(message: str) -> None: + if _debug_enabled(): + print(f"[DEBUG_FLYDSL_MXFP8_GEMM] {message}") + + +_SCALE_PACK_THREADS = 256 + + +def _compile_mx32_scale_pack_kernel( + dim: int, + qk: int, + source_colwise: bool, + stride0: int, + stride1: int, +): + """Build one fused raw-E8M0 -> HK-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. + """ + if dim % 64 != 0: + raise ValueError( + f"Scale outer dimension={dim} must be a multiple of 64" + ) + if qk % 4 != 0: + raise ValueError( + f"Scale K/32 dimension={qk} must be divisible by 4" + ) + + k128_tiles = qk // 4 + total_words = k128_tiles * dim + if total_words % _SCALE_PACK_THREADS != 0: + raise ValueError( + f"Packed scale words={total_words} must be divisible by " + f"{_SCALE_PACK_THREADS}" + ) + + # Select source addressing before FlyDSL captures the kernel. The emitted + # rowwise and columnwise binaries contain no runtime orientation branch. + if source_colwise: + 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) + ) + 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) + ) + + @flyc.kernel(known_block_size=[_SCALE_PACK_THREADS, 1, 1]) + def kernel_pack_mx32_scales(src: fx.Tensor, dst: fx.Tensor): + src_rsrc = buffer_ops.create_buffer_resource(src, max_size=True) + dst_rsrc = buffer_ops.create_buffer_resource(dst, max_size=True) + + linear = ( + fx.Index(fx.block_idx.x) * fx.Index(_SCALE_PACK_THREADS) + + fx.Index(gpu.thread_id("x")) + ) + k128 = linear // fx.Index(dim) + dst_row = linear % fx.Index(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 + + def load_scale_byte(group): + source_row = ( + tile * fx.Index(64) + + fx.Index(group * 16) + + row_within_16 + ) + value_i8 = buffer_ops.buffer_load( + src_rsrc, + _source_offset(source_k32, source_row), + vec_width=1, + dtype=T.i8, + ) + # Preserve the raw E8M0 byte when widening. Going through Uint8 + # avoids sign extension for scale bytes >= 0x80. + return fx.Int32(fx.Uint8(value_i8)) + + b0 = load_scale_byte(0) + b1 = load_scale_byte(1) + b2 = load_scale_byte(2) + b3 = load_scale_byte(3) + packed = b0 | (b1 << 8) | (b2 << 16) | (b3 << 24) + buffer_ops.buffer_store(packed, dst_rsrc, linear) + + @flyc.jit + def launch_pack_mx32_scales( + src: fx.Tensor, + dst: fx.Tensor, + stream: fx.Stream = fx.Stream(None), + ): + kernel_pack_mx32_scales(src, dst).launch( + grid=(total_words // _SCALE_PACK_THREADS, 1, 1), + block=(_SCALE_PACK_THREADS, 1, 1), + stream=stream, + ) + + return launch_pack_mx32_scales + + +@functools.lru_cache(maxsize=None) +def _cached_mx32_scale_pack_launch( + dim: int, + qk: int, + source_colwise: bool, + stride0: int, + stride1: int, +): + """Cache orientation-and-stride-specialized fused pack binaries.""" + return _compile_mx32_scale_pack_kernel( + dim, qk, source_colwise, stride0, stride1 + ) + + +def pack_mx32_scales_for_hk( + scales_u8: torch.Tensor, + *, + source_colwise: bool = False, + stream=None, +) -> torch.Tensor: + """Launch one fused GPU kernel producing HK MFMA-ready scale words. + + Input contracts: + * rowwise: ``[dim, K/32]`` + * columnwise: ``[K/32, dim]`` + + Output contract: + * ``[K/128, dim]`` ``torch.int32`` + """ + if scales_u8.dtype != torch.uint8: + raise TypeError( + f"MXFP8 scales must be torch.uint8 E8M0 bytes, got " + f"{scales_u8.dtype}" + ) + if scales_u8.ndim != 2: + raise ValueError( + f"MXFP8 scales must be rank 2, got {tuple(scales_u8.shape)}" + ) + if not scales_u8.is_cuda: + raise ValueError("MXFP8 scale packing requires a CUDA/ROCm tensor") + if any(stride <= 0 for stride in scales_u8.stride()): + raise ValueError( + f"MXFP8 scale packing requires positive strides, got " + f"{scales_u8.stride()}" + ) + + if source_colwise: + qk, dim = scales_u8.shape + else: + dim, qk = scales_u8.shape + + if qk % 4 != 0: + raise ValueError( + f"Scale K/32 dimension={qk} must be divisible by 4" + ) + if dim % 64 != 0: + raise ValueError( + f"Scale outer dimension={dim} must be a multiple of 64" + ) + + packed = torch.empty( + (qk // 4, dim), + dtype=torch.int32, + device=scales_u8.device, + ) + if stream is None: + 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, + ) + return packed + +def _encode_waitcnt(vmcnt=63, lgkmcnt=15): + """Encode the CDNA4/gfx950 ``S_WAITCNT`` SIMM16 operand. + + ``rocdl.s_waitcnt`` accepts the raw 16-bit immediate operand of the + 32-bit ``S_WAITCNT`` ISA instruction. On CDNA4, that SIMM16 field is: + + SIMM16[3:0] = vmcnt[3:0] + SIMM16[6:4] = expcnt[2:0] + SIMM16[11:8] = lgkmcnt[3:0] + SIMM16[15:14] = vmcnt[5:4] + + ``vmcnt`` is therefore one six-bit counter split across two noncontiguous + fields; bits [5:4] are placed in SIMM16[15:14], while bits [3:0] remain + in SIMM16[3:0]. + + A wait-counter field set to its maximum representable value is effectively + unconstrained: the instruction does not wait on that counter. This helper + always encodes ``expcnt=7`` and defaults to ``vmcnt=63`` and ``lgkmcnt=15``, + so callers specify only the counters on which they intend to wait. + + For example, ``_encode_waitcnt(lgkmcnt=0)`` returns ``0xC07F``, which the + assembler renders as ``s_waitcnt lgkmcnt(0)``. + See: https://llvm.org/docs/AMDGPU/gfx9_waitcnt.html + """ + if not 0 <= vmcnt <= 63: + raise ValueError(f"vmcnt must be in [0, 63], got {vmcnt}") + if not 0 <= lgkmcnt <= 15: + raise ValueError(f"lgkmcnt must be in [0, 15], got {lgkmcnt}") + + return ( + (7 << 4) # expcnt=7 -> SIMM16[6:4] (unconstrained) + | (vmcnt & 0x0F) # vmcnt[3:0] -> SIMM16[3:0] + | ((lgkmcnt & 0x0F) << 8) # lgkmcnt[3:0] -> SIMM16[11:8] + | ((vmcnt & 0x30) << 10) # vmcnt[5:4] -> SIMM16[15:14] + ) + + +# Keep the documented gfx950 encoding invariant executable and import-time cheap. +assert _encode_waitcnt(lgkmcnt=0) == 0xC07F + + +def _barrier(vmcnt=63, lgkmcnt=15): + if vmcnt != 63 or lgkmcnt != 15: + rocdl.s_waitcnt(_encode_waitcnt(vmcnt=vmcnt, lgkmcnt=lgkmcnt)) + rocdl.s_barrier() + +def _min(a, b): + return arith.select(a < b, a, b) + + +def _divmod(a, b): + return a // b, a % b + + +def _xcd_swizzle(num_pid_m, num_pid_n): + NUM_XCDS = 8 + WGM = 4 + NUM_CUS = 32 * NUM_XCDS + SWIZZLE_THRESHOLD = 4 * NUM_CUS + + wgid = fx.block_idx.x + num_wg = num_pid_m * num_pid_n + + # Simple row-major path. + simple_m, simple_n = _divmod(wgid, num_pid_n) + + # XCD-remapped grouped-M path. + intra_xcd, xcd = _divmod(wgid, NUM_XCDS) + wgid_remap = xcd * (num_wg // NUM_XCDS) + intra_xcd + num_wgid_in_group = WGM * num_pid_n + group_id, intra_group = _divmod(wgid_remap, num_wgid_in_group) + first_pid_m = group_id * WGM + group_size_m = _min(num_pid_m - first_pid_m, WGM) + pid_n, intra_group_m = _divmod(intra_group, group_size_m) + pid_m = first_pid_m + intra_group_m + + use_simple = (num_wg < SWIZZLE_THRESHOLD) | (num_wg % NUM_XCDS != 0) + return ( + arith.select(use_simple, simple_m, pid_m), + arith.select(use_simple, simple_n, pid_n), + ) + + +def _compile_kernel( + K: int, + a_fp8_dtype: torch.dtype, + b_fp8_dtype: torch.dtype, + output_dtype: torch.dtype, + layout: str, +): + """Build one compile-time-specialized TN, NN, or NT kernel. + + ``layout`` is a Python string consumed while constructing the FlyDSL IR. + It is not a runtime kernel argument. Each cache entry therefore contains + only the addressing, LDS reads, and scheduler directives for that layout. + """ + if layout not in ("TN", "NN", "NT"): + raise ValueError(f"Unsupported MXFP8 kernel layout: {layout}") + + a_transpose_read = layout == "NT" + b_transpose_read = layout in ("NN", "NT") + + BLOCK_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K + + fp8_input_types = { + torch.float8_e4m3fn: (fx.Float8E4M3FN, 0), + torch.float8_e5m2: (fx.Float8E5M2, 1), + } + try: + a_fx_dtype, a_matrix_format = fp8_input_types[a_fp8_dtype] + b_fx_dtype, b_matrix_format = fp8_input_types[b_fp8_dtype] + except KeyError as exc: + raise TypeError( + "FlyDSL MXFP8 input dtype must be torch.float8_e4m3fn or " + f"torch.float8_e5m2, got A={a_fp8_dtype}, B={b_fp8_dtype}" + ) from exc + + if output_dtype == torch.float16: + output_element_bytes = 2 + output_fx_dtype = fx.Float16 + elif output_dtype == torch.bfloat16: + output_element_bytes = 2 + output_fx_dtype = fx.BFloat16 + elif output_dtype == torch.float32: + output_element_bytes = 4 + output_fx_dtype = fx.Float32 + else: + raise TypeError( + "FlyDSL MXFP8 supports only float16, bfloat16, and float32 " + f"outputs, got {output_dtype}" + ) + + NUM_THREADS = 256 + WARP_SIZE = 64 + + SUBTILE_M = 64 + SUBTILE_N = 64 + + MFMA_M = 16 + MFMA_N = 16 + + SUBTILES_PER_WAVE = 4 + MFMA_M_PER_SUBTILE = SUBTILE_M // MFMA_M + MFMA_N_PER_SUBTILE = SUBTILE_N // MFMA_N + ACCS_PER_WAVE = SUBTILES_PER_WAVE * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + + ELEM_BYTES = 1 + VEC_BYTES = 16 + + LDS_ELEMS_A = BLOCK_M * BLOCK_K + LDS_ELEMS_B = BLOCK_N * BLOCK_K + LDS_BYTES_A = LDS_ELEMS_A * ELEM_BYTES + LDS_BYTES_B = LDS_ELEMS_B * ELEM_BYTES + + LOAD_PASSES_A = LDS_BYTES_A // (NUM_THREADS * VEC_BYTES) + LOAD_PASSES_B = LDS_BYTES_B // (NUM_THREADS * VEC_BYTES) + LOAD_PASSES_A_SUBTILE = LOAD_PASSES_A // 2 + LOAD_PASSES_B_SUBTILE = LOAD_PASSES_B // 2 + LOAD_PASSES_SCALES = 16 + + assert K % BLOCK_K == 0, f"K must be a multiple of {BLOCK_K}, got {K}" + NUM_K_TILES = K // BLOCK_K + assert NUM_K_TILES >= 4, f"K={K} gives {NUM_K_TILES} K128 tiles; the two-page pipeline needs at least 4" + + LDS_ELEMS_HALF = (BLOCK_M // 2) * BLOCK_K + LOAD_PASSES_HALF = LDS_ELEMS_HALF // (NUM_THREADS * VEC_BYTES) + assert LOAD_PASSES_HALF == LOAD_PASSES_A_SUBTILE == LOAD_PASSES_B_SUBTILE + + # Resolve every layout-dependent choice before FlyDSL captures kernel_gemm. + # These are ordinary Python callables/constants, so each cached layout emits + # only its selected addressing, fragment-read, and scheduler path. + Q0_SCHED_DSRD = 2 if a_transpose_read else 1 + PREFETCH_SCHED_DSRD = 4 if a_transpose_read else 2 + + if a_transpose_read: + 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)) + ) + + def _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, + ): + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + local_m_tile = ( + subtile_m_idx * fx.Index(SUBTILE_M) + + fx.Index(mi * MFMA_M) + - fx.Index(sm * (BLOCK_M // 2)) + ) + return load_transposed_frag_half( + lds_a[sm], + local_m_tile, + half, + ) + else: + def _a_leading_dim(c_m): + del c_m + return K + + 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 + ) + + def _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, + ): + 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) + return load_frag_half_at_byte_base( + lds_a[sm], + row_byte_base, + half, + ) + + if b_transpose_read: + 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)) + ) + + def _load_b_ni( + load_transposed_frag, + load_normal_b_frag, + lds_b, + sn, + ni, + reg_subtile_n_idx0, + lane_mod_16, + ): + del load_normal_b_frag, lane_mod_16 + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + local_n_tile = ( + subtile_n_idx * fx.Index(SUBTILE_N) + + fx.Index(ni * MFMA_N) + - fx.Index(sn * (BLOCK_N // 2)) + ) + return load_transposed_frag(lds_b[sn], local_n_tile) + else: + def _b_leading_dim(c_n): + del c_n + return K + + 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 + ) + + def _load_b_ni( + load_transposed_frag, + load_normal_b_frag, + lds_b, + sn, + ni, + reg_subtile_n_idx0, + 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 + ) + return load_normal_b_frag(lds_b, b_row_addr, sn) + + if (not a_transpose_read) or (not b_transpose_read): + def _normal_read_columns(lane_div_16, lane_mod_16): + reg_k_col0 = lane_div_16 * 16 + reg_k_col1 = 64 + lane_div_16 * 16 + _, col0 = swizzle_128(lane_mod_16, reg_k_col0) + _, col1 = swizzle_128(lane_mod_16, reg_k_col1) + return col0, col1 + else: + def _normal_read_columns(lane_div_16, lane_mod_16): + del lane_div_16, lane_mod_16 + return fx.Int32(0), fx.Int32(0) + + @fx.struct + class SharedStorage: + # Each logical 256x128 page is two independent 128x128 half-pages. + # The hot loop refills one 16-byte pass of one half-page at a time. + a0_0: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a0_1: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a1_0: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a1_1: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + b0_0: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + 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] + + @flyc.kernel(known_block_size=[NUM_THREADS, 1, 1]) + def kernel_gemm( + A: fx.Tensor, As: fx.Tensor, B: fx.Tensor, Bs: fx.Tensor, C: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32 + ): + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + lds_a0 = (lds.a0_0, lds.a0_1) + lds_a1 = (lds.a1_0, lds.a1_1) + lds_b0 = (lds.b0_0, lds.b0_1) + lds_b1 = (lds.b1_0, lds.b1_1) + + a_f8_ir_t = a_fx_dtype.ir_type + b_f8_ir_t = b_fx_dtype.ir_type + gA = make_fp8_buffer_tensor(A, a_f8_ir_t) + 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)) + as_rsrc = buffer_ops.create_buffer_resource(As, max_size=True) + bs_rsrc = buffer_ops.create_buffer_resource(Bs, max_size=True) + tx = gpu.thread_id("x") + + 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 = 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) + + # C can exceed the signed-i32 element/byte offset range for large M*N. + # Bias the buffer descriptor base once per CTA using an index/i64 GEP, + # then store with only tile-local i32 offsets. This keeps the hot store + # instruction form unchanged while avoiding i32 wrap in buffer_store(). + c_n_idx_for_base = fx.Index(c_n) + c_tile_base_elems = bx_m_idx * c_n_idx_for_base + by_n_idx + c_tile_base_bytes = c_tile_base_elems * fx.Index(output_element_bytes) + c_rsrc = buffer_ops.create_buffer_resource( + C, + max_size=True, + base_byte_offset=c_tile_base_bytes, + ) + + 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, + ) + + 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) + + 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): + packed = buffer_ops.buffer_load( + as_rsrc, + k128 * c_m_idx + bx_m_idx + row, + vec_width=1, + dtype=T.i32, + ) + return packed + + def load_b_scale_row(k128, row): + packed = buffer_ops.buffer_load( + bs_rsrc, + k128 * c_n_idx + by_n_idx + row, + vec_width=1, + dtype=T.i32, + ) + return packed + + 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_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_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(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_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 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 _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, 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), + ], + ( + f"v_mfma_scale_f32_16x16x128_f8f6f4 " + f"a[{acc_pin}:{acc_pin + 3}], " + f"$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}], " + f"$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 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), + ], + ( + f"v_mfma_scale_f32_16x16x128_f8f6f4 " + f"a[{dst_pin}:{dst_pin + 3}], " + f"$0, $1, " + f"a[{old_pin}:{old_pin + 3}], " + f"$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 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 + + 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 + for ii in range_constexpr(4): + row = row_base + fx.Index(ii) + c_idx = row * fx.Index(c_n) + col + value = Vec(acc)[ii] + if output_dtype != torch.float32: + value = value.to(output_fx_dtype) + buffer_ops.buffer_store(value, c_rsrc, 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, + 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) + + # 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) + + 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) + + 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) + + 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) + + 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_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_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_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_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) + + 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, next_scales_ready, refill_scales + + 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) + + 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) + + 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) + + 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) + + 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) + + rocdl.sched_barrier(0) + _barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + 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) + + 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) + + # 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, + 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, + a0_regs, + b0_regs, + scales0, + scales1, + ) + hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs, scales1) + 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) + + @flyc.jit + def launch_gemm( + A: fx.Tensor, + As: fx.Tensor, + B: fx.Tensor, + Bs: fx.Tensor, + C: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + # The integration only dispatches aligned shapes; no partial-tile masking exists. + grid_x = (c_m // BLOCK_M) * (c_n // BLOCK_N) + kernel_gemm( + A, + As, + B, + Bs, + C, + c_m, + c_n, + 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) + + return launch_gemm + +def do_gemm( + A: torch.Tensor, + As: torch.Tensor, + B: torch.Tensor, + Bs: torch.Tensor, + C: torch.Tensor, + stream=None, + *, + layout: str = "TN", +): + """Launch one cached compile-time MXFP8 layout specialization.""" + if layout == "TN": + M_runtime, K_runtime = A.shape + N_runtime, Kb_runtime = B.shape + elif layout == "NN": + M_runtime, K_runtime = A.shape + Kb_runtime, N_runtime = B.shape + elif layout == "NT": + K_runtime, M_runtime = A.shape + Kb_runtime, N_runtime = B.shape + else: + raise ValueError(f"Unsupported MXFP8 kernel layout: {layout}") + + assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" + 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}" + + if M_runtime % _BLOCK_M != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP8 {layout} GEMM requires M to be a multiple of " + f"{_BLOCK_M}, got M={M_runtime}" + ) + if N_runtime % _BLOCK_N != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP8 {layout} GEMM requires N to be a multiple of " + f"{_BLOCK_N}, got N={N_runtime}" + ) + if K_runtime % _BLOCK_K != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP8 {layout} GEMM requires K to be a multiple of " + f"{_BLOCK_K}, got K={K_runtime}" + ) + num_k_tiles = K_runtime // _BLOCK_K + if num_k_tiles < 4: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP8 {layout} GEMM requires at least 4 K{_BLOCK_K} " + f"tiles, got K={K_runtime} ({num_k_tiles} tiles)" + ) + + expected_as = (K_runtime // _BLOCK_K, M_runtime) + expected_bs = (K_runtime // _BLOCK_K, N_runtime) + 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)}" + ) + if C.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise TypeError( + "C dtype must be torch.float16, torch.bfloat16, or torch.float32, " + f"got {C.dtype}" + ) + + tensors = (A, As, B, Bs, C) + if any(t.device != A.device for t in tensors[1:]): + raise ValueError("A, B, packed scales, and C must be on the same device") + + if stream is None: + stream = torch.cuda.current_stream() + + # Preserve the exact flat descriptor contract used by the passing kernels. + A_arg = A.view(torch.uint8).contiguous().view(-1) + B_arg = B.view(torch.uint8).contiguous().view(-1) + As_arg = As.contiguous().view(-1) + Bs_arg = Bs.contiguous().view(-1) + C_arg = C.contiguous().view(-1) + + _cached_launch( + K_runtime, + A.dtype, + B.dtype, + C.dtype, + layout, + )( + A_arg, + As_arg, + B_arg, + Bs_arg, + C_arg, + M_runtime, + N_runtime, + stream=stream, + ) + + +@functools.lru_cache(maxsize=None) +def _cached_launch( + K: int, + a_fp8_dtype: torch.dtype, + b_fp8_dtype: torch.dtype, + output_dtype: torch.dtype, + layout: str, +): + """Cache independent TN/NN/NT binaries with no runtime layout argument.""" + return _compile_kernel( + K, + a_fp8_dtype, + b_fp8_dtype, + output_dtype, + layout, + ) + + +def _validate_common_payloads( + a: torch.Tensor, + b: torch.Tensor, + D: torch.Tensor, + *, + layout: str, +): + if a.ndim != 2 or b.ndim != 2: + raise ValueError( + f"FlyDSL MXFP8 {layout} expects rank-2 operands, got " + f"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 MXFP8 {layout} expects E4M3 or E5M2 payloads " + f"independently, got a={a.dtype} and b={b.dtype}" + ) + if a.device != b.device or D.device != a.device: + raise ValueError("A, B, and D must be on the same device") + if D.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise TypeError( + "FlyDSL MXFP8 output must be float16, bfloat16, or float32, " + f"got {D.dtype}" + ) + + +def mxfp8_matmul( + a: torch.Tensor, + a_scale: torch.Tensor, + b: torch.Tensor, + b_scale: torch.Tensor, + D: torch.Tensor, + stream=None, + *, + layout: str = "TN", +): + """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] + 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] + + 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. + """ + if layout not in ("TN", "NN", "NT"): + raise ValueError(f"Unsupported MXFP8 kernel layout: {layout}") + + _validate_common_payloads(a, b, D, layout=layout) + + if layout == "TN": + m, k = a.shape + n, kb = b.shape + elif layout == "NN": + m, k = a.shape + kb, n = b.shape + else: + k, m = a.shape + kb, n = b.shape + + if kb != k: + raise ValueError( + f"Incompatible MXFP8 {layout} operands: " + f"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)}") + if k % SCALE_GROUP_SIZE != 0: + raise ValueError( + f"K={k} must be divisible by MXFP8 scale group size " + f"{SCALE_GROUP_SIZE}" + ) + + if layout == "NT": + expected_a_scale = (k // SCALE_GROUP_SIZE, m) + else: + expected_a_scale = (m, k // SCALE_GROUP_SIZE) + + if layout == "TN": + expected_b_scale = (n, k // SCALE_GROUP_SIZE) + else: + expected_b_scale = (k // SCALE_GROUP_SIZE, n) + + if tuple(a_scale.shape) != expected_a_scale: + raise ValueError( + f"a_scale shape {tuple(a_scale.shape)} != expected " + f"{expected_a_scale} for {layout}" + ) + if tuple(b_scale.shape) != expected_b_scale: + raise ValueError( + f"b_scale shape {tuple(b_scale.shape)} != expected " + f"{expected_b_scale} for {layout}" + ) + if a_scale.dtype != torch.uint8 or b_scale.dtype != 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: + raise ValueError("A, B, scales, and D must be on the same device") + + if layout == "TN": + # TN selected backings already match the normal-read kernel contract: + # a [M,K], b [N,K] + a_kernel = a + b_kernel = b + a_scale_hk = pack_mx32_scales_for_hk( + a_scale, + source_colwise=False, + stream=stream, + ) + b_scale_hk = pack_mx32_scales_for_hk( + b_scale, + source_colwise=False, + stream=stream, + ) + elif layout == "NN": + a_kernel = a + b_kernel = b + a_scale_hk = pack_mx32_scales_for_hk( + a_scale, + source_colwise=False, + stream=stream, + ) + b_scale_hk = pack_mx32_scales_for_hk( + b_scale, + source_colwise=True, + stream=stream, + ) + else: + a_kernel = a + b_kernel = b + 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, + ) + + _debug( + f"{layout} kernel inputs: a={tuple(a_kernel.shape)}, " + f"b={tuple(b_kernel.shape)}, " + f"a_scale_hk={tuple(a_scale_hk.shape)}, " + f"b_scale_hk={tuple(b_scale_hk.shape)}, D={tuple(D.shape)}" + ) + + do_gemm( + a_kernel, + a_scale_hk, + b_kernel, + b_scale_hk, + D.view(m, n), + layout=layout, + stream=stream, + ) + return D + + +def mxfp8_matmul_nn(*args, **kwargs): + """Compatibility entry point for the common NN specialization.""" + kwargs["layout"] = "NN" + return mxfp8_matmul(*args, **kwargs) + + +def mxfp8_matmul_nt(*args, **kwargs): + """Compatibility entry point for the common NT specialization.""" + kwargs["layout"] = "NT" + return mxfp8_matmul(*args, **kwargs) + + +__all__ = [ + "BLOCK_M", + "BLOCK_N", + "BLOCK_K", + "SCALE_GROUP_SIZE", + "pack_mx32_scales_for_hk", + "do_gemm", + "mxfp8_matmul", + "mxfp8_matmul_nn", + "mxfp8_matmul_nt", +] diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/pf_gemm_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/pf_gemm_utils.py new file mode 100644 index 000000000..1bc109412 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/pf_gemm_utils.py @@ -0,0 +1,480 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""FlyDSL helpers for permute-free MoE grouped GEMM (Mega 8-wave / 32x32x16). + +Shared building blocks for ``pf_fwd.py`` and ``pf_dgrad.py``. Dense 4-wave BF16 +utilities live in ``fp16_gemm_utils.py``; the pipelined MMA loop lives in +``bf16_gemm.py`` as ``dense_mma_pipeline_bf16``. +""" + +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import fly as fly_dialect +from flydsl._mlir.dialects import llvm as _llvm +from flydsl.expr import buffer_ops +from flydsl.expr.arith import _to_raw as _raw +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec +from flydsl.expr.utils.arith import ArithValue +from flydsl.expr import arith, range_constexpr +from flydsl.expr.buffer_ops import _unwrap_value, buffer_store, create_buffer_resource + +from .bf16_gemm import BLOCK_K, dense_mma_pipeline_bf16 +from .fp16_gemm_utils import G2SLoader, ceildiv, make_bf16_buffer_tensor, swizzle_128 + + +def _inttoptr_lds(byte_addr): + """Integer byte address -> !llvm.ptr<3> (LDS). Parsed per call: the type is + bound to the current MLIRContext and cannot be cached across compiles.""" + return _llvm.inttoptr(ir.Type.parse("!llvm.ptr<3>"), _raw(fx.Int64(byte_addr))) + + +_gep = buffer_ops.get_element_ptr + + +def _lds_ptr_from_i32(addr_i32, byte_offset=0): + """Build an LDS pointer (ptr<3>) from an i32 byte address + optional static offset.""" + ptr = _inttoptr_lds(ArithValue(addr_i32).extui(T.i64)) + if byte_offset != 0: + ptr = _gep(ptr, static_byte_offset=byte_offset) + return ptr + + +def _packed_ds_read_tr16(base_ptr, byte_offsets): + n = len(byte_offsets) + v2i32 = ir.VectorType.get([2], ir.IntegerType.get_signless(32)) + struct_t = _llvm.StructType.get_literal([v2i32] * n) + asm = "\n".join(f"ds_read_b64_tr_b16 ${k}, ${n} offset:{byte_offsets[k]}" for k in range(n)) + constraints = ",".join(["=&v"] * n + ["v"] + ["~{memory}"]) + op = _llvm.InlineAsmOp( + res=struct_t, + operands_=[_raw(base_ptr)], + asm_string=asm, + constraints=constraints, + has_side_effects=True, + ) + return [Vec(_llvm.extractvalue(v2i32, op.result, [k])).bitcast(fx.BFloat16) for k in range(n)] + + +class _S2RLoaderBase: + """Shared ctor for LDS->register operand loaders: caches the per-lane id, + this wave's tile index, and the tile count.""" + + def __init__(self, wave_idx, n_tiles): + self.lane_id = fx.thread_idx.x % 64 + self.wave_idx = wave_idx + self.n_tiles = n_tiles + + +class _S2RLoaderBf16(_S2RLoaderBase): + """Shared skeleton for the bf16 operand loaders (cf. _MfmaBf16): n_tiles output + tiles, each a list of k-sub fragments. Subclasses supply the per-sub offset table + and _tile(), which holds the LDS address math -- transposed ds_read_tr_b16 or + swizzled buffer load (too different to share beyond this loop).""" + + def load(self, lds_src): + return [self._tile(lds_src, i) for i in range_constexpr(self.n_tiles)] + + +def _read_tr16_sub(base_i32, sub16, row_off): + """One tr16 sub-block (512 elems/block): packed double-read, the pair 128 bytes + apart, at sub16*512 + row_off, then assembled.""" + ptr = _lds_ptr_from_i32(base_i32 + (sub16 * 512 + row_off) * 2) + r0, r1 = _packed_ds_read_tr16(ptr, [0, 128]) + return r0.shuffle(r1, list(range(8))) + + +class S2RLoaderTrBf16(_S2RLoaderBf16): + """mfma_f32_32x32x16 operand via ds_read_tr_b16 transpose. Like S2RLoaderTr's + _K_BASE, _SUB lists the tr16 sub-block of each inst_k=16 mfma step (consecutive + here); its length is both the sub count and the per-tile block stride.""" + + _SUB = (0, 1, 2, 3) + + def _tile(self, lds_src, i): + m, kblk = self.lane_id % 32, self.lane_id // 32 + row_off = (m // 16) * 256 + kblk * 128 + (m % 16) * 4 + base_i32 = fx.Int32(fx.ptrtoint(lds_src.ptr)) + sub0 = (self.wave_idx * self.n_tiles + i) * len(self._SUB) + return [ + _read_tr16_sub(base_i32, sub0 + self._SUB[c], row_off) + for c in range_constexpr(len(self._SUB)) + ] + + +def _load8_bf16(lds_src, byte_off): + i8 = fx.recast_iter(fx.Uint8, lds_src.ptr) + p = fx.add_offset(i8, fx.make_int_tuple(byte_off)) + v = fx.make_view(p, fx.make_layout(16, 1)).load() + return v.bitcast(fx.BFloat16) + + +class S2RLoaderBf16(_S2RLoaderBf16): + """mfma_f32_32x32x16 operand (swizzled, non-transposed). Mirroring S2RLoaderTr, + _K_BASE lists the K-column (elems) of each inst_k=16 sub of a 32-row tile; its + length is the sub count -- no BLOCK_K needed.""" + + _K_BASE = (0, 16, 32, 48) + + def _tile(self, lds_src, i): + m, kblk = self.lane_id % 32, self.lane_id // 32 + row = self.wave_idx * (self.n_tiles * 32) + i * 32 + m + subs = [] + for c in range_constexpr(len(self._K_BASE)): + col_byte = (self._K_BASE[c] + kblk * 8) * 2 + _, cs = swizzle_128(row, col_byte) + subs.append(_load8_bf16(lds_src, row * 128 + cs)) + return subs + + +class _MfmaBf16: + """Grouped bf16 mfma: accumulate n_tiles_a x n_tiles_b output tiles. The k-sub + count is taken from each operand's fragment list (len(a[i])), so the atom's + (m, n, inst_k) is the only shape this class needs -- no BLOCK_K coupling.""" + + def __init__(self, n_tiles_a, n_tiles_b, m, n, inst_k): + self.atom = fx.make_mma_atom(fx.rocdl.MFMA(m, n, inst_k, fx.BFloat16)) + acc_len = m * n // 64 # f32 accum lanes per wave + self.accum_type = Vec.make_type(acc_len, fx.Float32) + self.zero_value = Vec.filled(acc_len, 0.0, fx.Float32) + self.n_tiles_a = n_tiles_a + self.n_tiles_b = n_tiles_b + + def idx(self, i, j): + return i * self.n_tiles_b + j + + def call(self, a, b, c): + assert len(a) == self.n_tiles_a + assert len(b) == self.n_tiles_b + for i in range_constexpr(self.n_tiles_a): + for j in range_constexpr(self.n_tiles_b): + acc = c[self.idx(i, j)] + for ks in range_constexpr(len(a[i])): + acc = fly_dialect.mma_atom_call_ssa( + [self.accum_type], self.atom, a[i][ks], b[j][ks], acc + ) + c[self.idx(i, j)] = acc + return c + + +class Mfma32x32x16(_MfmaBf16): + def __init__(self, n_tiles_a, n_tiles_b): + super().__init__(n_tiles_a, n_tiles_b, 32, 32, 16) + + +class StoreCBf16: + def __init__(self, C, c_rows, c_cols, out_ty, cache_modifier=0): + self.c_rows = c_rows + self.c_cols = c_cols + self.lane_id = fx.thread_idx.x % 64 + self.out_ty = out_ty + self.cache_modifier = cache_modifier + c_nbytes = c_rows * c_cols * 2 + gC = fx.rocdl.make_buffer_tensor(C, max_size=False, num_records_bytes=c_nbytes) + self.c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) + self.out_atom_1 = fx.make_copy_atom(fx.rocdl.BufferCopy16b(), out_ty) + self.reg_out_1 = fx.make_rmem_tensor(fx.make_layout(1, 1), out_ty) + self.c_rsrc = ( + create_buffer_resource(C, max_size=False, num_records_bytes=c_nbytes) + if cache_modifier + else None + ) + self.oob = fx.Int32(c_rows * c_cols) # out-of-bounds sink index + + def _store_masked(self, value, c_index, valid): + """Store one element to c_index (masked to the OOB sink when invalid).""" + idx = arith.select(valid, c_index, self.oob) + val = value.to(self.out_ty) + if self.cache_modifier: + buffer_store(val, self.c_rsrc, fx.Int32(idx), cache_modifier=self.cache_modifier) + else: + fx.memref_store_vec(Vec.filled(1, val, self.out_ty), self.reg_out_1) + fx.copy(self.out_atom_1, self.reg_out_1, fx.slice(self.c_div, (None, fx.Int32(idx)))) + + def store(self, c_frag, base_row, base_col): + n = self.lane_id % 32 + m_hi = (self.lane_id // 32) * 4 + col = base_col + n + col_valid = col < self.c_cols + for ti in range_constexpr(len(c_frag)): + acc = Vec(c_frag[ti]) + for r in range_constexpr(16): + row = base_row + ti * 32 + (r // 4) * 8 + m_hi + (r % 4) + self._store_masked(acc[r], row * self.c_cols + col, col_valid) + + def store16(self, c_frag, base_row, base_col): + n = self.lane_id % 16 + m_hi = (self.lane_id // 16) * 4 + col = base_col + n + col_valid = col < self.c_cols + for ti in range_constexpr(len(c_frag)): + acc = Vec(c_frag[ti]) + for r in range_constexpr(4): + row = base_row + ti * 16 + m_hi + r + self._store_masked(acc[r], row * self.c_cols + col, col_valid) + + def store_trans16(self, c_frag, group_idx, base_m, base_n, out_m, out_n): + n = self.lane_id % 16 + m_hi = (self.lane_id // 16) * 4 + glob_n = base_n + n + n_valid = glob_n < out_n + row_base = (group_idx * out_n + glob_n) * out_m + for ti in range_constexpr(len(c_frag)): + acc = Vec(c_frag[ti]) + for r in range_constexpr(4): + m = base_m + ti * 16 + m_hi + r + self._store_masked(acc[r], row_base + m, n_valid) + + +def compute_global_swizzle_bf16(lane_id, wave_id, K, n_rounds): + offsets = [] + n_waves = fx.block_dim.x // 64 + for r in range_constexpr(n_rounds): + row = lane_id // 8 + wave_id * 8 + r * (n_waves * 8) + col_byte = (lane_id % 8) * 16 + _, c = swizzle_128(row, col_byte) + offsets.append(row * K + c // 2) + return offsets + + +def compute_global_gather_swizzle_bf16(lane_id, wave_id, K, n_rounds, sorted_res, sorted_row_base): + """Per-lane global A offsets for a *gathering* grouped GEMM. + + Identical to :func:`compute_global_swizzle_bf16` except the tile row is redirected through a + gather index: the source row for tile row ``row`` is ``SORTED_IDS[sorted_row_base + row]`` + instead of the contiguous pool row. The bank swizzle (``c``) is still keyed on the *tile* row + so the LDS destination layout (and therefore the S2R transpose-read) is unchanged -- only the + global fetch address is redirected. The index loads happen once (K-invariant), so they are + amortized over the whole K-loop. + """ + offsets = [] + n_waves = fx.block_dim.x // 64 + for r in range_constexpr(n_rounds): + row = lane_id // 8 + wave_id * 8 + r * (n_waves * 8) + col_byte = (lane_id % 8) * 16 + _, c = swizzle_128(row, col_byte) + src_row = buffer_ops.buffer_load( + sorted_res, sorted_row_base + row, vec_width=1, dtype=T.i32 + ) + offsets.append(src_row * K + c // 2) + return offsets + + +def compute_global_swizzle_nn_bf16(lane_id, wave_id, c_n, n_steps): + offsets = [] + n_waves = fx.block_dim.x // 64 + kk = (lane_id % 32) // 2 + g = lane_id // 32 + n_in = g * 16 + (lane_id % 2) * 8 + for step in range_constexpr(n_steps): + idx = wave_id + step * n_waves + n_tile = idx // 4 + ks = idx % 4 + offsets.append((ks * 16 + kk) * c_n + n_tile * 32 + n_in) + return offsets + + +def make_value_attrs(waves_per_eu, agpr_alloc, fwg): + """Kernel value_attrs. agpr_alloc: 0 = compiler default; N>0 = force exactly + N AGPRs ("N,N"); -N = allow up to N ("0,N").""" + d = {"rocdl.waves_per_eu": waves_per_eu, "rocdl.flat_work_group_size": fwg} + if agpr_alloc != 0: + if agpr_alloc < 0: + alloc = f"0,{-agpr_alloc}" + else: + alloc = f"{agpr_alloc},{agpr_alloc}" + d["passthrough"] = [ + ["amdgpu-agpr-alloc", alloc], + ["amdgpu-mfma-vgpr-form", "false"], + ] + return d + + +def xcd_remap_pid(pid, total_pids, num_xcd): + """Remap the tile id so same-XCD workgroups gather into one contiguous + block, keeping each XCD's L2 reuse within that XCD. Bijection over + [0, total_pids); identity when num_xcd <= 1.""" + if num_xcd <= 1: + return pid + per_xcd = total_pids // num_xcd # floor + rem = total_pids - per_xcd * num_xcd + xcd = pid % num_xcd + local = pid // num_xcd + offset = xcd * per_xcd + arith.select(xcd < rem, xcd, rem) + return offset + local + + +def _i64(v): + # widen an i32 runtime value to i64 (avoids overflow in worst-case base offsets) + return ArithValue(arith.extsi(fx.T.i64(), _unwrap_value(v)), signed=True) + + +def _make_shared_storage(BLOCK_M, BLOCK_N): + a_lds_size = (BLOCK_M // 2) * BLOCK_K + b_lds_size = (BLOCK_N // 2) * BLOCK_K + + @fx.struct + class SharedStorage: + A_lds_cur_0: fx.Array[fx.BFloat16, a_lds_size, 16] + A_lds_cur_1: fx.Array[fx.BFloat16, a_lds_size, 16] + A_lds_next_0: fx.Array[fx.BFloat16, a_lds_size, 16] + A_lds_next_1: fx.Array[fx.BFloat16, a_lds_size, 16] + B_lds_cur_0: fx.Array[fx.BFloat16, b_lds_size, 16] + B_lds_cur_1: fx.Array[fx.BFloat16, b_lds_size, 16] + B_lds_next_0: fx.Array[fx.BFloat16, b_lds_size, 16] + B_lds_next_1: fx.Array[fx.BFloat16, b_lds_size, 16] + + return SharedStorage + + +def _gemm_bf16_nn_tn_tile_impl( + A, + B, + C, + c_m, + c_n, + lds, + block_m, + block_n, + *, + a_transpose, + K, + BLOCK_M, + BLOCK_N, + n_blocks=None, + GROUP_M=1, + num_xcd=8, + out_fp16=False, + nt_vmcnt=3, + b_group_base=None, + c_cache_modifier=0, +): + assert BLOCK_M >= 128 and BLOCK_N >= 256 and BLOCK_M % 128 == 0 and BLOCK_N % 256 == 0 + assert K % BLOCK_K == 0, f"bf16 NN/TN needs K % {BLOCK_K} == 0 (got K={K})" + N_TILES_A = BLOCK_M // 128 + N_TILES_B = BLOCK_N // 256 + LDS_BLOCK_M = BLOCK_M // 2 + LDS_BLOCK_N = BLOCK_N // 2 + N_LDS_STEPS_A = LDS_BLOCK_M // 64 + N_LDS_STEPS_B = LDS_BLOCK_N // 64 + N_LDS_ROUNDS = max(N_LDS_STEPS_A, N_LDS_STEPS_B) + + lane_id = fx.thread_idx.x % 64 + wave_id = fx.thread_idx.x // 64 + wave_m = wave_id // 4 + wave_n = wave_id % 4 + + if block_m is None: + num_pid_m = ceildiv(c_m, BLOCK_M) + pid = xcd_remap_pid(fx.block_idx.x, num_pid_m * n_blocks, num_xcd) + num_pid_in_group = GROUP_M * n_blocks + group_id = pid // num_pid_in_group + pid_in_group = pid % num_pid_in_group + first_pid_m = group_id * GROUP_M + remaining_m = num_pid_m - first_pid_m + group_size_m = arith.select(remaining_m < GROUP_M, remaining_m, fx.Int32(GROUP_M)) + block_m = first_pid_m + (pid_in_group % group_size_m) + block_n = pid_in_group // group_size_m + + if a_transpose: + A0_gl_offset = block_m * BLOCK_M + 0 + A1_gl_offset = block_m * BLOCK_M + LDS_BLOCK_M + a_k_step = BLOCK_K * c_m + else: + A0_gl_offset = (block_m * BLOCK_M) * K + A1_gl_offset = (block_m * BLOCK_M + LDS_BLOCK_M) * K + a_k_step = BLOCK_K + B0_gl_offset = block_n * BLOCK_N + 0 + B1_gl_offset = block_n * BLOCK_N + LDS_BLOCK_N + b_k_step = BLOCK_K * c_n + if b_group_base is not None: + B0_gl_offset = B0_gl_offset + b_group_base + B1_gl_offset = B1_gl_offset + b_group_base + + gA = make_bf16_buffer_tensor(A) + gB = make_bf16_buffer_tensor(B) + a_div = fx.logical_divide(gA, fx.make_layout(1, 1)) + b_div = fx.logical_divide(gB, fx.make_layout(1, 1)) + if a_transpose: + gl_off_a = compute_global_swizzle_nn_bf16(lane_id, wave_id, c_m, N_LDS_STEPS_A) + else: + gl_off_a = compute_global_swizzle_bf16(lane_id, wave_id, K, N_LDS_ROUNDS) + gl_off_b = compute_global_swizzle_nn_bf16(lane_id, wave_id, c_n, N_LDS_STEPS_B) + + mfma = Mfma32x32x16(N_TILES_A, N_TILES_B) + a_g2s = G2SLoader(a_div, gl_off_a, N_LDS_STEPS_A, fx.BFloat16.ir_type, wave_id) + b_g2s = G2SLoader(b_div, gl_off_b, N_LDS_STEPS_B, fx.BFloat16.ir_type, wave_id) + a_s2r = S2RLoaderTrBf16(wave_m, N_TILES_A) if a_transpose else S2RLoaderBf16(wave_m, N_TILES_A) + b_s2r = S2RLoaderTrBf16(wave_n, N_TILES_B) + _out_ty = fx.Float16 if out_fp16 else fx.BFloat16 + store_c = StoreCBf16(C, c_m, c_n, _out_ty, cache_modifier=c_cache_modifier) + + dense_mma_pipeline_bf16( + lds, + a_g2s, + b_g2s, + a_s2r, + b_s2r, + mfma, + store_c, + A0_gl_offset, + A1_gl_offset, + B0_gl_offset, + B1_gl_offset, + a_k_step, + b_k_step, + block_m, + block_n, + wave_m, + wave_n, + K, + BLOCK_M, + BLOCK_N, + nt_vmcnt, + ) + + +def gemm_bf16_nn_tile( + A, + B, + C, + c_m, + c_n, + lds, + block_m=None, + block_n=None, + *, + K, + BLOCK_M, + BLOCK_N, + n_blocks=None, + GROUP_M=1, + num_xcd=8, + out_fp16=False, + nt_vmcnt=3, + b_group_base=None, + c_cache_modifier=0, +): + _gemm_bf16_nn_tn_tile_impl( + A, + B, + C, + c_m, + c_n, + lds, + block_m, + block_n, + a_transpose=False, + K=K, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + n_blocks=n_blocks, + GROUP_M=GROUP_M, + num_xcd=num_xcd, + out_fp16=out_fp16, + nt_vmcnt=nt_vmcnt, + b_group_base=b_group_base, + c_cache_modifier=c_cache_modifier, + ) diff --git a/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_dgrad.py b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_dgrad.py new file mode 100644 index 000000000..c46f6f271 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_dgrad.py @@ -0,0 +1,377 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Permute-free MoE data-gradient (dgrad) grouped-GEMM: MegaMOE's fast bf16 NN GEMM. + +Backward companion to ``pf_fwd``. dgrad contracts the incoming grad against the +weight over the output-feature axis (NN layout):: + + dXrow[s, k] = sum_n grad[src(s), n] * W[e][n, k] + +There are two flavours, matching the TE permute-free contract (``index_a_by_route_pos``): + +* **route-read** (FC1 dgrad, ``gather=False``): ``grad`` is **dense route-ordered** + ``[num_routes, N]`` (or block-padded with only the dense head populated); row ``s`` is read + directly. ``dx`` is **block-padded route-ordered** ``[em_max, K]``. +* **gather** (FC2 dgrad, ``gather=True``): ``grad`` is **token-ordered** ``[num_recv, N]`` and each + route slot ``s`` gathers ``grad[SORTED[s]]`` (sentinel ``SORTED[s] == num_recv`` -> 0), writing + block-padded route-ordered ``dXrow[em_max, K]``. Mirrors the forward NT gather, one row map + redirecting the two LDS A half-tiles, only on the NN tile. + +The weight is bit-identical to the forward ``[E, N, K]`` (forward reads it NT, dgrad NN). + +Contract: + * ``grad_y`` [rows, N] bf16 incoming grad (rows = em_max route-read / num_recv gather) + * ``weight`` [E, N, K] bf16 per-expert weights (shared with forward) + * ``dx`` [em_max, K] bf16 block-padded route-ordered input grad per slot (in place) + * ``expert_ids`` [num_m_blocks] i32 expert id per BLOCK_M slot block + * ``num_tile_blocks`` [1] i32 real (non-padding) BLOCK_M block count (device) + * ``sorted_slot_ids`` [em_max] i32 gather index (gather=True); unused for route-read +""" + +from __future__ import annotations + +import functools + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.expr import arith +from flydsl.expr.buffer_ops import ( + buffer_load, + create_buffer_resource, + extract_base_index, +) +from flydsl.expr.typing import AddressSpace, PointerType + +from ..gemm.bf16_gemm import BLOCK_K, dense_mma_pipeline_bf16 +from ..gemm.fp16_gemm_utils import G2SLoader, ceildiv, make_bf16_buffer_tensor +from ..gemm.pf_gemm_utils import ( + Mfma32x32x16, + S2RLoaderBf16, + S2RLoaderTrBf16, + StoreCBf16, + _i64, + _make_shared_storage, + compute_global_gather_swizzle_bf16, + compute_global_swizzle_nn_bf16, + gemm_bf16_nn_tile, + make_value_attrs, + xcd_remap_pid, +) + +__all__ = ["compile_grouped_gemm_dgrad_bf16", "grouped_gemm_dgrad_bf16"] + + +def gemm_bf16_nn_gather_tile( + A, # flat [num_recv, N] grad buffer (gather source) + B, # weight [E, N, K] flat + C, # dx tile (already rebased to this row block) + c_n, + lds, + sorted_res, + sorted_row_base, + block_n, + *, + Kc, # contraction dim (forward intermediate feature N) + BLOCK_M, + BLOCK_N, + out_fp16=False, + nt_vmcnt=3, + b_group_base, +): + """One NN dgrad tile with the A (grad) rows *gathered* via ``sorted_slot_ids``. + + Mirrors the forward :func:`pf_fwd.gemm_bf16_nt_gather_tile` (same two-loader row + redirection, A base 0, store block_m 0) but on the NN B path (transpose-read B, ``b_k_step = + BLOCK_K * c_n``), contracting over ``Kc`` (= N). + """ + assert BLOCK_M >= 128 and BLOCK_N >= 256 and BLOCK_M % 128 == 0 and BLOCK_N % 256 == 0 + assert Kc % BLOCK_K == 0, f"NN gather needs N % {BLOCK_K} == 0 (got N={Kc})" + N_TILES_A = BLOCK_M // 128 + N_TILES_B = BLOCK_N // 256 + LDS_BLOCK_M = BLOCK_M // 2 + LDS_BLOCK_N = BLOCK_N // 2 + N_LDS_STEPS_A = LDS_BLOCK_M // 64 + N_LDS_STEPS_B = LDS_BLOCK_N // 64 + N_LDS_ROUNDS = max(N_LDS_STEPS_A, N_LDS_STEPS_B) + + lane_id = fx.thread_idx.x % 64 + wave_id = fx.thread_idx.x // 64 + wave_m = wave_id // 4 + wave_n = wave_id % 4 + + A0_gl_offset = fx.Int32(0) + A1_gl_offset = fx.Int32(0) + B0_gl_offset = block_n * BLOCK_N + b_group_base + B1_gl_offset = block_n * BLOCK_N + LDS_BLOCK_N + b_group_base + + gA = make_bf16_buffer_tensor(A) + gB = make_bf16_buffer_tensor(B) + a_div = fx.logical_divide(gA, fx.make_layout(1, 1)) + b_div = fx.logical_divide(gB, fx.make_layout(1, 1)) + + gl_off_a_lo = compute_global_gather_swizzle_bf16( + lane_id, wave_id, Kc, N_LDS_ROUNDS, sorted_res, sorted_row_base + ) + gl_off_a_hi = compute_global_gather_swizzle_bf16( + lane_id, wave_id, Kc, N_LDS_ROUNDS, sorted_res, sorted_row_base + fx.Int32(LDS_BLOCK_M) + ) + gl_off_b = compute_global_swizzle_nn_bf16(lane_id, wave_id, c_n, N_LDS_STEPS_B) + + mfma = Mfma32x32x16(N_TILES_A, N_TILES_B) + a_g2s = G2SLoader(a_div, gl_off_a_lo, N_LDS_STEPS_A, fx.BFloat16.ir_type, wave_id) + a_g2s_hi = G2SLoader(a_div, gl_off_a_hi, N_LDS_STEPS_A, fx.BFloat16.ir_type, wave_id) + b_g2s = G2SLoader(b_div, gl_off_b, N_LDS_STEPS_B, fx.BFloat16.ir_type, wave_id) + a_s2r = S2RLoaderBf16(wave_m, N_TILES_A) + b_s2r = S2RLoaderTrBf16(wave_n, N_TILES_B) + _out_ty = fx.Float16 if out_fp16 else fx.BFloat16 + store_c = StoreCBf16(C, fx.Int32(BLOCK_M), c_n, _out_ty) + + dense_mma_pipeline_bf16( + lds, + a_g2s, + b_g2s, + a_s2r, + b_s2r, + mfma, + store_c, + A0_gl_offset, + A1_gl_offset, + B0_gl_offset, + B1_gl_offset, + BLOCK_K, # a_k_step (contraction rides soffset) + BLOCK_K * c_n, # b_k_step (NN: B is [K, c_n] row-major) + fx.Int32(0), # store block_m: C is already rebased to this tile + block_n, + wave_m, + wave_n, + Kc, + BLOCK_M, + BLOCK_N, + nt_vmcnt, + a_g2s_hi=a_g2s_hi, + ) + + +def _dgrad_gather_body( + GY_flat, GY_tile, WEIGHT, DX_tile, lds, sorted_res, sorted_row_base, block_n, gbase, + *, N, Kout, BLOCK_M, BLOCK_N, out_fp16, nt_vmcnt, +): + """FC2 dgrad tile: gather the token-space grad rows via ``sorted_slot_ids`` (NN gather).""" + gemm_bf16_nn_gather_tile( + GY_flat, WEIGHT, DX_tile, fx.Int32(Kout), lds, sorted_res, sorted_row_base, block_n, + Kc=N, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, out_fp16=out_fp16, nt_vmcnt=nt_vmcnt, + b_group_base=gbase, + ) + + +def _dgrad_routeread_body( + GY_flat, GY_tile, WEIGHT, DX_tile, lds, sorted_res, sorted_row_base, block_n, gbase, + *, N, Kout, BLOCK_M, BLOCK_N, out_fp16, nt_vmcnt, +): + """FC1 dgrad tile: read the dense route-ordered grad directly (plain Mega NN tile).""" + gemm_bf16_nn_tile( + GY_tile, WEIGHT, DX_tile, fx.Int32(BLOCK_M), fx.Int32(Kout), lds, fx.Int32(0), block_n, + K=N, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, out_fp16=out_fp16, nt_vmcnt=nt_vmcnt, + b_group_base=gbase, + ) + + +@functools.lru_cache(maxsize=256) +def compile_grouped_gemm_dgrad_bf16( + N, # contraction dim (forward intermediate feature) + Kout, # output cols (forward hidden input feature) + BLOCK_M=256, + BLOCK_N=256, + GROUP_M=4, + num_xcd=1, + nt_vmcnt=3, + waves_per_eu=2, + agpr_alloc=0, + out_fp16=False, + gather=False, +): + """Compile (cached) the grouped BF16 NN dgrad launcher for one ``(N, Kout, tile)`` combo. + + ``gather=False`` (FC1 dgrad): ``grad`` is dense route-ordered, read per tile (plain Mega NN + tile). ``gather=True`` (FC2 dgrad): ``grad`` is token-ordered, gathered via ``SORTED`` into the + block-padded route-ordered output (NN gather tile). Both rebase the C tile in i64 to survive worst-case + pools; the grid front-loads via XCD swizzle over the real tile range. + """ + SharedStorage = _make_shared_storage(BLOCK_M, BLOCK_N) + # Compile-time tile selector (plain Python; resolved before the AST rewriter so the kernel + # body stays branch-free -- a device ``if gather`` would be lowered to real control flow). + tile_body = _dgrad_gather_body if gather else _dgrad_routeread_body + + @flyc.kernel(known_block_size=[512, 1, 1]) + def grouped_gemm_dgrad_k( + GRAD_Y: fx.Tensor, + WEIGHT: fx.Tensor, + DX: fx.Tensor, + TILE_TO_GROUP: fx.Tensor, + NUM_TILE_BLOCKS: fx.Int32, + SORTED: fx.Tensor, + A_ELEMS: fx.Int32, + ): + n_blocks = ceildiv(fx.Int32(Kout), BLOCK_N) + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + group_res = create_buffer_resource(TILE_TO_GROUP, max_size=True) + sorted_res = create_buffer_resource(SORTED, max_size=True) + # Real BLOCK_M tile count as a scalar (host-known, capture-safe -- no per-call device + # tensor, which a HIP graph capture forbids). Host-known int scalar. + real_tiles = NUM_TILE_BLOCKS + real_grid = real_tiles * n_blocks + + pool_ptr_ty = PointerType.get( + elem_ty=fx.BFloat16.ir_type, address_space=AddressSpace.Global, alignment=16 + ) + gy_base = fx.arith.ArithValue( + arith.index_cast(fx.T.i64(), extract_base_index(GRAD_Y)), signed=True + ) + dx_base = fx.arith.ArithValue( + arith.index_cast(fx.T.i64(), extract_base_index(DX)), signed=True + ) + # Flat 1D grad view (gather source; bounded to A_ELEMS so the sentinel row reads 0). + # Built unconditionally: the route-read tile simply ignores it. + GY_flat = fx.make_view(fx.inttoptr(pool_ptr_ty, gy_base), fx.make_layout(A_ELEMS, 1)) + + def _emit(): + pid = xcd_remap_pid(fx.block_idx.x, real_grid, num_xcd) + num_pid_m = real_tiles + num_pid_in_group = GROUP_M * n_blocks + group_id = pid // num_pid_in_group + pid_in_group = pid % num_pid_in_group + first_pid_m = group_id * GROUP_M + remaining_m = num_pid_m - first_pid_m + group_size_m = arith.select(remaining_m < GROUP_M, remaining_m, fx.Int32(GROUP_M)) + block_m = first_pid_m + (pid_in_group % group_size_m) + block_n = pid_in_group // group_size_m + g_idx = buffer_load(group_res, block_m, vec_width=1, dtype=fx.T.i32()) + # PF padding blocks mark expert_ids=-1; skip the full Mega pipeline. + if g_idx >= fx.Int32(0): + gbase = g_idx * fx.Int32(N) * fx.Int32(Kout) + sorted_row_base = block_m * fx.Int32(BLOCK_M) + + dx_byte_off = _i64(block_m * fx.Int32(BLOCK_M)) * _i64(fx.Int32(Kout)) * fx.Int64(2) + DX_tile = fx.make_view( + fx.inttoptr(pool_ptr_ty, dx_base + dx_byte_off), + fx.make_layout(fx.Int32(BLOCK_M) * fx.Int32(Kout), 1), + ) + # Route-read grad tile (rebased); ignored by the gather tile. + gy_byte_off = _i64(block_m * fx.Int32(BLOCK_M)) * _i64(fx.Int32(N)) * fx.Int64(2) + GY_tile = fx.make_view( + fx.inttoptr(pool_ptr_ty, gy_base + gy_byte_off), + fx.make_layout(fx.Int32(BLOCK_M) * fx.Int32(N), 1), + ) + + tile_body( + GY_flat, GY_tile, WEIGHT, DX_tile, lds, sorted_res, sorted_row_base, block_n, + gbase, N=N, Kout=Kout, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, out_fp16=out_fp16, + nt_vmcnt=nt_vmcnt, + ) + + if fx.block_idx.x < real_grid: + _emit() + + @flyc.jit + def launch(GRAD_Y, WEIGHT, DX, TILE_TO_GROUP, NUM_TILE_BLOCKS: fx.Int32, SORTED, A_ELEMS: fx.Int32, c_m: fx.Int32, stream: fx.Stream): + grid_x = ceildiv(c_m, BLOCK_M) * ceildiv(fx.Int32(Kout), BLOCK_N) + grouped_gemm_dgrad_k( + GRAD_Y, + WEIGHT, + DX, + TILE_TO_GROUP, + NUM_TILE_BLOCKS, + SORTED, + A_ELEMS, + value_attrs=make_value_attrs(waves_per_eu, agpr_alloc, "512,512"), + ).launch(grid=(grid_x, 1, 1), block=(512, 1, 1), stream=stream) + + return launch + + +def grouped_gemm_dgrad_bf16( + grad_y, # [rows, N] bf16 incoming grad (rows = em_max route-read / num_recv gather) + weight, # [E, N, K] bf16 per-expert weights (shared with forward) + dx, # [em_max, K] bf16 block-padded route-ordered input grad per slot (in place) + expert_ids, # [num_m_blocks] i32 + num_tile_blocks, # int real BLOCK_M block count (host scalar; capture-safe) + sorted_slot_ids=None, # [em_max] i32 gather index (required when gather=True) + *, + gather=False, + BLOCK_M=256, + BLOCK_N=256, + GROUP_M=4, + num_xcd=1, + nt_vmcnt=3, + waves_per_eu=2, + agpr_alloc=0, +): + """Host entry: grouped bf16 NN dgrad ``dXrow[s] = grad[src(s)] @ W[expert]``. + + ``gather=False`` reads ``grad`` at the route row (FC1 dgrad, ``grad`` rows == ``dx`` rows). + ``gather=True`` gathers ``grad[SORTED[s]]`` from a token-space buffer (FC2 dgrad); ``dx`` is + written over its full padded ``[em_max, K]`` extent (pad rows carry dead values). + """ + assert grad_y.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16 + assert dx.dtype == torch.bfloat16 + E, N, K = weight.shape + c_m = int(dx.shape[0]) + assert grad_y.shape[1] == N, f"grad_y N={grad_y.shape[1]} != weight N={N}" + assert dx.shape[1] == K, f"dx K={dx.shape[1]} != weight K={K}" + if not gather: + assert grad_y.shape[0] == c_m, ( + f"route-read dgrad expects grad_y rows == dx rows ({grad_y.shape[0]} != {c_m}); " + "did you mean gather=True (token-space grad)?" + ) + + grad_y = grad_y.contiguous() + if gather: + a_elems = int(grad_y.numel()) + if a_elems >= (1 << 31): + raise ValueError( + f"grad_y has {a_elems} elems (>= 2^31); the flat-view bound is int32. " + "Large-A support needs a per-lane i64 SRD rebase (TODO)." + ) + else: + # route-read rebases GRAD_Y per tile; the flat gather view (A_ELEMS) is unused. + a_elems = int(BLOCK_M) * int(N) + expert_ids_i32 = expert_ids.to(torch.int32) + if gather: + assert sorted_slot_ids is not None, "gather=True dgrad requires sorted_slot_ids" + sorted_arg = sorted_slot_ids.to(torch.int32) + else: + # route-read reads GRAD_Y at the route row; SORTED is unread. Reuse a live tensor to + # avoid a capture-illegal per-call allocation. + sorted_arg = expert_ids_i32 + + weight_flat = weight.reshape(E * N, K) + if not weight_flat.is_contiguous(): + weight_flat = weight_flat.contiguous() + weight_flat = weight_flat.view(-1) + launch = compile_grouped_gemm_dgrad_bf16( + N=N, + Kout=K, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + GROUP_M=GROUP_M, + num_xcd=num_xcd, + nt_vmcnt=int(nt_vmcnt), + waves_per_eu=int(waves_per_eu), + agpr_alloc=int(agpr_alloc), + gather=bool(gather), + ) + launch( + grad_y, + weight_flat, + dx, + expert_ids_i32, + int(num_tile_blocks), + sorted_arg, + a_elems, + c_m, + stream=torch.cuda.current_stream(), + ) + return dx diff --git a/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_fwd.py b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_fwd.py new file mode 100644 index 000000000..211a8978c --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_fwd.py @@ -0,0 +1,385 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Permute-free MoE forward gather grouped-GEMM: MegaMOE's fast bf16 GEMM + route-list gather. + +This is a port of the MegaMOE grouped bf16 GEMM (32x32x16 MFMA, 8-wave / 512-thread +workgroup, deep distance-2 3-buffer DMA ring, XCD swizzle with ``GROUP_M`` front-loading) +into Transformer Engine FlyDSL, with a *single* change: the A operand is fetched +through a per-row gather index instead of a contiguous pool. + +The gather is folded into MegaMOE's pipeline rather than implemented as a separate +pre-permute: we keep Mega's DMA/LDS/MFMA path verbatim and only redirect the A fetch +through ``sorted_slot_ids``, preserving Mega's throughput while keeping the permute-free +memory model (no pre-permutation, gather-on-demand). + +Contract (mirrors MegaMOE ``grouped_gemm_bf16_only`` + permute-free routing metadata): + * ``A`` [num_recv, K] bf16 received-token activations, UNPERMUTED (gather source) + * ``B`` [E, N, K] bf16 per-expert weights, NT (contiguous inner K) + * ``C`` [em_max, N] bf16 block-padded route-ordered output (expert-major, in place) + * ``sorted_slot_ids`` [em_max] i32 received-token row per padded slot (sentinel = num_recv) + * ``expert_ids`` [num_m_blocks] i32 expert id per ``BLOCK_M`` output block (padding tail: + ``-1``; those blocks early-exit in-kernel) + * ``num_tile_blocks`` [1] i32 real (non-padding) ``BLOCK_M`` block count (device) + +Gated activation and route-prob apply live in standalone Triton helpers, not here. +""" + +from __future__ import annotations + +import functools + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.expr import arith, range_constexpr +from flydsl.expr.buffer_ops import ( + buffer_load, + create_buffer_resource, + extract_base_index, +) +from flydsl.expr.typing import AddressSpace, PointerType + +from ..gemm.bf16_gemm import BLOCK_K, dense_mma_pipeline_bf16 +from ..gemm.fp16_gemm_utils import G2SLoader, ceildiv, make_bf16_buffer_tensor, swizzle_128 +from ..gemm.pf_gemm_utils import ( + Mfma32x32x16, + S2RLoaderBf16, + StoreCBf16, + _i64, + _make_shared_storage, + compute_global_gather_swizzle_bf16, + compute_global_swizzle_bf16, + make_value_attrs, + xcd_remap_pid, +) + +__all__ = ["compile_grouped_gemm_gather_bf16", "grouped_gemm_gather_bf16"] + + +def compute_global_identity_swizzle_bf16(lane_id, wave_id, K, n_rounds, sorted_row_base): + """Per-lane global A offsets reading the *route row directly* (identity gather). + + Same flat-buffer pipeline as :func:`compute_global_gather_swizzle_bf16` but the source row is + ``sorted_row_base + row`` computed arithmetically instead of loaded from a gather table. This + is the FC2 route-read (``index_a_by_route_pos``) path: it needs no ``sorted_slot_ids`` tensor, + so it stays inside a HIP graph capture (no per-call identity allocation) while reusing the + proven whole-buffer / base-0 gather tile (avoiding the per-tile A-rebase multi-block hazard). + """ + offsets = [] + n_waves = fx.block_dim.x // 64 + for r in range_constexpr(n_rounds): + row = lane_id // 8 + wave_id * 8 + r * (n_waves * 8) + col_byte = (lane_id % 8) * 16 + _, c = swizzle_128(row, col_byte) + offsets.append((sorted_row_base + row) * K + c // 2) + return offsets + + +def gemm_bf16_nt_gather_tile( + A, + B_T, + C, + c_m, + c_n, + lds, + sorted_res, + sorted_row_base, + block_n, + *, + K, + BLOCK_M, + BLOCK_N, + out_fp16=False, + nt_vmcnt=3, + b_group_base, + gather=True, +): + """One NT tile of the grouped GEMM with the A rows *gathered* via ``sorted_slot_ids``. + + Identical to MegaMOE's ``gemm_bf16_nt_tile`` except: + * the two LDS A half-tiles use *gather* swizzles (rows redirected through + ``sorted_slot_ids[sorted_row_base + tile_row]``), so the A base offset is 0 and only the + K-step rides the load ``soffset``; + * ``C`` is already rebased to this tile's row block by the caller, so the store block_m is 0. + """ + assert BLOCK_M >= 128 and BLOCK_N >= 256 and BLOCK_M % 128 == 0 and BLOCK_N % 256 == 0 + assert K % BLOCK_K == 0, f"bf16 NT gather needs K % {BLOCK_K} == 0 (got K={K})" + N_TILES_A = BLOCK_M // 128 + N_TILES_B = BLOCK_N // 256 + LDS_BLOCK_M = BLOCK_M // 2 + LDS_BLOCK_N = BLOCK_N // 2 + N_LDS_STEPS_A = LDS_BLOCK_M // 64 + N_LDS_STEPS_B = LDS_BLOCK_N // 64 + N_LDS_ROUNDS = max(N_LDS_STEPS_A, N_LDS_STEPS_B) + + lane_id = fx.thread_idx.x % 64 + wave_id = fx.thread_idx.x // 64 + wave_m = wave_id // 4 + wave_n = wave_id % 4 + + # A rows carried by the gather offsets -> global base is 0 (K rides soffset only). + A0_gl_offset = fx.Int32(0) + A1_gl_offset = fx.Int32(0) + B0_gl_offset = (block_n * BLOCK_N) * K + B1_gl_offset = (block_n * BLOCK_N + LDS_BLOCK_N) * K + if b_group_base is not None: + B0_gl_offset = B0_gl_offset + b_group_base + B1_gl_offset = B1_gl_offset + b_group_base + + # ``A`` MUST be a flat 1D buffer view (built by the caller): a linear gather offset + # (src_row*K + col) indexes row-major elements. A raw 2D tensor's logical_divide/slice + # indexes the outer (row) dim, so a flat offset runs off the end -> garbage. + gA = make_bf16_buffer_tensor(A) + gB = make_bf16_buffer_tensor(B_T) + a_div = fx.logical_divide(gA, fx.make_layout(1, 1)) + b_div = fx.logical_divide(gB, fx.make_layout(1, 1)) + + # Two gather swizzles: the lo half covers tile rows [0, LDS_BLOCK_M), the hi half + # [LDS_BLOCK_M, BLOCK_M); each redirects its tile row through sorted_slot_ids. The two + # halves gather independent rows, so they need distinct loaders (a_g2s / a_g2s_hi). + if gather: + gl_off_a_lo = compute_global_gather_swizzle_bf16( + lane_id, wave_id, K, N_LDS_ROUNDS, sorted_res, sorted_row_base + ) + gl_off_a_hi = compute_global_gather_swizzle_bf16( + lane_id, wave_id, K, N_LDS_ROUNDS, sorted_res, sorted_row_base + fx.Int32(LDS_BLOCK_M) + ) + else: + # FC2 route-read: identity (row = route position), no sorted_slot_ids load. + gl_off_a_lo = compute_global_identity_swizzle_bf16( + lane_id, wave_id, K, N_LDS_ROUNDS, sorted_row_base + ) + gl_off_a_hi = compute_global_identity_swizzle_bf16( + lane_id, wave_id, K, N_LDS_ROUNDS, sorted_row_base + fx.Int32(LDS_BLOCK_M) + ) + gl_off_b = compute_global_swizzle_bf16(lane_id, wave_id, K, N_LDS_ROUNDS) + + mfma = Mfma32x32x16(N_TILES_A, N_TILES_B) + a_g2s = G2SLoader(a_div, gl_off_a_lo, N_LDS_STEPS_A, fx.BFloat16.ir_type, wave_id) + a_g2s_hi = G2SLoader(a_div, gl_off_a_hi, N_LDS_STEPS_A, fx.BFloat16.ir_type, wave_id) + b_g2s = G2SLoader(b_div, gl_off_b, N_LDS_STEPS_B, fx.BFloat16.ir_type, wave_id) + a_s2r = S2RLoaderBf16(wave_m, N_TILES_A) + b_s2r = S2RLoaderBf16(wave_n, N_TILES_B) + _out_ty = fx.Float16 if out_fp16 else fx.BFloat16 + store_c = StoreCBf16(C, c_m, c_n, _out_ty) + + dense_mma_pipeline_bf16( + lds, + a_g2s, + b_g2s, + a_s2r, + b_s2r, + mfma, + store_c, + A0_gl_offset, + A1_gl_offset, + B0_gl_offset, + B1_gl_offset, + BLOCK_K, + BLOCK_K, + fx.Int32(0), # store block_m: C is already rebased to this tile + block_n, + wave_m, + wave_n, + K, + BLOCK_M, + BLOCK_N, + nt_vmcnt, + a_g2s_hi=a_g2s_hi, + ) + + +@functools.lru_cache(maxsize=256) +def compile_grouped_gemm_gather_bf16( + K, + BLOCK_M=256, + BLOCK_N=256, + GROUP_M=4, + num_xcd=1, + nt_vmcnt=3, # gfx950 G2S LDS hazard: vmcnt>=4 races (nondeterministic); 3 is det + waves_per_eu=2, + agpr_alloc=0, + out_fp16=False, + gather=True, +): + """Compile (cached) the gathering grouped BF16 GEMM launcher for one ``(K, tile)`` combo. + + Grid is over-launched to the padded output pool; each block early-exits past the real tile + range (``num_tile_blocks``) or when ``expert_ids[block_m] < 0`` (PF padding tail). Mirrors + MegaMOE's ``compile_grouped_gemm_bf16`` (NT) exactly apart from the gather A fetch, the extra + ``SORTED`` argument, and the padding-block guard. Returns the flyc launch callable. + + The FC2 route-read (``index_a_by_route_pos``) case reuses this same gathering kernel with an + identity ``SORTED`` (``SORTED[s] = s``), so no separate no-gather tile is needed. + """ + SharedStorage = _make_shared_storage(BLOCK_M, BLOCK_N) + + @flyc.kernel(known_block_size=[512, 1, 1]) + def grouped_gemm_gather_k( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + TILE_TO_GROUP: fx.Tensor, + NUM_TILE_BLOCKS: fx.Int32, + SORTED: fx.Tensor, + A_ELEMS: fx.Int32, + c_n: fx.Int32, + ): + n_blocks = ceildiv(c_n, BLOCK_N) + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + group_res = create_buffer_resource(TILE_TO_GROUP, max_size=True) + sorted_res = create_buffer_resource(SORTED, max_size=True) + # Real (non-padding) BLOCK_M tile count as a scalar (host-known, capture-safe -- avoids a + # per-call device tensor that a HIP graph capture forbids). Host-known int scalar. + real_tiles = NUM_TILE_BLOCKS + # XCD-swizzle over the REAL tile range only (front-loaded); swizzling the full padded + # pool scatters real tiles -> ~2x slower. + real_grid = real_tiles * n_blocks + + pool_ptr_ty = PointerType.get( + elem_ty=fx.BFloat16.ir_type, address_space=AddressSpace.Global, alignment=16 + ) + # Flat 1D view of the whole received-token buffer, bounded to A_ELEMS elements. The + # gather offsets index this row-major buffer directly; the buffer resource clamps the + # padding sentinel (src_row == num_recv) to an OOB read of 0. (int32 element count: + # holds up to 2^31 elems; larger A needs a per-lane i64 rebase like Mega's fp8 path.) + a_base = fx.arith.ArithValue(arith.index_cast(fx.T.i64(), extract_base_index(A)), signed=True) + A_flat = fx.make_view(fx.inttoptr(pool_ptr_ty, a_base), fx.make_layout(A_ELEMS, 1)) + + def _emit(): + pid = xcd_remap_pid(fx.block_idx.x, real_grid, num_xcd) + num_pid_m = real_tiles + num_pid_in_group = GROUP_M * n_blocks + group_id = pid // num_pid_in_group + pid_in_group = pid % num_pid_in_group + first_pid_m = group_id * GROUP_M + remaining_m = num_pid_m - first_pid_m + group_size_m = arith.select(remaining_m < GROUP_M, remaining_m, fx.Int32(GROUP_M)) + block_m = first_pid_m + (pid_in_group % group_size_m) + block_n = pid_in_group // group_size_m + g_idx = buffer_load(group_res, block_m, vec_width=1, dtype=fx.T.i32()) + # PF padding blocks mark expert_ids=-1; skip the full Mega pipeline. + if g_idx >= fx.Int32(0): + gbase = g_idx * fx.Int32(K) * c_n + # Worst-case pool (cap*N > 2^31): rebase C per tile in int64, int32 in-resource + # offset. Mirrors the fused nt path. + c_byte_off = _i64(block_m * fx.Int32(BLOCK_M)) * _i64(c_n) * fx.Int64(2) + c_base = fx.arith.ArithValue(arith.index_cast(fx.T.i64(), extract_base_index(C)), signed=True) + C_tile = fx.make_view( + fx.inttoptr(pool_ptr_ty, c_base + c_byte_off), + fx.make_layout(fx.Int32(BLOCK_M) * c_n, 1), + ) + sorted_row_base = block_m * fx.Int32(BLOCK_M) + gemm_bf16_nt_gather_tile( + A_flat, + B, + C_tile, + fx.Int32(BLOCK_M), + c_n, + lds, + sorted_res, + sorted_row_base, + block_n, + K=K, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + out_fp16=out_fp16, + nt_vmcnt=nt_vmcnt, + b_group_base=gbase, + gather=gather, + ) + + if fx.block_idx.x < real_grid: + _emit() + + @flyc.jit + def launch(A, B, C, TILE_TO_GROUP, NUM_TILE_BLOCKS: fx.Int32, SORTED, A_ELEMS: fx.Int32, c_m: fx.Int32, c_n: fx.Int32, stream: fx.Stream): + grid_x = ceildiv(c_m, BLOCK_M) * ceildiv(c_n, BLOCK_N) + grouped_gemm_gather_k( + A, + B, + C, + TILE_TO_GROUP, + NUM_TILE_BLOCKS, + SORTED, + A_ELEMS, + c_n, + value_attrs=make_value_attrs(waves_per_eu, agpr_alloc, "512,512"), + ).launch(grid=(grid_x, 1, 1), block=(512, 1, 1), stream=stream) + + return launch + + +def grouped_gemm_gather_bf16( + A, # [num_recv, K] bf16 received-token activations (UNPERMUTED gather source) + weight, # [E, N, K] bf16 per-expert B (NT) + output, # [em_max, N] bf16 block-padded route-ordered C (in place) + expert_ids, # [num_m_blocks] i32 expert per BLOCK_M output block + num_tile_blocks, # int real BLOCK_M block count (host scalar; capture-safe) + sorted_slot_ids, # [em_max] i32 received-token row per padded slot (sentinel = num_recv) + *, + BLOCK_M=256, + BLOCK_N=256, + GROUP_M=4, + num_xcd=1, + nt_vmcnt=3, + waves_per_eu=2, + agpr_alloc=0, + gather=True, +): + """Host entry: grouped bf16 NT GEMM. With ``gather=True`` (FC1) ``C[pos] = A[SORTED[pos]] + @ B[expert]^T``; with ``gather=False`` (FC2 route-read) ``A`` is **block-padded route-ordered** + ``[em_max, K]`` read at the route slot (``sorted_slot_ids`` unused, may be a dummy). + + ``output`` is written in place over its full padded ``[em_max, N]`` extent (padding rows carry + dead values, ignored by downstream stages keyed on the same routing metadata). ``c_m`` is the + padded slot count (``output.shape[0]``); the grid self-bounds to ``num_tile_blocks``. + """ + assert A.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16 + assert output.dtype == torch.bfloat16 + c_m = int(output.shape[0]) + E, N, K = weight.shape + assert K == A.shape[1], f"weight K={K} != A K={A.shape[1]}" + assert output.shape[1] == N, f"output N={output.shape[1]} != weight N={N}" + + A = A.contiguous() + a_elems = int(A.numel()) + if a_elems >= (1 << 31): + raise ValueError( + f"A has {a_elems} elems (>= 2^31); the flat-view gather bound is int32. " + "Large-A support needs a per-lane i64 SRD rebase (TODO)." + ) + weight_flat = weight.reshape(E * N, K).contiguous().view(-1) + expert_ids_i32 = expert_ids.to(torch.int32) + if gather: + sorted_arg = sorted_slot_ids.to(torch.int32) + else: + # FC2 route-read: the kernel synthesizes the identity index on-device, so SORTED is + # unread. Reuse a live tensor as the (unused) arg to avoid a capture-illegal allocation. + sorted_arg = expert_ids_i32 + launch = compile_grouped_gemm_gather_bf16( + K=K, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + GROUP_M=GROUP_M, + num_xcd=num_xcd, + nt_vmcnt=int(nt_vmcnt), + waves_per_eu=int(waves_per_eu), + agpr_alloc=int(agpr_alloc), + gather=bool(gather), + ) + launch( + A, + weight_flat, + output, + expert_ids_i32, + int(num_tile_blocks), + sorted_arg, + a_elems, + c_m, + N, + stream=torch.cuda.current_stream(), + ) + return output diff --git a/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_wgrad.py b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_wgrad.py new file mode 100644 index 000000000..1275f060f --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_wgrad.py @@ -0,0 +1,671 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Permute-free MoE weight-gradient (wgrad) grouped GEMM in FlyDSL. + +Contract: the gradient operand is the *block-padded* ``[em_max, N]`` buffer +and ``SORTED`` maps each padded slot to a received-token row:: + + dW[e][n, k] = sum_{routed slot s of e, valid} grad[grad_base[e] + local(s), n] + * x[SORTED[s], k] + +where ``grad_base[e] = block_start[e] * block_size_m`` is the expert's block-padded first-slot +row, ``local(s)`` is the slot's offset within expert ``e``'s block-padded range, and padding +slots (``SORTED[s] == num_recv_tokens``) are masked to zero. The grad walk is a plain +contiguous row scan (no ``SORTED`` indirection); ``SORTED`` is only consulted for the ``x`` +gather token and the padding mask. + +The contraction tile is staged through LDS and transposed on-read: + + 1. Coalesced fill: each 32-slot contraction step loads ``grad[slot, n_feat]`` and + ``x[token(slot), k_feat]`` into LDS as ``[slot(row), feature(col)]`` tiles with + wide vector loads along the contiguous feature axis. + 2. Hardware transpose-read: ``ds_read_tr16_b64`` reads the ``[slot, feature]`` tile + transposed into the MFMA A/B fragment layout ``[feature, slot]`` -- so the token + slot becomes the matrix-core contraction axis with no VGPR shuffle and no strided + global gather. + +A workgroup of ``warps_n x warps_k`` warps computes one ``block_n x block_k`` output +tile. All warps **cooperatively fill** the shared LDS contraction tile once per step +(amortizing the global gather), then each warp owns a ``(block_n/warps_n) x +(block_k/warps_k)`` sub-tile of 16x16 MFMA atoms, transpose-reading its own feature +columns out of the shared tile. ``warps_n = warps_k = 1`` reduces to the single-warp v2. + +Fixed configuration (no runtime toggles): + - bf16 inputs, bf16 output (FC1: block-padded route-ordered grad + token-gathered ``x``) + - optional ``accumulate``: overwrite (default) or read-modify-write into ``dW`` + - DMA + XOR chunk swizzle fill, 3-stage LDS pipeline, LLVM DMA alias scopes +""" + +from __future__ import annotations + +import functools + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm as _llvm +from flydsl._mlir.dialects import scf +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import arith, buffer_ops, const_expr, gpu, ptrtoint, range_constexpr, rocdl, vector +from flydsl.expr.typing import T +from flydsl.runtime.device import get_rocm_arch +from flydsl.utils.smem_allocator import SmemAllocator + +from ..tensor_shim import ptr_rsrc + +__all__ = ["compile_moe_wgrad_v2", "WGRAD_BLOCK_M"] + +# MFMA atom dims (16x16x32 bf16). +WMMA_M = 16 # output rows -> N (grad feature) +WMMA_N = 16 # output cols -> K (x feature) +WMMA_K = 32 # contraction -> token-slot tile +C_FRAG = 4 # f32 values per lane for the C (dW) fragment +WARP_SIZE = 64 + +# Coalesced fill vector width (bf16 elems). 8 bf16 = 16 B = one global_load_dwordx4. +FILL_V = 8 + +WGRAD_BLOCK_M = 32 # contraction (slot) step; matches the align block_size +NUM_BUF = 3 # triple-buffered DMA pipeline (pipe_stages == 3) + + +@functools.lru_cache(maxsize=None) +def compile_moe_wgrad_v2( + *, + block_n: int = 64, + block_k: int = 64, + warps_n: int = 1, + warps_k: int = 1, + accumulate: bool = False, + swap_gather: bool = False, + out_dtype: str = "bf16", +): + # ``out_dtype`` selects the ``dW`` element type: ``"bf16"`` (default) truncates the f32 MFMA + # accumulator on store; ``"fp32"`` stores/reads-modify-writes the f32 accumulator directly. + # The latter lets an fp32 ``main_grad`` sink accumulate in-kernel (no bf16 scratch + fold). + if out_dtype not in ("bf16", "fp32"): + raise ValueError(f"moe_wgrad out_dtype must be 'bf16' or 'fp32', got {out_dtype!r}") + # ``swap_gather`` mirrors the two operands' gather/contiguous roles for the FC2 wgrad: + # FC1 walks ``GRAD`` (the [em_max, N] block-padded gradient) contiguously and token-gathers + # ``X`` (the [num_recv, K] activation) -> ``dW[E, N, K]``. FC2 instead token-gathers ``GRAD`` + # (the [num_recv, N] token-space dL/dFC2out) and walks ``X`` (the [em_max, K] block-padded + # activation) contiguously, so the same kernel emits the native ``dW2[E, H, F]`` layout with + # no transpose. Only the per-operand ``clamp_row`` and which operand's resource is bounded to + # ``[num_recv, feat]`` differ. + if block_n % (warps_n * WMMA_M) != 0 or block_k % (warps_k * WMMA_N) != 0: + raise ValueError("block_n/block_k must be multiples of warps_*16") + + n_threads = warps_n * warps_k * WARP_SIZE + if (WGRAD_BLOCK_M * block_n) % (n_threads * FILL_V) != 0: + raise ValueError("32*block_n must be a multiple of n_threads*FILL_V") + if (WGRAD_BLOCK_M * block_k) % (n_threads * FILL_V) != 0: + raise ValueError("32*block_k must be a multiple of n_threads*FILL_V") + + gpu_arch = get_rocm_arch() + WN = block_n // warps_n # per-warp grad-feature span + WK = block_k // warps_k # per-warp x-feature span + M_STEPS = WN // WMMA_M # grad-feature atoms per warp (MFMA-M) + N_STEPS = WK // WMMA_N # x-feature atoms per warp (MFMA-N) + NACC = M_STEPS * N_STEPS + + # DMA fill writes each lane's 16B contiguously into LDS (no per-row pad possible), + # so the swizzle path uses an un-padded stride and breaks bank conflicts with an + # XOR chunk-swizzle instead (see the fill + transpose-read below). + SG = block_n # grad LDS row stride (bf16 elems) + SX = block_k # x LDS row stride + # Swizzle granule = FILL_V bf16 (one 16B DMA unit). The XOR maps the feature chunk + # index with the slot so consecutive contraction slots land on distinct banks. + CPR_G_SWZ = block_n // FILL_V # feature chunks per grad row + CPR_X_SWZ = block_k // FILL_V # feature chunks per x row + + G_TILE_ELEMS = WGRAD_BLOCK_M * SG + X_TILE_ELEMS = WGRAD_BLOCK_M * SX + G_FILLS = (WGRAD_BLOCK_M * block_n) // (n_threads * FILL_V) + X_FILLS = (WGRAD_BLOCK_M * block_k) // (n_threads * FILL_V) + + out_is_f32 = out_dtype == "fp32" + KERNEL_NAME = ( + f"moe_wgrad_routelist_bf16_{block_n}x{block_k}_w{warps_n}x{warps_k}" + f"{'_o32' if out_is_f32 else ''}{'_acc' if accumulate else ''}" + f"{'_sg' if swap_gather else ''}_dsz_s3_dsa_v2" + ) + + # LDS allocation: NUM_BUF-buffered grad tile + x tile (2 bytes/bf16). + allocator = SmemAllocator(None, arch=gpu_arch, global_sym_name="smem") + g_lds_off = allocator._align(allocator.ptr, 16) + allocator.ptr = g_lds_off + G_TILE_ELEMS * 2 * NUM_BUF + x_lds_off = allocator._align(allocator.ptr, 16) + allocator.ptr = x_lds_off + X_TILE_ELEMS * 2 * NUM_BUF + # The DMA path backs LDS with a raw llvm addrspace(3) global (buffer_load_lds needs a + # real global for M0; the memref allocator does not work). Same offsets/size. + LDS_TOTAL_BYTES = allocator._align(allocator.ptr, 128) + LDS_SYM = KERNEL_NAME + "_lds" + + @flyc.kernel(known_block_size=[n_threads, 1, 1]) + def wgrad_kernel( + dW: fx.Pointer, # [E, N, K] bf16 output + X: fx.Pointer, # [num_recv_tokens, K] bf16 (received-token activations) + GRAD: fx.Pointer, # [em_max, N] bf16 (block-padded per-slot gradient) + SORTED: fx.Pointer, # [padded] i32 received-token row per slot (sentinel = num_recv_tokens) + BLOCK_START: fx.Pointer, # [E] i32 (block units) + BLOCKS_PER_EXPERT: fx.Pointer, # [E] i32 + GRAD_BASE: fx.Pointer, # [E] i32 (block-padded first-slot row = block_start[e] * block_size_m) + N: fx.Int32, + K: fx.Int32, + num_recv_tokens: fx.Int32, + ): + bf16 = T.bf16 + c0 = arith.constant(0, index=True) + + dW_rsrc = ptr_rsrc(dW) + x_rsrc = ptr_rsrc(X) + grad_rsrc = ptr_rsrc(GRAD) + sorted_rsrc = ptr_rsrc(SORTED) + bstart_rsrc = ptr_rsrc(BLOCK_START) + bpe_rsrc = ptr_rsrc(BLOCKS_PER_EXPERT) + gbase_rsrc = ptr_rsrc(GRAD_BASE) + + # DMA path: raw addrspace(3) global backs LDS; reads + DMA both GEP off it. + smem_raw_ptr = _llvm.mlir_addressof(ir.Type.parse("!llvm.ptr<3>"), LDS_SYM) + + tid = fx.Int32(gpu.thread_id("x")) + n_tile = fx.Int32(gpu.block_id("x")) # along N (grad feature) + k_tile = fx.Int32(gpu.block_id("y")) # along K (x feature) + expert = fx.Int32(gpu.block_id("z")) # expert id + + wid = tid // WARP_SIZE + lane = tid % WARP_SIZE + wn_id = wid // warps_k # warp row (grad feature) + wk_id = wid % warps_k # warp col (x feature) + lane_n = lane % 16 # MFMA C col (k_feat) + lane_m_base = lane // 16 # 0..3 + tr_k_group = (lane % 16) // 4 # 0..3 + tr_col_sub = lane % 4 # 0..3 + + warp_n_base = wn_id * fx.Int32(WN) # grad-feature col base of this warp + warp_k_base = wk_id * fx.Int32(WK) # x-feature col base of this warp + + # One alias scope per (operand, ring slot): the grad tile and the x tile each own + # NUM_BUF disjoint LDS byte ranges. A fill of slot r and a read of slot r share a + # scope so their RAW dependency survives; every other pair (different slot, or the + # other operand) is marked noalias, which is what lets SIInsertWaitcnts keep older + # in-flight fills streaming instead of draining vmcnt to 0 before each transpose read. + _ALIAS_DOMAIN = '#llvm.alias_scope_domain' + _SCOPE_IDS = tuple( + [f"g{r}" for r in range(NUM_BUF)] + [f"x{r}" for r in range(NUM_BUF)] + ) + + def _scope_attr(ids): + inner = ", ".join( + f'#llvm.alias_scope' + for sid in ids + ) + return ir.Attribute.parse(f"[{inner}]") + + _MY_SCOPE = {sid: _scope_attr((sid,)) for sid in _SCOPE_IDS} + _NOALIAS_SCOPE = { + sid: _scope_attr(tuple(o for o in _SCOPE_IDS if o != sid)) + for sid in _SCOPE_IDS + } + + def _g_sid(slot): + return f"g{slot % NUM_BUF}" + + def _x_sid(slot): + return f"x{slot % NUM_BUF}" + + def _scope_kw(sid): + """alias/noalias metadata kwargs for one (operand, ring slot), or {}.""" + if sid is None: + return {} + return { + "alias_scopes": _MY_SCOPE[sid], + "noalias_scopes": _NOALIAS_SCOPE[sid], + } + + N_idx = arith.index_cast(T.index, N) + K_idx = arith.index_cast(T.index, K) + nrecv_idx = arith.index_cast(T.index, num_recv_tokens) + + # DMA fill cannot mask padding slots in registers, so bound the *token-gathered* + # operand's resource to its real [num_recv, feat] extent: the sentinel token + # (== num_recv) and any pipeline overrun then read out-of-bounds -> hardware + # returns 0. A zeroed gathered column makes the padding slot's outer product + # ``gathered (x) 0 == 0``, so the paired contiguous-walk operand may safely read + # garbage (clamped to row 0) for those slots. FC1 gathers ``x`` (stride K); FC2 + # (``swap_gather``) gathers ``grad`` (stride N). + if swap_gather: + grad_addr_i64 = arith.index_cast(T.i64, ptrtoint(GRAD)) + grad_nrec_bytes = nrecv_idx * N_idx * arith.index(2) + grad_rsrc = buffer_ops.create_buffer_resource_from_addr( + grad_addr_i64, num_records_bytes=grad_nrec_bytes + ) + else: + x_addr_i64 = arith.index_cast(T.i64, ptrtoint(X)) + x_nrec_bytes = nrecv_idx * K_idx * arith.index(2) + x_rsrc = buffer_ops.create_buffer_resource_from_addr( + x_addr_i64, num_records_bytes=x_nrec_bytes + ) + + n_block_base = n_tile * block_n + k_block_base = k_tile * block_k + n_base_idx = arith.index_cast(T.index, n_block_base) + k_base_idx = arith.index_cast(T.index, k_block_base) + + # Per-expert routed-slot range. ``base_slot`` is the wgrad-align slot offset into + # ``SORTED`` (holds the received-token row for the ``x`` gather); ``grad_base_idx`` is + # expert ``e``'s block-padded first-slot row into the ``[em_max, N]`` grad buffer. + bstart = buffer_load_i32(bstart_rsrc, expert) + nblocks = buffer_load_i32(bpe_rsrc, expert) + gbase = buffer_load_i32(gbase_rsrc, expert) + base_slot = arith.index_cast(T.index, bstart) * arith.index(WGRAD_BLOCK_M) + num_slots = arith.index_cast(T.index, nblocks) * arith.index(WGRAD_BLOCK_M) + grad_base_idx = arith.index_cast(T.index, gbase) + + CPR_G = block_n // FILL_V + CPR_X = block_k // FILL_V + + def _tile_idx(i, cpr): + chunk = tid + fx.Int32(i * n_threads) + slot_i32 = chunk // fx.Int32(cpr) + feat_i32 = (chunk % fx.Int32(cpr)) * fx.Int32(FILL_V) + slot_idx = arith.index_cast(T.index, slot_i32) + feat_idx = arith.index_cast(T.index, feat_i32) + return slot_idx, feat_idx, slot_i32, feat_i32 + + # ``load_slot_ids`` issues the (indirect) ``sorted`` index loads. These are + # prefetched *two* steps ahead and carried across the loop as iter_args. + def load_slot_ids(s_base_idx): + g_ids = [ + buffer_load_i32_idx(sorted_rsrc, base_slot + s_base_idx + _tile_idx(i, CPR_G)[0]) + for i in range_constexpr(G_FILLS) + ] + x_ids = [ + buffer_load_i32_idx(sorted_rsrc, base_slot + s_base_idx + _tile_idx(i, CPR_X)[0]) + for i in range_constexpr(X_FILLS) + ] + return g_ids, x_ids + + def _dma_one(rsrc, ids, slot_base_idx, cpr, feat_base_idx, dim_idx, lds_off, + buf_byte, n_fills, clamp_row, sid=None): + # Issue the global->LDS DMA for one operand: each lane streams FILL_V bf16 + # from ``global[row, swizzled_feat]`` straight into contiguous LDS (no VGPR + # staging). The LDS destination base is wave-uniform (readfirstlane); the + # hardware spreads lane L to base + L*16B, reconstructing the contiguous + # ``phys_linear * FILL_V`` layout the swizzled transpose read expects. + for i in range_constexpr(n_fills): + phys = tid + fx.Int32(i * n_threads) + slot = phys // fx.Int32(cpr) + chunk = phys % fx.Int32(cpr) + slot_idx = arith.index_cast(T.index, slot) + token = arith.index_cast(T.index, ids[i]) + in_range = arith.cmpi( + arith.CmpIPredicate.ult, slot_base_idx + slot_idx, num_slots + ) + valid = arith.andi( + in_range, arith.cmpi(arith.CmpIPredicate.ult, token, nrecv_idx) + ) + if const_expr(clamp_row): + # grad: contiguous route walk; clamp overrun to row 0 for fault safety + # (its padding contribution is cancelled by the zeroed x column). + row_idx = valid.select(grad_base_idx + slot_base_idx + slot_idx, c0) + else: + # x: gather by received-token; sentinel/OOB row -> hardware 0. + row_idx = token + swz = slot & fx.Int32(cpr - 1) + glob_feat = (chunk ^ swz) * fx.Int32(FILL_V) + glob_feat_idx = arith.index_cast(T.index, glob_feat) + voff_elem = row_idx * dim_idx + feat_base_idx + glob_feat_idx + voff_byte = arith.index_cast(T.i32, voff_elem * arith.index(2)) + lds_perlane = fx.Int32(lds_off) + buf_byte + phys * fx.Int32(FILL_V * 2) + lds_base = rocdl.readfirstlane(T.i32, lds_perlane) + rocdl.raw_ptr_buffer_load_lds( + rsrc, _gep_lds(smem_raw_ptr, lds_base), fx.Int32(FILL_V * 2), + voff_byte, fx.Int32(0), fx.Int32(0), fx.Int32(1), + **_scope_kw(sid), + ) + + def dma_fill(g_ids, x_ids, slot_base_idx, g_buf_byte, x_buf_byte, wbuf=None): + # FC1 (swap_gather=False): block-padded grad route walk (clamp_row=True) + token-gather + # ``x`` (clamp_row=False). FC2 (swap_gather=True) flips both: token-gather ``grad`` + # (clamp_row=False) + block-padded activation walk on ``x`` (clamp_row=True). + # ``wbuf`` is the Python ring-slot index this tile is being staged into (for alias scopes). + g_sid = _g_sid(wbuf) if wbuf is not None else None + x_sid = _x_sid(wbuf) if wbuf is not None else None + _dma_one( + grad_rsrc, g_ids, slot_base_idx, CPR_G_SWZ, n_base_idx, N_idx, + g_lds_off, g_buf_byte, G_FILLS, clamp_row=not swap_gather, sid=g_sid, + ) + _dma_one( + x_rsrc, x_ids, slot_base_idx, CPR_X_SWZ, k_base_idx, K_idx, + x_lds_off, x_buf_byte, X_FILLS, clamp_row=swap_gather, sid=x_sid, + ) + + def _dma_barrier(keep=0): + # DMA lands on vmcnt (global load); drain it before the workgroup barrier so + # all waves observe the freshly-staged LDS tile. ``keep`` leaves that many + # vmem ops in flight (graduated wait) -- for the 3-stage pipeline this keeps the + # just-issued tile's DMA streaming across the barrier so it overlaps the next + # iteration's MFMA too (only the tile read next is fully drained). lgkmcnt(0) + # retires this wave's ds_reads before the buffer is recycled NUM_BUF steps on. + asm = f"s_waitcnt vmcnt({keep}) lgkmcnt(0)\ns_barrier" + _llvm.InlineAsmOp( + res=None, operands_=[], asm_string=asm, + constraints="", has_side_effects=True, is_align_stack=False, + ) + + def compute(accs, g_buf_byte, x_buf_byte, dma_prefetch=None, slot=None): + # A fragments are read one-per-mi to keep VGPR pressure low. B fragments are + # loop-invariant across mi, fetched once and interleaved with the mi==0 MFMAs + # so their ds_read latency overlaps compute. The MFMA region runs at raised + # priority so the matrix pipe stays fed while reads are in flight. + g_read_sid = _g_sid(slot) if slot is not None else None + x_read_sid = _x_sid(slot) if slot is not None else None + + def read_a(mi): + return _tr_read_frag_swz( + smem_raw_ptr, g_lds_off, SG, CPR_G_SWZ, warp_n_base, mi * WMMA_M, + lane_m_base, tr_k_group, tr_col_sub, g_buf_byte, + alias_kw=_scope_kw(g_read_sid), + ) + + def read_b(nj): + return _tr_read_frag_swz( + smem_raw_ptr, x_lds_off, SX, CPR_X_SWZ, warp_k_base, nj * WMMA_N, + lane_m_base, tr_k_group, tr_col_sub, x_buf_byte, + alias_kw=_scope_kw(x_read_sid), + ) + + new_accs = [None] * NACC + b_frags = [None] * N_STEPS + + if const_expr(dma_prefetch is not None): + # Burst *all* current-tile transpose reads first, then issue the next-tile + # DMA. Keeping the reads ahead of the DMA in program order stops the compiler + # from planting an ``s_waitcnt vmcnt(0)`` in front of the first ds_read; + # the DMA then streams under the MFMA burst instead. + a_frags = [read_a(mi) for mi in range_constexpr(M_STEPS)] + for nj in range_constexpr(N_STEPS): + b_frags[nj] = read_b(nj) + rocdl.sched_barrier(0) + dma_prefetch() + rocdl.sched_barrier(0) + rocdl.s_setprio(1) + for mi in range_constexpr(M_STEPS): + for nj in range_constexpr(N_STEPS): + idx = mi * N_STEPS + nj + new_accs[idx] = rocdl.mfma_f32_16x16x32_bf16( + T.vec(C_FRAG, T.f32), + [a_frags[mi], b_frags[nj], accs[idx], 0, 0, 0], + ) + rocdl.sched_barrier(0) + rocdl.s_setprio(0) + return new_accs + + a_next = read_a(0) + b_frags[0] = read_b(0) # first B needed for the very first MFMA + rocdl.s_setprio(1) + for mi in range_constexpr(M_STEPS): + a_cur = a_next + if const_expr(mi + 1 < M_STEPS): + a_next = read_a(mi + 1) # prefetch next A while MFMA-ing current + for nj in range_constexpr(N_STEPS): + # Prefetch the next B fragment during the first mi only; its ds_read + # then overlaps this step's MFMA and all later mi reuse the resident frag. + if const_expr(mi == 0 and nj + 1 < N_STEPS): + b_frags[nj + 1] = read_b(nj + 1) + idx = mi * N_STEPS + nj + new_accs[idx] = rocdl.mfma_f32_16x16x32_bf16( + T.vec(C_FRAG, T.f32), [a_cur, b_frags[nj], accs[idx], 0, 0, 0] + ) + rocdl.sched_barrier(0) + rocdl.s_setprio(0) + return new_accs + + acc_init = [ + arith.constant_vector(0.0, T.vec(C_FRAG, T.f32)) + for _ in range(NACC) + ] + acc_iter_args = NACC + + # NUM_BUF-stage ping-pong pipeline: MFMA the current LDS buffer while the next + # step's global gather is in flight. Prefetch distance D == NUM_BUF - 1: stage D + # tiles up front and keep the most-recently-issued tile's DMA in flight across the + # barrier (graduated ``vmcnt``) so it overlaps two MFMA steps instead of one. + D = NUM_BUF - 1 + PER_TILE_DMA = G_FILLS + X_FILLS + # Graduated ``vmcnt`` kept in flight across each sub-tile's barrier: leaves the + # just-issued tile's DMA streaming (only the tile read next is fully drained). + KEEP = (NUM_BUF - 2) * PER_TILE_DMA + acc_ty = T.vec(C_FRAG, T.f32) + + # Static NUM_BUF-buffer ring, distance-D, slot-loop unrolled by NUM_BUF. Sub-tile j + # always reads buffer j and prefetches tile (base+j+D) into buffer (j+D) % NUM_BUF -- + # both Python constants, so every LDS byte offset is compile-time constant and the + # backend can prove read(buf j) never aliases the in-flight DMA write(buf (j+D)). + def g_byte(j): + return fx.Int32((j % NUM_BUF) * G_TILE_ELEMS * 2) + + def x_byte(j): + return fx.Int32((j % NUM_BUF) * X_TILE_ELEMS * 2) + + ID_SET = G_FILLS + X_FILLS + + # Slot ids are carried NUM_BUF-ahead so their ~500-cycle SORTED-load latency is + # retired before ``dma_fill`` consumes them. + def pack_idsets(sets): + out = [] + for g_ids, x_ids in sets: + out += g_ids + x_ids + return out + + def unpack_idsets(args, base): + sets = [] + for s in range_constexpr(NUM_BUF): + off = base + s * ID_SET + g_ids = [args[off + i] for i in range(G_FILLS)] + x_ids = [args[off + G_FILLS + i] for i in range(X_FILLS)] + sets.append((g_ids, x_ids)) + return sets + + def mk_prefetch(g_ids, x_ids, pf_slot, wbuf): + return lambda: dma_fill( + g_ids, x_ids, pf_slot, g_byte(wbuf), x_byte(wbuf), wbuf=wbuf + ) + + # Prologue: stage tiles 0..D-1 into buffers 0..D-1 (distance-D), then drain fully. + for j in range_constexpr(D): + j_slot = arith.index(j * WMMA_K) + gidj, xidj = load_slot_ids(j_slot) + dma_fill(gidj, xidj, j_slot, g_byte(j), x_byte(j), wbuf=j) + _dma_barrier() + + # Initial carried id-sets: iteration 0 prefetches tiles D..D+NUM_BUF-1. + init_sets = [ + load_slot_ids(arith.index((D + j) * WMMA_K)) + for j in range_constexpr(NUM_BUF) + ] + + ntiles = fx.Int32(arith.index_cast(T.i32, num_slots)) // fx.Int32(WMMA_K) + group_end = arith.index_cast( + T.index, (ntiles // fx.Int32(NUM_BUF)) * fx.Int32(NUM_BUF * WMMA_K) + ) + step_big = arith.index(NUM_BUF * WMMA_K) + + loop = scf.ForOp( + c0, group_end, step_big, iter_args=acc_init + pack_idsets(init_sets) + ) + with ir.InsertionPoint(loop.body): + s_base = loop.induction_variable + accs = [loop.body.arguments[1 + i] for i in range(NACC)] + cur_sets = unpack_idsets(loop.body.arguments, 1 + acc_iter_args) + + # sub-tile j: read buffer j (tile s_base+j), prefetch tile s_base+j+D into + # buffer (j+D) % NUM_BUF using its carried id-set. + for j in range_constexpr(NUM_BUF): + g_ids_j, x_ids_j = cur_sets[j] + pf_slot = s_base + arith.index((j + D) * WMMA_K) + accs = compute( + accs, g_byte(j), x_byte(j), + dma_prefetch=mk_prefetch(g_ids_j, x_ids_j, pf_slot, j + D), + slot=j, + ) + _dma_barrier(KEEP) + + nxt_sets = [ + load_slot_ids(s_base + arith.index((NUM_BUF + D + j) * WMMA_K)) + for j in range_constexpr(NUM_BUF) + ] + scf.YieldOp(accs + pack_idsets(nxt_sets)) + + accs = [loop.results[i] for i in range(NACC)] + + # Tail: rem in 0..NUM_BUF-1 leftover tiles, already prefetched into buffers + # 0..rem-1 by the last group. Drain any in-flight DMA, then consume read-only. + _dma_barrier() + rem = ntiles - (ntiles // fx.Int32(NUM_BUF)) * fx.Int32(NUM_BUF) + + def _tail(j, accs): + if const_expr(j >= NUM_BUF - 1): + return accs + has_j = arith.cmpi(arith.CmpIPredicate.ugt, rem, fx.Int32(j)) + tif = scf.IfOp(has_j, results_=[acc_ty] * NACC, has_else=True) + with ir.InsertionPoint(tif.then_block): + a2 = compute(accs, g_byte(j), x_byte(j), dma_prefetch=None, slot=j) + scf.YieldOp(_tail(j + 1, a2)) + with ir.InsertionPoint(tif.else_block): + scf.YieldOp(accs) + return [tif.results[i] for i in range(NACC)] + + accs = _tail(0, accs) + + # Epilogue: C[m=n_feat, n=k_feat], lane holds 4 rows. Each dW element is owned by + # exactly one workgroup (grid = N x K x E over disjoint output tiles), so when + # ``accumulate`` is set the read-modify-write into the destination is race-free. + # ``out_ty`` is the ``dW`` element type: bf16 truncates the f32 accumulator on store; + # fp32 stores it directly (and accumulates without the bf16 round-trip). + out_ty = T.f32 if out_is_f32 else bf16 + E_NK_row = arith.index_cast(T.index, expert) * N_idx * K_idx + for mi in range_constexpr(M_STEPS): + for nj in range_constexpr(N_STEPS): + acc_idx = mi * N_STEPS + nj + acc = accs[acc_idx] + c_n = k_block_base + warp_k_base + fx.Int32(nj * WMMA_N) + lane_n + c_n_idx = arith.index_cast(T.index, c_n) + k_ok = arith.cmpi(arith.CmpIPredicate.ult, c_n_idx, K_idx) + for ii in range_constexpr(C_FRAG): + n_out = ( + n_block_base + warp_n_base + fx.Int32(mi * WMMA_M) + + lane_m_base * C_FRAG + fx.Int32(ii) + ) + n_out_idx = arith.index_cast(T.index, n_out) + in_bounds = arith.andi( + arith.cmpi(arith.CmpIPredicate.ult, n_out_idx, N_idx), k_ok + ) + store_if = scf.IfOp(in_bounds, results_=[], has_else=False) + with ir.InsertionPoint(store_if.then_block): + val = vector.extract( + acc, static_position=[ii], dynamic_position=[] + ) # f32 accumulator + out_off = E_NK_row + n_out_idx * K_idx + c_n_idx + if const_expr(accumulate): + prev = buffer_ops.buffer_load( + dW_rsrc, out_off, vec_width=1, dtype=out_ty + ) + prev_f32 = prev if const_expr(out_is_f32) else arith.extf(T.f32, prev) + val = arith.addf(val, prev_f32) + store_val = val if const_expr(out_is_f32) else arith.truncf(bf16, val) + buffer_ops.buffer_store(store_val, dW_rsrc, out_off) + scf.YieldOp([]) + + @flyc.jit + def launch_wgrad( + dW: fx.Pointer, + X: fx.Pointer, + GRAD: fx.Pointer, + SORTED: fx.Pointer, + BLOCK_START: fx.Pointer, + BLOCKS_PER_EXPERT: fx.Pointer, + GRAD_BASE: fx.Pointer, + N: fx.Int32, + K: fx.Int32, + num_recv_tokens: fx.Int32, + num_experts: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + _llvm.GlobalOp( + global_type=ir.Type.parse(f"!llvm.array<{LDS_TOTAL_BYTES} x i8>"), + sym_name=LDS_SYM, + linkage=ir.Attribute.parse("#llvm.linkage"), + addr_space=3, + alignment=1024, + ) + gx = (N + block_n - 1) // block_n + gy = (K + block_k - 1) // block_k + gz = num_experts + wgrad_kernel._func.__name__ = KERNEL_NAME + wgrad_kernel( + dW, X, GRAD, SORTED, BLOCK_START, BLOCKS_PER_EXPERT, GRAD_BASE, + N, K, num_recv_tokens, + ).launch(grid=(gx, gy, gz), block=(n_threads, 1, 1), stream=stream) + + return launch_wgrad + + +def _gep_lds(base_ptr, byte_i32): + """GEP an LDS !llvm.ptr<3> by a runtime i8 byte offset off a real LDS base pointer. + + ``buffer_load_lds`` derives its M0 write base from the LDS pointer, which the backend + only lowers correctly when the pointer is a GEP off a genuine addrspace(3) global. + We route the reads through the same base so alias analysis ties the DMA writes to the + transpose reads. + """ + return _llvm.getelementptr( + ir.Type.parse("!llvm.ptr<3>"), rocdl._to_ir(base_ptr), [rocdl._to_ir(byte_i32)], + [-(2 ** 31)], T.i8, None, + ) + + +def _tr_read_frag_swz( + smem_base, lds_off, stride, cpr, warp_col_base, col_const, + lane_m_base, tr_k_group, tr_col_sub, buf_byte, alias_kw=None, +): + """Swizzled transpose-read for the DMA fill (un-padded LDS). + + The DMA writes each contraction tile contiguously (stride == feature span), so bank + conflicts on the transpose read are broken by an XOR *chunk* swizzle: the physical + feature chunk of logical ``(slot, feat)`` is ``(feat // FILL_V) XOR (slot & (cpr - 1))`` + -- the *same* map the fill applies to the global gather column, so the read lands on + exactly the element the DMA staged. + """ + col_run = warp_col_base + tr_col_sub * fx.Int32(4) + feat_log = col_run + fx.Int32(col_const) # logical feature column + chunk = feat_log // fx.Int32(FILL_V) + within = feat_log % fx.Int32(FILL_V) + row_lo = lane_m_base * fx.Int32(8) + tr_k_group # contraction slot (0..31) + + def _read(slot): + swz = slot & fx.Int32(cpr - 1) + phys_chunk = chunk ^ swz + phys_feat = phys_chunk * fx.Int32(FILL_V) + within + elem = slot * fx.Int32(stride) + phys_feat + byte = elem * fx.Int32(2) + fx.Int32(lds_off) + buf_byte + raw = rocdl.ds_read_tr16_b64( + T.vec(4, T.bf16), _gep_lds(smem_base, byte), **(alias_kw or {}) + ).result + return fx.Vector(raw, (4,), fx.BFloat16) + + lo = _read(row_lo) + hi = _read(row_lo + fx.Int32(4)) + return lo.shuffle(hi, [0, 1, 2, 3, 4, 5, 6, 7]) + + +def buffer_load_i32(rsrc, off_i32): + return buffer_ops.buffer_load(rsrc, off_i32, vec_width=1, dtype=T.i32) + + +def buffer_load_i32_idx(rsrc, off_idx): + return buffer_ops.buffer_load(rsrc, off_idx, vec_width=1, dtype=T.i32) diff --git a/transformer_engine/pytorch/flydsl_kernels/tensor_shim.py b/transformer_engine/pytorch/flydsl_kernels/tensor_shim.py new file mode 100644 index 000000000..ccea9cb4d --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/tensor_shim.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +import torch +import flydsl.compiler as flyc +from flydsl.expr import arith, buffer_ops, ptrtoint +from flydsl.expr.typing import T + + +def ptr_rsrc(ptr): + """Convert an fx.Pointer kernel arg to a buffer resource for buffer_load/store.""" + addr_i64 = arith.index_cast(T.i64, ptrtoint(ptr)) + return buffer_ops.create_buffer_resource_from_addr(addr_i64) + + +def ptr_arg(t: torch.Tensor): + """Wrap a torch.Tensor as an fx.Pointer (PointerJitArg) for kernel launch.""" + import flydsl.expr as fx + + type_name = type(t).__name__ + module_name = type(t).__module__ + if type_name == "FakeTensor" or "fake_tensor" in module_name: + return flyc.from_c_void_p(fx.Uint8, 0) + return flyc.from_c_void_p(fx.Uint8, t.data_ptr()) + + +def _run_compiled(exe, *args): + """First call: ``flyc.compile(exe, *args)`` compiles **and** executes the kernel. + Subsequent calls: fast dispatch via the cached ``CompiledFunction``. + """ + cf = getattr(exe, "_cf", None) + if cf is None: + cf = flyc.compile(exe, *args) + exe._cf = cf + else: + cf(*args) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index f534da5c3..d3b7df98c 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -70,6 +70,11 @@ if IS_HIP_EXTENSION: from transformer_engine.pytorch.triton_kernels.grouped_gemm import general_grouped_gemm_triton + from transformer_engine.pytorch.moe import ( + is_permute_free_grouped_gemm_enabled, + permute_free_grouped_gemm_forward, + permute_free_grouped_gemm_backward, + ) import os __all__ = ["GroupedLinear"] @@ -405,6 +410,7 @@ def forward( ctx, inp: torch.Tensor, m_splits: torch.Tensor, + dispatched_probs: Optional[torch.Tensor], non_tensor_args: Tuple, *weights_and_biases, ) -> Tuple[torch.Tensor, list]: @@ -437,7 +443,12 @@ def forward( m_splits_tensor, actual_m_splits, unpad_output, + routing_metadata, + grouped_weight_param, ) = non_tensor_args + # The gated-activation fusion hint rides on the routing metadata (single source of + # truth); ``None`` on every non-permute-free / non-FC1 path. + perm_free_activation = getattr(routing_metadata, "activation", None) if fp8: backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override else: @@ -448,11 +459,34 @@ def forward( # Check if Triton kernel should be used use_grouped_gemm_triton = IS_HIP_EXTENSION and os.getenv("NVTE_USE_GROUPED_GEMM_TRITON", "0") == "1" and not fp8 and not fuse_wgrad_accumulation + # Permute-free gather GEMM (fwd + dgrad gather, wgrad on default path). + # Training is supported: fp8 / bias are excluded. ``fuse_wgrad_accumulation`` is only + # supported together with a single grouped weight param, where the permute-free backward + # accumulates wgrad directly into the grouped param's ``main_grad``. + use_perm_free_grouped_gemm = ( + IS_HIP_EXTENSION + and is_permute_free_grouped_gemm_enabled() + and routing_metadata is not None + and not fp8 + and (not fuse_wgrad_accumulation or grouped_weight_param is not None) + and activation_dtype == torch.bfloat16 + and not use_bias + ) + if use_perm_free_grouped_gemm and use_grouped_gemm_triton: + raise RuntimeError( + "NVTE_PERMUTE_FREE_GROUPED_GEMM and NVTE_USE_GROUPED_GEMM_TRITON cannot both be enabled." + ) + num_gemms = len(m_splits) weights = weights_and_biases[:num_gemms] biases = weights_and_biases[num_gemms:] device = inp.device - weight_requires_grad = weights[0].requires_grad + # Grouped weights expose detached per-expert views; the grouped param carries requires_grad. + weight_requires_grad = ( + grouped_weight_param.requires_grad + if grouped_weight_param is not None + else weights[0].requires_grad + ) # Configure quantizers if save_original_input and isinstance(input_quantizers[0], Float8Quantizer): @@ -503,7 +537,20 @@ def forward( # Initialize input tensors in_features = weights[0].size(-1) - if inp.size(-1) != in_features: + perm_free_route_space = ( + getattr(routing_metadata, "route_space", False) if routing_metadata is not None else False + ) + # FC2 with gated activation consumes FC1's raw 2F [gate|up] buffer (width 2F). + expect_in_features = ( + 2 * in_features + if ( + use_perm_free_grouped_gemm + and perm_free_route_space + and perm_free_activation is not None + ) + else in_features + ) + if inp.size(-1) != expect_in_features: raise ValueError( f"Input tensor (shape={tuple(inp.size())}) is not compatible with " f"weight tensor (shape={tuple(weights[0].size())})" @@ -546,7 +593,9 @@ def forward( # Convert splits to list of ints for compatibility with split functions m_splits = m_splits.tolist() - inp_view = inp.reshape(-1, in_features) + # Use ``expect_in_features`` (2F on the gated FC2 route-space path, else in_features) so a + # raw 2F [gate|up] preact is not mangled into F-wide rows by the reshape. + inp_view = inp.reshape(-1, expect_in_features) inputmats: list if fp8 and not debug: # Disable bulk allocation when CPU offloading is active: offloading skips small @@ -565,6 +614,8 @@ def forward( ) elif use_grouped_gemm_triton: inputmats = [cast_if_needed(inp_view, activation_dtype)] + elif use_perm_free_grouped_gemm: + inputmats = [cast_if_needed(inp_view, activation_dtype)] else: inputmats = torch.split(cast_if_needed(inp_view, activation_dtype), m_splits) @@ -597,12 +648,15 @@ def forward( if fp8 and activation_dtype == torch.float32: bias_dtype = torch.bfloat16 # FP8 GEMM only supports BF16/FP16 bias biases = [cast_if_needed(bias, bias_dtype) for bias in biases] if use_bias else biases - # Initialize output tensor - out = torch.empty( - [sum(m_splits), weights_fp8[0].size(0)], - dtype=activation_dtype, - device=device, - ) + # Initialize output tensor. The permute-free path allocates its own worst-case padded + # [T * min(topk, E), out_features] output inside permute_free_grouped_gemm_bf16 (valid rows are + # the dense route range [0, num_routes); the tail is inert zero padding). + if not use_perm_free_grouped_gemm: + out = torch.empty( + [sum(m_splits), weights_fp8[0].size(0)], + dtype=activation_dtype, + device=device, + ) # Choose whether to use split accumulator use_split_accumulator = _2X_ACC_FPROP @@ -612,25 +666,37 @@ def forward( use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator # Perform GEMM - if use_grouped_gemm_triton: + if use_perm_free_grouped_gemm: + # FC1 emits raw 2F [gate|up]; the ``activation`` hint on the metadata is consumed on + # FC2, which applies the gated activation in a standalone pass and then runs a plain + # GEMM (the fused-prologue path regressed throughput). Route probs ride with FC2 too. + out = permute_free_grouped_gemm_forward( + inputmats[0], + weights_fp8, + routing_metadata, + activation=perm_free_activation if perm_free_route_space else None, + dispatched_probs=dispatched_probs if perm_free_route_space else None, + ) + elif 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, - ) + if not use_perm_free_grouped_gemm: + 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 @@ -681,21 +747,51 @@ def forward( if backward_override == "high_precision" and inp.requires_grad else [None] * num_gemms ) + # Permute-free FC2 backward needs the route probs when the forward fused them. tensors_to_save, tensor_objects = prepare_for_saving( *inputmats, *weights_fp8, *saved_weights, *biases, + dispatched_probs, ) ctx.save_for_backward(*tensors_to_save) ctx.tensor_objects = tensor_objects + ctx.perm_free_fc2_activation = ( + perm_free_activation + if ( + use_perm_free_grouped_gemm + and getattr(routing_metadata, "route_space", False) + ) + else None + ) ctx.grad_input_quantizers = grad_input_quantizers ctx.grad_output_quantizers = grad_output_quantizers ctx.grad_weight_quantizers = grad_weight_quantizers - ctx.weights_requires_grad = weights[0].requires_grad - if fuse_wgrad_accumulation and ctx.weights_requires_grad: + ctx.weights_requires_grad = weight_requires_grad + # Permute-free backward routes the whole [E, out, in] wgrad straight to the grouped + # param, rather than through the detached per-expert views (which carry no grad edge). + # With fuse_wgrad_accumulation it accumulates into ``main_grad``; otherwise it + # accumulates into the grouped param's autograd ``.grad`` (standard accumulation). + ctx.perm_free_grouped = ( + use_perm_free_grouped_gemm + and grouped_weight_param is not None + and ctx.weights_requires_grad + ) + ctx.grouped_fuse_wgrad = fuse_wgrad_accumulation + if ctx.perm_free_grouped: + ctx.grouped_weight_ref = weakref.ref(grouped_weight_param) + if ctx.grouped_fuse_wgrad: + ctx.grouped_overwrite_main_grad = getattr( + grouped_weight_param, "overwrite_main_grad", False + ) + if hasattr(grouped_weight_param, "__fsdp_param__"): + ctx.grouped_main_grad_func = grouped_weight_param.get_main_grad + else: + ctx.grouped_main_grad_func = lambda: grouped_weight_param.main_grad + elif fuse_wgrad_accumulation and ctx.weights_requires_grad: # Keep weakrefs to weights to preserve attributes like main_grad # when we need to modify the weight python objects ctx.origin_weight_refs = [weakref.ref(w) for w in weights] @@ -742,6 +838,8 @@ def forward( ctx.input_quantizers = input_quantizers ctx.use_grouped_gemm_triton = use_grouped_gemm_triton ctx.num_input_tensors = len(inputmats) + ctx.use_perm_free_grouped_gemm = use_perm_free_grouped_gemm + ctx.routing_metadata = routing_metadata if use_perm_free_grouped_gemm else None # backward overrides if backward_override is not None: @@ -756,7 +854,10 @@ def forward( ctx.grad_output_quantizers = [None] * num_gemms ctx.reduce_and_update_bwd_fp8_tensors = False - # [*, in_features] -> [*, out_features] except first dimension changes for SP + # [*, in_features] -> [*, out_features], or worst-case padded [T * min(topk, E), out_features] + # (permute-free route-list path; valid rows are the dense route range [0, num_routes)). + if use_perm_free_grouped_gemm: + return out, new_workspaces return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces @staticmethod @@ -961,6 +1062,99 @@ def backward( weights = saved_tensors[num_inputs: num_inputs + N] saved_weights = saved_tensors[num_inputs + N : num_inputs + 2 * N] biases = saved_tensors[num_inputs + 2 * N : num_inputs + 3 * N] + dispatched_probs = saved_tensors[num_inputs + 3 * N] + + # Permute-free gather GEMM: both dgrad and wgrad gather along the + # contraction axis inside a single Triton kernel. + if getattr(ctx, "use_perm_free_grouped_gemm", False): + # Single grouped weight: the GEMM only ever sees detached per-expert views, so + # the whole [E, out, in] wgrad goes straight to the grouped param (the positional + # autograd return is dead). Point the kernel at the destination accumulator up + # front so it folds the wgrad in directly -- no scratch dW and no separate + # add/copy. Both FC1 and FC2 (route_space, via swap_gather) emit [E, out, in] + # directly into the buffer. + grouped = getattr(ctx, "perm_free_grouped", False) and ctx.weights_requires_grad + wgrad_out = None + wgrad_accumulate = False + if grouped: + grouped_weight = ctx.grouped_weight_ref() + assert ( + grouped_weight is not None + ), "grouped weight was removed before its wgrad could be applied" + gw_shape = tuple(grouped_weight.shape) + if ctx.grouped_fuse_wgrad: + wgrad_out = ctx.grouped_main_grad_func().view(gw_shape) + wgrad_accumulate = not ctx.grouped_overwrite_main_grad + else: + # Accumulate into the grouped param's autograd ``.grad``; allocate it on + # the first backward (overwrite) and add in place thereafter so grad + # accumulation across microbatches still works. + if grouped_weight.grad is None: + grouped_weight.grad = torch.empty( + gw_shape, dtype=grouped_weight.dtype, device=grad_output.device + ) + wgrad_accumulate = False + else: + wgrad_accumulate = True + wgrad_out = grouped_weight.grad + + # The wrapper decides FC1 vs FC2 dgrad/wgrad from the routing metadata. + pf_result = permute_free_grouped_gemm_backward( + grad_output, + routing=ctx.routing_metadata, + weights=weights, + num_gemms=ctx.num_gemms, + hidden_states=inputmats[0], + requires_dgrad=ctx.requires_dgrad, + requires_wgrad=ctx.weights_requires_grad, + dispatched_probs=dispatched_probs, + fc2_activation=getattr(ctx, "perm_free_fc2_activation", None), + wgrad_out=wgrad_out, + wgrad_accumulate=wgrad_accumulate, + ) + if getattr(ctx, "perm_free_grouped", False) and pf_result.wgrad_stacked is not None: + grouped_weight = ctx.grouped_weight_ref() + assert ( + grouped_weight is not None + ), "grouped weight was removed before its wgrad could be applied" + if pf_result.wgrad_applied: + # FC1: the kernel already folded the wgrad into main_grad / .grad. + if ctx.grouped_fuse_wgrad and hasattr( + grouped_weight, "grad_added_to_main_grad" + ): + grouped_weight.grad_added_to_main_grad = True + else: + # Fresh wgrad tensor: sink into main_grad / .grad. + dW = pf_result.wgrad_stacked + if ctx.grouped_fuse_wgrad: + main_grad = ctx.grouped_main_grad_func().view(dW.shape) + if ctx.grouped_overwrite_main_grad: + main_grad.copy_(dW) + else: + main_grad.add_(dW) + if hasattr(grouped_weight, "grad_added_to_main_grad"): + grouped_weight.grad_added_to_main_grad = True + elif grouped_weight.grad is None: + grouped_weight.grad = dW + else: + grouped_weight.grad.add_(dW) + wgrad_list = [None] * ctx.num_gemms + else: + # Fall back to returning positional per-expert wgrad (separate leaf params). + # Split the [E, out, in] gradient into zero-copy per-expert views. + wgrad_list = ( + list(pf_result.wgrad_stacked) + if pf_result.wgrad_stacked is not None + else [None] * ctx.num_gemms + ) + return ( + pf_result.dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, + None, # m_splits + pf_result.grad_probs, # dispatched_probs + None, # non_tensor_args + *wgrad_list, + *([None] * ctx.num_gemms), + ) # Restore from weakrefs to get original weight python objects # (preserves attributes like main_grad, grad_added_to_main_grad, etc.) @@ -1270,6 +1464,7 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): return ( dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, None, # m_splits + None, # dispatched_probs None, # non_tensor_args *wgrad_list, *grad_biases, @@ -1346,6 +1541,19 @@ class GroupedLinear(TransformerEngineBaseModule): GroupedLinear doesn't really handle the TP communications inside. The ``tp_size`` and ``parallel_mode`` are used to determine the shapes of weights and biases. The TP communication should be handled in the dispatch and combine stages of MoE models. + + Permute-free MoE (ROCm, bf16) + ----------------------------- + When ``NVTE_PERMUTE_FREE_GROUPED_GEMM=1``, pass a ``permute_free_metadata`` + (:class:`PermuteFreeMetadata`, carrying the boolean ``routing_map`` + ``[num_recv_tokens, num_local_experts]`` + a ``route_space`` direction) instead of + permuting activations before this module. The caller must skip ``moe_permute``. FC1 + (``route_space=False``) takes ``[num_recv_tokens, in_features]`` and produces the + worst-case padded ``[T * min(topk, E), out_features]`` route buffer (valid rows are the dense + route range ``[0, num_routes)``; the tail is inert zero padding); FC2 + (``route_space=True``) takes the route-ordered ``[T * min(topk, E), in_features]`` and fuses the + scatter back to token order, returning ``[num_recv_tokens, out_features]``. Requires + ``bias=False`` and bf16. The router-weight combine happens upstream (at the activation). """ def __init__( @@ -1767,6 +1975,8 @@ def forward( m_splits_tensor: Optional[torch.Tensor] = None, actual_m_splits: Optional[List[int]] = None, unpad_output: bool = False, + permute_free_metadata: Optional["PermuteFreeMetadata"] = None, + dispatched_probs: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: """ Apply the linear transformation to the input. @@ -1798,6 +2008,28 @@ def forward( When True, unpad the GEMM output from sum(m_splits) to sum(actual_m_splits) rows before returning. Used by the ROCm fused-pad-cast-transpose path; ignored on CUDA. + permute_free_metadata : PermuteFreeMetadata, optional + Route-list routing metadata (carrying the boolean ``routing_map`` + ``[num_recv_tokens, num_local_experts]`` plus a ``route_space`` + direction). When set with ``NVTE_PERMUTE_FREE_GROUPED_GEMM=1``, run + the route-list GEMM on unpermuted bf16 activations. ``route_space= + False`` (FC1) gathers per expert into the worst-case padded + ``[T * min(topk, E), out_features]`` route buffer (valid rows are the dense + route range ``[0, num_routes)``; the tail is inert zero padding); + ``route_space=True`` (FC2) reads route-ordered input, applies the + standalone gated activation (when ``activation`` is set), runs a plain + GEMM, and scatter-combines back to token order, returning + ``[num_recv_tokens, out_features]``. TE builds/caches the expert-sorted + alignment buffers on the metadata. The ``activation`` hint + (``"silu"`` / ``"gelu"``) is carried on this object for the FC2 + direction: FC1 (``route_space=False``) emits raw ``2F`` ``[gate | up]``; + FC2 recompute applies ``act(gate) * up`` (and optionally route probs) + before the GEMM. + dispatched_probs : torch.Tensor, optional + ``[num_recv_tokens, num_local_experts]`` gating probabilities. On the FC2 + permute-free path, multiplied into the gated activation during the + standalone recompute pass; its gradient is returned to the router through + autograd. Must be a leaf/differentiable tensor for training. """ debug = self.is_debug_iter() is_grad_enabled = torch.is_grad_enabled() @@ -1836,6 +2068,12 @@ def forward( # Preprocess input tensor if isinstance(inp, QuantizedTensorStorage): raise TypeError("GroupedLinear doesn't support input tensor in FP8.") + + # The permute-free path is driven entirely by the PermuteFreeMetadata (which carries + # the boolean routing_map + the route_space direction). It is built once by the + # caller and shared across FC1/FC2 to avoid a duplicate align build; TE builds/caches + # the align buffers on the object. + routing_metadata = permute_free_metadata inp = self.prepare_forward(inp, num_gemms=self.num_gemms) try: @@ -1897,9 +2135,13 @@ def forward( m_splits_tensor, actual_m_splits, unpad_output, + routing_metadata, + # Grouped weight param (single_grouped_weight): permute-free backward accumulates + # wgrad into its ``main_grad`` instead of returning it through the detached views. + getattr(self, "weight", None) if self.single_grouped_weight else None, ) out, new_workspaces = linear_fn( - *autograd_ctx, inp, m_splits, non_tensor_args, *weight_tensors, *bias_tensors + *autograd_ctx, inp, m_splits, dispatched_probs, non_tensor_args, *weight_tensors, *bias_tensors ) if cache_weight: diff --git a/transformer_engine/pytorch/moe/__init__.py b/transformer_engine/pytorch/moe/__init__.py new file mode 100644 index 000000000..e4313ae45 --- /dev/null +++ b/transformer_engine/pytorch/moe/__init__.py @@ -0,0 +1,41 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""Permute-free MoE grouped GEMM (FlyDSL) integration for PyTorch.""" + +from .permute_free_grouped_gemm import ( + MoERoutingMetadata, + PermuteFreeBackwardResult, + PermuteFreeMetadata, + is_permute_free_grouped_gemm_enabled, + permute_free_grouped_gemm_backward, + permute_free_grouped_gemm_bf16, + permute_free_grouped_gemm_bf16_dgrad, + permute_free_grouped_gemm_bf16_fc2, + permute_free_grouped_gemm_bf16_fc2_dgrad, + permute_free_grouped_gemm_bf16_fc2_wgrad, + permute_free_grouped_gemm_bf16_wgrad, + permute_free_grouped_gemm_forward, + permute_free_gated_act_bwd, + permute_free_gated_act_fwd, + prepare_moe_align, +) + +__all__ = [ + "MoERoutingMetadata", + "PermuteFreeBackwardResult", + "PermuteFreeMetadata", + "is_permute_free_grouped_gemm_enabled", + "permute_free_grouped_gemm_backward", + "permute_free_grouped_gemm_bf16", + "permute_free_grouped_gemm_bf16_dgrad", + "permute_free_grouped_gemm_bf16_fc2", + "permute_free_grouped_gemm_bf16_fc2_dgrad", + "permute_free_grouped_gemm_bf16_fc2_wgrad", + "permute_free_grouped_gemm_bf16_wgrad", + "permute_free_grouped_gemm_forward", + "permute_free_gated_act_bwd", + "permute_free_gated_act_fwd", + "prepare_moe_align", +] diff --git a/transformer_engine/pytorch/moe/moe_routing.py b/transformer_engine/pytorch/moe/moe_routing.py new file mode 100644 index 000000000..4f06127c9 --- /dev/null +++ b/transformer_engine/pytorch/moe/moe_routing.py @@ -0,0 +1,161 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""MoE routing metadata for permute-free grouped GEMM.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +import torch + + +@dataclass +class WgradAlign: + """Lazily-built block-``CONTRACT_M`` align buffers for the route-list wgrad kernel. + + Held in a *mutable* container so that a metadata and its ``dataclasses.replace`` copy + (e.g. the FC2 ``route_space=True`` view) share one instance by reference. The wgrad align + is a pure function of ``routing_map`` (identical for FC1 and FC2), but it is built lazily in + the backward -- after the ``replace`` that shares the fwd align has already happened. Sharing + this holder lets whichever backward runs first build the buffers once; the other reuses them, + avoiding a duplicate align build (and a second live copy). + """ + + sorted_slot_ids: Optional[torch.Tensor] = None + block_start: Optional[torch.Tensor] = None + blocks_per_expert: Optional[torch.Tensor] = None + block_size: Optional[int] = None + + +@dataclass +class MoERoutingMetadata: + """Routing tensors for the route-list gather-in-GEMM MoE path. + + Parameters + ---------- + routing_map: + Boolean mask ``[num_recv_tokens, num_local_experts]``, True where a received token + feeds a local expert. ``num_routes = routing_map.sum()``. + num_experts: + Local expert count on this rank. Optional -- defaults to ``routing_map.size(1)``. + topk: + Upstream router top-k -- the host-known maximum number of local experts any received + token can feed. Optional. When provided, ``prepare_moe_align`` tightens the (still + sync-free) static over-allocation of the block-padded route buffers from the dense + ``num_recv_tokens * num_experts`` bound down to + ``num_recv_tokens * min(topk, num_experts)``, shrinking ``em_max`` (and the zero-init + it costs) by up to ``num_experts / topk``. Leave ``None`` to keep the dense bound. + + The remaining fields are lazily populated by ``prepare_moe_align`` and cached for + FC1 fwd/dgrad reuse (all in the expert-sorted, block-padded route-list layout): + + The align buffers are built sync-free and over-allocated to static, shape-derived + upper bounds; the real extents are carried as device scalars so no host sync is + needed to construct them. + + sorted_slot_ids: + ``[T * min(topk, E)]`` received-token row to gather for each block-padded position (sentinel + ``num_recv_tokens`` for padding and for the over-allocated tail). + expert_ids: + ``[blocks_max]`` local expert owning each ``BLOCK_SIZE_M`` block (``-1`` past the + real block count; those blocks are never visited). + slot_expert_ids: + ``[T * min(topk, E)]`` per-slot local expert id, derived from ``expert_ids`` and + ``block_size_m`` (``expert_ids[slot // block_size_m]``). Populated by + ``prepare_moe_align`` for the standalone gated-activation kernels. + num_tokens_post_padded: + ``[1]`` device scalar = real ``em`` (block-padded route count). Bounds the kernel. + block_start: + ``[num_experts]`` per-expert first block index (block units). Expert ``e``'s block-padded + slots start at ``block_start[e] * block_size_m``. + token_routes / token_route_count: + Token -> its block-padded slot positions, built sync-free for the contention-free + gather-combine that replaces the atomic scatter in the token combine (FC2 fwd) and the + FC1 input-grad reduction (FC1 dgrad). ``token_routes`` is + ``[num_recv_tokens, min(topk, num_experts)]`` (int32); for token ``t`` the first + ``token_route_count[t]`` entries are its padded slots (expert-ascending), the rest + are unused padding. + block_size_m: + ``BLOCK_SIZE_M`` used to build the fwd/dgrad align buffers. + wgrad_align: + Shared, lazily-built block-``CONTRACT_M`` align buffers for the route-list wgrad kernel + (see :class:`WgradAlign`). The holder is shared by reference across a metadata and its + ``dataclasses.replace`` copy, so FC1 and FC2 build the (identical) wgrad align only once. + route_counts / route_within: + Cached block-size-independent scan (per-expert counts and within-expert ranks) + shared by the fwd/dgrad and wgrad align builds. + """ + + routing_map: torch.Tensor + num_experts: Optional[int] = None + topk: Optional[int] = None + sorted_slot_ids: Optional[torch.Tensor] = None + expert_ids: Optional[torch.Tensor] = None + slot_expert_ids: Optional[torch.Tensor] = None + num_tokens_post_padded: Optional[torch.Tensor] = None + block_start: Optional[torch.Tensor] = None + token_routes: Optional[torch.Tensor] = None + token_route_count: Optional[torch.Tensor] = None + block_size_m: Optional[int] = None + # Shared mutable holder so a metadata and its ``dataclasses.replace`` copy (FC2 route-space + # view) reuse one wgrad align build. ``replace`` copies this reference, and ``__post_init__`` + # only creates a fresh holder when the field is genuinely absent, preserving the shared one. + wgrad_align: Optional[WgradAlign] = None + # Block-size-independent scan (per-expert counts + within-expert ranks), shared by the + # fwd/dgrad and wgrad align builds so it is computed once per routing map. + route_counts: Optional[torch.Tensor] = None + route_within: Optional[torch.Tensor] = None + + def __post_init__(self): + # num_experts is redundant with the routing map width (one column per local + # expert), so the caller can pass just ``routing_map``. + if self.num_experts is None: + self.num_experts = int(self.routing_map.size(1)) + # Fresh holder only for a genuinely new metadata; a ``replace`` copy passes the existing + # (possibly already-populated) holder through so FC1 and FC2 share the same wgrad align. + if self.wgrad_align is None: + self.wgrad_align = WgradAlign() + + @property + def num_recv_tokens(self) -> int: + """Number of received tokens (rows of ``routing_map`` / the activation buffer).""" + return int(self.routing_map.size(0)) + + +@dataclass +class PermuteFreeMetadata(MoERoutingMetadata): + """Routing metadata for the permute-free grouped GEMM, tagged with a direction. + + Extends :class:`MoERoutingMetadata` + + - ``route_space=False`` (FC1): the input lives in **token-ordered** + ``[num_recv_tokens, in]``. The forward *gathers* per expert (``index_a_by_route_pos= + False``) into the **block-padded route-ordered** ``[em_max, out]`` buffer; the dgrad + combines the input gradient back to token rows (contention-free gather-combine). + - ``route_space=True`` (FC2): the input is already **block-padded route-ordered** + ``[em_max, in]`` (FC1's output). The forward reads by route slot + (``index_a_by_route_pos=True``) and combines each token's routes back to + **token-ordered** ``[num_recv_tokens, out]`` (contention-free gather-combine); the dgrad + gathers the token-ordered grad back into the block-padded route-ordered buffer. + + The align buffers are identical for both directions, so a single built metadata can be + reused for FC1 and FC2 (e.g. via ``dataclasses.replace(meta, route_space=True)``), + avoiding a duplicate align build. + + Activation hint (optional): + + activation: + Gated activation for the FC2 standalone pass -- ``"silu"`` or ``"gelu"``. + ``None`` leaves activation to the caller. FC1 emits raw ``2F``; this hint is + consumed on the FC2 direction (``route_space=True``) to run + :func:`permute_free_gated_act_fwd` before the plain FC2 GEMM. + + (The per-route gating probabilities are *not* carried here: they need a gradient, so they + are passed as a separate autograd tensor argument to the module rather than as metadata.) + """ + + route_space: bool = False + activation: Optional[str] = None \ No newline at end of file diff --git a/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py b/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py new file mode 100644 index 000000000..275d8dd75 --- /dev/null +++ b/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py @@ -0,0 +1,1132 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""Permute-free route-list grouped GEMM for MoE (bf16). + +Layout vocabulary: + * **token-ordered** ``[num_recv_tokens, F]`` — one row per received token. + * **dense route-ordered** ``[num_routes, F]`` — one row per ``(token, expert)`` route, no padding. + * **block-padded route-ordered** ``[em_max, F]`` — expert-sorted routes with per-expert + ``block_m`` padding (sync-free over-allocated upper bound). + +TE builds expert-sorted ``sorted_slot_ids`` (token row per route slot) plus per-expert +``block_start``, then runs gather-in-GEMM. FC1 fwd output is block-padded route-ordered +(valid route slots are ``[0, num_routes)``; tail is inert padding). FC1 dgrad returns +token-ordered ``dA = [num_recv_tokens, in_features]`` via gather-combine. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch + +from .moe_routing import MoERoutingMetadata, PermuteFreeMetadata +from .pf_helper_kernels import ( + fused_gated_act_prob_bwd, + fused_gated_act_prob_fwd, + route_gather_combine, + route_list_align, + route_list_scan, +) + +__all__ = [ + "MoERoutingMetadata", + "PermuteFreeMetadata", + "permute_free_grouped_gemm_bf16", + "permute_free_grouped_gemm_bf16_dgrad", + "permute_free_grouped_gemm_bf16_wgrad", + "permute_free_grouped_gemm_forward", + "permute_free_grouped_gemm_backward", + "PermuteFreeBackwardResult", + "prepare_moe_align", + "is_permute_free_grouped_gemm_enabled", +] + +_WGRAD_CONTRACT_M = 32 + +# Minimum gather/dgrad align block_m (128x256 MFMA floor). +_FLYDSL_MIN_BLOCK_M = 128 +_FLYDSL_FWD_BLOCK_M = 256 + + +def _get_flydsl_fwd(): + """Return the permute-free FlyDSL gather-GEMM launcher.""" + from .pf_fwd_wrapper import flydsl_moe_fwd + + return flydsl_moe_fwd + + +def _expand_expert_ids_per_slot( + expert_ids: torch.Tensor, + block_m: int, + routes_max: int, +) -> torch.Tensor: + """Broadcast block-level ``expert_ids`` to per-slot ids ``[routes_max]``. + + ``prepare_moe_align`` lays routes out expert-by-expert, padding each expert's count up to + a multiple of ``block_m``. ``expert_ids[b]`` records which expert owns M-block ``b`` (an + expert with more than ``block_m`` routes spans several consecutive blocks, all with the same + id). FlyDSL GEMM reads ``expert_ids`` at block granularity; the standalone gated-activation + kernels index row-by-row and need ``expert[slot]`` to look up ``dispatched_probs[token, e]``. + + This is a lookup, not ``expert = slot // block_m``: ``slot // block_m`` is the block index, + then ``expert_ids[block_idx]`` is the owner assigned during align. + """ + block_m = int(block_m) + pos = torch.arange(routes_max, device=expert_ids.device, dtype=torch.int64) + # Map each padded slot to its M-block; clamp covers the static over-allocated tail. + block_idx = (pos // block_m).clamp(max=expert_ids.numel() - 1) + return expert_ids[block_idx].to(torch.int32) + + +def _expert_per_route(routing: MoERoutingMetadata, routes_max: int) -> torch.Tensor: + """Per-route local expert id ``[routes_max]`` from the route-list block metadata.""" + if ( + routing.expert_ids is None + or routing.block_size_m is None + or routing.block_size_m <= 0 + ): + raise ValueError("_expert_per_route requires prepared routing align buffers.") + if ( + routing.slot_expert_ids is not None + and routing.slot_expert_ids.shape[0] == routes_max + ): + return routing.slot_expert_ids + slot_expert_ids = _expand_expert_ids_per_slot( + routing.expert_ids, routing.block_size_m, routes_max + ) + routing.slot_expert_ids = slot_expert_ids + return slot_expert_ids + + +def _get_flydsl_wgrad(): + """Return the FC1 FlyDSL wgrad launcher.""" + from .pf_wgrad_wrapper import flydsl_moe_wgrad_autotuned + + return flydsl_moe_wgrad_autotuned + + +def _pf_moe_fwd( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + routing: MoERoutingMetadata, + *, + num_recv_tokens: int, + block_m: int, + index_a_by_route_pos: bool = False, + dgrad: bool = False, +) -> None: + """Run the route-list gather-GEMM (forward or dgrad) via the FlyDSL permute-free kernels.""" + launch = _get_flydsl_fwd() + block_m = int(block_m) + launch( + A, + B, + C, + routing.sorted_slot_ids, + routing.expert_ids, + num_recv_tokens=num_recv_tokens, + block_m=block_m, + index_a_by_route_pos=index_a_by_route_pos, + dgrad=dgrad, + ) + + +def _fwd_align_block_size_m(num_tokens: int) -> int: + """Forward align ``block_m`` (128 or 256); shared by FC1 fwd/dgrad and FC2 fwd.""" + return _FLYDSL_FWD_BLOCK_M if num_tokens >= _FLYDSL_MIN_BLOCK_M else _FLYDSL_MIN_BLOCK_M + + +def _ensure_fwd_align( + routing: MoERoutingMetadata, + A: torch.Tensor, + B: torch.Tensor, +) -> int: + """Return ``routing.block_size_m``, building fwd align buffers when missing.""" + if routing.sorted_slot_ids is not None and routing.block_size_m is not None: + return int(routing.block_size_m) + block_m = _fwd_align_block_size_m(routing.num_recv_tokens) + prepare_moe_align(routing, block_m) + return block_m + + +def is_permute_free_grouped_gemm_enabled() -> bool: + from torch.utils.cpp_extension import IS_HIP_EXTENSION + + return IS_HIP_EXTENSION and os.getenv("NVTE_PERMUTE_FREE_GROUPED_GEMM", "0") == "1" + + +def moe_align_route_list( + routing_map: torch.Tensor, + *, + num_experts: int, + block_size: int, + scan=None, + topk: Optional[int] = None, + build_inverse_map: bool = False, +) -> Tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + Optional[torch.Tensor], + Optional[torch.Tensor], +]: + """Build the expert-sorted route-list buffers from a boolean routing map. + + Parameters + ---------- + routing_map: + Boolean ``[num_recv_tokens, num_experts]``; True where a received token feeds a + local expert. + num_experts: + Local expert count. + block_size: + ``BLOCK_SIZE_M`` (fwd/dgrad) or ``CONTRACT_M`` (wgrad) the layout is padded to. + scan: + Optional ``(counts, within)`` from :func:`route_list_scan`, shared across block + sizes so the block-independent scan is only paid once per routing map. + topk: + Host-known upper bound (the router top-k) on the number of experts any token routes + to. When provided, tightens the static over-allocation from ``T * num_experts`` to + ``T * min(topk, num_experts)`` (still sync-free). + + Returns + ------- + ``(sorted_slot_ids, expert_ids, num_tokens_post_padded, block_start, blocks_per_expert, + token_routes, token_route_count)``. Index tensors are ``int32``; + ``num_tokens_post_padded`` (``[1]``) is a device scalar. The last two are the token->routes + inverse map, ``None`` unless ``build_inverse_map``. + """ + return route_list_align( + routing_map, + num_experts=num_experts, + block_size=block_size, + scan=scan, + topk=topk, + build_inverse_map=build_inverse_map, + ) + + +def _ensure_route_scan(metadata: MoERoutingMetadata): + """Compute + cache the block-independent scan (counts, within) once per routing map.""" + if metadata.route_counts is None or metadata.route_within is None: + metadata.route_counts, metadata.route_within = route_list_scan( + metadata.routing_map, num_experts=metadata.num_experts + ) + return metadata.route_counts, metadata.route_within + + +def prepare_moe_align(metadata: MoERoutingMetadata, block_m: int) -> MoERoutingMetadata: + """Build and cache the fwd/dgrad route-list align buffers on ``metadata`` (sync-free).""" + if ( + metadata.sorted_slot_ids is not None + and metadata.block_size_m == block_m + and metadata.token_routes is not None + ): + if metadata.slot_expert_ids is None: + metadata.slot_expert_ids = _expand_expert_ids_per_slot( + metadata.expert_ids, + block_m, + metadata.sorted_slot_ids.shape[0], + ) + return metadata + + counts, within = _ensure_route_scan(metadata) + + ( + sorted_slot_ids, + expert_ids, + num_tokens_post_padded, + block_start, + _blocks_per_expert, + token_routes, + token_route_count, + ) = moe_align_route_list( + metadata.routing_map, + num_experts=metadata.num_experts, + block_size=block_m, + scan=(counts, within), + topk=metadata.topk, + build_inverse_map=True, + ) + metadata.sorted_slot_ids = sorted_slot_ids # [T * min(topk, E)] int32: block-padded slot -> token + metadata.expert_ids = expert_ids # [blocks_max] int32: expert owning each block (-1 past end) + metadata.slot_expert_ids = _expand_expert_ids_per_slot( + expert_ids, block_m, sorted_slot_ids.shape[0] + ) + metadata.num_tokens_post_padded = num_tokens_post_padded # [1] int32 device scalar: padded extent + metadata.block_start = block_start # [E] int32: first block index of each expert (block units) + metadata.block_size_m = block_m # int: BLOCK_SIZE_M the layout is padded to + metadata.token_routes = token_routes # [T, min(topk, E)] int32: token -> padded route slot indices + metadata.token_route_count = token_route_count # [T] int32: number of routes per token + return metadata + + +def _prepare_wgrad_align( + metadata: MoERoutingMetadata, contract_m: int +) -> MoERoutingMetadata: + """Build and cache the block-``contract_m`` align buffers for the wgrad kernel. + + The buffers live in the shared :class:`WgradAlign` holder, so the FC1 metadata and its + ``dataclasses.replace`` FC2 route-space copy reuse a single build (whichever backward runs + first fills the holder; the other hits the cache below). + """ + cache = metadata.wgrad_align + if cache.sorted_slot_ids is not None and cache.block_size == contract_m: + return metadata + + ( + sorted_slot_ids, + _expert_ids, + _num_tokens_post_padded, + block_start, + blocks_per_expert, + _token_routes, + _token_route_count, + ) = moe_align_route_list( + metadata.routing_map, + num_experts=metadata.num_experts, + block_size=contract_m, + scan=_ensure_route_scan(metadata), + topk=metadata.topk, + ) + cache.sorted_slot_ids = sorted_slot_ids + cache.block_start = block_start + cache.blocks_per_expert = blocks_per_expert + cache.block_size = contract_m + return metadata + + +def _try_view_grouped(weights: list[torch.Tensor]) -> Optional[torch.Tensor]: + """Return a zero-copy ``[E, out, in]`` view when the per-expert tensors are contiguous, + uniformly-strided, sequential slices of **one shared storage** (e.g. a single grouped weight + buffer). + + Returns ``None`` when the layout is not a single contiguous block, so the caller can fall + back to ``torch.stack``. Note that sequential ``data_ptr``s are not sufficient: separate + parameters can be allocated back-to-back yet own distinct storages, so ``as_strided`` from + the first tensor would overrun its storage. We therefore require a single shared storage. + """ + w0 = weights[0] + if not w0.is_contiguous(): + return None + out, in_ = w0.shape + stride = out * in_ + itemsize = w0.element_size() + base_ptr = w0.data_ptr() + storage_ptr = w0.untyped_storage().data_ptr() + for i, w in enumerate(weights): + if ( + w.shape != w0.shape + or w.dtype != w0.dtype + or not w.is_contiguous() + or w.untyped_storage().data_ptr() != storage_ptr + or w.data_ptr() != base_ptr + i * stride * itemsize + ): + return None + # Guard: the shared storage must actually span all E experts from w0's offset. + needed = (w0.storage_offset() + len(weights) * stride) * itemsize + if w0.untyped_storage().nbytes() < needed: + return None + # All experts form one contiguous [E, out, in] block starting at w0's storage offset. + return torch.as_strided(w0, (len(weights), out, in_), (stride, in_, 1)) + + +def _stack_expert_weights(weights: torch.Tensor | list[torch.Tensor]) -> torch.Tensor: + if isinstance(weights, torch.Tensor): + if weights.dim() != 3: + raise ValueError( + f"Stacked expert weights must be 3D [num_experts, out, in], got {weights.shape}." + ) + return weights + if not weights: + raise ValueError("At least one expert weight tensor is required.") + grouped = _try_view_grouped(weights) + return grouped if grouped is not None else torch.stack(weights, dim=0) + + +def permute_free_grouped_gemm_bf16( + hidden_states: torch.Tensor, + weights: torch.Tensor | list[torch.Tensor], + routing: MoERoutingMetadata, +) -> torch.Tensor: + """Route-list gather-in-GEMM (FC1 forward) for bf16 MoE. + + Computes ``C[slot] = A[sorted_slot_ids[slot]] @ W[expert(slot)]^T`` for each block-padded + slot, writing directly into the ``[em_max, out_features]`` block-padded output. Gated + activation is **not** fused here; FC1 emits the raw ``2F`` ``[gate | up]`` pre-activation + and the standalone gated-act helpers apply ``act(gate) * up [* prob]`` on FC2. + + Parameters + ---------- + hidden_states: + Received activations ``[num_recv_tokens, in_features]``, bf16, contiguous. + weights: + Expert weights ``[num_experts, out_features, in_features]`` or list of ``[out, in]``. + routing: + ``MoERoutingMetadata`` carrying (or able to build) the route-list align buffers. + + Returns + ------- + torch.Tensor + Block-padded ``[em_max, out_features]``, bf16 (expert-contiguous, each expert padded to + ``block_size``). The valid rows are the block-padded slots whose + ``sorted_slot_ids[slot] < num_recv_tokens``; the padding slots carry inert dead values. + Consumers locate each expert's rows via the routing metadata: + + - ``block_start[e]`` (block units): expert ``e``'s rows occupy padded slots + ``[block_start[e] * block_size, ...)``, with within-rank offset matching each route. + - ``num_tokens_post_padded = sum_e ceil(counts[e] / block_size) * block_size``: the + block-padded route count (each expert's row count rounded up to ``block_size``). + """ + if hidden_states.dtype != torch.bfloat16: + raise TypeError( + f"permute_free_grouped_gemm_bf16 requires bf16 input, got {hidden_states.dtype}." + ) + if not hidden_states.is_contiguous(): + hidden_states = hidden_states.contiguous() + + weights_stacked = _stack_expert_weights(weights) + if weights_stacked.stride(-1) != 1: + weights_stacked = weights_stacked.contiguous() + + num_recv_tokens, in_features = hidden_states.shape + num_experts, out_features, in_k = weights_stacked.shape + if in_k != in_features: + raise ValueError( + f"Weight in_features ({in_k}) does not match hidden_states ({in_features})." + ) + if num_experts != routing.num_experts: + raise ValueError( + f"num_experts mismatch: weights have {num_experts}, routing has {routing.num_experts}." + ) + + block_size_m = _ensure_fwd_align(routing, hidden_states, weights_stacked) + + # Worst-case (sync-free) allocation: size the output to the block-padded upper bound + # em_max = sorted_slot_ids.shape[0], which is derived purely from shapes (num_recv_tokens, + # num_experts, block_size) and is always >= num_routes. This avoids the device->host sync + # (.item()) that a dense route-ordered [num_routes, N] allocation would require. + em_max = routing.sorted_slot_ids.shape[0] + output = torch.empty( + (em_max, out_features), + dtype=torch.bfloat16, + device=hidden_states.device, + ) + + _pf_moe_fwd( + hidden_states, + weights_stacked, + output, + routing, + num_recv_tokens=num_recv_tokens, + block_m=block_size_m, + index_a_by_route_pos=False, + ) + return output + + +def permute_free_gated_act_bwd( + grad_output: torch.Tensor, + preact: torch.Tensor, + routing: MoERoutingMetadata, + *, + activation: str, + dispatched_probs: Optional[torch.Tensor] = None, + return_fc2_input: bool = False, +): + """Activation (+ route-prob) backward for the fused permute-free FC1 epilogue. + + Reconstructs the raw ``2F`` GEMM-output gradient ``dpre = [d_gate | d_up]`` (and, when the + forward fused the route prob, the prob gradient) from the saved ``[T * min(topk, E), 2F]`` + pre-activation. The returned ``dpre`` feeds the *unchanged* route-list dgrad / wgrad, so the + FC1 backward stays symmetric with the non-activation path; the caller owns the autograd + boundary and routes ``dprob`` back to ``dispatched_probs``. + + Parameters + ---------- + grad_output: + ``[T * min(topk, E), F]`` grad wrt the fused FC1 output (route/padded layout). + preact: + ``[T * min(topk, E), 2F]`` raw ``[gate | up]`` pre-activation saved by the forward. + dispatched_probs: + ``[num_recv_tokens, E]`` gating probs (or ``None`` for a silu/gelu-only fusion). + + ``return_fc2_input`` additionally re-materialises the ``F``-wide forward activation + ``fc2_input = act(gate)*up[*prob]`` inside this pass (one extra multiply + store from values + already in registers) so the FC2 wgrad can consume it directly instead of re-streaming the + ``2F`` preact through a separate recompute -- eliminating a full ``2F`` HBM read. + + Returns + ------- + ``(dpre, dprob)`` by default, or ``(dpre, dprob, fc2_input)`` when ``return_fc2_input`` -- + ``dpre`` is ``[T * min(topk, E), 2F]`` bf16; ``dprob`` matches ``dispatched_probs`` (or + ``None`` when no probs were fused); ``fc2_input`` is ``[em_max, F]`` bf16. + """ + # Block-padded canonical layout: grad_out (FC2 dgrad) and preact both live in the [em_max] + # padded slot order, and the kernel indexes grad_out/preact/dpre by the same row as + # token/expert -- so we must feed the padded slot arrays (sorted_slot_ids + slot expert), + # not the dense route arrays, or the padded valid slots beyond routes_max never get a dpre + # row (and each route pairs a correct token with the wrong padded grad row). + em_max = int(routing.sorted_slot_ids.shape[0]) + token = routing.sorted_slot_ids.to(torch.int32) + expert = _expert_per_route(routing, em_max) + # ``num_tokens_post_padded`` is a device scalar bounding the real padded extent, so the tail + # programs exit early (sync-free) instead of streaming the padding through HBM. + dpre, dprob, fc2_input = fused_gated_act_prob_bwd( + grad_output.contiguous(), + preact, + token, + expert, + num_recv_tokens=routing.num_recv_tokens, + activation=activation, + dispatched_probs=dispatched_probs, + num_routes_bound=routing.num_tokens_post_padded, + emit_act=return_fc2_input, + ) + if return_fc2_input: + return dpre, dprob, fc2_input + return dpre, dprob + + +def permute_free_gated_act_fwd( + preact: torch.Tensor, + routing: MoERoutingMetadata, + *, + activation: str, + dispatched_probs: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Apply the fused gated activation ``act(gate) * up * prob`` to the ``2F`` pre-activation. + + Maps the raw ``[T * min(topk, E), 2F]`` ``[gate | up]`` GEMM output to the ``F``-wide FC2 + input in the block-padded slot layout. Used both by the forward (to form the FC2 input) and, + on the backward, to reconstruct that ``F``-wide activation just-in-time for the FC2 wgrad -- + so the backward can checkpoint only the ``2F`` pre-activation and feed the *unchanged* + full-speed wgrad kernel a transient buffer (freed right after) instead of persisting the + activation across the fwd/bwd boundary. Uses the same per-route routing arrays as + :func:`permute_free_gated_act_bwd`. + + Returns + ------- + torch.Tensor + ``act``, shape ``[em_max, F]``, bf16 (block-padded slot layout, matching the FC1 + forward output and the layout the wgrad route-reads). + """ + # Block-padded canonical layout: emit one row per padded slot (keyed by sorted_slot_ids) so + # the activation lines up with the [em_max] grad the wgrad route-reads. Padding slots + # (token sentinel >= num_recv_tokens) are masked to zero by the kernel. + em_max = int(routing.sorted_slot_ids.shape[0]) + token = routing.sorted_slot_ids.to(torch.int32) + expert = _expert_per_route(routing, em_max) + return fused_gated_act_prob_fwd( + preact, + token, + expert, + num_recv_tokens=routing.num_recv_tokens, + activation=activation, + dispatched_probs=dispatched_probs, + num_routes_bound=routing.num_tokens_post_padded, + ) + + +def permute_free_grouped_gemm_bf16_dgrad( + grad_output: torch.Tensor, + weights: torch.Tensor | list[torch.Tensor], + routing: MoERoutingMetadata, +) -> torch.Tensor: + """Route-list gather-in-GEMM dgrad (FC1 backward wrt input). + + ``grad_output`` is the **dense route-ordered** gradient ``[num_routes, out_features]``. The + kernel contracts over ``out_features`` to produce per-route ``dX`` in a block-padded + route-ordered buffer, then gather-combine sums back to **token-ordered** ``dA``. + + Returns + ------- + torch.Tensor + ``dA``, shape ``[num_recv_tokens, in_features]``, bf16. + """ + if grad_output.dtype != torch.bfloat16: + raise TypeError( + f"permute_free_grouped_gemm_bf16_dgrad requires bf16 grad, got {grad_output.dtype}." + ) + if grad_output.dim() != 2: + raise ValueError( + f"grad_output must be dense route-ordered [num_routes, out_features], got {grad_output.shape}." + ) + if not grad_output.is_contiguous(): + grad_output = grad_output.contiguous() + + weights_stacked = _stack_expert_weights(weights) + num_experts, out_features, in_features = weights_stacked.shape + if grad_output.shape[-1] != out_features: + raise ValueError( + f"grad_output out_features ({grad_output.shape[-1]}) " + f"does not match weights ({out_features})." + ) + if num_experts != routing.num_experts: + raise ValueError( + f"num_experts mismatch: weights have {num_experts}, routing has {routing.num_experts}." + ) + + if weights_stacked.stride(-1) != 1: + weights_stacked = weights_stacked.contiguous() + + block_size_m = _ensure_fwd_align(routing, grad_output, weights_stacked) + + # Two-stage dgrad: (1) store per-route dX into block-padded route-ordered [em_max, in], + # then (2) token-parallel gather-combine back to token-ordered dA. Replaces atomic scatter. + em_max = routing.sorted_slot_ids.shape[0] + route_buf = torch.empty( + (em_max, in_features), + dtype=torch.bfloat16, + device=grad_output.device, + ) + _pf_moe_fwd( + grad_output, + weights_stacked, + route_buf, + routing, + num_recv_tokens=routing.num_recv_tokens, + block_m=block_size_m, + index_a_by_route_pos=True, + dgrad=True, + ) + return route_gather_combine( + route_buf, + routing.token_routes, + routing.token_route_count, + routing.num_recv_tokens, + out_dtype=torch.bfloat16, + ) + + +def permute_free_grouped_gemm_bf16_wgrad( + hidden_states: torch.Tensor, + grad_output: torch.Tensor, + weights_shape, + routing: MoERoutingMetadata, + *, + out: Optional[torch.Tensor] = None, + accumulate: bool = False, + swap_gather: bool = False, +) -> torch.Tensor: + """Route-list fused weight-gradient (FC1 backward wrt weights), FlyDSL-only. + + Computes ``dW[e] = sum_{slot in e} grad[slot]^T @ A[sorted_slot_ids[slot]]`` by gathering + the activation operand (received-token row) and reading the block-padded grad row (base + ``block_start[e] * block_size_m`` + within-rank) in the FlyDSL kernel. + + Parameters + ---------- + hidden_states: + Received activations ``[num_recv_tokens, in_features]``, bf16. + grad_output: + Compact per-route gradient ``[num_routes, out_features]``, bf16. + weights_shape: + ``(num_experts, out_features, in_features)``. + out: + Optional ``[E, out, in]`` destination (bf16 or fp32) to fold the wgrad into directly + (``+=`` if ``accumulate`` else ``=``). An fp32 sink (e.g. ``main_grad``) accumulates + in-kernel via the fp32-store kernel variant -- no bf16 scratch + separate fold. + accumulate: + With ``out``: add into it rather than overwrite. + swap_gather: + FC2 wgrad mode. Gathers ``grad_output`` (token-space ``[num_recv, out_features]``) and + walks ``hidden_states`` (block-padded ``[em_max, in_features]``) contiguously, emitting the + native ``[E, out_features, in_features]`` layout directly (no transpose). + + Returns + ------- + torch.Tensor + The wgrad ``[num_experts, out_features, in_features]`` -- ``out`` when supplied, else a + freshly allocated bf16 tensor. + """ + if hidden_states.dtype != torch.bfloat16 or grad_output.dtype != torch.bfloat16: + raise TypeError("permute_free_grouped_gemm_bf16_wgrad requires bf16 inputs.") + if grad_output.dim() != 2: + raise ValueError( + f"grad_output must be dense route-ordered [num_routes, out_features], got {grad_output.shape}." + ) + + num_experts, out_features, in_features = (int(v) for v in weights_shape) + if grad_output.shape[-1] != out_features: + raise ValueError( + f"grad_output out_features ({grad_output.shape[-1]}) " + f"does not match weights ({out_features})." + ) + if hidden_states.shape[-1] != in_features: + raise ValueError( + f"hidden_states in_features ({hidden_states.shape[-1]}) " + f"does not match weights ({in_features})." + ) + if num_experts != routing.num_experts: + raise ValueError( + f"num_experts mismatch: weights have {num_experts}, routing has {routing.num_experts}." + ) + + x = hidden_states + if not x.is_contiguous(): + x = x.contiguous() + if not grad_output.is_contiguous(): + grad_output = grad_output.contiguous() + + contract_m = _WGRAD_CONTRACT_M + # The grad operand lives in the forward's block-padded [em_max] slot layout (its row for + # route (e, w) is block_start[e]*block_size_m + w). Ensure the forward align is present so we + # can pass that padded per-expert base to the kernel; routes are contiguous within an expert, + # so base + within-rank indexes the correct padded grad row. + if routing.block_start is None or routing.block_size_m is None: + stub_w = torch.empty( + num_experts, out_features, in_features, + device=hidden_states.device, dtype=torch.bfloat16, + ) + _ensure_fwd_align(routing, hidden_states, stub_w) + routing = _prepare_wgrad_align(routing, contract_m) + grad_base = (routing.block_start.to(torch.int64) * int(routing.block_size_m)).to(torch.int32) + + flydsl_wgrad = _get_flydsl_wgrad() + + if out is not None: + assert tuple(out.shape) == (num_experts, out_features, in_features), ( + f"wgrad out shape {tuple(out.shape)} != {(num_experts, out_features, in_features)}" + ) + if out.dtype not in (torch.bfloat16, torch.float32): + raise TypeError( + f"FlyDSL wgrad requires bf16 or fp32 out, got {out.dtype}." + ) + flydsl_wgrad( + x, + grad_output, + out, + routing.wgrad_align.sorted_slot_ids, + routing.wgrad_align.block_start, + routing.wgrad_align.blocks_per_expert, + grad_base, + num_recv_tokens=routing.num_recv_tokens, + accumulate=bool(accumulate), + swap_gather=bool(swap_gather), + ) + return out + + dW = torch.zeros( + (num_experts, out_features, in_features), + dtype=torch.bfloat16, + device=hidden_states.device, + ) + flydsl_wgrad( + x, + grad_output, + dW, + routing.wgrad_align.sorted_slot_ids, + routing.wgrad_align.block_start, + routing.wgrad_align.blocks_per_expert, + grad_base, + num_recv_tokens=routing.num_recv_tokens, + accumulate=bool(accumulate), + swap_gather=bool(swap_gather), + ) + return dW + + +def permute_free_grouped_gemm_bf16_fc2( + fc2_input: torch.Tensor, + weights: torch.Tensor | list[torch.Tensor], + routing: MoERoutingMetadata, +) -> torch.Tensor: + """Route-list FC2 forward with gather-combine to token-ordered output (bf16 MoE). + + ``fc2_input`` is **block-padded route-ordered** ``F``-wide FC1 activation (or a transient + buffer rebuilt from the saved ``2F`` preact via :func:`permute_free_gated_act_fwd`). For each + route slot the kernel reads its row (``index_a_by_route_pos=True``) and writes + block-padded route-ordered per-route GEMM output; gather-combine then sums each token's + routes into **token-ordered** ``[num_recv_tokens, out]``. + + Parameters + ---------- + fc2_input: + Block-padded route-ordered activations ``[em_max, in_features(F)]``, bf16. + weights: + Expert weights ``[num_experts, out_features(H), in_features(F)]`` (W2) or list of ``[H, F]``. + routing: + ``MoERoutingMetadata`` carrying (or able to build) the route-list align buffers. + + Returns + ------- + torch.Tensor + ``[num_recv_tokens, out_features]``, bf16 (**token-ordered**); ready for the cross-rank + combine. No separate unpermute is needed. + """ + if fc2_input.dtype != torch.bfloat16: + raise TypeError( + f"permute_free_grouped_gemm_bf16_fc2 requires bf16 input, got {fc2_input.dtype}." + ) + if not fc2_input.is_contiguous(): + fc2_input = fc2_input.contiguous() + + weights_stacked = _stack_expert_weights(weights) + if weights_stacked.stride(-1) != 1: + weights_stacked = weights_stacked.contiguous() + + num_experts, out_features, in_features = weights_stacked.shape + if fc2_input.shape[-1] != in_features: + raise ValueError( + f"fc2_input in_features ({fc2_input.shape[-1]}) does not match weights " + f"({in_features})." + ) + if num_experts != routing.num_experts: + raise ValueError( + f"num_experts mismatch: weights have {num_experts}, routing has {routing.num_experts}." + ) + + block_size_m = _ensure_fwd_align(routing, fc2_input, weights_stacked) + + # Two-stage combine: (1) block-padded route-ordered per-route GEMM into [em_max, out], + # then (2) gather-combine to token-ordered output. Replaces atomic scatter-to-token. + em_max = routing.sorted_slot_ids.shape[0] + route_buf = torch.empty( + (em_max, out_features), + dtype=torch.bfloat16, + device=fc2_input.device, + ) + _pf_moe_fwd( + fc2_input, + weights_stacked, + route_buf, + routing, + num_recv_tokens=routing.num_recv_tokens, + block_m=block_size_m, + index_a_by_route_pos=True, + ) + return route_gather_combine( + route_buf, + routing.token_routes, + routing.token_route_count, + routing.num_recv_tokens, + out_dtype=torch.bfloat16, + ) + + +def permute_free_grouped_gemm_bf16_fc2_dgrad( + grad_output: torch.Tensor, + weights: torch.Tensor | list[torch.Tensor], + routing: MoERoutingMetadata, +) -> torch.Tensor: + """FC2 dgrad: gather **token-ordered** grad back into **block-padded route-ordered** dX. + + ``grad_output`` is **token-ordered** ``[num_recv_tokens, out_features]`` (grad of FC2's + gather-combine output). For each route slot the kernel gathers its token's grad row + (``index_a_by_route_pos=False``) and computes ``grad[token] @ W2[e]^T``, writing + block-padded route-ordered ``d(fc2_input)`` ``[em_max, in_features]``. + + Returns + ------- + torch.Tensor + ``[em_max, in_features]``, bf16 (**block-padded route-ordered**; valid slots are the + dense route range ``[0, num_routes)``). + """ + if grad_output.dtype != torch.bfloat16: + raise TypeError( + f"permute_free_grouped_gemm_bf16_fc2_dgrad requires bf16 grad, got " + f"{grad_output.dtype}." + ) + if grad_output.dim() != 2: + raise ValueError( + f"grad_output must be [num_recv_tokens, out_features], got {grad_output.shape}." + ) + if not grad_output.is_contiguous(): + grad_output = grad_output.contiguous() + + weights_stacked = _stack_expert_weights(weights) + num_experts, out_features, in_features = weights_stacked.shape + if grad_output.shape[-1] != out_features: + raise ValueError( + f"grad_output out_features ({grad_output.shape[-1]}) does not match weights " + f"({out_features})." + ) + if num_experts != routing.num_experts: + raise ValueError( + f"num_experts mismatch: weights have {num_experts}, routing has {routing.num_experts}." + ) + + if weights_stacked.stride(-1) != 1: + weights_stacked = weights_stacked.contiguous() + + block_size_m = _ensure_fwd_align(routing, grad_output, weights_stacked) + + # Block-padded route-ordered output; the gather-GEMM writes only the dense route range + # [0, num_routes) and never visits the tail (bounded by num_tokens_post_padded, sentinel-masked). + # Left uninitialized (torch.empty) to skip the em_max*N zero-init; consumers read only the + # dense route range via the routing metadata, so the garbage tail is never observed. + em_max = routing.sorted_slot_ids.shape[0] + dx = torch.empty( + (em_max, in_features), + dtype=torch.bfloat16, + device=grad_output.device, + ) + _pf_moe_fwd( + grad_output, + weights_stacked, + dx, + routing, + num_recv_tokens=grad_output.shape[0], + block_m=block_size_m, + index_a_by_route_pos=False, + dgrad=True, + ) + return dx + + +def permute_free_grouped_gemm_bf16_fc2_wgrad( + fc2_input: torch.Tensor, + grad_output: torch.Tensor, + weights_shape, + routing: MoERoutingMetadata, + *, + out: Optional[torch.Tensor] = None, + accumulate: bool = False, + preact: Optional[torch.Tensor] = None, + dispatched_probs: Optional[torch.Tensor] = None, + activation: Optional[str] = None, +) -> torch.Tensor: + """FC2 wgrad: ``dW2[e] = sum_{route in e} grad[token(route)]^T @ fc2_input[route]``. + + The two operands play mirror roles to FC1: ``grad_output`` is token-space and must be + token-gathered, ``fc2_input`` is route-ordered and read contiguously. Which one the kernel + gathers decides the output orientation: + + * **FC2 wgrad** uses ``swap_gather``: the FlyDSL kernel token-gathers ``grad_output`` and + walks the block-padded ``fc2_input`` contiguously, writing the native ``[E, H, F]`` layout + directly (no transpose, and ``out``/``accumulate`` fold straight into a bf16 destination). + + Parameters + ---------- + fc2_input: + Route-ordered activations ``[T * min(topk, E), in_features(F)]``, bf16. + grad_output: + Token-space gradient ``[num_recv_tokens, out_features(H)]``, bf16. + weights_shape: + ``(num_experts, out_features(H), in_features(F))`` -- the W2 shape. + out / accumulate: + Optional ``[E, H, F]`` destination; ``+=`` if ``accumulate`` else ``=``. Returned when given. + preact / dispatched_probs / activation: + Recompute-from-preact mode. When ``preact`` (``[T * min(topk, E), 2F]`` raw ``[gate | up]``) + is given, ``fc2_input`` is ignored and the ``F``-wide activation ``act(gate)*up*prob`` is + re-materialised here into a transient buffer (freed after the wgrad), so the backward can + checkpoint only the ``2F`` preact instead of the ``F``-wide activation. ``activation`` picks + the nonlinearity (``"silu"``/``"gelu"``); ``dispatched_probs`` (``[num_recv, E]``) is the + fused route-prob table (omit when the forward did not fuse probs). The wgrad kernel itself + is unchanged and runs at full stored-act speed. + """ + num_experts, out_features, in_features = (int(v) for v in weights_shape) + + # Recompute-from-preact: rebuild the transient activation for the (unchanged) wgrad. + if preact is not None: + if activation is None: + raise ValueError("permute_free FC2 wgrad recompute requires an activation.") + fc2_input = permute_free_gated_act_fwd( + preact, routing, activation=activation, dispatched_probs=dispatched_probs + ) + + # FC2 wgrad via swap_gather: gather grad_output [num_recv, H] (N=H) and walk the block-padded + # fc2_input [em_max, F] (K=F) contiguously -> native [E, H, F], no transpose. + weights_shape = (num_experts, out_features, in_features) # [E, H, F] + if out is not None and out.dtype in (torch.bfloat16, torch.float32): + # bf16 or fp32 sink: the kernel writes/accumulates straight into it (fp32 via the + # fp32-store variant), so no scratch dW and no separate fold pass. + return permute_free_grouped_gemm_bf16_wgrad( + fc2_input, + grad_output, + weights_shape, + routing, + out=out, + accumulate=accumulate, + swap_gather=True, + ) + dW = permute_free_grouped_gemm_bf16_wgrad( + fc2_input, + grad_output, + weights_shape, + routing, + swap_gather=True, + ) # [E, H, F] bf16 + if out is None: + return dW + # Exotic sink dtype the kernel can't write directly: fold the bf16 result here. + if accumulate: + out.add_(dW) + else: + out.copy_(dW) + return out + + +# --------------------------------------------------------------------------- +# Direction-aware dispatch: pick the FC1/FC2 fwd/bwd kernels from the routing +# metadata so callers (e.g. GroupedLinear) don't have to branch on route_space / +# activation themselves. +# --------------------------------------------------------------------------- +@dataclass +class PermuteFreeBackwardResult: + """Gradients from :func:`permute_free_grouped_gemm_backward`. + + ``dgrad`` / ``wgrad_stacked`` are ``None`` when the corresponding gradient was not requested. + ``wgrad_stacked`` is the weight gradient as a single ``[E, out, in]`` tensor (the natural + kernel output): accumulate it directly into a grouped param's ``main_grad``, or split it into + per-expert views (``list(wgrad_stacked)``) for the positional autograd return. ``grad_probs`` + is the route-prob gradient from the standalone FC2 gated-activation backward (``None`` + otherwise). + """ + + dgrad: Optional[torch.Tensor] = None + grad_probs: Optional[torch.Tensor] = None + wgrad_stacked: Optional[torch.Tensor] = None + # True when the wgrad was written directly into the caller-provided ``wgrad_out``; + # False when ``wgrad_stacked`` is a fresh tensor the caller still needs to sink. + wgrad_applied: bool = False + + +def permute_free_grouped_gemm_forward( + hidden_states: torch.Tensor, + weights: torch.Tensor | list[torch.Tensor], + routing: MoERoutingMetadata, + *, + activation: Optional[str] = None, + dispatched_probs: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Dispatch the permute-free grouped-GEMM forward from the routing direction + fusion hints. + + ``routing.route_space`` selects the direction: + + - ``True`` (FC2): route-ordered ``2F`` pre-activation (FC1 output). The gated activation + ``act(gate)*up[*prob]`` is applied in a standalone pass into an ``F``-wide transient, then + a plain gather-GEMM + gather-combine -> ``[num_recv, out]``. + - ``False`` (FC1): gather-in-GEMM -> padded ``[em_max, out]`` raw ``[gate | up]`` (``2F``). + """ + if getattr(routing, "route_space", False): + fc2_input = hidden_states + if activation is not None: + fc2_input = permute_free_gated_act_fwd( + hidden_states, + routing, + activation=activation, + dispatched_probs=dispatched_probs, + ) + return permute_free_grouped_gemm_bf16_fc2(fc2_input, weights, routing) + return permute_free_grouped_gemm_bf16(hidden_states, weights, routing) + + +def permute_free_grouped_gemm_backward( + grad_output: torch.Tensor, + *, + routing: MoERoutingMetadata, + weights: list[torch.Tensor], + num_gemms: int, + hidden_states: Optional[torch.Tensor] = None, + requires_dgrad: bool = False, + requires_wgrad: bool = False, + dispatched_probs: Optional[torch.Tensor] = None, + fc2_activation: Optional[str] = None, + wgrad_out: Optional[torch.Tensor] = None, + wgrad_accumulate: bool = False, +) -> PermuteFreeBackwardResult: + """Dispatch the permute-free grouped-GEMM backward (mirror of the forward dispatch). + + ``routing.route_space`` picks the FC2 vs FC1 dgrad/wgrad kernels. On the FC2 path, + ``fc2_activation`` runs the standalone gated-activation backward after ``fc2_dgrad``; + ``hidden_states`` is the saved ``2F`` pre-activation for act-bwd / wgrad recompute. + + ``wgrad_out`` (optional ``[E, out, in]`` accumulator) folds the wgrad straight into the + caller's buffer (``+=`` if ``wgrad_accumulate`` else ``=``) instead of returning a fresh + tensor. When used, ``result.wgrad_applied`` is ``True`` and ``result.wgrad_stacked`` is + that same buffer. + """ + grad_output = grad_output.contiguous() + route_space = getattr(routing, "route_space", False) + dgrad = None + grad_probs = None + wgrad_stacked = None + wgrad_applied = False + + if route_space: + # FC2: grad is token-space [num_recv, out]. dgrad gathers back to the route buffer, + # then the standalone gated-activation backward maps dL/dF -> dL/d(2F) (+ dprob). + recompute = fc2_activation is not None and hidden_states is not None + # When both grads are wanted with a gated activation, fuse the wgrad's activation + # recompute into the act-bwd (which is already streaming the 2F preact) so the wgrad + # can reuse the F-wide fc2_input instead of re-reading the 2F preact a second time. + fused_fc2_input = None + if requires_dgrad: + weights_stacked = _stack_expert_weights(weights) # [E, H, F], zero-copy when grouped + dgrad_f = permute_free_grouped_gemm_bf16_fc2_dgrad( + grad_output, weights_stacked, routing + ) + if recompute: + if requires_wgrad: + dgrad, grad_probs, fused_fc2_input = permute_free_gated_act_bwd( + dgrad_f, + hidden_states, + routing, + activation=fc2_activation, + dispatched_probs=dispatched_probs, + return_fc2_input=True, + ) + else: + dgrad, grad_probs = permute_free_gated_act_bwd( + dgrad_f, + hidden_states, + routing, + activation=fc2_activation, + dispatched_probs=dispatched_probs, + ) + if not (dispatched_probs is not None and dispatched_probs.requires_grad): + grad_probs = None + else: + dgrad = dgrad_f + if requires_wgrad: + weights_shape = (num_gemms, weights[0].size(0), weights[0].size(1)) + if fused_fc2_input is not None: + # Reuse the F-wide activation emitted by the act-bwd above: a plain stored-act + # wgrad, no second pass over the 2F preact. + dW = permute_free_grouped_gemm_bf16_fc2_wgrad( + fused_fc2_input, grad_output, weights_shape, routing, + out=wgrad_out, accumulate=wgrad_accumulate, + ) # [E, H, F] + else: + dW = permute_free_grouped_gemm_bf16_fc2_wgrad( + hidden_states, grad_output, weights_shape, routing, + out=wgrad_out, accumulate=wgrad_accumulate, + preact=hidden_states if recompute else None, + dispatched_probs=dispatched_probs if recompute else None, + activation=fc2_activation if recompute else None, + ) # [E, H, F] + wgrad_stacked = dW + wgrad_applied = wgrad_out is not None + else: + # FC1: grad is the padded [T * min(topk, E), 2F] route buffer (from FC2 act-bwd); dgrad + # scatters the input grad back to received-token rows. No FC1 activation backward. + if requires_dgrad: + weights_stacked = _stack_expert_weights(weights) # [E, N, H], zero-copy when grouped + dgrad = permute_free_grouped_gemm_bf16_dgrad(grad_output, weights_stacked, routing) + if requires_wgrad: + weights_shape = (num_gemms, weights[0].size(0), weights[0].size(1)) + # bf16 or fp32 ``main_grad`` sink both accumulate in-kernel (fp32 via the fp32-store + # variant); no scratch dW + separate fold. + dW = permute_free_grouped_gemm_bf16_wgrad( + hidden_states, grad_output, weights_shape, routing, + out=wgrad_out, accumulate=wgrad_accumulate, + ) # [E, N, H] + wgrad_stacked = dW + wgrad_applied = wgrad_out is not None + + return PermuteFreeBackwardResult( + dgrad=dgrad, grad_probs=grad_probs, wgrad_stacked=wgrad_stacked, + wgrad_applied=wgrad_applied, + ) diff --git a/transformer_engine/pytorch/moe/pf_fwd_wrapper.py b/transformer_engine/pytorch/moe/pf_fwd_wrapper.py new file mode 100644 index 000000000..ca26af741 --- /dev/null +++ b/transformer_engine/pytorch/moe/pf_fwd_wrapper.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. + +"""FlyDSL permute-free MoE route-list forward gather-GEMM op wrapper. + +Mirrors the route-list forward gather-GEMM contract so the same routing +metadata (``sorted_slot_ids`` / ``expert_ids`` / ``block_start``) can be reused verbatim. +Writes the block-padded ``[em_max, WIDTH_N]`` slot output in place. + +Forward gather-GEMM is FlyDSL-only (MegaMOE-ported plain GEMM). +""" + +from __future__ import annotations + +import torch + +__all__ = [ + "flydsl_moe_fwd", +] + +# MegaMOE hand-tuned bf16 grouped-GEMM tile geometry (offline-swept in primus-turbo's +# bench_mega_moe: BLOCK_N=256, GROUP_M=4 fwd / GROUP_M=8 FC1 NN dgrad, num_xcd=1). +# ``block_m`` comes from the route-list align (128 or 256); N-tile and grouping are fixed here. +_PF_BLOCK_N = 256 +_PF_GROUP_M = 4 +_PF_DGRAD_FC1_GROUP_M = 8 # bench_mega_moe grouped_gemm_combine L1-dgrad (NN) sweep +_PF_NUM_XCD = 1 + + +def _run_gather_gemm( + A, B, C, sorted_slot_ids, expert_ids, *, num_recv_tokens, block_m, + transpose_b, index_a_by_route_pos, +): + """Dispatch fwd/dgrad to the permute-free FlyDSL gather-GEMM kernels. + + Covers FC1 gather (``index_a_by_route_pos=False``), FC2 route-read + (``index_a_by_route_pos=True``), and dgrad (``transpose_b``). Uses MegaMOE's fixed tile + geometry (32x32x16 MFMA, ``BLOCK_N=256``, ``GROUP_M=4``, ``num_xcd=1``). + """ + block_m = int(block_m) + em_max = int(sorted_slot_ids.shape[0]) + if em_max % block_m != 0: + raise ValueError( + f"permute-free gather-GEMM expects em_max ({em_max}) divisible by block_m ({block_m}); " + "check routing align block_size_m." + ) + num_tile_blocks = em_max // block_m + expert_ids_i32 = expert_ids.to(torch.int32) + if expert_ids_i32.numel() != num_tile_blocks: + raise ValueError( + f"permute-free gather-GEMM expects expert_ids length {num_tile_blocks} " + f"(em_max={em_max}, BLOCK_M={block_m}), got {expert_ids_i32.numel()}; " + f"routing align block_size_m must match block_m={block_m}" + ) + + if transpose_b: + # dgrad: NN GEMM contracting the incoming grad against the forward weight [E, N, K] + # route-read from block-padded route-ordered grad; FC2 dgrad gathers token-ordered + # grad [num_recv, N] into block-padded route-ordered dX [em_max, K]. + from ..flydsl_kernels.permute_free_grouped_gemm.pf_dgrad import grouped_gemm_dgrad_bf16 + + dgrad_group_m = _PF_DGRAD_FC1_GROUP_M if index_a_by_route_pos else _PF_GROUP_M + grouped_gemm_dgrad_bf16( + A, B, C, expert_ids_i32, num_tile_blocks, sorted_slot_ids, + gather=not index_a_by_route_pos, + BLOCK_M=block_m, BLOCK_N=_PF_BLOCK_N, GROUP_M=dgrad_group_m, num_xcd=_PF_NUM_XCD, + ) + return + + from ..flydsl_kernels.permute_free_grouped_gemm.pf_fwd import grouped_gemm_gather_bf16 + + grouped_gemm_gather_bf16( + A, B, C, expert_ids_i32, num_tile_blocks, sorted_slot_ids, + gather=not index_a_by_route_pos, + BLOCK_M=block_m, BLOCK_N=_PF_BLOCK_N, GROUP_M=_PF_GROUP_M, num_xcd=_PF_NUM_XCD, + ) + + +def flydsl_moe_fwd( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + sorted_slot_ids: torch.Tensor, + expert_ids: torch.Tensor, + *, + num_recv_tokens: int, + block_m: int, + index_a_by_route_pos: bool = False, + dgrad: bool = False, +) -> None: + """Route-list gather-GEMM, writing block-padded route-ordered ``C[em_max, WIDTH_N]`` in place. + + ``A`` is token-ordered ``[num_recv, K]`` gathered via ``sorted_slot_ids`` when + ``index_a_by_route_pos=False`` (FC1), or block-padded route-ordered read by route slot when + ``index_a_by_route_pos=True`` (FC2). ``B`` is ``[num_experts, N_OUT, K]`` (contiguous inner + ``K``). Set ``dgrad=True`` for the data-gradient path (same ``B`` layout). Gated activation + and route-prob apply are **not** fused here; use the standalone helpers in + :mod:`pf_helper_kernels`. + """ + assert A.dtype == B.dtype == C.dtype == torch.bfloat16 + assert A.stride(1) == 1, "A must be contiguous along the contraction (K)" + assert B.stride(2) == 1, "B must be [E, N, K] with contiguous inner K" + + _run_gather_gemm( + A, B, C, sorted_slot_ids, expert_ids, + num_recv_tokens=num_recv_tokens, block_m=block_m, + transpose_b=dgrad, index_a_by_route_pos=index_a_by_route_pos, + ) diff --git a/transformer_engine/pytorch/moe/pf_helper_kernels.py b/transformer_engine/pytorch/moe/pf_helper_kernels.py new file mode 100644 index 000000000..8ce1bc1ae --- /dev/null +++ b/transformer_engine/pytorch/moe/pf_helper_kernels.py @@ -0,0 +1,850 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""Route-list MoE Triton helpers (bf16): align/scan, gated-act, gather-combine. + +Grouped GEMM for FC1/FC2 forward, backward (dgrad), and wgrad is FlyDSL-only +(``pf_fwd_wrapper``, ``pf_wgrad_wrapper``). This module retains Triton kernels for +routing metadata construction, gated-activation recompute/bwd, and the token-order +gather-combine pass that follows the block-padded route-ordered GEMM outputs. +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import triton +import triton.language as tl + + +# --- Fused gated-activation epilogue helpers (exp2-based, no libdevice) --- +@triton.jit +def _tanh(x): + # tanh via exp2 (1.44269504089 = 1/ln2): tanh(x) = 2*sigmoid(2x) - 1. + return 2.0 / (1.0 + tl.exp2(-2.0 * x * 1.44269504089)) - 1.0 + + +@triton.jit +def _silu(x): + # x * sigmoid(x), sigmoid via exp2. + return x / (1.0 + tl.exp2(-(x * 1.44269504089))) + + +@triton.jit +def _gelu_tanh(x): + # tanh approximation of GELU (matches PyTorch approximate='tanh'). + inner = 0.7978845608028654 * (x + 0.044715 * x * x * x) + return 0.5 * x * (1.0 + _tanh(inner)) + + +# Activation selector for the fused epilogue. Keep the set small and explicit; add a helper +# above and an entry here to extend. Passed to the kernel as a compile-time int so each +# activation specializes to its own instance (no runtime branch in the hot path). +_ACT_IDS = {"silu": 0, "gelu": 1} + + +@triton.jit +def _apply_activation(x, ACTIVATION: tl.constexpr): + if ACTIVATION == 0: + return _silu(x) + else: + return _gelu_tanh(x) + + +# --- Activation derivatives (for the fused gated-activation backward) --- +@triton.jit +def _silu_grad(x): + # d/dx [x*sigmoid(x)] = sigmoid(x) * (1 + x*(1 - sigmoid(x))). + s = 1.0 / (1.0 + tl.exp2(-(x * 1.44269504089))) + return s + x * s * (1.0 - s) + + +@triton.jit +def _gelu_tanh_grad(x): + # d/dx [0.5*x*(1+tanh(inner))], inner = c*(x + 0.044715 x^3), c = sqrt(2/pi). + inner = 0.7978845608028654 * (x + 0.044715 * x * x * x) + t = _tanh(inner) + dinner = 0.7978845608028654 * (1.0 + 3.0 * 0.044715 * x * x) + return 0.5 * (1.0 + t) + 0.5 * x * (1.0 - t * t) * dinner + + +@triton.jit +def _apply_activation_grad(x, ACTIVATION: tl.constexpr): + if ACTIVATION == 0: + return _silu_grad(x) + else: + return _gelu_tanh_grad(x) + + +# --- Fused activation + derivative (for the gated-activation backward) --- +# The backward needs *both* act(x) and act'(x) on the same input. Computing them via +# the separate helpers above evaluates the transcendental twice (sigmoid for silu, +# tanh for gelu). These return the pair from a single transcendental -- the backward +# kernel is VALU-bound on exp2, so sharing it is the dominant win. +@triton.jit +def _silu_and_grad(x): + # s = sigmoid(x); silu = x*s; d/dx silu = s + x*s*(1-s). + s = 1.0 / (1.0 + tl.exp2(-(x * 1.44269504089))) + act = x * s + return act, s + act * (1.0 - s) + + +@triton.jit +def _gelu_tanh_and_grad(x): + inner = 0.7978845608028654 * (x + 0.044715 * x * x * x) + t = _tanh(inner) + dinner = 0.7978845608028654 * (1.0 + 3.0 * 0.044715 * x * x) + return 0.5 * x * (1.0 + t), 0.5 * (1.0 + t) + 0.5 * x * (1.0 - t * t) * dinner + + +@triton.jit +def _apply_activation_and_grad(x, ACTIVATION: tl.constexpr): + if ACTIVATION == 0: + return _silu_and_grad(x) + else: + return _gelu_tanh_and_grad(x) + + +def _gated_act_bwd_autotune_configs() -> list: + """Tile/warp configs for the gated-act-bwd autotuner. + + Each program owns a ``BLOCK_M`` x ``BLOCK_H`` route/feature tile. The kernel is + latency-bound (the per-route work is tiny and the memory unit is never stalled), so + the sweep brackets the route-tile height ``BLOCK_M`` (more routes/block = more loads + in flight to hide latency) against the feature width ``BLOCK_H`` and the warp count, + trading memory-level parallelism against VGPR/occupancy. + """ + return [ + triton.Config({"BLOCK_M": bm, "BLOCK_H": bh}, num_warps=w) + for bm, bh, w in ( + # ``num_warps`` such that BLOCK_H/(warps*64) >= 8 gives 128-bit (dwordx4) + # global loads/stores; the (1,1024,2) / (2,1024,2) points below hit that, + # while the wider-warp points keep dwordx2 -- the tuner picks per shape. + (1, 1024, 4), + (1, 1024, 2), + (2, 1024, 2), + (2, 1024, 4), + (2, 512, 4), + (4, 1024, 8), + (8, 512, 4), + (16, 256, 4), + (16, 512, 8), + (32, 256, 8), + (32, 128, 4), + (8, 256, 4), + (4, 512, 4), + (64, 128, 8), + ) + ] + + +@triton.autotune(configs=_gated_act_bwd_autotune_configs(), key=["F", "HAS_PROBS"]) +@triton.jit +def _gated_act_prob_bwd_kernel( + grad_out_ptr, # [T * min(topk, E), F] grad wrt the fused FC1 output (= act(g)*u*prob) + preact_ptr, # [T * min(topk, E), 2F] raw pre-activation [gate | up] + probs_ptr, # [num_recv_tokens, E] + token_ptr, # [routes_max] route -> received-token row + expert_ptr, # [routes_max] route -> local expert + dpre_ptr, # [T * min(topk, E), 2F] out: grad wrt the raw 2F GEMM output + grad_probs_ptr, # [num_recv_tokens, E] out (fp32) + act_out_ptr, # [T * min(topk, E), F] out: fused fc2_input = act(g)*u*prob (EMIT_ACT only) + nbound_ptr, # [1] int32 device scalar: dynamic upper bound on route slots + num_recv_tokens, + F, + stride_gom, + stride_goh, + stride_prem, + stride_preh, + stride_pm, + stride_pe, + stride_dprem, + stride_dpreh, + stride_gpm, + stride_gpe, + stride_aom, + stride_aoh, + ACTIVATION: tl.constexpr, + HAS_PROBS: tl.constexpr, + EMIT_ACT: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_H: tl.constexpr, +): + """Backward of the fused gated-activation (+ optional route-prob) FC1 epilogue. + + Each program owns a ``BLOCK_M`` x ``BLOCK_H`` tile of routes x gate features. Given + ``grad_out`` (F), the saved pre-activation ``[gate | up]`` (2F) and the per-route prob + ``p``, it emits the 2F grad wrt the raw GEMM output (``dpre = [d_gate | d_up]``) and, + when probs are fused, reduces ``dp = sum_f grad_out * act(gate) * up`` per route and + scatters it to ``grad_dispatched_probs[token, expert]``. + + ``act(gate)`` and ``act'(gate)`` are produced together from a single transcendental + (the kernel is VALU-heavy on ``exp2``). Each ``(token, expert)`` cell is written by + exactly one route, so the prob-grad scatter is a plain store (no atomics). Padded / + over-allocated routes carry the ``token == num_recv_tokens`` sentinel and are masked + out of the prob load/store; their ``dpre`` rows are inert (ignored downstream). + + When ``EMIT_ACT`` is set the kernel *also* stores the ``F``-wide forward activation + ``fc2_input = act(gate)*up[*prob]`` to ``act_out_ptr`` -- the same buffer the FC2 wgrad + would otherwise recompute in a second full pass over the ``2F`` preact. Since ``act(gate)``, + ``up`` and ``prob`` are already live in registers here, this is one extra multiply + store + and removes the redundant ``2F`` HBM read (see :func:`fused_gated_act_prob_fwd`). + + The block-padded route buffers are statically over-allocated to the worst-case + ``routes_max = T * topk`` (sync-free shape bound), but the real routes occupy only the + dense route-ordered head ``[0, num_routes)`` -- under expert parallelism this can be ~topk*E_local/E + times smaller. ``nbound_ptr`` carries the actual (block-padded) route extent as a device + scalar so tail programs beyond it exit before touching HBM, instead of grinding through + the padding at full memory bandwidth. + """ + pid = tl.program_id(axis=0) + num_routes_bound = tl.load(nbound_ptr) + if pid * BLOCK_M >= num_routes_bound: + return + r_offs = pid * BLOCK_M + tl.arange(0, BLOCK_M) + r_mask = r_offs < num_routes_bound + token = tl.load(token_ptr + r_offs, mask=r_mask, other=num_recv_tokens).to(tl.int64) + valid = token < num_recv_tokens + expert = tl.load(expert_ptr + r_offs, mask=r_mask, other=0).to(tl.int64) + if HAS_PROBS: + prob = tl.load( + probs_ptr + token * stride_pm + expert * stride_pe, mask=valid, other=0.0 + ).to(tl.float32) + + dp_acc = tl.zeros((BLOCK_M,), dtype=tl.float32) + for h0 in range(0, F, BLOCK_H): + offs = h0 + tl.arange(0, BLOCK_H) + hmask = offs < F + m = r_mask[:, None] & hmask[None, :] + prow = r_offs[:, None] * stride_prem + g = tl.load( + preact_ptr + prow + offs[None, :] * stride_preh, mask=m, other=0.0 + ).to(tl.float32) + u = tl.load( + preact_ptr + prow + (offs[None, :] + F) * stride_preh, mask=m, other=0.0 + ).to(tl.float32) + go = tl.load( + grad_out_ptr + r_offs[:, None] * stride_gom + offs[None, :] * stride_goh, + mask=m, + other=0.0, + ).to(tl.float32) + + act_g, dact_g = _apply_activation_and_grad(g, ACTIVATION) + if HAS_PROBS: + dp_acc += tl.sum(go * act_g * u, axis=1) + da = go * prob[:, None] + else: + da = go + d_up = da * act_g + d_gate = da * u * dact_g + drow = r_offs[:, None] * stride_dprem + tl.store( + dpre_ptr + drow + offs[None, :] * stride_dpreh, + d_gate.to(dpre_ptr.dtype.element_ty), + mask=m, + ) + tl.store( + dpre_ptr + drow + (offs[None, :] + F) * stride_dpreh, + d_up.to(dpre_ptr.dtype.element_ty), + mask=m, + ) + if EMIT_ACT: + # Re-emit the F-wide forward activation from values already in registers, so the + # FC2 wgrad can consume it directly instead of re-streaming the 2F preact. + fi = act_g * u + if HAS_PROBS: + fi = fi * prob[:, None] + tl.store( + act_out_ptr + r_offs[:, None] * stride_aom + offs[None, :] * stride_aoh, + fi.to(act_out_ptr.dtype.element_ty), + mask=m, + ) + + if HAS_PROBS: + tl.store( + grad_probs_ptr + token * stride_gpm + expert * stride_gpe, + dp_acc, + mask=valid, + ) + + +def fused_gated_act_prob_bwd( + grad_out: torch.Tensor, + preact: torch.Tensor, + token: torch.Tensor, + expert: torch.Tensor, + *, + num_recv_tokens: int, + activation: str, + dispatched_probs: Optional[torch.Tensor] = None, + grad_probs_shape: Optional[torch.Size] = None, + num_routes_bound: Optional[torch.Tensor] = None, + emit_act: bool = False, +): + """Backward of the fused gated-activation (+ route-prob) FC1 epilogue. + + Parameters + ---------- + grad_out: + ``[T * min(topk, E), F]`` grad wrt the fused FC1 output (route/padded layout). + preact: + ``[T * min(topk, E), 2F]`` raw pre-activation ``[gate | up]`` saved by the forward. + token, expert: + ``[routes_max]`` per-route received-token row / local-expert id (int32). + dispatched_probs: + ``[num_recv_tokens, E]`` gating probs, or ``None`` if the forward did not fuse the + route-prob multiply. When given, its gradient is returned. + num_routes_bound: + Optional ``[1]`` int32 device scalar giving a (block-padded) upper bound on the number + of *real* route slots. The route buffers are statically sized to the worst case + ``routes_max = T * topk``, but under expert parallelism only a small dense head is + populated; passing the actual extent (e.g. ``num_tokens_post_padded`` from the routing + metadata) lets tail programs exit early instead of streaming the padding through HBM. + When ``None`` the full static ``routes_max`` is used (no early exit). + emit_act: + When ``True`` the kernel additionally re-materialises the ``F``-wide forward activation + ``fc2_input = act(gate)*up[*prob]`` (``[T * min(topk, E), F]`` bf16) from the values it + already computes, so the FC2 wgrad can consume it directly instead of re-reading the ``2F`` + preact in a separate recompute pass. Returned as the third element (``None`` otherwise). + + Returns + ------- + (dpre, grad_probs, fc2_input) + ``dpre`` is ``[T * min(topk, E), 2F]`` (bf16), grad wrt the raw GEMM output; ``grad_probs`` is + ``[num_recv_tokens, E]`` (matching ``dispatched_probs.dtype``) or ``None``; ``fc2_input`` is + ``[T * min(topk, E), F]`` (bf16) when ``emit_act`` else ``None``. + """ + em_max, F = grad_out.shape + if preact.shape[0] != em_max or preact.shape[1] != 2 * F: + raise ValueError( + f"preact must be [T * min(topk, E), 2F]={ (em_max, 2 * F) }, got {tuple(preact.shape)}." + ) + routes_max = int(token.shape[0]) + has_probs = dispatched_probs is not None + + dpre = torch.empty((em_max, 2 * F), dtype=torch.bfloat16, device=grad_out.device) + if emit_act: + act_out = torch.empty((em_max, F), dtype=torch.bfloat16, device=grad_out.device) + stride_aom, stride_aoh = act_out.stride(0), act_out.stride(1) + else: + act_out = None + stride_aom = stride_aoh = 0 + if has_probs: + grad_probs = torch.zeros(dispatched_probs.shape, dtype=torch.float32, device=grad_out.device) + probs_ptr = dispatched_probs + stride_pm, stride_pe = dispatched_probs.stride(0), dispatched_probs.stride(1) + stride_gpm, stride_gpe = grad_probs.stride(0), grad_probs.stride(1) + else: + grad_probs = None + probs_ptr = grad_out # unused + stride_pm = stride_pe = stride_gpm = stride_gpe = 0 + + act_id = _ACT_IDS[activation] + if num_routes_bound is None: + nbound = torch.tensor([routes_max], dtype=torch.int32, device=grad_out.device) + else: + nbound = num_routes_bound + grid = lambda meta: (triton.cdiv(routes_max, meta["BLOCK_M"]),) # noqa: E731 + _gated_act_prob_bwd_kernel[grid]( + grad_out, + preact, + probs_ptr, + token, + expert, + dpre, + grad_probs if has_probs else grad_out, + act_out if emit_act else dpre, + nbound, + num_recv_tokens, + F, + grad_out.stride(0), + grad_out.stride(1), + preact.stride(0), + preact.stride(1), + stride_pm, + stride_pe, + dpre.stride(0), + dpre.stride(1), + stride_gpm, + stride_gpe, + stride_aom, + stride_aoh, + ACTIVATION=act_id, + HAS_PROBS=has_probs, + EMIT_ACT=emit_act, + ) + if has_probs: + grad_probs = grad_probs.to(dispatched_probs.dtype) + return dpre, grad_probs, act_out + + +@triton.autotune(configs=_gated_act_bwd_autotune_configs(), key=["F", "HAS_PROBS"]) +@triton.jit +def _gated_act_prob_fwd_kernel( + preact_ptr, # [routes_max, 2F] raw pre-activation [gate | up] + probs_ptr, # [num_recv_tokens, E] + token_ptr, # [routes_max] route -> received-token row + expert_ptr, # [routes_max] route -> local expert + act_ptr, # [routes_max, F] out: act(gate) * up * prob + nbound_ptr, # [1] int32 device scalar: dynamic upper bound on route slots + num_recv_tokens, + F, + stride_prem, + stride_preh, + stride_pm, + stride_pe, + stride_am, + stride_ah, + ACTIVATION: tl.constexpr, + HAS_PROBS: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_H: tl.constexpr, +): + """Recompute the fused gated-activation FC1 output ``act(gate) * up * prob`` from preact. + + This is the *forward* counterpart of :func:`_gated_act_prob_bwd_kernel`: given the saved + ``2F`` pre-activation ``[gate | up]`` and the per-route prob it re-materialises the + ``F``-wide activation that the forward fused into the FC1 epilogue, so the backward can + checkpoint only the ``2F`` preact (never the ``F``-wide act) and rebuild it just-in-time + for the FC2 wgrad. Padded / over-allocated routes carry the ``token == num_recv_tokens`` + sentinel and get ``prob == 0`` (their act rows are ignored downstream). ``nbound_ptr`` + bounds tail programs to the real block-padded route extent (sync-free, EP-friendly early exit). + """ + pid = tl.program_id(axis=0) + num_routes_bound = tl.load(nbound_ptr) + if pid * BLOCK_M >= num_routes_bound: + return + r_offs = pid * BLOCK_M + tl.arange(0, BLOCK_M) + r_mask = r_offs < num_routes_bound + if HAS_PROBS: + token = tl.load(token_ptr + r_offs, mask=r_mask, other=num_recv_tokens).to(tl.int64) + valid = token < num_recv_tokens + expert = tl.load(expert_ptr + r_offs, mask=r_mask, other=0).to(tl.int64) + prob = tl.load( + probs_ptr + token * stride_pm + expert * stride_pe, mask=valid, other=0.0 + ).to(tl.float32) + + for h0 in range(0, F, BLOCK_H): + offs = h0 + tl.arange(0, BLOCK_H) + hmask = offs < F + m = r_mask[:, None] & hmask[None, :] + prow = r_offs[:, None] * stride_prem + g = tl.load( + preact_ptr + prow + offs[None, :] * stride_preh, mask=m, other=0.0 + ).to(tl.float32) + u = tl.load( + preact_ptr + prow + (offs[None, :] + F) * stride_preh, mask=m, other=0.0 + ).to(tl.float32) + act = _apply_activation(g, ACTIVATION) * u + if HAS_PROBS: + act = act * prob[:, None] + tl.store( + act_ptr + r_offs[:, None] * stride_am + offs[None, :] * stride_ah, + act.to(act_ptr.dtype.element_ty), + mask=m, + ) + + +def fused_gated_act_prob_fwd( + preact: torch.Tensor, + token: torch.Tensor, + expert: torch.Tensor, + *, + num_recv_tokens: int, + activation: str, + dispatched_probs: Optional[torch.Tensor] = None, + num_routes_bound: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Recompute the fused gated-activation FC1 output from the saved pre-activation. + + Re-materialises ``act = act(gate) * up * prob`` (``[routes_max, F]``, bf16) from the + ``[routes_max, 2F]`` preact ``[gate | up]`` -- the inverse of checkpointing the ``F``-wide + activation. Feeding this transient buffer to the *unchanged* FC2 wgrad keeps that kernel at + full (stored-act) speed while only the ``2F`` preact is persisted across the fwd/bwd + boundary. See :func:`fused_gated_act_prob_bwd` for the argument contract (token / expert / + ``num_routes_bound`` are the same per-route routing arrays). + """ + routes_max, two_f = preact.shape + if two_f % 2 != 0: + raise ValueError(f"preact must be [routes_max, 2F], got a {two_f}-wide last dim.") + F = two_f // 2 + has_probs = dispatched_probs is not None + + act = torch.empty((routes_max, F), dtype=torch.bfloat16, device=preact.device) + if has_probs: + probs_ptr = dispatched_probs + stride_pm, stride_pe = dispatched_probs.stride(0), dispatched_probs.stride(1) + else: + probs_ptr = preact # unused + stride_pm = stride_pe = 0 + + act_id = _ACT_IDS[activation] + if num_routes_bound is None: + nbound = torch.tensor([routes_max], dtype=torch.int32, device=preact.device) + else: + nbound = num_routes_bound + grid = lambda meta: (triton.cdiv(routes_max, meta["BLOCK_M"]),) # noqa: E731 + _gated_act_prob_fwd_kernel[grid]( + preact, + probs_ptr, + token, + expert, + act, + nbound, + num_recv_tokens, + F, + preact.stride(0), + preact.stride(1), + stride_pm, + stride_pe, + act.stride(0), + act.stride(1), + ACTIVATION=act_id, + HAS_PROBS=has_probs, + ) + return act + + + +@triton.jit +def _gather_combine_kernel( + src_ptr, # block-padded [em_max, N] (padded slot order; padding rows carry dead values) + token_routes_ptr, # [T, MAXK] int32: block-padded slot positions per token + token_count_ptr, # [T] int32: number of routes for each token + out_ptr, # [T, N] out + N, + stride_sm, + stride_sn, + stride_om, + stride_on, + MAXK: tl.constexpr, + BLOCK_N: tl.constexpr, + compute_type: tl.constexpr, +): + t = tl.program_id(axis=0) + pid_n = tl.program_id(axis=1) + cnt = tl.load(token_count_ptr + t) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + n_mask = offs_n < N + acc = tl.zeros((BLOCK_N,), dtype=tl.float32) + # Sum the token's route rows (expert-ascending -> deterministic). Columns >= cnt are + # unused padding and skipped, so no padded/garbage slot is ever gathered (this masked + # padded->token reduction is the sole output-side masking point in the block-padded path). + for j in range(0, MAXK): + if j < cnt: + r = tl.load(token_routes_ptr + t * MAXK + j).to(tl.int64) + v = tl.load(src_ptr + r * stride_sm + offs_n * stride_sn, mask=n_mask, other=0.0) + acc += v.to(tl.float32) + tl.store(out_ptr + t * stride_om + offs_n * stride_on, acc.to(compute_type), mask=n_mask) + + +def route_gather_combine( + src: torch.Tensor, + token_routes: torch.Tensor, + token_route_count: torch.Tensor, + num_recv_tokens: int, + *, + out_dtype: torch.dtype = torch.bfloat16, + block_n: Optional[int] = None, +) -> torch.Tensor: + """Combine per-route rows into per-token rows via a contention-free gather. + + Parameters + ---------- + src: + Compact per-route buffer ``[T * min(topk, E), N]`` (route order; only ``[0, num_routes)`` valid). + token_routes / token_route_count: + Token->routes inverse map from the align place kernel (``[T, MAXK]`` route positions + and the per-token route count). + num_recv_tokens: + Number of output token rows ``T``. + block_n: + N-tile width (``BLOCK_N``). ``None`` (default) picks ``min(next_pow2(N), 4096)``: wide + tiles issue larger contiguous per-route loads and cut redundant index loads. + + Returns + ------- + torch.Tensor + ``[num_recv_tokens, N]`` in ``out_dtype`` -- each row is the fp32 sum of its token's + route rows, cast once. No atomics, no host sync. + """ + if src.dim() != 2: + raise ValueError(f"src must be [T * min(topk, E), N], got {tuple(src.shape)}.") + if not src.is_contiguous(): + src = src.contiguous() + n = src.shape[1] + if block_n is None: + block_n = min(triton.next_power_of_2(n), 4096) + maxk = int(token_routes.shape[1]) + out = torch.empty((num_recv_tokens, n), dtype=out_dtype, device=src.device) + compute_type = tl.bfloat16 if out_dtype == torch.bfloat16 else tl.float32 + grid = (num_recv_tokens, triton.cdiv(n, block_n)) + _gather_combine_kernel[grid]( + src, + token_routes, + token_route_count, + out, + n, + src.stride(0), + src.stride(1), + out.stride(0), + out.stride(1), + MAXK=maxk, + BLOCK_N=block_n, + compute_type=compute_type, + ) + return out + + +@triton.jit +def _counts_within_kernel( + routing_map_ptr, # [T, E] (bool/int8), True where token t feeds local expert e + within_ptr, # [E, T] int32 out: exclusive within-expert rank of each routed cell + counts_ptr, # [E] int32 out: routed-token count per expert + T, + stride_t, + stride_e, + BLOCK_T: tl.constexpr, +): + e = tl.program_id(axis=0) + offs = tl.arange(0, BLOCK_T) + mask = offs < T + vals = tl.load(routing_map_ptr + offs * stride_t + e * stride_e, mask=mask, other=0).to( + tl.int32 + ) + # Exclusive prefix sum over tokens => within-expert rank; total => per-expert count. + incl = tl.cumsum(vals, axis=0) + excl = incl - vals + tl.store(within_ptr + e * T + offs, excl, mask=mask) + tl.store(counts_ptr + e, tl.sum(vals, axis=0)) + + +@triton.jit +def _expert_meta_kernel( + counts_ptr, # [E] int32 + blocks_per_expert_ptr, # [E] int32 out + block_start_ptr, # [E] int32 out (block units) + expert_ids_ptr, # [blocks_max] int32 out (expert owning each block, -1 past the end) + ntpp_ptr, # [1] int32 out: block-padded token extent + E, + blocks_max, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_E: tl.constexpr, + BLOCK_B: tl.constexpr, +): + offs_e = tl.arange(0, BLOCK_E) + mask_e = offs_e < E + counts = tl.load(counts_ptr + offs_e, mask=mask_e, other=0) + blocks_per_expert = (counts + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + cblocks = tl.cumsum(blocks_per_expert, axis=0) # inclusive prefix over experts + # block_start is over padded block token counts + block_start = cblocks - blocks_per_expert + total_blocks = tl.max(tl.where(mask_e, cblocks, 0), axis=0) + + pid = tl.program_id(axis=0) + if pid == 0: + tl.store(blocks_per_expert_ptr + offs_e, blocks_per_expert, mask=mask_e) + tl.store(block_start_ptr + offs_e, block_start, mask=mask_e) + tl.store(ntpp_ptr, total_blocks * BLOCK_SIZE_M) + + # expert_ids[b] = #{e : cblocks[e] <= b} then -1 for blocks past the real extent. + offs_b = pid * BLOCK_B + tl.arange(0, BLOCK_B) + mask_b = offs_b < blocks_max + cblocks_valid = tl.where(mask_e, cblocks, (1 << 30)) + le = (cblocks_valid[None, :] <= offs_b[:, None]).to(tl.int32) + expert_ids = tl.sum(le, axis=1) + expert_ids = tl.where(offs_b < total_blocks, expert_ids, -1) + tl.store(expert_ids_ptr + offs_b, expert_ids, mask=mask_b) + + +@triton.jit +def _route_list_place_kernel( + routing_map_ptr, # [T, E] (bool/int8), True where token t feeds local expert e + within_ptr, # [E, T] int32, exclusive within-expert rank of each routed cell + block_start_ptr, # [E] int32: first block index of each expert (block units) + sorted_slot_ids_ptr, # [T * min(topk, E)] int32, sentinel-init (T) + token_routes_ptr, # [T, MAXK] int32 out: token->route positions (only if BUILD_INVERSE) + token_count_ptr, # [T] int32 out: routes per token (only if BUILD_INVERSE) + T, + E, + stride_t, + stride_e, + BLOCK_SIZE_M: tl.constexpr, + BUILD_INVERSE: tl.constexpr, + MAXK: tl.constexpr, +): + # One row of the map per program; place the (few) routed cells deterministically into their + # block-padded slot ``bs*BLOCK_SIZE_M + within[e, t]``. The token->routes inverse map (used by + # the contention-free gather-combine) is emitted here too when ``BUILD_INVERSE`` -- folding + # what was a second per-token kernel launch into this one (expert-ascending, no atomics). + t = tl.program_id(axis=0) + if t >= T: + return + j = tl.zeros((), dtype=tl.int32) + for e in range(0, E): + is_routed = tl.load(routing_map_ptr + t * stride_t + e * stride_e) + if is_routed != 0: + w = tl.load(within_ptr + e * T + t) + bs = tl.load(block_start_ptr + e) + slot = bs * BLOCK_SIZE_M + w + tl.store(sorted_slot_ids_ptr + slot, t) + if BUILD_INVERSE: + # Block-padded canonical layout: every intermediate route tensor (the GEMM + # output, the FC2 route-read input, the activation) is indexed by the padded + # slot ``bs*BLOCK_SIZE_M + w``. The token->routes inverse map therefore stores the + # padded slot so the final padded->token gather-combine reads the correct rows. + tl.store(token_routes_ptr + t * MAXK + j, slot) + j += 1 + if BUILD_INVERSE: + tl.store(token_count_ptr + t, j) + + +def route_list_scan( + routing_map: torch.Tensor, + *, + num_experts: int, +): + """Block-size-independent scan: per-expert token counts + within-expert ranks. + + Returned ``(counts [E], within [E, T])`` can be reused across multiple block sizes + (e.g. the fwd/dgrad ``BLOCK_SIZE_M`` and the wgrad ``CONTRACT_M``), so the scan is + only paid once per routing map. Pass them to :func:`route_list_align` via ``scan=``. + """ + device = routing_map.device + T = int(routing_map.size(0)) + E = int(num_experts) + within = torch.empty((E, T), dtype=torch.int32, device=device) + counts = torch.empty((E,), dtype=torch.int32, device=device) + _counts_within_kernel[(E,)]( + routing_map, + within, + counts, + T, + routing_map.stride(0), + routing_map.stride(1), + BLOCK_T=triton.next_power_of_2(max(T, 1)), + ) + return counts, within + + +def route_list_align( + routing_map: torch.Tensor, + *, + num_experts: int, + block_size: int, + scan=None, + topk: int | None = None, + build_inverse_map: bool = False, +): + """Sync-free fused build of the route-list align buffers. + + Parameters + ---------- + scan: + Optional ``(counts, within)`` from :func:`route_list_scan` for this ``routing_map``. + When supplied the block-independent scan kernel is skipped (shared across block + sizes); otherwise it is computed here. + topk: + Host-known upper bound on the number of experts any token routes to (the router + top-k). When provided, the static over-allocation bound is tightened from the dense + ``T * num_experts`` to ``T * min(topk, num_experts)`` -- still + sync-free, but shrinking the padded buffers by ``num_experts / topk``. + build_inverse_map: + When True, the place kernel also emits the token->routes inverse map (used by the + contention-free gather-combine in FC2 fwd / FC1 dgrad) in the same launch, so no + separate inverse-map kernel is needed. Block-independent, so build it on the fwd align + only. + + Returns + ------- + ``(sorted_slot_ids, expert_ids, num_tokens_post_padded, block_start, blocks_per_expert, + token_routes, token_route_count)`` -- index tensors ``int32``; ``num_tokens_post_padded`` + (``[1]``) is a device scalar. ``token_routes`` (``[T, min(topk, E)]``) and + ``token_route_count`` (``[T]``) are ``None`` unless ``build_inverse_map``. + """ + if routing_map.dtype != torch.bool: + routing_map = routing_map.bool() + routing_map = routing_map.contiguous() + device = routing_map.device + T = int(routing_map.size(0)) + E = int(num_experts) + + # Static (sync-free) upper bounds from shapes only. Each token routes to at most + # ``min(topk, E)`` experts, so that tightens the dense ``T * E`` bound. + max_per_token = E if topk is None else min(int(topk), E) + routes_max = T * max_per_token + blocks_max = (routes_max + block_size - 1) // block_size + E + em_max = blocks_max * block_size + + # Per-expert count + exclusive within-expert rank (one program per expert). Reused + # across block sizes when the caller passes a precomputed scan. + if scan is None: + counts, within = route_list_scan(routing_map, num_experts=E) + else: + counts, within = scan + + # Per-expert placement metadata + per-block expert ids (single launch). + blocks_per_expert = torch.empty((E,), dtype=torch.int32, device=device) # [E]: ceil(count[e]/block_m) + block_start = torch.empty((E,), dtype=torch.int32, device=device) # [E]: first M-block index of expert e + expert_ids = torch.empty((blocks_max,), dtype=torch.int32, device=device) # [blocks_max]: owner of M-block b (-1 past end) + num_tokens_post_padded = torch.empty((1,), dtype=torch.int32, device=device) # [1] device scalar: real em (padded route count) + block_b = 256 + # From per-expert route counts, derive the expert-sorted block layout: how many M-blocks + # each expert needs, where each expert's slots start (block_start), which expert owns each + # M-block (expert_ids), and the real padded route extent (num_tokens_post_padded). + _expert_meta_kernel[(triton.cdiv(blocks_max, block_b),)]( + counts, + blocks_per_expert, + block_start, + expert_ids, + num_tokens_post_padded, + E, + blocks_max, + BLOCK_SIZE_M=block_size, + BLOCK_E=triton.next_power_of_2(max(E, 1)), + BLOCK_B=block_b, + ) + + # Optional token->routes inverse map, emitted by the same place kernel. Width is the + # tightened per-token bound; zero-init so any unused tail column is a safe (in-range) + # index (the gather masks columns >= count anyway). + if build_inverse_map: + maxk = max(max_per_token, 1) + token_routes = torch.zeros((T, maxk), dtype=torch.int32, device=device) + token_route_count = torch.empty((T,), dtype=torch.int32, device=device) + else: + maxk = 1 + token_routes = torch.empty((1,), dtype=torch.int32, device=device) # unused stub + token_route_count = token_routes + + # One program per token: for each routed (t, e), write sorted_slot_ids[slot] = t where + # slot = block_start[e] * block_m + within[e, t]. Unwritten slots stay at sentinel T. + # Optionally emits token_routes / token_route_count (token -> padded slot indices). + sorted_slot_ids = torch.full((em_max,), T, dtype=torch.int32, device=device) + _route_list_place_kernel[(T,)]( + routing_map, + within, + block_start, + sorted_slot_ids, + token_routes, + token_route_count, + T, + E, + routing_map.stride(0), + routing_map.stride(1), + BLOCK_SIZE_M=block_size, + BUILD_INVERSE=build_inverse_map, + MAXK=maxk, + ) + + return ( + sorted_slot_ids, + expert_ids, + num_tokens_post_padded, + block_start, + blocks_per_expert, + token_routes if build_inverse_map else None, + token_route_count if build_inverse_map else None, + ) diff --git a/transformer_engine/pytorch/moe/pf_wgrad_wrapper.py b/transformer_engine/pytorch/moe/pf_wgrad_wrapper.py new file mode 100644 index 000000000..9758cc5fd --- /dev/null +++ b/transformer_engine/pytorch/moe/pf_wgrad_wrapper.py @@ -0,0 +1,242 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. + +"""FlyDSL permute-free MoE weight-gradient (wgrad) op wrapper. + +Mirrors the route-list wgrad contract so the same routing metadata +(``sorted_slot_ids`` holding the received-token row per slot, plus ``block_start`` / +``blocks_per_expert`` and the per-expert block-padded grad base ``grad_base``) can be reused +verbatim. + +Fixed kernel configuration (matches ``pf_wgrad.py``): FC1 route-list wgrad with block-padded +``grad`` + token-gathered ``x``, bf16 into ``dw`` (overwrite or accumulate), DMA + XOR +chunk swizzle fill, 3-stage LDS pipeline. Only the workgroup tile geometry is selectable +/ autotuned. +""" + +from __future__ import annotations + +import torch + +from ..flydsl_kernels.permute_free_grouped_gemm.pf_wgrad import ( + WGRAD_BLOCK_M, + compile_moe_wgrad_v2, +) +from ..flydsl_kernels.tensor_shim import _run_compiled, ptr_arg + +__all__ = ["flydsl_moe_wgrad", "flydsl_moe_wgrad_autotuned", "WGRAD_BLOCK_M"] + + +def flydsl_moe_wgrad( + x: torch.Tensor, + grad: torch.Tensor, + dw: torch.Tensor, + sorted_slot_ids: torch.Tensor, + block_start: torch.Tensor, + blocks_per_expert: torch.Tensor, + grad_base: torch.Tensor, + *, + num_recv_tokens: int, + block_n: int = 128, + block_k: int = 128, + warps_n: int = 2, + warps_k: int = 2, + accumulate: bool = False, + swap_gather: bool = False, +) -> None: + """Compute FC1 grouped wgrad ``grad[slot]^T @ x[token(slot)]`` into ``dw``, per expert. + + See the permute-free wgrad API for the argument contract: ``x`` is + ``[num_recv_tokens, K]`` (gathered by received-token row), ``grad`` is the block-padded + ``[em_max, N]`` per-slot gradient, ``sorted_slot_ids`` maps each block-padded slot to its + received-token row (sentinel ``num_recv_tokens`` for padding), and ``grad_base[e]`` is the + block-padded first-slot row of expert ``e`` (``block_start[e] * block_size_m``). + + ``swap_gather`` selects the FC2 wgrad variant: ``grad`` is the token-space ``[num_recv, N]`` + operand (gathered) and ``x`` is the block-padded ``[em_max, K]`` operand (contiguous walk), + emitting the native ``dw[E, N, K]`` orientation with no transpose. + + ``block_n``/``block_k`` and ``warps_n``/``warps_k`` select the workgroup tile; the + defaults (``128x128`` over ``2x2`` warps) are a strong general config on CDNA4. + """ + num_experts, N, K = dw.shape + + assert x.dtype == grad.dtype == torch.bfloat16 + assert dw.dtype in (torch.bfloat16, torch.float32) + assert x.is_contiguous() and grad.is_contiguous() and dw.is_contiguous() + + out_dtype = "fp32" if dw.dtype == torch.float32 else "bf16" + exe = compile_moe_wgrad_v2( + block_n=int(block_n), + block_k=int(block_k), + warps_n=int(warps_n), + warps_k=int(warps_k), + accumulate=bool(accumulate), + swap_gather=bool(swap_gather), + out_dtype=out_dtype, + ) + + _run_compiled( + exe, + ptr_arg(dw), + ptr_arg(x), + ptr_arg(grad), + ptr_arg(sorted_slot_ids), + ptr_arg(block_start), + ptr_arg(blocks_per_expert), + ptr_arg(grad_base), + int(N), + int(K), + int(num_recv_tokens), + int(num_experts), + torch.cuda.current_stream(), + ) + + +# Tile configs the autotuner sweeps (block_n, block_k, warps_n, warps_k). +_AUTOTUNE_TILES = [ + (64, 64, 1, 1), + (128, 64, 1, 1), + (128, 128, 2, 2), + (128, 128, 4, 2), + (256, 128, 4, 2), + (256, 256, 4, 4), + (128, 256, 2, 4), + (256, 128, 2, 2), + (256, 256, 2, 4), + (256, 256, 4, 2), +] + +_wgrad_autotuner = None + + +def _wgrad_run( + dw, + x, + grad, + sorted_slot_ids, + block_start, + blocks_per_expert, + grad_base, + N, + K, + num_recv_tokens, + num_experts, + block_n=128, + block_k=128, + warps_n=2, + warps_k=2, + accumulate=False, + swap_gather=False, +): + """Dispatch target for the FlyDSL autotuner: compile (lru-cached) + launch one tile.""" + exe = compile_moe_wgrad_v2( + block_n=int(block_n), + block_k=int(block_k), + warps_n=int(warps_n), + warps_k=int(warps_k), + accumulate=bool(accumulate), + swap_gather=bool(swap_gather), + ) + _run_compiled( + exe, + ptr_arg(dw), + ptr_arg(x), + ptr_arg(grad), + ptr_arg(sorted_slot_ids), + ptr_arg(block_start), + ptr_arg(blocks_per_expert), + ptr_arg(grad_base), + int(N), + int(K), + int(num_recv_tokens), + int(num_experts), + torch.cuda.current_stream(), + ) + + +def _get_autotuner(warmup=10, rep=30): + """Build the shape-keyed Autotuner lazily (one instance, disk-cached results).""" + global _wgrad_autotuner + if _wgrad_autotuner is None: + from flydsl.autotune import Autotuner, Config + + configs = [ + Config(block_n=bn, block_k=bk, warps_n=wn, warps_k=wk) + for (bn, bk, wn, wk) in _AUTOTUNE_TILES + ] + _wgrad_autotuner = Autotuner( + _wgrad_run, + configs, + key=["x", "N", "num_experts"], + warmup=warmup, + rep=rep, + ) + return _wgrad_autotuner + + +def _select_wgrad_config( + x, grad, sorted_slot_ids, block_start, blocks_per_expert, grad_base, + N, K, num_recv_tokens, num_experts, +): + """Return the autotuned ``(block_n, block_k, warps_n, warps_k)`` for this problem.""" + tuner = _get_autotuner() + scratch = torch.empty(num_experts, N, K, device=x.device, dtype=torch.bfloat16) + args = ( + scratch, x, grad, sorted_slot_ids, block_start, blocks_per_expert, grad_base, + int(N), int(K), int(num_recv_tokens), int(num_experts), + ) + key = tuner._make_key(args, {}) + if key not in tuner.cache: + tuner(*args) + cfg = tuner.cache[key].kwargs + return cfg["block_n"], cfg["block_k"], cfg["warps_n"], cfg["warps_k"] + + +def flydsl_moe_wgrad_autotuned( + x: torch.Tensor, + grad: torch.Tensor, + dw: torch.Tensor, + sorted_slot_ids: torch.Tensor, + block_start: torch.Tensor, + blocks_per_expert: torch.Tensor, + grad_base: torch.Tensor, + *, + num_recv_tokens: int, + accumulate: bool = False, + swap_gather: bool = False, +) -> None: + """Shape-autotuned variant of :func:`flydsl_moe_wgrad`. + + First call for a given ``(x.shape, N, num_experts)`` benchmarks every tile in + ``_AUTOTUNE_TILES`` and caches the fastest (in-memory + on disk under + ``~/.flydsl/autotune/``). The tile geometry is independent of ``swap_gather`` (the two + variants share the same memory-access shape), so the sweep runs on the default variant and + the fastest tile is reused for the ``swap_gather`` launch. + """ + num_experts, N, K = dw.shape + + assert x.dtype == grad.dtype == torch.bfloat16 + assert dw.dtype in (torch.bfloat16, torch.float32) + assert x.is_contiguous() and grad.is_contiguous() and dw.is_contiguous() + + block_n, block_k, warps_n, warps_k = _select_wgrad_config( + x, grad, sorted_slot_ids, block_start, blocks_per_expert, grad_base, + N, K, num_recv_tokens, num_experts, + ) + flydsl_moe_wgrad( + x, + grad, + dw, + sorted_slot_ids, + block_start, + blocks_per_expert, + grad_base, + num_recv_tokens=int(num_recv_tokens), + block_n=block_n, + block_k=block_k, + warps_n=warps_n, + warps_k=warps_k, + accumulate=bool(accumulate), + swap_gather=bool(swap_gather), + )