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