From a6b354c4c7584e98cb38e449ab5efed953d4ac5a Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 26 Aug 2026 17:07:32 +0000 Subject: [PATCH] Add gfx950 MXFP8 CK grouped GEMM and split backends per architecture Bump the composable_kernel pin to 763da730b. That commit removes the gfx950 guard in UniversalGemmBasePolicy that disabled transposed LDS loads for K warp tiles above 64, so the gfx1250 WMMA/TDM instantiation now fails to compile under a gfx950 device pass. CK gates TDM to gfx1250 in its own build, and there is no gfx950 warp tile that makes the TDM pipeline compile, so the backends have to be separated rather than reshaped. Split ck_mx_grouped_gemm.cpp into a shared templated body (ck_mx_grouped_gemm_impl.h) plus one thin translation unit per architecture, following the ck_grouped_gemm_fp16_{nn,nt,tn,tt} pattern. Each per-arch source is built by its own object library pinned with HIP_ARCHITECTURES and gated on CMAKE_HIP_ARCHITECTURES, as gemm/kittens does for its CDNA3/CDNA4 kernels, and exports NVTE_HAVE_MX_GROUPED_GEMM_* so the dispatcher can guard the switch case like KITTENS_HAVE_*. Add the gfx950 backend on the CompAsync pipeline, with shape, warp tile and epilogue taken from KernelTypesMxGemmCompAsync in CK's test/ck_tile/grouped_gemm_mx test suite. This needs a device port of ck_tile::preShuffleScaleBuffer_gfx950, which upstream provides only as a host loop over host memory. Pass TilesPacked_ to CShuffleEpilogueProblem. It defaults to false in the header while the block GEMM emits contiguous MRepeat/NRepeat, which sent correct values through the wrong LDS distribution and produced a permuted output. Only CShuffleEpilogue reads it; TdmEpilogue does not, which is why the gfx1250 path was unaffected. Verified on gfx950: te.GroupedLinear MXFP8 forward through the CK path matches the hipBLASLt baseline exactly. The gfx1250 path is not compiled in a gfx950-only build and remains unverified against this CK pin. Co-Authored-By: Claude Opus 5 (1M context) --- 3rdparty/composable_kernel | 2 +- test_ck_mxfp8_grouped_linear.py | 396 +++++++++++ transformer_engine/common/CMakeLists.txt | 41 ++ .../ck_grouped_gemm/ck_mx_grouped_gemm.cpp | 456 +----------- .../ck_mx_grouped_gemm_gfx1250.cpp | 18 + .../ck_mx_grouped_gemm_gfx950.cpp | 18 + .../ck_grouped_gemm/ck_mx_grouped_gemm_impl.h | 649 ++++++++++++++++++ 7 files changed, 1142 insertions(+), 438 deletions(-) create mode 100644 test_ck_mxfp8_grouped_linear.py create mode 100644 transformer_engine/common/gemm/ck_grouped_gemm/ck_mx_grouped_gemm_gfx1250.cpp create mode 100644 transformer_engine/common/gemm/ck_grouped_gemm/ck_mx_grouped_gemm_gfx950.cpp create mode 100644 transformer_engine/common/gemm/ck_grouped_gemm/ck_mx_grouped_gemm_impl.h diff --git a/3rdparty/composable_kernel b/3rdparty/composable_kernel index 0d18f4fc05..763da730b2 160000 --- a/3rdparty/composable_kernel +++ b/3rdparty/composable_kernel @@ -1 +1 @@ -Subproject commit 0d18f4fc05a31890e5ee365cfc15d82e1ba94669 +Subproject commit 763da730b297c14f2a88f595755a049e8c066340 diff --git a/test_ck_mxfp8_grouped_linear.py b/test_ck_mxfp8_grouped_linear.py new file mode 100644 index 0000000000..f37aa6a0ca --- /dev/null +++ b/test_ck_mxfp8_grouped_linear.py @@ -0,0 +1,396 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# License for AMD contributions = MIT. See LICENSE for more information + +"""Forward-pass check for the CK MXFP8 grouped GEMM behind te.GroupedLinear. + +The CK grouped GEMM is opt-in. Without NVTE_USE_CK_GROUPED_GEMM=1 the MXFP8 +grouped path falls through to hipBLASLt, so a test that merely runs GroupedLinear +proves nothing about CK. This script therefore launches two clean worker +processes -- Transformer Engine is imported only after the backend environment is +established -- and compares them: + + ck : NVTE_USE_CK_GROUPED_GEMM=1 + baseline : no backend flag (hipBLASLt) + +Selection notes (transformer_engine/common/gemm/rocm_gemm.cu): + + kittens_grouped_mxfp8_enabled() == use_hk || (use_cutlass && !use_ck) + +so NVTE_USE_CK_GROUPED_GEMM=1 on its own both enables CK and keeps HipKittens +out of the way -- provided NVTE_USE_HIPKITTENS_GROUPED_GEMM is not also set. + +Reaching the CK entry point is not the same as CK serving the GEMM: it declines +(returns false) on unsupported shapes or insufficient workspace and the caller +quietly uses hipBLASLt instead. Both workers run with +NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK=1 so that decision is visible, and the +ck worker is treated as failed if the fallback warning appears. + +Usage: + + python test_ck_mxfp8_grouped_linear.py + python test_ck_mxfp8_grouped_linear.py --num-gemms 8 --hidden 2048 --ffn 4096 +""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +import tempfile +from pathlib import Path + +# Emitted by cublaslt_gemm.cu when handled_by_ck is false. +FALLBACK_MARKER = "Fallback to cuBLAS grouped GEMM." + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--num-gemms", type=int, default=4, help="Number of experts.") + parser.add_argument("--hidden", type=int, default=1024, help="in_features (K).") + parser.add_argument("--ffn", type=int, default=1024, help="out_features (N).") + parser.add_argument( + "--tokens-per-expert", + type=int, + default=256, + help="Rows per group (M). Must keep every split MXFP8-aligned.", + ) + parser.add_argument("--seed", type=int, default=1234) + parser.add_argument( + "--atol", + type=float, + default=None, + help="Absolute tolerance vs the hipBLASLt baseline. Default scales with K.", + ) + parser.add_argument( + "--rtol", + type=float, + default=0.05, + help="Relative tolerance vs the hipBLASLt baseline.", + ) + # Worker-only arguments. + parser.add_argument("--worker-backend", choices=("ck", "baseline"), default=None) + parser.add_argument("--summary-json", type=Path, default=None) + parser.add_argument("--output-pt", type=Path, default=None) + return parser.parse_args() + + +def validate_args(args: argparse.Namespace) -> None: + # ScaleBlockSize is 32 and the gfx1250 scale preshuffle needs KScale % 4 == 0, + # i.e. K % 128 == 0. gfx950 packs scales in pairs, needing only K % 64 == 0. + # Require the stricter of the two so one shape exercises both backends. + if args.hidden % 128 != 0: + raise SystemExit(f"--hidden must be a multiple of 128, got {args.hidden}") + if args.ffn % 128 != 0: + raise SystemExit(f"--ffn must be a multiple of 128, got {args.ffn}") + if args.tokens_per_expert % 32 != 0: + raise SystemExit( + f"--tokens-per-expert must be a multiple of 32, got {args.tokens_per_expert}" + ) + if args.num_gemms < 1: + raise SystemExit("--num-gemms must be >= 1") + + +# --------------------------------------------------------------------------- +# Worker +# --------------------------------------------------------------------------- + + +def run_worker(args: argparse.Namespace) -> None: + # Backend environment is already set by the parent before this process + # imports torch or Transformer Engine. + import torch + + import transformer_engine.pytorch as te + from transformer_engine.common.recipe import Format, MXFP8BlockScaling + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA/ROCm device is not available") + + backend = args.worker_backend + assert backend is not None + if backend == "ck" and os.environ.get("NVTE_USE_CK_GROUPED_GEMM") != "1": + raise RuntimeError("ck worker was launched without NVTE_USE_CK_GROUPED_GEMM=1") + if backend == "baseline" and os.environ.get("NVTE_USE_CK_GROUPED_GEMM"): + raise RuntimeError("baseline worker unexpectedly has NVTE_USE_CK_GROUPED_GEMM set") + + # Same seed in both workers so the two GroupedLinear instances hold identical + # weights and see identical inputs; any output difference is the GEMM backend. + torch.manual_seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + torch.set_grad_enabled(False) + + device = torch.device("cuda") + dtype = torch.bfloat16 + + layer = te.GroupedLinear( + args.num_gemms, + args.hidden, + args.ffn, + bias=False, + params_dtype=dtype, + device=device, + ).eval() + + m_splits = [args.tokens_per_expert] * args.num_gemms + inp = torch.randn(sum(m_splits), args.hidden, device=device, dtype=dtype) + + recipe = MXFP8BlockScaling(fp8_format=Format.E4M3) + + def fp8_context(): + autocast = getattr(te, "autocast", None) + if autocast is not None: + try: + return autocast(enabled=True, recipe=recipe) + except TypeError: + pass + fp8_autocast = getattr(te, "fp8_autocast", None) + if fp8_autocast is None: + raise RuntimeError( + "Transformer Engine exposes neither te.autocast nor te.fp8_autocast" + ) + return fp8_autocast(enabled=True, fp8_recipe=recipe) + + with fp8_context(): + out = layer(inp, m_splits) + torch.cuda.synchronize() + + if args.output_pt is not None: + torch.save(out.detach().float().cpu(), args.output_pt) + + summary = { + "backend": backend, + "num_gemms": args.num_gemms, + "hidden": args.hidden, + "ffn": args.ffn, + "m_splits": m_splits, + "out_shape": list(out.shape), + "out_dtype": str(out.dtype), + "has_nan": bool(torch.isnan(out).any().item()), + "has_inf": bool(torch.isinf(out).any().item()), + "abs_mean": float(out.detach().float().abs().mean().item()), + "NVTE_ROCM_ENABLE_MXFP8": os.environ.get("NVTE_ROCM_ENABLE_MXFP8"), + "NVTE_USE_CK_GROUPED_GEMM": os.environ.get("NVTE_USE_CK_GROUPED_GEMM"), + "NVTE_USE_HIPKITTENS_GROUPED_GEMM": os.environ.get( + "NVTE_USE_HIPKITTENS_GROUPED_GEMM" + ), + } + if args.summary_json is not None: + args.summary_json.write_text(json.dumps(summary, indent=2)) + print(json.dumps(summary, indent=2)) + + +# --------------------------------------------------------------------------- +# Parent +# --------------------------------------------------------------------------- + + +def worker_command(args: argparse.Namespace, backend: str, summary_json: Path, + output_pt: Path) -> list[str]: + return [ + sys.executable, + str(Path(__file__).resolve()), + "--num-gemms", str(args.num_gemms), + "--hidden", str(args.hidden), + "--ffn", str(args.ffn), + "--tokens-per-expert", str(args.tokens_per_expert), + "--seed", str(args.seed), + "--worker-backend", backend, + "--summary-json", str(summary_json), + "--output-pt", str(output_pt), + ] + + +def run_backend(args: argparse.Namespace, backend: str, workdir: Path) -> dict: + summary_json = workdir / f"{backend}_summary.json" + output_pt = workdir / f"{backend}_out.pt" + + env = os.environ.copy() + # MXFP8 is off by default on ROCm; "1" enables it (quantization.py:177). + # "2" would also make it the default recipe, which we do not need -- the + # worker passes MXFP8BlockScaling explicitly. + env.setdefault("NVTE_ROCM_ENABLE_MXFP8", "1") + # Make the CK-declined-and-fell-back case observable in both workers. + env["NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK"] = "1" + # HipKittens takes precedence over CK when this is set; keep it out of the way. + env.pop("NVTE_USE_HIPKITTENS_GROUPED_GEMM", None) + env.pop("NVTE_USE_CUTLASS_GROUPED_GEMM", None) + if backend == "ck": + env["NVTE_USE_CK_GROUPED_GEMM"] = "1" + else: + env.pop("NVTE_USE_CK_GROUPED_GEMM", None) + + proc = subprocess.run( + worker_command(args, backend, summary_json, output_pt), + env=env, + capture_output=True, + text=True, + ) + combined = proc.stdout + proc.stderr + print(f"----- {backend} worker (exit {proc.returncode}) -----") + print(combined.strip()) + + result = { + "backend": backend, + "returncode": proc.returncode, + "fell_back": FALLBACK_MARKER in combined, + "output_pt": output_pt if output_pt.exists() else None, + "summary": json.loads(summary_json.read_text()) if summary_json.exists() else None, + } + return result + + +def main() -> None: + args = parse_args() + + if args.worker_backend is not None: + run_worker(args) + return + + validate_args(args) + args.m_splits_resolved = [args.tokens_per_expert] * args.num_gemms + + import torch # parent only needs this for the comparison + + failures: list[str] = [] + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + ck = run_backend(args, "ck", workdir) + baseline = run_backend(args, "baseline", workdir) + + for res in (ck, baseline): + if res["returncode"] != 0: + failures.append(f"{res['backend']} worker exited {res['returncode']}") + elif res["summary"] is None: + failures.append(f"{res['backend']} worker produced no summary") + else: + if res["summary"]["has_nan"]: + failures.append(f"{res['backend']} output contains NaN") + if res["summary"]["has_inf"]: + failures.append(f"{res['backend']} output contains Inf") + + # The point of the test: CK must actually have served the GEMM. + if ck["fell_back"]: + failures.append( + "CK declined the grouped GEMM and fell back to hipBLASLt " + f'(saw "{FALLBACK_MARKER}")' + ) + + if ck["output_pt"] and baseline["output_pt"]: + ck_out = torch.load(ck["output_pt"]) + base_out = torch.load(baseline["output_pt"]) + if ck_out.shape != base_out.shape: + failures.append(f"shape mismatch {ck_out.shape} vs {base_out.shape}") + else: + atol = args.atol + if atol is None: + # MXFP8 quantization error grows with the reduction length. + atol = 0.05 * (args.hidden ** 0.5) + diff = (ck_out - base_out).abs() + max_abs = float(diff.max().item()) + denom = base_out.abs().clamp_min(1e-6) + max_rel = float((diff / denom).max().item()) + close = torch.allclose(ck_out, base_out, atol=atol, rtol=args.rtol) + print( + f"\nck vs baseline: max_abs={max_abs:.6g} max_rel={max_rel:.6g} " + f"(atol={atol:.6g} rtol={args.rtol:.6g})" + ) + # Diagnostics for the "right values, wrong places" failure mode: + # a permuted output keeps the multiset of magnitudes intact, so + # abs_mean matches while the elementwise diff is large. + if not close: + ck_sorted = ck_out.flatten().sort().values + base_sorted = base_out.flatten().sort().values + perm_max = float((ck_sorted - base_sorted).abs().max().item()) + print(f" sorted-values max_abs={perm_max:.6g} " + "(small => same values, wrong positions)") + if ck_out.shape[0] == ck_out.shape[1]: + t_max = float((ck_out - base_out.T).abs().max().item()) + print(f" vs baseline transposed: max_abs={t_max:.6g}") + + # Per-group diff, and whether group g of ck matches some + # other group of baseline (i.e. the groups got reordered). + offs = [0] + for s in args.m_splits_resolved: + offs.append(offs[-1] + s) + print(" per-group max_abs (ck[g] vs base[g]):", end=" ") + for g in range(len(args.m_splits_resolved)): + blk_c = ck_out[offs[g]:offs[g + 1]] + blk_b = base_out[offs[g]:offs[g + 1]] + print(f"g{g}={float((blk_c - blk_b).abs().max().item()):.4g}", end=" ") + print() + for g in range(len(args.m_splits_resolved)): + blk_c = ck_out[offs[g]:offs[g + 1]] + matches = [ + h for h in range(len(args.m_splits_resolved)) + if args.m_splits_resolved[h] == args.m_splits_resolved[g] + and torch.equal(blk_c, base_out[offs[h]:offs[h + 1]]) + ] + if matches != [g]: + print(f" ck group {g} exactly matches baseline group(s) {matches}") + + # Row-level: does every ck row appear somewhere in baseline? + row_perm_exact = torch.equal( + ck_out.sort(dim=0).values, base_out.sort(dim=0).values + ) + col_perm_exact = torch.equal( + ck_out.sort(dim=1).values, base_out.sort(dim=1).values + ) + print(f" columnwise-sorted equal: {row_perm_exact} " + f"rowwise-sorted equal: {col_perm_exact}") + + # Tile-local rearrangement: TransposeC flips the layout each + # warp/block tile writes, which is a permutation that is + # neither a row nor a column nor a whole-matrix transpose. + H, W = ck_out.shape + for T in (16, 32, 64, 128): + if H % T or W % T: + continue + a = ck_out.reshape(H // T, T, W // T, T) + b = base_out.reshape(H // T, T, W // T, T) + intra = float((a - b.permute(0, 3, 2, 1)).abs().max().item()) + line = f" tile {T}x{T}: intra-tile-transpose max_abs={intra:.6g}" + if H == W: + grid = float((a - b.permute(2, 1, 0, 3)).abs().max().item()) + line += f" tile-grid-transpose max_abs={grid:.6g}" + print(line) + + # Shift/shear: a wrong leading dimension moves data by a + # constant offset in the flat buffer, preserving the multiset + # while matching none of the structured permutations above. + flat_c = ck_out.flatten() + flat_b = base_out.flatten() + first = flat_c[0] + cand = (flat_b == first).nonzero().flatten().tolist()[:8] + print(f" ck[0,0]={float(first):.6g} occurs in baseline at flat idx {cand}") + for k in cand: + if torch.equal(flat_c, torch.roll(flat_b, -k)): + print(f" EXACT: ck == baseline rolled by {k} elements " + f"({k // W} rows + {k % W} cols)") + break + # Row-level shift. + row0 = ck_out[0] + row_matches = [ + r for r in range(H) if torch.equal(base_out[r], row0) + ][:8] + print(f" ck row 0 matches baseline rows {row_matches}") + if not close: + failures.append( + f"ck and baseline outputs differ beyond tolerance " + f"(max_abs={max_abs:.6g}, max_rel={max_rel:.6g})" + ) + + print() + if failures: + for f in failures: + print(f"FAIL: {f}") + sys.exit(1) + print("PASS: CK MXFP8 grouped GEMM served te.GroupedLinear forward " + "and matched the hipBLASLt baseline.") + + +if __name__ == "__main__": + main() diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index dd4730e475..b92f86f389 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -737,6 +737,47 @@ else() # USE_ROCM target_compile_definitions(transformer_engine PUBLIC USE_HIPKITTENS_GEMM) list(APPEND transformer_engine_LINKER_LIBS kittens_gemm) endif() + + # MXFP8 grouped GEMM backends are architecture specific: CK's TDM pipeline only + # compiles for gfx1250, so it cannot be built for every architecture in + # CMAKE_HIP_ARCHITECTURES like the rest of the library. Compile each backend + # into its own object library pinned to its architecture, as gemm/kittens does + # for its CDNA3/CDNA4 kernels, and tell the dispatcher which ones exist. + set(_mx_grouped_gemm_arch_objs "") + set(_mx_grouped_gemm_have_defs "") + + function(mx_grouped_gemm_add_arch GFX SOURCE) + if(NOT "${GFX}" IN_LIST CMAKE_HIP_ARCHITECTURES) + return() + endif() + string(TOUPPER "${GFX}" _gfx_upper) + set(_target "mx_grouped_gemm_${GFX}") + + add_library(${_target} OBJECT ${SOURCE}) + set_source_files_properties(${SOURCE} PROPERTIES LANGUAGE HIP) + set_target_properties(${_target} PROPERTIES + HIP_ARCHITECTURES "${GFX}" POSITION_INDEPENDENT_CODE ON) + target_include_directories(${_target} PRIVATE + $) + target_compile_definitions(${_target} PRIVATE + $) + target_compile_options(${_target} PRIVATE -fno-gpu-rdc) + target_link_libraries(${_target} PRIVATE hip::host hip::device) + + set(_mx_grouped_gemm_arch_objs ${_mx_grouped_gemm_arch_objs} + $ PARENT_SCOPE) + set(_mx_grouped_gemm_have_defs ${_mx_grouped_gemm_have_defs} + NVTE_HAVE_MX_GROUPED_GEMM_${_gfx_upper} PARENT_SCOPE) + endfunction() + + mx_grouped_gemm_add_arch(gfx950 + gemm/ck_grouped_gemm/ck_mx_grouped_gemm_gfx950_hip.cpp) + mx_grouped_gemm_add_arch(gfx1250 + gemm/ck_grouped_gemm/ck_mx_grouped_gemm_gfx1250_hip.cpp) + + target_sources(transformer_engine PRIVATE ${_mx_grouped_gemm_arch_objs}) + target_compile_definitions(transformer_engine PRIVATE ${_mx_grouped_gemm_have_defs}) + target_link_libraries(transformer_engine PUBLIC ${transformer_engine_LINKER_LIBS}) endif() diff --git a/transformer_engine/common/gemm/ck_grouped_gemm/ck_mx_grouped_gemm.cpp b/transformer_engine/common/gemm/ck_grouped_gemm/ck_mx_grouped_gemm.cpp index 72d1935914..c4f4825f8b 100644 --- a/transformer_engine/common/gemm/ck_grouped_gemm/ck_mx_grouped_gemm.cpp +++ b/transformer_engine/common/gemm/ck_grouped_gemm/ck_mx_grouped_gemm.cpp @@ -5,282 +5,11 @@ ************************************************************************/ #include "ck_grouped_gemm_common.h" +#include "ck_mx_grouped_gemm_impl.h" namespace transformer_engine { namespace grouped_gemm { -using mx_grouped_gemm_kargs = ck_tile::MxGroupedGemmHostArgs<>; - -static constexpr ck_tile::index_t ScaleBlockSize = 32; - -enum struct MxGemmPipelineType { - CompTDMV1, - CompTDMV2 -}; - -template -struct MxGemmPipelineTypeSelector; - -template -struct MxGemmPipelineTypeSelector { - using base_pipeline = ck_tile::BaseGemmPipelineAgBgCrCompTDM; - using pipeline = ck_tile::GemmPipelineAgBgCrCompTDMV1; - static constexpr auto GetName() { return "GemmPipelineAgBgCrCompTDMV1"; } -}; - -template -struct MxGemmPipelineTypeSelector { - using base_pipeline = ck_tile::BaseGemmPipelineAgBgCrCompTDM; - using pipeline = ck_tile::GemmPipelineAgBgCrCompTDMV2; - static constexpr auto GetName() { return "GemmPipelineAgBgCrCompTDMV2"; } -}; - -struct GroupedGemKernelParam_Wmma { - static const bool kPadM = false; - static const bool kPadN = false; - static const bool kPadK = false; - static const int kBlockPerCu = 1; - static const ck_tile::index_t M_Tile = 64; - static const ck_tile::index_t N_Tile = 64; - static const ck_tile::index_t K_Tile = 128; - static const ck_tile::index_t M_Warp = 2; - static const ck_tile::index_t N_Warp = 2; - static const ck_tile::index_t K_Warp = 1; - static const ck_tile::index_t M_Warp_Tile = 32; - static const ck_tile::index_t N_Warp_Tile = 32; - static constexpr ck_tile::index_t K_Warp_Tile = 128; -}; - -// gfx1250 scale preshuffle. -// -// Unlike the existing MXFP8 GEMM scale swizzle defined in: -// transformer_engine/common/swizzle/swizzle.cu -// -// CK gfx1250 WMMA kernels expect scales in the layout below: -// -// Input scales are logically [MN, KScale] -// -// The output layout groups KScale into tiles of 4 (= 128 / ScaleBlockSize) -// and additionally blocks M into chunks of 32 rows: -// -// [MN, KScale] -// -> [MN/32, KScale/4, 32, 4] -// -// For A scales, rows=M and output_rows is M padded to M_Warp_Tile. -// For B scales, rows=N and output_rows is currently N. -template -__global__ void preshuffle_scale_gfx1250_kernel(const ScaleType *__restrict__ src, - ScaleType *__restrict__ dst, - int actual_rows, - int output_rows, - int KScale) { - static_assert(ScaleBlockSize == 32 && sizeof(ScaleType) == 1, - "gfx1250 scale preshuffle only supports 8-bit scale with ScaleBlockSize=32"); - constexpr int MPerXdlops = 16; - constexpr int KPerXdlops = 128; - constexpr int MNPack = 2; - constexpr int KPack = 1; - constexpr int MNStep = MPerXdlops; // 16 - constexpr int KStep = KPerXdlops / ScaleBlockSize; // 4 - const int K0 = KScale / (KPack * KStep); - const int linear = blockIdx.x * blockDim.x + threadIdx.x; - const int total = output_rows * KScale; - if (linear >= total) { - return; - } - const int mn = linear / KScale; - const int k = linear % KScale; - const int iMNRepeat = mn / (MNStep * MNPack); - const int tempmn = mn % (MNStep * MNPack); - const int iKRepeat = k / (KStep * KPack); - const int tempk = k % (KStep * KPack); - const int outputIndex = - (iMNRepeat * MNPack * MNStep) * (KStep * KPack * K0) + - (iKRepeat * KStep * KPack) * (MNStep * MNPack) + - tempmn * (KStep * KPack) + - tempk; - ScaleType value{}; - if (mn < actual_rows) { - if constexpr (KStride) { - value = src[mn * KScale + k]; - } else { - value = src[k * actual_rows + mn]; - } - } - dst[outputIndex] = value; -} - -template -void preShuffleScaleBuffer_gfx1250(const ScaleType *src, - ScaleType *dst, - int actual_rows, - int output_rows, - int KScale, - hipStream_t stream) { - constexpr int KPerXdlops = 128; - constexpr int KStep = KPerXdlops / ScaleBlockSize; // 4 - if (KScale % KStep != 0) { - NVTE_ERROR("preshuffle_scale_gfx1250: KScale must be a multiple of 4, " - "i.e. original K must be a multiple of 128 for ScaleBlockSize=32."); - } - const int total = output_rows * KScale; - constexpr int block_size = 256; - const int grid_size = (total + block_size - 1) / block_size; - hipLaunchKernelGGL((preshuffle_scale_gfx1250_kernel), - dim3(grid_size), - dim3(block_size), - 0, - stream, - src, - dst, - actual_rows, - output_rows, - KScale); - NVTE_CHECK_CUDA(hipGetLastError()); -} - -template -bool invoke_mx_grouped_gemm(const std::vector &descs, - const GroupedGemmRunContext &ctx, - const ck_tile::stream_config &stream_cfg, - bool warn_fallback) { - // Check hardware WMMA support for the warp tile. - static constexpr bool has_wmma_support = - ck_tile::has_wmma_traits_v; - - NVTE_CHECK(has_wmma_support, - "ck_tile_mx_grouped_gemm: unsupported gfx125 WMMA traits for " - "AType/BType/AccType with warp tile shape ", - MXFP8GemmConfig::M_Warp_Tile, "x", - MXFP8GemmConfig::N_Warp_Tile, "x", - MXFP8GemmConfig::K_Warp_Tile); - - using CLayout = RowMajor; - constexpr bool preshuffle = false; - constexpr bool DoubleSmemBuffer = true; // TDM pipeline requires double smem buffer - constexpr bool TransposeC = - std::is_same_v && - MXFP8GemmConfig::M_Warp_Tile == MXFP8GemmConfig::N_Warp_Tile; - static constexpr bool StructuredSparsity = false; - static constexpr bool NumWaveGroup = 1; - constexpr ck_tile::index_t TileParitionerGroupNum = 8; - constexpr ck_tile::index_t TileParitionerM01 = 4; - using GemmShape = - ck_tile::TileGemmShape, - ck_tile::sequence, - ck_tile::sequence>; - using TilePartitioner = ck_tile:: - GemmSpatiallyLocalTilePartitioner; - TRANSFORMER_ENGINE_SWITCH_CONDITION(ctx.transA, kTransA, { - using ALayout = std::conditional_t; - TRANSFORMER_ENGINE_SWITCH_CONDITION(ctx.transB, kTransB, { - using BLayout = std::conditional_t; - using GemmUniversalTraits = ck_tile::TileGemmUniversalTraits; - using UniversalGemmProblem = - ck_tile::MxGemmPipelineProblem; - /* Make pipeline selective. */ - using GemmPipeline = - typename MxGemmPipelineTypeSelector< - PipelineType, - UniversalGemmProblem>::pipeline; - - using GemmEpilogue = ck_tile::TdmEpilogue< - ck_tile::CShuffleEpilogueProblem, // DsDataType - float, - CType, - ck_tile::tuple<>, // DsLayout - CLayout, - ck_tile::element_wise::PassThrough, - TilePartitioner::MPerBlock, - TilePartitioner::NPerBlock, - MXFP8GemmConfig::M_Warp, - MXFP8GemmConfig::N_Warp, - MXFP8GemmConfig::M_Warp_Tile, - MXFP8GemmConfig::N_Warp_Tile, - MXFP8GemmConfig::K_Warp_Tile, - UniversalGemmProblem::TransposeC, - 1, /* kNumWaveGroups_ */ - false, /* FixedVectorSize_ */ - 1, /* VectorSizeC_ */ - 1, /* BlockedXDLN_PerWarp_ */ - DoubleSmemBuffer, /* DoubleSmemBuffer */ - AType, /* AType_ */ - BType /* BType_ */>>; - using Kernel = ck_tile::MxGroupedGemmKernel; - - if (!has_sufficient_workspace(ctx)) { - return false; - } - - auto kargs = Kernel::MakeKargs(descs); - if (!Kernel::IsSupportedArgument(kargs)) { - if (warn_fallback) { - NVTE_WARN("ck_tile_mx_grouped_gemm: CK_Tile kernel arguments not supported for this config. " - "Falling back."); - } - return false; - } - const dim3 blocks = Kernel::BlockSize(); - const dim3 grids = Kernel::GridSize(kargs); - NVTE_CHECK_CUDA(hipMemcpyAsync(ctx.workspace, - kargs.data(), - kargs.size() * sizeof(typename decltype(kargs)::value_type), - hipMemcpyHostToDevice, - ctx.stream)); - ck_tile::ignore = ck_tile::launch_kernel( - stream_cfg, ck_tile::make_kernel( - Kernel{}, grids, blocks, 0, - ck_tile::cast_pointer_to_constant_address_space(ctx.workspace), - kargs.size())); - return true; - }); - }); - return false; -} - bool ck_tile_mx_grouped_gemm(const NVTETensor *A, const NVTETensor *B, NVTETensor *D, @@ -294,13 +23,6 @@ bool ck_tile_mx_grouped_gemm(const NVTETensor *A, const bool warn_fallback = getenv("NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK", false); - if (detect_gpu_arch() != GPUArch::GFX1250) { - if (warn_fallback) { - NVTE_WARN("ck_tile_mx_grouped_gemm: only supported on gfx1250. Falling back."); - } - return false; - } - if (group_num <= 0) { return true; } @@ -357,9 +79,6 @@ bool ck_tile_mx_grouped_gemm(const NVTETensor *A, NVTE_CHECK(is_fp8_dtype(a_dtype), "ck_tile_mx_grouped_gemm: A dtype must be FP8"); NVTE_CHECK(is_fp8_dtype(b_dtype), "ck_tile_mx_grouped_gemm: B dtype must be FP8"); - using AScaleType = ck_tile::e8m0_t; - using BScaleType = ck_tile::e8m0_t; - void *ws_ptr = nullptr; size_t ws_bytes = 0; if (workspace) { @@ -385,162 +104,25 @@ bool ck_tile_mx_grouped_gemm(const NVTETensor *A, .accumulate = false, }; - const ck_tile::stream_config s{ctx.stream}; - - std::vector descs; - descs.reserve(group_num); - - NVTE_CHECK(ctx.workspace != nullptr, - "ck_tile_mx_grouped_gemm: workspace is required for shuffled MXFP8 scales."); - - // Carve regions from the end of the workspace for mxfp8 scales. - // Layout: [CK kargs workspace ... | a_scales (i) | b_scales (i) | ... | a_scales (group_num-1) | b_scales (group_num-1)] - constexpr size_t kScaleWorkspaceAlign = 256; - uint8_t *scale_workspace_base = reinterpret_cast(ctx.workspace); - size_t scale_workspace_end = - (ctx.workspace_bytes / kScaleWorkspaceAlign) * kScaleWorkspaceAlign; - - for (int i = 0; i < group_num; i++) { - const transformer_engine::Tensor *const A_te = - transformer_engine::convertNVTETensorCheck(ctx.A[i]); - const transformer_engine::Tensor *const B_te = - transformer_engine::convertNVTETensorCheck(ctx.B[i]); - transformer_engine::Tensor *D_te = - transformer_engine::convertNVTETensorCheck(ctx.D[i]); - - const auto &a = ctx.use_a_columnwise_data ? A_te->columnwise_data : A_te->data; - const auto &b = ctx.use_b_columnwise_data ? B_te->columnwise_data : B_te->data; - const auto &d = D_te->data; - const auto &a_scales = - ctx.use_a_columnwise_data ? A_te->columnwise_scale_inv : A_te->scale_inv; - const auto &b_scales = - ctx.use_b_columnwise_data ? B_te->columnwise_scale_inv : B_te->scale_inv; - - int64_t Ad0 = 0, Ad1 = 0, Bd0 = 0, Bd1 = 0, Dd0 = 0, Dd1 = 0; - - if (!get_flat_2d_dims(*A_te, Ad0, Ad1)) { - NVTE_ERROR("ck_tile_mx_grouped_gemm: expected rank>=2 for normalized A in group ", i); - } - - if (!get_flat_2d_dims(*B_te, Bd0, Bd1)) { - NVTE_ERROR("ck_tile_mx_grouped_gemm: expected rank>=2 for normalized B in group ", i); - } - - if (!get_flat_2d_dims(*D_te, Dd0, Dd1)) { - NVTE_ERROR("ck_tile_mx_grouped_gemm: expected rank>=2 for normalized D in group ", i); - } - if (a.dptr == nullptr || b.dptr == nullptr || a_scales.dptr == nullptr || - b_scales.dptr == nullptr) { - NVTE_ERROR("ck_tile_mx_grouped_gemm: effective A/B data or scale_inv is missing."); - } - if (a_scales.shape.size() != 2 || b_scales.shape.size() != 2) { - NVTE_ERROR("ck_tile_mx_grouped_gemm: expected effective A/B scale_inv tensors to be rank-2."); - } - - const size_t M = ctx.transA ? Ad1 : Ad0; - const size_t K = ctx.transA ? Ad0 : Ad1; - const size_t N = ctx.transB ? Bd0 : Bd1; - const size_t Kb = ctx.transB ? Bd1 : Bd0; - if (K % ScaleBlockSize != 0) { - NVTE_ERROR("ck_tile_mx_grouped_gemm: K must be a multiple of ScaleBlockSize for MX GEMM", i); - } - const int KScale = static_cast(K / ScaleBlockSize); - if (Kb != K) { - NVTE_ERROR("ck_tile_mx_grouped_gemm: K mismatch between A and B in group ", i, - ". op(A)=", M, "x", K, ", op(B)=", Kb, "x", N); - } - if (Dd0 != M || Dd1 != N) { - NVTE_ERROR("ck_tile_mx_grouped_gemm: D shape mismatch in group ", i, - ". D=", Dd0, "x", Dd1, ", expected=", M, "x", N); - } - - const ck_tile::index_t stride_A = static_cast(Ad1); - const ck_tile::index_t stride_B = static_cast(Bd1); - const ck_tile::index_t stride_E = static_cast(Dd1); - - // Pre-shuffle scale buffers for the hardware. - const int a_scale_actual_rows = static_cast(M); - const int a_scale_output_rows = - ck_tile::integer_least_multiple( - static_cast(M), - static_cast(GroupedGemKernelParam_Wmma::M_Warp_Tile)); - const int b_scale_actual_rows = static_cast(N); - const int b_scale_output_rows = static_cast(N); - const size_t a_scale_shuffled_bytes = - static_cast(a_scale_output_rows) * - static_cast(KScale) * - sizeof(AScaleType); - const size_t b_scale_shuffled_bytes = - static_cast(b_scale_output_rows) * - static_cast(KScale) * - sizeof(BScaleType); - const size_t scale_pair_bytes = - a_scale_shuffled_bytes + b_scale_shuffled_bytes; - scale_workspace_end = - (scale_workspace_end / kScaleWorkspaceAlign) * kScaleWorkspaceAlign; - - NVTE_CHECK(scale_workspace_end >= scale_pair_bytes, - "ck_tile_mx_grouped_gemm: insufficient workspace for shuffled MXFP8 scales. " - "Need current group scale bytes=", scale_pair_bytes, - ", available workspace bytes=", scale_workspace_end, - ". Increase the grouped GEMM workspace size."); - - scale_workspace_end -= scale_pair_bytes; - uint8_t *scale_pair_ptr = scale_workspace_base + scale_workspace_end; - - void *a_scale_shuffled_ptr = scale_pair_ptr; - void *b_scale_shuffled_ptr = scale_pair_ptr + a_scale_shuffled_bytes; - - // CK expects canonical pre-shuffled scale buffers laid out as - // A: [M, KScale] and B: [N, KScale], independent of A/B data layouts. - // TE rowwise MXFP8 scale_inv is [rows, KScale] and can be read with - // KStride=true. TE columnwise_scale_inv is [KScale, rows] and must be - // read with KStride=false before writing CK's canonical shuffled layout. - if (ctx.use_a_columnwise_data) { - preShuffleScaleBuffer_gfx1250( - reinterpret_cast(a_scales.dptr), - reinterpret_cast(a_scale_shuffled_ptr), - a_scale_actual_rows, a_scale_output_rows, KScale, stream); - } else { - preShuffleScaleBuffer_gfx1250( - reinterpret_cast(a_scales.dptr), - reinterpret_cast(a_scale_shuffled_ptr), - a_scale_actual_rows, a_scale_output_rows, KScale, stream); - } - - if (ctx.use_b_columnwise_data) { - preShuffleScaleBuffer_gfx1250( - reinterpret_cast(b_scales.dptr), - reinterpret_cast(b_scale_shuffled_ptr), - b_scale_actual_rows, b_scale_output_rows, KScale, stream); - } else { - preShuffleScaleBuffer_gfx1250( - reinterpret_cast(b_scales.dptr), - reinterpret_cast(b_scale_shuffled_ptr), - b_scale_actual_rows, b_scale_output_rows, KScale, stream); - } - descs.emplace_back(mx_grouped_gemm_kargs( - a.dptr, a_scale_shuffled_ptr, b.dptr, b_scale_shuffled_ptr, - {/*ds_ptr*/}, d.dptr, 1, // kbatch - M, N, K, stride_A, stride_B, {/*stride_Ds*/}, stride_E)); + // Dispatch to per-architecture translation unit. Each backend is only built + // when its architecture is targeted, so guard on the availability macro the + // build defines for it. + switch (ctx.arch) { +#ifdef NVTE_HAVE_MX_GROUPED_GEMM_GFX1250 + case GPUArch::GFX1250: + return ck_tile_mx_grouped_gemm_dispatch_gfx1250(a_dtype, b_dtype, d_dtype, ctx); +#endif +#ifdef NVTE_HAVE_MX_GROUPED_GEMM_GFX950 + case GPUArch::GFX950: + return ck_tile_mx_grouped_gemm_dispatch_gfx950(a_dtype, b_dtype, d_dtype, ctx); +#endif + default: + if (warn_fallback) { + NVTE_WARN("ck_tile_mx_grouped_gemm: no MX grouped GEMM kernel built for this " + "architecture. Falling back."); + } + return false; } - ctx.workspace_bytes = scale_workspace_end; - - // Invoke the GEMM. - bool ok = false; - TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY(a_dtype, a_te_type, { - using AType = typename TETypeToCKType::type; - TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY(b_dtype, b_te_type, { - using BType = typename TETypeToCKType::type; - TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(d_dtype, d_te_type, { - using CType = typename TETypeToCKType::type; - ok = invoke_mx_grouped_gemm(descs, ctx, s, warn_fallback); - }); // NOLINT(*) - }); // NOLINT(*) - }); // NOLINT(*) - return ok; } } // namespace grouped_gemm diff --git a/transformer_engine/common/gemm/ck_grouped_gemm/ck_mx_grouped_gemm_gfx1250.cpp b/transformer_engine/common/gemm/ck_grouped_gemm/ck_mx_grouped_gemm_gfx1250.cpp new file mode 100644 index 0000000000..62b8732d60 --- /dev/null +++ b/transformer_engine/common/gemm/ck_grouped_gemm/ck_mx_grouped_gemm_gfx1250.cpp @@ -0,0 +1,18 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * License for AMD contributions = MIT. See LICENSE for more information + ************************************************************************/ + +#include "ck_mx_grouped_gemm_impl.h" + +namespace transformer_engine { +namespace grouped_gemm { + +bool ck_tile_mx_grouped_gemm_dispatch_gfx1250(DType a_dtype, DType b_dtype, DType d_dtype, + const GroupedGemmRunContext& ctx) { + return ck_tile_mx_grouped_gemm_impl(a_dtype, b_dtype, d_dtype, ctx); +} + +} // namespace grouped_gemm +} // namespace transformer_engine diff --git a/transformer_engine/common/gemm/ck_grouped_gemm/ck_mx_grouped_gemm_gfx950.cpp b/transformer_engine/common/gemm/ck_grouped_gemm/ck_mx_grouped_gemm_gfx950.cpp new file mode 100644 index 0000000000..d5da5398dd --- /dev/null +++ b/transformer_engine/common/gemm/ck_grouped_gemm/ck_mx_grouped_gemm_gfx950.cpp @@ -0,0 +1,18 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * License for AMD contributions = MIT. See LICENSE for more information + ************************************************************************/ + +#include "ck_mx_grouped_gemm_impl.h" + +namespace transformer_engine { +namespace grouped_gemm { + +bool ck_tile_mx_grouped_gemm_dispatch_gfx950(DType a_dtype, DType b_dtype, DType d_dtype, + const GroupedGemmRunContext& ctx) { + return ck_tile_mx_grouped_gemm_impl(a_dtype, b_dtype, d_dtype, ctx); +} + +} // namespace grouped_gemm +} // namespace transformer_engine diff --git a/transformer_engine/common/gemm/ck_grouped_gemm/ck_mx_grouped_gemm_impl.h b/transformer_engine/common/gemm/ck_grouped_gemm/ck_mx_grouped_gemm_impl.h new file mode 100644 index 0000000000..b5b92ffb62 --- /dev/null +++ b/transformer_engine/common/gemm/ck_grouped_gemm/ck_mx_grouped_gemm_impl.h @@ -0,0 +1,649 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * License for AMD contributions = MIT. See LICENSE for more information + ************************************************************************/ + +#pragma once + +#include "ck_grouped_gemm_common.h" + +namespace transformer_engine { +namespace grouped_gemm { + +using mx_grouped_gemm_kargs = ck_tile::MxGroupedGemmHostArgs<>; + +static constexpr ck_tile::index_t ScaleBlockSize = 32; + +enum struct MxGemmPipelineType { + CompTDMV1, + CompTDMV2, + CompAsync +}; + +template +struct MxGemmPipelineTypeSelector; + +template +struct MxGemmPipelineTypeSelector { + using base_pipeline = ck_tile::BaseGemmPipelineAgBgCrCompTDM; + using pipeline = ck_tile::GemmPipelineAgBgCrCompTDMV1; + static constexpr auto GetName() { return "GemmPipelineAgBgCrCompTDMV1"; } +}; + +template +struct MxGemmPipelineTypeSelector { + using base_pipeline = ck_tile::BaseGemmPipelineAgBgCrCompTDM; + using pipeline = ck_tile::GemmPipelineAgBgCrCompTDMV2; + static constexpr auto GetName() { return "GemmPipelineAgBgCrCompTDMV2"; } +}; + +template +struct MxGemmPipelineTypeSelector { + using base_pipeline = ck_tile::BaseGemmPipelineAgBgCrCompAsync; + using pipeline = ck_tile::GemmPipelineAgBgCrCompAsync; + static constexpr auto GetName() { return "GemmPipelineAgBgCrCompAsync"; } +}; + +// The TDM pipelines drive the epilogue directly; the async pipeline writes C +// through LDS and needs the CShuffle epilogue. +template +struct MxGemmEpilogueTypeSelector; + +template +struct MxGemmEpilogueTypeSelector { + using epilogue = ck_tile::TdmEpilogue; +}; + +template +struct MxGemmEpilogueTypeSelector { + using epilogue = ck_tile::TdmEpilogue; +}; + +template +struct MxGemmEpilogueTypeSelector { + using epilogue = ck_tile::CShuffleEpilogue; +}; + +// gfx1250 scale preshuffle. +// +// Unlike the existing MXFP8 GEMM scale swizzle defined in: +// transformer_engine/common/swizzle/swizzle.cu +// +// CK gfx1250 WMMA kernels expect scales in the layout below: +// +// Input scales are logically [MN, KScale] +// +// The output layout groups KScale into tiles of 4 (= 128 / ScaleBlockSize) +// and additionally blocks M into chunks of 32 rows: +// +// [MN, KScale] +// -> [MN/32, KScale/4, 32, 4] +// +// For A scales, rows=M and output_rows is M padded to M_Warp_Tile. +// For B scales, rows=N and output_rows is currently N. +template +__global__ void preshuffle_scale_gfx1250_kernel(const ScaleType *__restrict__ src, + ScaleType *__restrict__ dst, + int actual_rows, + int output_rows, + int KScale) { + static_assert(ScaleBlockSize == 32 && sizeof(ScaleType) == 1, + "gfx1250 scale preshuffle only supports 8-bit scale with ScaleBlockSize=32"); + constexpr int MPerXdlops = 16; + constexpr int KPerXdlops = 128; + constexpr int MNPack = 2; + constexpr int KPack = 1; + constexpr int MNStep = MPerXdlops; // 16 + constexpr int KStep = KPerXdlops / ScaleBlockSize; // 4 + const int K0 = KScale / (KPack * KStep); + const int linear = blockIdx.x * blockDim.x + threadIdx.x; + const int total = output_rows * KScale; + if (linear >= total) { + return; + } + const int mn = linear / KScale; + const int k = linear % KScale; + const int iMNRepeat = mn / (MNStep * MNPack); + const int tempmn = mn % (MNStep * MNPack); + const int iKRepeat = k / (KStep * KPack); + const int tempk = k % (KStep * KPack); + const int outputIndex = + (iMNRepeat * MNPack * MNStep) * (KStep * KPack * K0) + + (iKRepeat * KStep * KPack) * (MNStep * MNPack) + + tempmn * (KStep * KPack) + + tempk; + ScaleType value{}; + if (mn < actual_rows) { + if constexpr (KStride) { + value = src[mn * KScale + k]; + } else { + value = src[k * actual_rows + mn]; + } + } + dst[outputIndex] = value; +} + +template +void preShuffleScaleBuffer_gfx1250(const ScaleType *src, + ScaleType *dst, + int actual_rows, + int output_rows, + int KScale, + hipStream_t stream) { + constexpr int KPerXdlops = 128; + constexpr int KStep = KPerXdlops / ScaleBlockSize; // 4 + if (KScale % KStep != 0) { + NVTE_ERROR("preshuffle_scale_gfx1250: KScale must be a multiple of 4, " + "i.e. original K must be a multiple of 128 for ScaleBlockSize=32."); + } + const int total = output_rows * KScale; + constexpr int block_size = 256; + const int grid_size = (total + block_size - 1) / block_size; + hipLaunchKernelGGL((preshuffle_scale_gfx1250_kernel), + dim3(grid_size), + dim3(block_size), + 0, + stream, + src, + dst, + actual_rows, + output_rows, + KScale); + NVTE_CHECK_CUDA(hipGetLastError()); +} + +// gfx950 scale preshuffle. +// +// Device port of ck_tile::preShuffleScaleBuffer_gfx950 (ck_tile/host/mx_processing.hpp), +// which is a host-only loop over host memory. The index math below is copied from +// it verbatim; keep the two in sync. One thread handles one (packed_mn, packed_k) +// pair and writes the MNPack * KPack scales belonging to it. +// +// KStride selects how the source is read, matching the gfx1250 kernel: +// true -> src is [MN, KScale] (TE rowwise scale_inv) +// false -> src is [KScale, MN] (TE columnwise_scale_inv) +template +__global__ void preshuffle_scale_gfx950_kernel(const ScaleType *__restrict__ src, + ScaleType *__restrict__ dst, + int actual_rows, + int output_rows, + int KScale) { + static_assert(sizeof(ScaleType) == 1, + "gfx950 scale preshuffle only supports 8-bit scale types"); + constexpr int NumScalesPerDword = 4 / sizeof(ScaleType); + const int MN_packed = output_rows / MNPack; + const int K_packed = KScale / KPack; + const int linear = blockIdx.x * blockDim.x + threadIdx.x; + if (linear >= MN_packed * K_packed) { + return; + } + const int packed_mn = linear / K_packed; + const int packed_k = linear % K_packed; + const int mn_lane = packed_mn % XdlMNThread; + const int mn_group = packed_mn / XdlMNThread; + const int k_lane = packed_k % XdlKThread; + const int k_group = packed_k / XdlKThread; + for (int ik = 0; ik < KPack; ik++) { + for (int imn = 0; imn < MNPack; imn++) { + const int byteIdx = ik * MNPack + imn; + const int orig_mn = mn_group * XdlMNThread * MNPack + imn * XdlMNThread + mn_lane; + const int orig_k = k_group * XdlKThread * KPack + ik * XdlKThread + k_lane; + ScaleType value{}; + if (orig_mn < actual_rows) { + if constexpr (KStride) { + value = src[orig_k + static_cast(orig_mn) * KScale]; + } else { + value = src[orig_mn + static_cast(orig_k) * actual_rows]; + } + } + const int64_t outputIndex = + byteIdx + static_cast(mn_lane) * NumScalesPerDword + + static_cast(packed_k) * XdlMNThread * NumScalesPerDword + + static_cast(mn_group) * XdlMNThread * NumScalesPerDword * K_packed; + dst[outputIndex] = value; + } + } +} + +template +void preShuffleScaleBuffer_gfx950(const ScaleType *src, + ScaleType *dst, + int actual_rows, + int output_rows, + int KScale, + hipStream_t stream) { + if (output_rows % MNPack != 0 || KScale % KPack != 0) { + NVTE_ERROR("preshuffle_scale_gfx950: output_rows must be a multiple of ", MNPack, + " and KScale a multiple of ", KPack, "."); + } + const int total = (output_rows / MNPack) * (KScale / KPack); + constexpr int block_size = 256; + const int grid_size = (total + block_size - 1) / block_size; + hipLaunchKernelGGL((preshuffle_scale_gfx950_kernel), + dim3(grid_size), + dim3(block_size), + 0, + stream, + src, + dst, + actual_rows, + output_rows, + KScale); + NVTE_CHECK_CUDA(hipGetLastError()); +} + +template struct MxTileCfg; + +template <> struct MxTileCfg { + static const bool kPadM = false; + static const bool kPadN = false; + static const bool kPadK = false; + static const int kBlockPerCu = 1; + static const ck_tile::index_t M_Tile = 64; + static const ck_tile::index_t N_Tile = 64; + static const ck_tile::index_t K_Tile = 128; + static const ck_tile::index_t M_Warp = 2; + static const ck_tile::index_t N_Warp = 2; + static const ck_tile::index_t K_Warp = 1; + static const ck_tile::index_t M_Warp_Tile = 32; + static const ck_tile::index_t N_Warp_Tile = 32; + static constexpr ck_tile::index_t K_Warp_Tile = 128; + static constexpr MxGemmPipelineType PipelineType = MxGemmPipelineType::CompTDMV1; + // WMMA writes C transposed when the warp tile is square and C is RowMajor. + static constexpr bool TransposeC = true; + // A scales are padded to the warp tile; the WMMA swizzle blocks M by 32 rows. + static constexpr ck_tile::index_t ScalePadM = M_Warp_Tile; +}; + +// gfx950 MFMA. Shape and pipeline follow KernelTypesMxGemmCompAsync in CK's +// test/ck_tile/grouped_gemm_mx/test_mx_grouped_gemm_pipeline_kernel_types.hpp; +// K_Warp_Tile is what get_k_warp_tile() there yields for a 16-wide MFMA warp tile. +template <> struct MxTileCfg { + static const bool kPadM = false; + static const bool kPadN = false; + static const bool kPadK = false; + static const int kBlockPerCu = 1; + static const ck_tile::index_t M_Tile = 64; + static const ck_tile::index_t N_Tile = 64; + static const ck_tile::index_t K_Tile = 256; + static const ck_tile::index_t M_Warp = 2; + static const ck_tile::index_t N_Warp = 2; + static const ck_tile::index_t K_Warp = 1; + static const ck_tile::index_t M_Warp_Tile = 16; + static const ck_tile::index_t N_Warp_Tile = 16; + static constexpr ck_tile::index_t K_Warp_Tile = 128; + static constexpr MxGemmPipelineType PipelineType = MxGemmPipelineType::CompAsync; + static constexpr bool TransposeC = false; + // CK's test pads A scales to the block tile before the MFMA swizzle. + static constexpr ck_tile::index_t ScalePadM = M_Tile; +}; + +// XdlPack factors for the gfx950 scale swizzle, derived exactly as CK's test does: +// pack two iterations together whenever the per-warp iteration count is even. +template +inline constexpr ck_tile::index_t MxXdlPackEff = + (IterPerWarp >= 2 && IterPerWarp % 2 == 0) ? 2 : 1; + +template +bool invoke_mx_grouped_gemm(const std::vector &descs, + const GroupedGemmRunContext &ctx, + const ck_tile::stream_config &stream_cfg, + bool warn_fallback) { + + using Cfg = MxTileCfg; + + // Check hardware WMMA support for the warp tile. gfx950 uses MFMA, not WMMA, + // so the gfx125 traits do not apply there. + if constexpr (Arch == GPUArch::GFX1250) { + static constexpr bool has_wmma_support = + ck_tile::has_wmma_traits_v; + + NVTE_CHECK(has_wmma_support, + "ck_tile_mx_grouped_gemm: unsupported gfx125 WMMA traits for " + "AType/BType/AccType with warp tile shape ", + Cfg::M_Warp_Tile, "x", + Cfg::N_Warp_Tile, "x", + Cfg::K_Warp_Tile); + } + + using CLayout = RowMajor; + constexpr bool preshuffle = false; + // Both the TDM and async pipelines static_assert on this being true. + constexpr bool DoubleSmemBuffer = true; + constexpr bool TransposeC = Cfg::TransposeC; + static constexpr bool StructuredSparsity = false; + static constexpr bool NumWaveGroup = 1; + constexpr ck_tile::index_t TileParitionerGroupNum = 8; + constexpr ck_tile::index_t TileParitionerM01 = 4; + using GemmShape = + ck_tile::TileGemmShape, + ck_tile::sequence, + ck_tile::sequence>; + using TilePartitioner = ck_tile:: + GemmSpatiallyLocalTilePartitioner; + TRANSFORMER_ENGINE_SWITCH_CONDITION(ctx.transA, kTransA, { + using ALayout = std::conditional_t; + TRANSFORMER_ENGINE_SWITCH_CONDITION(ctx.transB, kTransB, { + using BLayout = std::conditional_t; + using GemmUniversalTraits = ck_tile::TileGemmUniversalTraits; + using UniversalGemmProblem = + ck_tile::MxGemmPipelineProblem; + /* Make pipeline selective. */ + using GemmPipeline = + typename MxGemmPipelineTypeSelector< + Cfg::PipelineType, + UniversalGemmProblem>::pipeline; + + using GemmEpilogueProblem = + ck_tile::CShuffleEpilogueProblem, // DsDataType + float, + CType, + ck_tile::tuple<>, // DsLayout + CLayout, + ck_tile::element_wise::PassThrough, + TilePartitioner::MPerBlock, + TilePartitioner::NPerBlock, + Cfg::M_Warp, + Cfg::N_Warp, + Cfg::M_Warp_Tile, + Cfg::N_Warp_Tile, + Cfg::K_Warp_Tile, + UniversalGemmProblem::TransposeC, + 1, /* kNumWaveGroups_ */ + false, /* FixedVectorSize_ */ + 1, /* VectorSizeC_ */ + 1, /* BlockedXDLN_PerWarp_ */ + DoubleSmemBuffer, /* DoubleSmemBuffer */ + AType, /* AType_ */ + BType, /* BType_ */ + // TilesPacked_: the block GEMM emits contiguous + // MRepeat/NRepeat because the scales are packed. + // Only CShuffleEpilogue reads this; TdmEpilogue + // ignores it, so gfx1250 is unaffected. + !preshuffle>; + using GemmEpilogue = + typename MxGemmEpilogueTypeSelector::epilogue; + using Kernel = ck_tile::MxGroupedGemmKernel; + + if (!has_sufficient_workspace(ctx)) { + return false; + } + + auto kargs = Kernel::MakeKargs(descs); + if (!Kernel::IsSupportedArgument(kargs)) { + if (warn_fallback) { + NVTE_WARN("ck_tile_mx_grouped_gemm: CK_Tile kernel arguments not supported for this config. " + "Falling back."); + } + return false; + } + const dim3 blocks = Kernel::BlockSize(); + const dim3 grids = Kernel::GridSize(kargs); + NVTE_CHECK_CUDA(hipMemcpyAsync(ctx.workspace, + kargs.data(), + kargs.size() * sizeof(typename decltype(kargs)::value_type), + hipMemcpyHostToDevice, + ctx.stream)); + ck_tile::ignore = ck_tile::launch_kernel( + stream_cfg, ck_tile::make_kernel( + Kernel{}, grids, blocks, 0, + ck_tile::cast_pointer_to_constant_address_space(ctx.workspace), + kargs.size())); + return true; + }); + }); + return false; +} + +template +bool ck_tile_mx_grouped_gemm_impl(DType a_dtype, DType b_dtype, DType d_dtype, + const GroupedGemmRunContext& ctx_in) { + using Cfg = MxTileCfg; + + using AScaleType = ck_tile::e8m0_t; + using BScaleType = ck_tile::e8m0_t; + + // gfx950 scale swizzle parameters (unused on gfx1250, whose swizzle is fixed). + constexpr ck_tile::index_t MIterPerWarp = Cfg::M_Tile / (Cfg::M_Warp * Cfg::M_Warp_Tile); + constexpr ck_tile::index_t NIterPerWarp = Cfg::N_Tile / (Cfg::N_Warp * Cfg::N_Warp_Tile); + constexpr ck_tile::index_t KIterPerWarp = Cfg::K_Tile / Cfg::K_Warp_Tile; + constexpr ck_tile::index_t MXdlPack = MxXdlPackEff; + constexpr ck_tile::index_t NXdlPack = MxXdlPackEff; + constexpr ck_tile::index_t KXdlPack = MxXdlPackEff; + constexpr ck_tile::index_t XdlMNThread = Cfg::M_Warp_Tile; + constexpr ck_tile::index_t XdlKThread = 64 / XdlMNThread; + + // The scale workspace is carved out of the tail of the caller's workspace, so + // this routine shrinks workspace_bytes for the kernel launch below. + GroupedGemmRunContext ctx = ctx_in; + + const bool warn_fallback = + getenv("NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK", false); + + const ck_tile::stream_config s{ctx.stream}; + + std::vector descs; + descs.reserve(ctx.group_num); + + NVTE_CHECK(ctx.workspace != nullptr, + "ck_tile_mx_grouped_gemm: workspace is required for shuffled MXFP8 scales."); + + // Carve regions from the end of the workspace for mxfp8 scales. + // Layout: [CK kargs workspace ... | a_scales (i) | b_scales (i) | ... | a_scales (group_num-1) | b_scales (group_num-1)] + constexpr size_t kScaleWorkspaceAlign = 256; + uint8_t *scale_workspace_base = reinterpret_cast(ctx.workspace); + size_t scale_workspace_end = + (ctx.workspace_bytes / kScaleWorkspaceAlign) * kScaleWorkspaceAlign; + + for (int i = 0; i < ctx.group_num; i++) { + const transformer_engine::Tensor *const A_te = + transformer_engine::convertNVTETensorCheck(ctx.A[i]); + const transformer_engine::Tensor *const B_te = + transformer_engine::convertNVTETensorCheck(ctx.B[i]); + transformer_engine::Tensor *D_te = + transformer_engine::convertNVTETensorCheck(ctx.D[i]); + + const auto &a = ctx.use_a_columnwise_data ? A_te->columnwise_data : A_te->data; + const auto &b = ctx.use_b_columnwise_data ? B_te->columnwise_data : B_te->data; + const auto &d = D_te->data; + const auto &a_scales = + ctx.use_a_columnwise_data ? A_te->columnwise_scale_inv : A_te->scale_inv; + const auto &b_scales = + ctx.use_b_columnwise_data ? B_te->columnwise_scale_inv : B_te->scale_inv; + + int64_t Ad0 = 0, Ad1 = 0, Bd0 = 0, Bd1 = 0, Dd0 = 0, Dd1 = 0; + + if (!get_flat_2d_dims(*A_te, Ad0, Ad1)) { + NVTE_ERROR("ck_tile_mx_grouped_gemm: expected rank>=2 for normalized A in group ", i); + } + + if (!get_flat_2d_dims(*B_te, Bd0, Bd1)) { + NVTE_ERROR("ck_tile_mx_grouped_gemm: expected rank>=2 for normalized B in group ", i); + } + + if (!get_flat_2d_dims(*D_te, Dd0, Dd1)) { + NVTE_ERROR("ck_tile_mx_grouped_gemm: expected rank>=2 for normalized D in group ", i); + } + if (a.dptr == nullptr || b.dptr == nullptr || a_scales.dptr == nullptr || + b_scales.dptr == nullptr) { + NVTE_ERROR("ck_tile_mx_grouped_gemm: effective A/B data or scale_inv is missing."); + } + if (a_scales.shape.size() != 2 || b_scales.shape.size() != 2) { + NVTE_ERROR("ck_tile_mx_grouped_gemm: expected effective A/B scale_inv tensors to be rank-2."); + } + + const size_t M = ctx.transA ? Ad1 : Ad0; + const size_t K = ctx.transA ? Ad0 : Ad1; + const size_t N = ctx.transB ? Bd0 : Bd1; + const size_t Kb = ctx.transB ? Bd1 : Bd0; + if (K % ScaleBlockSize != 0) { + NVTE_ERROR("ck_tile_mx_grouped_gemm: K must be a multiple of ScaleBlockSize for MX GEMM", i); + } + const int KScale = static_cast(K / ScaleBlockSize); + if (Kb != K) { + NVTE_ERROR("ck_tile_mx_grouped_gemm: K mismatch between A and B in group ", i, + ". op(A)=", M, "x", K, ", op(B)=", Kb, "x", N); + } + if (Dd0 != M || Dd1 != N) { + NVTE_ERROR("ck_tile_mx_grouped_gemm: D shape mismatch in group ", i, + ". D=", Dd0, "x", Dd1, ", expected=", M, "x", N); + } + + const ck_tile::index_t stride_A = static_cast(Ad1); + const ck_tile::index_t stride_B = static_cast(Bd1); + const ck_tile::index_t stride_E = static_cast(Dd1); + + // Pre-shuffle scale buffers for the hardware. + const int a_scale_actual_rows = static_cast(M); + const int a_scale_output_rows = + ck_tile::integer_least_multiple( + static_cast(M), + static_cast(Cfg::ScalePadM)); + const int b_scale_actual_rows = static_cast(N); + const int b_scale_output_rows = static_cast(N); + // gfx1250 writes one scale per input element; the gfx950 swizzle packs + // MNPack x KPack scales per dword group, so its buffer is sized in packs. + const auto shuffled_scale_elems = [&](int rows, ck_tile::index_t mn_pack) -> size_t { + if constexpr (Arch == GPUArch::GFX1250) { + return static_cast(rows) * static_cast(KScale); + } else { + return static_cast(rows / mn_pack * 2) * + static_cast(KScale / KXdlPack * 2); + } + }; + const size_t a_scale_shuffled_bytes = + shuffled_scale_elems(a_scale_output_rows, MXdlPack) * sizeof(AScaleType); + const size_t b_scale_shuffled_bytes = + shuffled_scale_elems(b_scale_output_rows, NXdlPack) * sizeof(BScaleType); + const size_t scale_pair_bytes = + a_scale_shuffled_bytes + b_scale_shuffled_bytes; + scale_workspace_end = + (scale_workspace_end / kScaleWorkspaceAlign) * kScaleWorkspaceAlign; + + NVTE_CHECK(scale_workspace_end >= scale_pair_bytes, + "ck_tile_mx_grouped_gemm: insufficient workspace for shuffled MXFP8 scales. " + "Need current group scale bytes=", scale_pair_bytes, + ", available workspace bytes=", scale_workspace_end, + ". Increase the grouped GEMM workspace size."); + + scale_workspace_end -= scale_pair_bytes; + uint8_t *scale_pair_ptr = scale_workspace_base + scale_workspace_end; + + void *a_scale_shuffled_ptr = scale_pair_ptr; + void *b_scale_shuffled_ptr = scale_pair_ptr + a_scale_shuffled_bytes; + + // CK expects canonical pre-shuffled scale buffers laid out as + // A: [M, KScale] and B: [N, KScale], independent of A/B data layouts. + // TE rowwise MXFP8 scale_inv is [rows, KScale] and can be read with + // KStride=true. TE columnwise_scale_inv is [KScale, rows] and must be + // read with KStride=false before writing CK's canonical shuffled layout. + TRANSFORMER_ENGINE_SWITCH_CONDITION(ctx.use_a_columnwise_data, kAColwise, { + if constexpr (Arch == GPUArch::GFX1250) { + preShuffleScaleBuffer_gfx1250( + reinterpret_cast(a_scales.dptr), + reinterpret_cast(a_scale_shuffled_ptr), + a_scale_actual_rows, a_scale_output_rows, KScale, ctx.stream); + } else { + preShuffleScaleBuffer_gfx950( + reinterpret_cast(a_scales.dptr), + reinterpret_cast(a_scale_shuffled_ptr), + a_scale_actual_rows, a_scale_output_rows, KScale, ctx.stream); + } + }); + + TRANSFORMER_ENGINE_SWITCH_CONDITION(ctx.use_b_columnwise_data, kBColwise, { + if constexpr (Arch == GPUArch::GFX1250) { + preShuffleScaleBuffer_gfx1250( + reinterpret_cast(b_scales.dptr), + reinterpret_cast(b_scale_shuffled_ptr), + b_scale_actual_rows, b_scale_output_rows, KScale, ctx.stream); + } else { + preShuffleScaleBuffer_gfx950( + reinterpret_cast(b_scales.dptr), + reinterpret_cast(b_scale_shuffled_ptr), + b_scale_actual_rows, b_scale_output_rows, KScale, ctx.stream); + } + }); + descs.emplace_back(mx_grouped_gemm_kargs( + a.dptr, a_scale_shuffled_ptr, b.dptr, b_scale_shuffled_ptr, + {/*ds_ptr*/}, d.dptr, 1, // kbatch + M, N, K, stride_A, stride_B, {/*stride_Ds*/}, stride_E)); + } + ctx.workspace_bytes = scale_workspace_end; + + // Invoke the GEMM. + bool ok = false; + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY(a_dtype, a_te_type, { + using AType = typename TETypeToCKType::type; + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY(b_dtype, b_te_type, { + using BType = typename TETypeToCKType::type; + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(d_dtype, d_te_type, { + using CType = typename TETypeToCKType::type; + ok = invoke_mx_grouped_gemm(descs, ctx, s, warn_fallback); + }); // NOLINT(*) + }); // NOLINT(*) + }); // NOLINT(*) + return ok; +} + +// Per-architecture dispatch function signature. +// Each architecture file implements one of these. The TDM pipeline is gfx1250 +// only, so each file is built for its own architecture and the shared template +// above is never instantiated for an architecture it does not support. +bool ck_tile_mx_grouped_gemm_dispatch_gfx1250(DType a_dtype, DType b_dtype, DType d_dtype, + const GroupedGemmRunContext& ctx); +bool ck_tile_mx_grouped_gemm_dispatch_gfx950(DType a_dtype, DType b_dtype, DType d_dtype, + const GroupedGemmRunContext& ctx); + +} // namespace grouped_gemm +} // namespace transformer_engine \ No newline at end of file