From 9c14bbf81129ac19c54d274d2120b67368defa96 Mon Sep 17 00:00:00 2001 From: Yihua Wei Date: Mon, 24 Aug 2026 15:54:34 -0700 Subject: [PATCH 1/4] add sm120 matmul support --- python/cudnn/gemm/frost/__init__.py | 2 +- python/cudnn/gemm/frost/compiler.py | 38 +- python/cudnn/gemm/frost/graph_analyzer.py | 6 +- python/cudnn/gemm/frost/kernel_registry.py | 75 +- .../frost/kernel_templates/sm120_matmul.py | 935 ++++++++++++++++++ python/cudnn/gemm/frost/tile_config.py | 61 +- test/python/gemm/frost/test_sm120_matmul.py | 680 +++++++++++++ 7 files changed, 1776 insertions(+), 21 deletions(-) create mode 100644 python/cudnn/gemm/frost/kernel_templates/sm120_matmul.py create mode 100644 test/python/gemm/frost/test_sm120_matmul.py diff --git a/python/cudnn/gemm/frost/__init__.py b/python/cudnn/gemm/frost/__init__.py index e4f0cf4d9..5ca7012c2 100644 --- a/python/cudnn/gemm/frost/__init__.py +++ b/python/cudnn/gemm/frost/__init__.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""cudnn.gemm.frost: JIT fused sm100 GEMM kernels from cuDNN graphs via the CuTe DSL. +"""cudnn.gemm.frost: JIT fused GEMM kernels from cuDNN graphs via the CuTe DSL. User code uses the plain cuDNN frontend API. The analyzer reads the python IR (``graph.nodes``) directly; :class:`cudnn.gemm.frost.engine.FrostGemmEngine` diff --git a/python/cudnn/gemm/frost/compiler.py b/python/cudnn/gemm/frost/compiler.py index 5c3f9a05b..bf5e50749 100644 --- a/python/cudnn/gemm/frost/compiler.py +++ b/python/cudnn/gemm/frost/compiler.py @@ -451,11 +451,24 @@ def _epi_tile_cols(config: TileConfig, cta_group: int) -> int: return cols +# Widest epilogue chunk in ELEMENTS per pipeline family. The tcgen05 drains +# take the global cap; sm120's transposed-STG epilogue stores at most the +# 8-column fragment row run each compute warp owns (template: 8 % _STG_V == 0). +_EPI_CHUNK_ELEMS_BY_PIPELINE = { + "sm100": MAX_EPI_CHUNK_ELEMS, + "sm103": MAX_EPI_CHUNK_ELEMS, + "sm107": MAX_EPI_CHUNK_ELEMS, + "sm120": 8, +} + + def _epi_vec_bytes(chain: FusionChain, config: TileConfig, cta_group: int) -> int: """The epilogue chunk width the kernel is rendered with: the chain-derived width additionally clamped so it divides every power-of-2 subtile span of - this config's N-tile (see ``_compute_output_vec_bytes``).""" - return _compute_output_vec_bytes(chain, tile_cols=_epi_tile_cols(config, cta_group)) + this config's N-tile (see ``_compute_output_vec_bytes``), capped at the + pipeline's widest store run (:data:`_EPI_CHUNK_ELEMS_BY_PIPELINE`).""" + vec = _compute_output_vec_bytes(chain, tile_cols=_epi_tile_cols(config, cta_group)) + return min(vec, _EPI_CHUNK_ELEMS_BY_PIPELINE[config.pipeline] * DTYPE_BYTES[chain.output_dtype]) def _mainloop_chain_zero_preserving(ops) -> bool: @@ -2724,6 +2737,10 @@ def _check_input_alignment(chain: FusionChain) -> None: _EPI_N_BASE = 32 # drain width when the epilogue is already hidden behind the MMA _EPI_N_MAX = 64 # per-lane fp32 registers the drain can hold +# Families whose templates the compiler renders with the TMA-store epilogue; +# the sm120 warp kernel always takes its transposed-STG path. +_TMA_STORE_EPI_PIPELINES = ("sm100", "sm103", "sm107") + def _epi_n(cfg, cta_group: int, out_dt: str) -> int: cols = _epi_tile_cols(cfg, cta_group) @@ -2771,7 +2788,10 @@ def _use_tma_store_epi(chain, cfg, vec_bytes_epi: int, cta_group: int) -> bool: - out dtype ∈ {bf16, fp16}: the drain widens to epi_n only for 2-byte output. - M-major output: 16B-aligned M (16x256b TMEM-load + stmatrix.trans + tma_store), and not MoE: the six MoE templates carry only the N-major TMA store. + - a :data:`_TMA_STORE_EPI_PIPELINES` family: sm120 renders transposed-STG only. """ + if cfg.pipeline not in _TMA_STORE_EPI_PIPELINES: + return False if chain.is_multi_gemm: # No multi-accumulator hook in the TMA-store path → STG only. return False @@ -2969,6 +2989,13 @@ def probe_supported( _check_executable(chain) if chain.has_moe or chain.has_block_scale: return # specialized paths validate at compile + if config is DEFAULT_CONFIG: + # The default probe geometry is an sm100 config; when the auto path + # would build another family (the active GPU is outside sm100's SM + # range, e.g. SM 12.x), probe that family and its execution strategy. + from .kernel_registry import preferred_strategy + + config, cta_group = preferred_strategy(chain, config, cta_group) if chain.is_multi_gemm: from .kernel_registry import select_template @@ -2982,7 +3009,7 @@ def probe_supported( _check_supported(chain, config) from .kernel_registry import select_template as _sel_tmpl - _arch_reason = _sel_tmpl(chain, config, cta_group).active_reject(config) + _arch_reason = _sel_tmpl(chain, config, cta_group).active_reject(config, chain) if _arch_reason is not None: raise NotImplementedError(_arch_reason) _check_dtype_config_compat(chain, config, cta_group) @@ -3044,11 +3071,12 @@ def jit_from_cudnn_graph( f"→ {tmpl.file}. Use cta_group=1." ) # Plain-matmul (pipeline × input/acc dtype combo [× GPU for the rare - # special-case combos]) gate, then the template family's active-GPU gate. + # special-case combos]) gate, then the template family's active-GPU gate + # and its chain-level scope (e.g. the sm120 TN/N-major-output contract). _check_supported(chain, config) from .kernel_registry import select_template as _sel_tmpl - _arch_reason = _sel_tmpl(chain, config, cta_group).active_reject(config) + _arch_reason = _sel_tmpl(chain, config, cta_group).active_reject(config, chain) if _arch_reason is not None: raise NotImplementedError(_arch_reason) _check_dtype_config_compat(chain, config, cta_group) diff --git a/python/cudnn/gemm/frost/graph_analyzer.py b/python/cudnn/gemm/frost/graph_analyzer.py index 7690a9bc1..234065d1b 100644 --- a/python/cudnn/gemm/frost/graph_analyzer.py +++ b/python/cudnn/gemm/frost/graph_analyzer.py @@ -563,8 +563,8 @@ def build_gemm_plan(graph: cudnn.pygraph): if not _graph_has_gemm(graph): raise ValueError("cudnn.gemm.frost: graph has no matmul / moe_grouped_matmul node; nothing to compile") from .compiler import jit_from_cudnn_graph - from .kernel_registry import preferred_pipeline - from .tile_config import as_pipeline, select_config + from .kernel_registry import preferred_strategy + from .tile_config import select_config chain = analyze(graph) tile_m = chain.matmul.M @@ -582,7 +582,7 @@ def build_gemm_plan(graph: cudnn.pygraph): b_n_major=chain.matmul.b_major == "n", b_elem_bytes=DTYPE_BYTES[chain.matmul.b_dtype], ) - config = as_pipeline(config, preferred_pipeline(chain)) + config, cta_group = preferred_strategy(chain, config, cta_group) return jit_from_cudnn_graph(graph, config=config, cta_group=cta_group) diff --git a/python/cudnn/gemm/frost/kernel_registry.py b/python/cudnn/gemm/frost/kernel_registry.py index 5c5b13a29..3d03998fe 100644 --- a/python/cudnn/gemm/frost/kernel_registry.py +++ b/python/cudnn/gemm/frost/kernel_registry.py @@ -27,7 +27,7 @@ from enum import Enum from .fusion_ir import BINARY_OPS, UNARY_OPS, FusionChain -from .tile_config import CATALOG, TileConfig, config_class_for_pipeline +from .tile_config import CATALOG, TileConfig, as_pipeline, config_class_for_pipeline def _pipeline_from_file(template_file: str) -> str: @@ -46,6 +46,7 @@ def _pipeline_from_file(template_file: str) -> str: "sm100": ((100, 120),), "sm103": ((103, 110),), "sm107": ((107, 110),), + "sm120": ((100, 130),), } # Pointwise ops a mainloop-fusion template can transform in SMEM. @@ -198,6 +199,13 @@ def _bs_key(a: str, sfa: str, b: str, sfb: str, kblk: int) -> tuple: "sm107": { GraphType.BLOCK_SCALE_MATMUL: _BLOCK_SCALE_CASES, }, + # sm120 carries ONLY the graph types it has templates for. No block-scale + # row until an sm120 block-scale template lands: a row here without its + # MMA_GPU_ARCH_SPECIAL_CASES narrowing would accept descriptor-less combos + # (see test_gpu_gated_cases_are_narrowed_everywhere). + "sm120": { + GraphType.MATMUL: _MATMUL_CASES, + }, } # The ONE home for checks that need template SM family × mma dtype × ACTUAL @@ -386,11 +394,15 @@ def multi_mma_m_reject(self, config: TileConfig) -> str | None: ) return None - def active_reject(self, config: TileConfig) -> str | None: + def active_reject(self, config: TileConfig, chain: FusionChain | None = None) -> str | None: """The gates a JIT path applies once it has picked this template: the active GPU's SM range, then capabilities a pure-geometry config can ask - for that this template does not implement.""" - return self.arch_active_reject() or self.multi_mma_m_reject(config) + for that this template does not implement, and — when the caller passes + the ``chain`` — the template-specific scope (:meth:`_extra_reject`).""" + reason = self.arch_active_reject() or self.multi_mma_m_reject(config) + if reason is None and chain is not None: + reason = self._extra_reject(chain, config) + return reason def candidate_configs(self, chain: FusionChain) -> tuple[TileConfig, ...]: """Catalog geometries this template accepts for ``chain`` — by @@ -411,6 +423,33 @@ def _extra_reject(self, chain: FusionChain, config: TileConfig) -> str | None: return None +class Sm120KernelTemplate(KernelTemplate): + """The sm120 warp-MMA template. Its v1 scope is enforced by render-time + asserts in the template source; encoding it here lets the funnel (and the + jit paths, via ``active_reject``) reject cleanly instead of faulting + mid-render: TN GEMM (K-major A and B), N-major non-fp4 output, and an + epilogue that stores whole (n, n+1) accumulator pairs.""" + + def _extra_reject(self, chain: FusionChain, config: TileConfig) -> str | None: + mm = chain.matmul + if mm.a_major != "k" or mm.b_major != "k": + return f"{self.file} supports only K-major A and B (TN GEMM); " f"got A {mm.a_major}-major, B {mm.b_major}-major" + if chain.out_major != "n": + return f"{self.file} supports only an N-major output" + if chain.output_dtype == "fp4_e2m1": + return f"{self.file} does not support fp4 output" + from . import compiler as C + from .dtypes import DTYPE_BYTES + + try: + vec = C._epi_vec_bytes(chain, config, self.cta_group) + except ValueError as e: + return str(e) + if vec < 2 * DTYPE_BYTES[chain.output_dtype]: + return f"{self.file} stores whole (n, n+1) accumulator pairs; the output " f"layout admits only {vec}-byte epilogue chunks" + return None + + # Registry — one entry per template file (20 today). A geometry config expands # across these via `candidates`. cta_group / mainloop live HERE. @@ -423,13 +462,14 @@ def _mm( graph_type: GraphType = GraphType.MATMUL, supports_multi_gemm: bool = True, supports_multi_mma_m: bool = True, + template_cls: "type[KernelTemplate] | None" = None, ) -> KernelTemplate: pipeline = _pipeline_from_file(file) if pipeline not in PIPELINE_ARCH_RANGES: raise KeyError( f"template {file!r}: pipeline family {pipeline!r} has no SM-range entry in " f"PIPELINE_ARCH_RANGES — add one when introducing a new family" ) - cls = MainloopKernelTemplate if mainloop else KernelTemplate + cls = template_cls or (MainloopKernelTemplate if mainloop else KernelTemplate) return cls( file=file, pipeline=pipeline, @@ -525,13 +565,24 @@ def _mm( cta_group=2, graph_type=GraphType.MOE_BLOCK_SCALE, ), + # sm120 (consumer Blackwell) warp-MMA matmul: no clusters, CLC persistent + # scheduler, single-GEMM only (no per-GEMM operand indexing in the compute + # warps). The v1 scope gates (TN, N-major output) live on Sm120KernelTemplate. + _mm( + "sm120_matmul.py", + cta_group=1, + supports_multi_gemm=False, + template_cls=Sm120KernelTemplate, + ), ) # Pipeline families the AUTO path (``tile_config.select_config``) may build # with, best first. sm103 is deliberately absent: its 384-byte K-tile is outside # select_config's geometry ladder, so it stays an explicit-config pipeline. -_AUTO_PIPELINE_ORDER: tuple[str, ...] = ("sm107", "sm100") +# sm120 is last: it serves the GPUs (SM 12.x) the tcgen05 families cannot, and +# never outranks them where both run. +_AUTO_PIPELINE_ORDER: tuple[str, ...] = ("sm107", "sm100", "sm120") def preferred_pipeline(chain: FusionChain) -> str: @@ -546,6 +597,18 @@ def preferred_pipeline(chain: FusionChain) -> str: return _AUTO_PIPELINE_ORDER[-1] +def preferred_strategy(chain: FusionChain, config: TileConfig, cta_group: int) -> tuple[TileConfig, int]: + """Re-target an auto pick at the family :func:`preferred_pipeline` chooses: + the geometry crosses via ``as_pipeline``, and ``cta_group`` survives only + if the family has a template for it (else the smallest it does have — + sm120 is warp-scoped MMA, 1-CTA only).""" + pipeline = preferred_pipeline(chain) + if pipeline == config.pipeline: + return config, cta_group + groups = {t.cta_group for t in TEMPLATES if t.pipeline == pipeline} + return as_pipeline(config, pipeline), cta_group if cta_group in groups else min(groups) + + def select_template( chain: FusionChain, config: TileConfig, diff --git a/python/cudnn/gemm/frost/kernel_templates/sm120_matmul.py b/python/cudnn/gemm/frost/kernel_templates/sm120_matmul.py new file mode 100644 index 000000000..d5b8df84e --- /dev/null +++ b/python/cudnn/gemm/frost/kernel_templates/sm120_matmul.py @@ -0,0 +1,935 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""sm120 (GeForce/consumer Blackwell, CC 12.0) GEMM kernel: persistent + CLC +dynamic scheduler (2-stage ring) + warp-level MMA. +""" + +from __future__ import annotations + +from functools import lru_cache +from typing import Callable + +import cutlass.experimental.primitives as nvvm +import cutlass.experimental.cuda.tensor_map as _tma +import cutlass._mlir_helpers.vector as _cvec +from cutlass import apply_swizzle as _apply_smem_swizzle +import cutlass +import cutlass.cute as cute +from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_tensor +from cutlass.cute.runtime import make_fake_stream +from cuda.bindings import driver as _cuda +from cutlass.cute.arch import clc as cute_clc + +# @@INJECT_TILE_CONSTANTS@@ + + +CLC_SCHED_STAGES = 2 + +# Programmatic Dependent Launch (PDL, sm_90+; supported on sm_120). +USE_PDL = True + +# Double-buffer for the TMA-store epilogue path. +EPI_SMEM_STAGES = 2 + +# Named barrier id for cross-warp sync of the 8 compute warps around TMA stores. +EPI_SYNC_BAR_ID = 1 + +# Compute-warp grid over the CTA tile (warp_row x warp_col). +NUM_COMPUTE_WARPS = 8 +WARPS_M = 4 +WARPS_N = 2 + +TMA_WARP_ID = 8 +SCHEDULER_WARP_ID = 9 +NUM_WARPS = 12 + +# CLC-ring consumers: every compute warp + the TMA producer + the scheduler +# itself each arrive once (elected) per consumed response slot. +NUM_CLC_CONSUMER_WARPS = NUM_COMPUTE_WARPS + 2 + +EPI_REG_COUNT = 232 +PROD_REG_COUNT = 24 + +# --------------------------------------------------------------------------- +# Geometry derived from the injected tile constants (all plain Python ints — +# resolved at render/import time, traced as constants). +# --------------------------------------------------------------------------- + +_ELEM_BITS = ab_dtype.width +_ELEM_BYTES = _ELEM_BITS // 8 +_ELEMS_16B = 16 // _ELEM_BYTES +# One k-block = 32 bytes of K = the K extent of one mma.sync (k16 for 16-bit, +# k32 for 8-bit operands) — fort's UNIT_MATRIX_{A,B} column span. +_K_BLK_ELEMS = (32 * 8) // _ELEM_BITS +_NUM_K_BLOCKS = (cta_tile_mnk[2] * _ELEM_BITS) // (32 * 8) +_CTA_K_ELEMS = cta_tile_mnk[2] + +_WARP_TILE_M = cta_tile_mnk[0] // WARPS_M +_WARP_TILE_N = cta_tile_mnk[1] // WARPS_N +_M_FRAGS = _WARP_TILE_M // 16 +_N_FRAGS = _WARP_TILE_N // 8 +_N_FRAG_PAIRS = _N_FRAGS // 2 +_ACC_REGS = _M_FRAGS * _N_FRAGS * 4 + +_EPI_N = epi_tile_mn[1] + +# SMEM K-row swizzle: the TMA s{128,64,32}b pattern == cutlass.Swizzle(b, 4, 3) +# with b = log2(row_bytes / 16). ldmatrix addresses below apply the same XOR +# (fort: swizzled_bank_id = bank ^ ((bank / 8) % SWIZZLE_SCALE)). +_AB_SW_BBITS = (ab_smem_swizzle_bytes // 16).bit_length() - 1 +_AB_SWIZZLE = cutlass.Swizzle(_AB_SW_BBITS, 4, 3) +# Epilogue staging tile swizzle — matches the s64b TMA-store descriptor. +_EPI_SWIZZLE = cutlass.Swizzle(2, 4, 3) + +# ---- Transposed STG epilogue staging (fort "Sheet3" scheme) ----------------- +_STG_EPI_LANE_QUAD = 4 # one STS.128 = 4 x 32-bit acc regs per lane +_STG_EPI_PAD = 4 # 16B skew after each 128-element batch (sheet's X cells) +_STG_EPI_BATCH_STRIDE = 32 * _STG_EPI_LANE_QUAD + _STG_EPI_PAD # 132 +_STG_EPI_GROUP_FRAGS = 4 # fragments (= STS batches) per 32-column group +_STG_EPI_WARP_ELEMS = _STG_EPI_GROUP_FRAGS * _STG_EPI_BATCH_STRIDE # 528 +_STG_EPI_NGRP = (_N_FRAGS + _STG_EPI_GROUP_FRAGS - 1) // _STG_EPI_GROUP_FRAGS +_STG_V = (vec_bytes_epi * 8) // cd_dtype.width + +if not use_tma_store_epi: + _STG_EPI_BYTES = 4 * _STG_EPI_WARP_ELEMS * NUM_COMPUTE_WARPS + _AB_STAGE_BYTES = (cta_tile_mnk[0] + cta_tile_mnk[1]) * _CTA_K_ELEMS * _ELEM_BYTES + 16 + ab_stages = ab_stages - -(-_STG_EPI_BYTES // _AB_STAGE_BYTES) + assert ab_stages >= 1, "transposed STG epilogue: staging stream cannot be funded from the AB pipeline" + +# ---- v1 scope guards (fail at render/import, not at runtime) --------------- +assert cluster_shape_mnk == (1, 1, 1), "sm120 has no thread-block clusters (CC 12.0): cluster_shape must be (1,1,1)" +assert threads_per_cta == NUM_WARPS * 32, f"sm120 template is a fixed 12-warp kernel (384 threads), got {threads_per_cta}" +assert not multicast_a and not multicast_b, "sm120 has no TMA multicast (no clusters)" +assert not a_is_m_major and not b_is_n_major, "sm120 template v1 supports K-major A and K-major B only (TN GEMM)" +assert not cd_out_is_m_major, "sm120 template v1 supports N-major output only" +assert num_gemms == 1 and num_a_operands == 1 and num_b_operands == 1, "sm120 template v1 is single-GEMM only" +assert cd_fake_n_div == 1, "sm120 template v1 does not support fp4 output" +assert cta_tile_mnk[0] % WARPS_M == 0 and _WARP_TILE_M % 16 == 0, f"cta_tile_m={cta_tile_mnk[0]} must be a multiple of {WARPS_M * 16}" +assert cta_tile_mnk[1] % WARPS_N == 0 and _WARP_TILE_N % 8 == 0, f"cta_tile_n={cta_tile_mnk[1]} must be a multiple of {WARPS_N * 8}" +assert ab_smem_swizzle_bytes == _CTA_K_ELEMS * _ELEM_BYTES, "SMEM K-row width must equal the swizzle span" +assert _NUM_K_BLOCKS * _K_BLK_ELEMS == _CTA_K_ELEMS, "cta_tile_k must be a multiple of 32 bytes" +if use_tma_store_epi: + assert vec_bytes_epi * 8 == 2 * cd_dtype.width, "sm120 TMA-store epilogue drains one (n, n+1) accumulator pair per thread" +else: + assert _STG_V >= 2 and _STG_V % 2 == 0 and 8 % _STG_V == 0, "transposed STG epilogue: the store vector must be whole pairs tiling the 8-column row run" +if use_tma_store_epi: + assert cta_tile_mnk[1] % _EPI_N == 0, "TMA-store epilogue needs cta_tile_n to be a whole number of epi subtiles" + +# --------------------------------------------------------------------------- +# The warp MMA instruction, resolved from the injected MMA dtypes. +# sm120 tensor cores are warp-scoped: mma.sync.aligned.m16n8k16 (16-bit A/B) +# or .m16n8k32 (8-bit A/B), row.col (both operands K-major), fp32/s32 acc — +# fort emits the same instruction pair per XMMA (bf16mma_fp32_16x16x16). +# --------------------------------------------------------------------------- + +_PTX_AB_TAG = { + cutlass.BFloat16: "bf16", + cutlass.Float16: "f16", + cutlass.Float8E4M3FN: "e4m3", + cutlass.Float8E5M2: "e5m2", + cutlass.Int8: "s8", +} +assert mma_a_dtype in _PTX_AB_TAG and mma_b_dtype in _PTX_AB_TAG, f"unsupported sm120 MMA input dtypes: {mma_a_dtype} x {mma_b_dtype}" +_MMA_SHAPE = "m16n8k16" if _ELEM_BITS == 16 else "m16n8k32" +_MMA_C_TAG = "f32" if mma_c_dtype == cutlass.Float32 else "s32" +_MMA_PTX = ( + f"mma.sync.aligned.{_MMA_SHAPE}.row.col" + f".{_MMA_C_TAG}.{_PTX_AB_TAG[mma_a_dtype]}.{_PTX_AB_TAG[mma_b_dtype]}.{_MMA_C_TAG} " + "{$0,$1,$2,$3}, {$4,$5,$6,$7}, {$8,$9}, {$10,$11,$12,$13};" +) + + +@cute.jit +def _mma_16x8_k32b(a0, a1, a2, a3, b0, b1, c0, c1, c2, c3): + """One warp-wide mma.sync on a (16, 8, 32-byte-K) fragment. + + A carrier: 4x b32 regs (ldmatrix.x4 of a [16 x 32B] K-major SMEM region). + B carrier: 2x b32 regs. D/C: 4 accumulator regs (f32 or s32). + """ + return cute.arch.inline_ptx( + _MMA_PTX, + write_only_types=[mma_c_dtype, mma_c_dtype, mma_c_dtype, mma_c_dtype], + read_only_args=[a0, a1, a2, a3, b0, b1, c0, c1, c2, c3], + ) + + +@cute.jit +def _auto_swizzle_w(m, n, k, nt_n): + """N-super-block width for the tile rasterization, resolved per launch. + + ``tile_swizzle_n > 0`` pins it. Otherwise: the walk keeps one operand slice + resident and re-reads the other every super-block, so block along the SHORTER + problem side. Once that side outgrows what L2 can hold onto while C streams + through it, keeping it is no longer free -- fall back to the widest N block the + budget does cover. + """ + if cutlass.const_expr(tile_swizzle_n > 0): + return tile_swizzle_n + budget = cutlass.Int64(swizzle_l2_budget_bytes) + row_bytes = (cutlass.Int64(ab_dtype.width) * k) // 8 + cap = cutlass.max(budget // (row_bytes * cgrp_tile_mnk[1]), cutlass.Int64(1)) + w = cutlass.min(cutlass.Int64(nt_n), cap) + if cutlass.min(m, n) * row_bytes <= budget and m <= n: + w = cutlass.Int64(1) + return cutlass.Int32(w) + + +def _l2_swizzle_tile(raw_m, raw_n, nt_m, nt_n, swizzle_w): + """N-direction super-block rasterization of the (m, n) tile coord, for + L2 reuse. Applied identically to the launch-grid coords and to every CLC + response, so a stolen CTA id lands on the same logical tile the canceled + CTA would have computed (fort's ``swizzle()`` plays the same role). + ``swizzle_w == 1`` falls out of the math as the identity mapping. + """ + t = raw_n * nt_m + raw_m + blk = nt_m * swizzle_w + sb = t // blk + off = t - sb * blk + base_n = sb * swizzle_w + cur_S = cutlass.min(cutlass.Int32(swizzle_w), nt_n - base_n) + log_m = off // cur_S + log_n = base_n + off - log_m * cur_S + return log_m, log_n + + +@cute.kernel +def _kernel( + m: cutlass.Int64, + n: cutlass.Int64, + k: cutlass.Int64, + # @@INJECT_KERNEL_AB_DESC_PARAMS@@ + # @@INJECT_KERNEL_TAP_PARAMS@@ + # @@INJECT_KERNEL_REDUCTION_STRIDE_PARAMS@@ + # @@INJECT_KERNEL_AUX_PARAMS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_KERNEL_TMA_C_PARAMS@@ + # @@TMA_STORE_ONLY:END@@ +) -> None: + # @@INJECT_AB_DESC_LISTS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_TMA_C_LISTS@@ + tma_c_desc = tma_c_descs[0] + # @@TMA_STORE_ONLY:END@@ + + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + tidx = cute.arch.thread_idx()[0] + bidx = cute.arch.block_idx()[0] + bidy = cute.arch.block_idx()[1] + bidz = cute.arch.block_idx()[2] + gridx = cute.arch.grid_dim()[0] + gridy = cute.arch.grid_dim()[1] + + if warp_idx == TMA_WARP_ID: + for _i in cutlass.range_constexpr(num_a_operands): + nvvm.prefetch_tensormap(tma_a_descs[_i].get_ptr()) + for _j in cutlass.range_constexpr(num_b_operands): + nvvm.prefetch_tensormap(tma_b_descs[_j].get_ptr()) + + # @@TMA_STORE_ONLY:BEGIN@@ + nvvm.prefetch_tensormap(tma_c_desc.get_ptr()) + # @@TMA_STORE_ONLY:END@@ + + # First tile from the launch grid (grid == tile grid); later tiles come + # from canceled-CTA ids delivered through the CLC response ring. + swizzle_w = _auto_swizzle_w(m, n, k, gridy) + init_tile_m, init_tile_n = _l2_swizzle_tile(bidx, bidy, gridx, gridy, swizzle_w) + init_tile_l = bidz + + ab_full_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) + ab_empty_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) + + _clc_response_raw = cutlass.Array(cutlass.Int128, CLC_SCHED_STAGES, space=cutlass.AddressSpace.smem, alignment=16) + clc_response_ptr_base = cute.make_ptr( + cutlass.Int128, + _clc_response_raw.data_ptr(), + mem_space=cute.AddressSpace.smem, + ) + clc_full_mbar_ptr = cutlass.Array(cutlass.Int64, CLC_SCHED_STAGES, space=cutlass.AddressSpace.smem, alignment=8) + clc_empty_mbar_ptr = cutlass.Array(cutlass.Int64, CLC_SCHED_STAGES, space=cutlass.AddressSpace.smem, alignment=8) + clc_full_mbar_cute_base = cute.make_ptr( + cutlass.Int64, + clc_full_mbar_ptr.data_ptr(), + mem_space=cute.AddressSpace.smem, + ) + + sA_elems = cta_tile_mnk[0] * cta_tile_mnk[2] + sB_elems = cta_tile_mnk[1] * cta_tile_mnk[2] + smem_a_list = [ + cutlass.Array( + ab_dtype, + sA_elems * ab_stages, + space=cutlass.AddressSpace.smem, + alignment=1024, + ) + for _ in range(num_a_operands) + ] + smem_b_list = [ + cutlass.Array( + ab_dtype, + sB_elems * ab_stages, + space=cutlass.AddressSpace.smem, + alignment=1024, + ) + for _ in range(num_b_operands) + ] + + # @@TMA_STORE_ONLY:BEGIN@@ + epi_subtile_elems = cta_tile_mnk[0] * epi_tile_mn[1] + smem_d_ptr = cutlass.Array( + cd_dtype, + epi_subtile_elems * EPI_SMEM_STAGES, + space=cutlass.AddressSpace.smem, + alignment=1024, + ) + # @@TMA_STORE_ONLY:END@@ + + # @@STG_ONLY:BEGIN@@ + # Per-compute-warp staging stream for the transposed STG epilogue (raw + # accumulator dtype; 4 batches x (128 elems + 16B pad) = 528 elems, one + # 32-column group of one m-frag at a time). Slices are warp-private, so + # the round trip only needs bar.warp syncs — no CTA barrier. + smem_stg_epi = cutlass.Array( + mma_c_dtype, + _STG_EPI_WARP_ELEMS * NUM_COMPUTE_WARPS, + space=cutlass.AddressSpace.smem, + alignment=1024, + ) + # @@STG_ONLY:END@@ + + # ab full: one producer-elected arrive_expect_tx per stage. + # ab empty: one elected arrive per compute warp per stage (fort inits this + # to GROUPS_M * WARPS_PER_GROUP = the 8 math warps). + # clc full: tx-count armed by the scheduler; completed by the response. + # clc empty: one elected arrive per consumer warp per slot. + if warp_idx == 0: + if nvvm.elect_sync(): + for i in range(ab_stages): + nvvm.mbarrier_init(ab_full_mbar_ptr.subview(i), 1) + nvvm.mbarrier_init(ab_empty_mbar_ptr.subview(i), NUM_COMPUTE_WARPS) + for i in range(CLC_SCHED_STAGES): + nvvm.mbarrier_init(clc_full_mbar_ptr.subview(i), 1) + nvvm.mbarrier_init(clc_empty_mbar_ptr.subview(i), NUM_CLC_CONSUMER_WARPS) + nvvm.fence_mbarrier_init() + nvvm.barrier_cta_sync(0) + + sA_bytes = sA_elems * (ab_dtype.width // 8) + sB_bytes = sB_elems * (ab_dtype.width // 8) + num_tma_copy_bytes = num_a_operands * sA_bytes + num_b_operands * sB_bytes + + # @@INJECT_TAP_PTRS@@ + + VEC_BYTES = vec_bytes_epi + vsize = (VEC_BYTES * 8) // cd_dtype.width + + M = m + N = n + num_k_tiles = cute.ceil_div(k, cta_tile_mnk[2]) + + # -- CLC scheduler warp --------------------------------------------------- + # fort's scheduler warp: wait empty(slot) -> arm 16 tx bytes -> try_cancel + # into the slot -> wait full(slot) -> read validity -> arrive empty. No + # cluster: every CTA is its own leader and the response is CTA-local. + if warp_idx == SCHEDULER_WARP_ID: + nvvm.setmaxregister(PROD_REG_COUNT, nvvm.SetMaxRegisterAction.DECREASE) + sched_iter = cutlass.Int32(0) + clc_empty_phase = cutlass.Int32(1) + clc_full_phase = cutlass.Int32(0) + is_valid_sched = cutlass.Int32(1) + while is_valid_sched != 0: + stage = sched_iter % CLC_SCHED_STAGES + if stage == 0 and sched_iter != 0: + clc_empty_phase = clc_empty_phase ^ 1 + clc_full_phase = clc_full_phase ^ 1 + + while not nvvm.mbarrier_try_wait_parity(clc_empty_mbar_ptr.subview(stage), clc_empty_phase, time_limit=10_000_000): + pass + + if nvvm.elect_sync(): + nvvm.mbarrier_arrive_expect_tx(clc_full_mbar_ptr.subview(stage), 16) + cute_clc.issue_clc_query( + clc_full_mbar_cute_base + stage, + clc_response_ptr_base + stage, + multicast=False, + ) + + while not nvvm.mbarrier_try_wait_parity(clc_full_mbar_ptr.subview(stage), clc_full_phase, time_limit=10_000_000): + pass + + _m_idx, _n_idx, _l_idx, vld = cute_clc.clc_response(clc_response_ptr_base + stage) + cute.arch.fence_proxy("async.shared", space="cta") + is_valid_sched = vld + + nvvm.bar_warp_sync(0xFFFFFFFF) + if nvvm.elect_sync(): + nvvm.mbarrier_arrive(clc_empty_mbar_ptr.subview(stage)) + + sched_iter += 1 + + # -- TMA producer warp ---------------------------------------------------- + if warp_idx == TMA_WARP_ID: + nvvm.setmaxregister(PROD_REG_COUNT, nvvm.SetMaxRegisterAction.DECREASE) + if cutlass.const_expr(USE_PDL): + if nvvm.elect_sync(): + nvvm.griddepcontrol("wait") + ab_empty_phase_bit = cutlass.Int32(1) + ab_iter = cutlass.Int32(0) + tile_m = init_tile_m + tile_n = init_tile_n + tile_l = init_tile_l + tile_iter = cutlass.Int32(0) + is_valid = cutlass.Int32(1) + clc_full_phase_tma = cutlass.Int32(0) + while is_valid != 0: + coord_m = tile_m * cgrp_tile_mnk[0] + coord_n = tile_n * cgrp_tile_mnk[1] + if cutlass.const_expr(matmul_a_batch == 1): + tile_l_a = cutlass.Int32(0) + else: + tile_l_a = tile_l + if cutlass.const_expr(matmul_b_batch == 1): + tile_l_b = cutlass.Int32(0) + else: + tile_l_b = tile_l + + for k_tile_idx in range(num_k_tiles): + stage = ab_iter % ab_stages + if stage == 0 and ab_iter != 0: + ab_empty_phase_bit = ab_empty_phase_bit ^ 1 + + while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(stage), ab_empty_phase_bit, time_limit=10_000_000): + pass + + coord_k = k_tile_idx * cta_tile_mnk[2] + if nvvm.elect_sync(): + nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes) + # K-major A: TMA box [K_tile, cta_m] at (k, m, l); OOB rows/cols + # are hardware zero-filled (K tails contribute 0 to the MMA). + for _ai in cutlass.range_constexpr(num_a_operands): + nvvm.cp_async_bulk_tensor_shared_cta_global( + smem_a_list[_ai].subview(sA_elems * stage), + tma_a_descs[_ai].get_ptr(), + (coord_k, coord_m, tile_l_a), + ab_full_mbar_ptr.subview(stage), + ) + # K-major B: TMA box [K_tile, cta_n] at (k, n, l). + for _bj in cutlass.range_constexpr(num_b_operands): + nvvm.cp_async_bulk_tensor_shared_cta_global( + smem_b_list[_bj].subview(sB_elems * stage), + tma_b_descs[_bj].get_ptr(), + (coord_k, coord_n, tile_l_b), + ab_full_mbar_ptr.subview(stage), + ) + ab_iter += 1 + + consumer_stage = tile_iter % CLC_SCHED_STAGES + if consumer_stage == 0 and tile_iter != 0: + clc_full_phase_tma = clc_full_phase_tma ^ 1 + while not nvvm.mbarrier_try_wait_parity( + clc_full_mbar_ptr.subview(consumer_stage), + clc_full_phase_tma, + time_limit=10_000_000, + ): + pass + m_idx, n_idx, l_idx, vld = cute_clc.clc_response(clc_response_ptr_base + consumer_stage) + cute.arch.fence_proxy("async.shared", space="cta") + is_valid = vld + tile_m, tile_n = _l2_swizzle_tile(m_idx, n_idx, gridx, gridy, swizzle_w) + tile_l = l_idx + nvvm.bar_warp_sync(0xFFFFFFFF) + if nvvm.elect_sync(): + nvvm.mbarrier_arrive(clc_empty_mbar_ptr.subview(consumer_stage)) + tile_iter += 1 + + # Drain: wait until the compute warps have consumed the final stage so + # the producer never exits with a stage it would have re-armed pending. + tail_stage = ab_iter % ab_stages + tail_phase = ab_empty_phase_bit + if tail_stage == 0 and ab_iter != 0: + tail_phase = tail_phase ^ 1 + for _ in range(ab_stages - 1): + tail_stage = tail_stage + 1 + if tail_stage == ab_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 + if nvvm.elect_sync(): + while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + + # -- Compute warps: mma.sync mainloop + epilogue -------------------------- + if warp_idx < NUM_COMPUTE_WARPS: + nvvm.setmaxregister(EPI_REG_COUNT, nvvm.SetMaxRegisterAction.INCREASE) + if cutlass.const_expr(USE_PDL): + nvvm.griddepcontrol("wait") + + lane = tidx % 32 + lane_div4 = lane // 4 + lane_mod4 = lane % 4 + warp_row = warp_idx % WARPS_M + warp_col = warp_idx // WARPS_M + + # ldmatrix lane->address maps (see PTX ldmatrix; addresses are 16B rows). + # A x4 tile order = (rows 0-7, rows 8-15) x (16B col 0, 16B col 1) — + # matching the a0..a3 fragment order of mma.sync (fort Lds_tile_8). + a_ldm_row = (lane % 8) + 8 * ((lane // 8) % 2) + a_ldm_col16 = lane // 16 + # B x4 covers TWO 8-col n-frags: (n rows 0-7, n rows 8-15) each split + # over (16B col 0, 16B col 1) -> regs (b0,b1) frag0 + (b0,b1) frag1 + # (fort Lds_tile_10). + b_ldm_pair_row = (lane % 8) + 8 * (lane // 16) + b_ldm_pair_col16 = (lane // 8) % 2 + # B x2 tail: one n-frag (rows 0-7 x two 16B cols; lanes 16-31 unused). + b_ldm_tail_row = lane % 8 + b_ldm_tail_col16 = (lane // 8) % 2 + + acc = cutlass.Array(mma_c_dtype, _ACC_REGS, alignment=16) + + ab_full_phase_bit = cutlass.Int32(0) + ab_iter = cutlass.Int32(0) + # @@TMA_STORE_ONLY:BEGIN@@ + epi_stage_idx = cutlass.Int32(EPI_SMEM_STAGES - 1) + # @@TMA_STORE_ONLY:END@@ + tile_m = init_tile_m + tile_n = init_tile_n + tile_l = init_tile_l + tile_iter = cutlass.Int32(0) + is_valid = cutlass.Int32(1) + clc_full_phase_epi = cutlass.Int32(0) + while is_valid != 0: + coord_m = tile_m * cgrp_tile_mnk[0] + coord_n = tile_n * cgrp_tile_mnk[1] + + for _z in cutlass.range_constexpr(_ACC_REGS): + acc[_z] = mma_c_dtype(0) + + for k_tile_idx in range(num_k_tiles): + stage = ab_iter % ab_stages + if stage == 0 and ab_iter != 0: + ab_full_phase_bit = ab_full_phase_bit ^ 1 + + while not nvvm.mbarrier_try_wait_parity(ab_full_mbar_ptr.subview(stage), ab_full_phase_bit, time_limit=10_000_000): + pass + + sA_ptr = smem_a_list[0].subview(sA_elems * stage).data_ptr() + sB_ptr = smem_b_list[0].subview(sB_elems * stage).data_ptr() + + for k_blk in cutlass.range_constexpr(_NUM_K_BLOCKS): + kb_base = k_blk * _K_BLK_ELEMS + a_frags = [] + for mf in cutlass.range_constexpr(_M_FRAGS): + a_row = warp_row * _WARP_TILE_M + mf * 16 + a_ldm_row + a_off = a_row * _CTA_K_ELEMS + kb_base + a_ldm_col16 * _ELEMS_16B + a_frags.append( + nvvm.ldmatrix( + _apply_smem_swizzle(sA_ptr + a_off, _AB_SWIZZLE), + 4, + nvvm.MMALayout.ROW, + ) + ) + b_frags = [] + for npair in cutlass.range_constexpr(_N_FRAG_PAIRS): + b_row = warp_col * _WARP_TILE_N + npair * 16 + b_ldm_pair_row + b_off = b_row * _CTA_K_ELEMS + kb_base + b_ldm_pair_col16 * _ELEMS_16B + bv = nvvm.ldmatrix( + _apply_smem_swizzle(sB_ptr + b_off, _AB_SWIZZLE), + 4, + nvvm.MMALayout.ROW, + ) + b_frags.append((bv[0], bv[1])) + b_frags.append((bv[2], bv[3])) + if cutlass.const_expr(_N_FRAGS % 2 == 1): + b_row = warp_col * _WARP_TILE_N + (_N_FRAGS - 1) * 8 + b_ldm_tail_row + b_off = b_row * _CTA_K_ELEMS + kb_base + b_ldm_tail_col16 * _ELEMS_16B + bt = nvvm.ldmatrix( + _apply_smem_swizzle(sB_ptr + b_off, _AB_SWIZZLE), + 2, + nvvm.MMALayout.ROW, + ) + b_frags.append((bt[0], bt[1])) + + for mf in cutlass.range_constexpr(_M_FRAGS): + av = a_frags[mf] + for nf in cutlass.range_constexpr(_N_FRAGS): + b0, b1 = b_frags[nf] + _o = (mf * _N_FRAGS + nf) * 4 + acc[_o:4] = _mma_16x8_k32b( + av[0], + av[1], + av[2], + av[3], + b0, + b1, + acc[_o + 0], + acc[_o + 1], + acc[_o + 2], + acc[_o + 3], + ) + + # Stage fully consumed by this warp (ldmatrix is synchronous). + nvvm.bar_warp_sync(0xFFFFFFFF) + if nvvm.elect_sync(): + nvvm.mbarrier_arrive(ab_empty_mbar_ptr.subview(stage)) + ab_iter += 1 + + # -- Epilogue: accumulators are already in registers ------------------ + + # @@INJECT_AUX_VIEWS@@ + + # @@TMA_STORE_ONLY:BEGIN@@ + for subtile_idx in cutlass.range_constexpr(cta_tile_mnk[1] // _EPI_N): + subtile_col_offset = subtile_idx * _EPI_N + epi_stage_idx = (epi_stage_idx + 1) % EPI_SMEM_STAGES + smem_subtile_ptr = smem_d_ptr.subview(epi_stage_idx * epi_subtile_elems) + col = coord_n + subtile_col_offset + + for mf in cutlass.range_constexpr(_M_FRAGS): + for half in cutlass.range_constexpr(2): + row_in_cta = warp_row * _WARP_TILE_M + mf * 16 + half * 8 + lane_div4 + row = coord_m + row_in_cta + for nf in cutlass.range_constexpr(_N_FRAGS): + # Only fragments inside this subtile's N span write — + # warp-uniform predicate (warp_col is warp-uniform). + frag_rel_col = warp_col * _WARP_TILE_N + nf * 8 - subtile_col_offset + if frag_rel_col >= 0: + if frag_rel_col < _EPI_N: + _o = (mf * _N_FRAGS + nf) * 4 + half * 2 + _pair = acc[_o:2] + if cutlass.const_expr(acc_widen_to_fp32): + # INT8 int32 accumulate -> widen to fp32; `+ 0.0` + # forces a fresh fp32 register so int32->fp32 isn't + # folded into an invalid int32->fp8 cast. + _pf = _pair.to(cutlass.Float32) + vec_f32 = _pf + cutlass.full_like(_pf, 0.0) + else: + vec_f32 = _pair + col_j = col + frag_rel_col + lane_mod4 * 2 + linear_idx = tile_l * out_stride_l_0 + row * out_stride_m_0 + col_j * out_stride_n_0 + + # @@INJECT_EPILOGUE@@ + + _e_off = row_in_cta * _EPI_N + frag_rel_col + lane_mod4 * 2 + (smem_subtile_ptr.data_ptr() + _e_off).store_swizzled( + vec_out, + alignment=vec_bytes_epi, + swizzle=_EPI_SWIZZLE, + ) + + cute.arch.fence_view_async_shared() + nvvm.barrier_cta_sync( + barrier_id=EPI_SYNC_BAR_ID, + thread_count=NUM_COMPUTE_WARPS * 32, + ) + + if warp_idx == 0: + if nvvm.elect_sync(): + nvvm.cp_async_bulk_tensor_global_shared_cta( + tma_c_desc.get_ptr(), + smem_subtile_ptr, + (col, coord_m, tile_l), + ) + nvvm.cp_async_bulk_commit_group() + nvvm.cp_async_bulk_wait_group(EPI_SMEM_STAGES - 1, read=True) + + nvvm.barrier_cta_sync( + barrier_id=EPI_SYNC_BAR_ID, + thread_count=NUM_COMPUTE_WARPS * 32, + ) + # @@TMA_STORE_ONLY:END@@ + + # @@STG_ONLY:BEGIN@@ + _stg_stage = smem_stg_epi.subview(warp_idx * _STG_EPI_WARP_ELEMS) + for mf in cutlass.range_constexpr(_M_FRAGS): + for grp in cutlass.range_constexpr(_STG_EPI_NGRP): + _nf0 = grp * _STG_EPI_GROUP_FRAGS + _grp_frags = min(_STG_EPI_GROUP_FRAGS, _N_FRAGS - _nf0) + # -- STS_128: reg-index-order dump, one batch per fragment -- + for b in cutlass.range_constexpr(_grp_frags): + _o = (mf * _N_FRAGS + _nf0 + b) * 4 + _s_off = b * _STG_EPI_BATCH_STRIDE + lane * _STG_EPI_LANE_QUAD + (_stg_stage.data_ptr() + _s_off).store(acc[_o:4], alignment=16) + nvvm.bar_warp_sync(0xFFFFFFFF) + # -- LDS: 16 contiguous elems = both row-halves of one frag -- + _seg = (_stg_stage.data_ptr() + lane_mod4 * _STG_EPI_BATCH_STRIDE + lane_div4 * 16).load(alignment=16, count=16) + # Short tail group: trailing lanes own no fragment there + # (True at trace time for full groups — no guard emitted). + _lane_active = True if _grp_frags == _STG_EPI_GROUP_FRAGS else lane_mod4 < _grp_frags + if _lane_active: + for half in cutlass.range_constexpr(2): + row_in_cta = warp_row * _WARP_TILE_M + mf * 16 + half * 8 + lane_div4 + row = coord_m + row_in_cta + if row < M: + _row = cutlass.Array(mma_c_dtype, 8, alignment=16) + for sj in cutlass.range_constexpr(4): + _row[2 * sj] = _seg[4 * sj + 2 * half] + _row[2 * sj + 1] = _seg[4 * sj + 2 * half + 1] + for sv in cutlass.range_constexpr(8 // _STG_V): + col = coord_n + warp_col * _WARP_TILE_N + (_nf0 + lane_mod4) * 8 + sv * _STG_V + col_j = col + if col_j + vsize <= N: + # NB: Array slices are [start:COUNT], not + # [start:stop] (matches acc[_o:2] above). + _vec = _row[sv * _STG_V : _STG_V] + if cutlass.const_expr(acc_widen_to_fp32): + _pf = _vec.to(cutlass.Float32) + vec_f32 = _pf + cutlass.full_like(_pf, 0.0) + else: + vec_f32 = _vec + linear_idx = tile_l * out_stride_l_0 + row * out_stride_m_0 + col_j * out_stride_n_0 + + # @@INJECT_STG_VEC_BINDINGS@@ + + # @@INJECT_EPILOGUE@@ + nvvm.bar_warp_sync(0xFFFFFFFF) + # @@STG_ONLY:END@@ + + consumer_stage = tile_iter % CLC_SCHED_STAGES + if consumer_stage == 0 and tile_iter != 0: + clc_full_phase_epi = clc_full_phase_epi ^ 1 + while not nvvm.mbarrier_try_wait_parity( + clc_full_mbar_ptr.subview(consumer_stage), + clc_full_phase_epi, + time_limit=10_000_000, + ): + pass + m_idx, n_idx, l_idx, vld = cute_clc.clc_response(clc_response_ptr_base + consumer_stage) + cute.arch.fence_proxy("async.shared", space="cta") + is_valid = vld + tile_m, tile_n = _l2_swizzle_tile(m_idx, n_idx, gridx, gridy, swizzle_w) + tile_l = l_idx + nvvm.bar_warp_sync(0xFFFFFFFF) + if nvvm.elect_sync(): + nvvm.mbarrier_arrive(clc_empty_mbar_ptr.subview(consumer_stage)) + + tile_iter += 1 + + # No more tiles for this CTA: all its global A/B reads have been issued + # (fort fires launch_dependent_grids at the same point). + if cutlass.const_expr(USE_PDL): + if warp_idx == 0: + if nvvm.elect_sync(): + nvvm.griddepcontrol("launch_dependents") + + # @@TMA_STORE_ONLY:BEGIN@@ + if warp_idx == 0: + nvvm.cp_async_bulk_wait_group(0, read=True) + # @@TMA_STORE_ONLY:END@@ + + # -- Unused donor warps --------------------------------------------------- + if warp_idx > SCHEDULER_WARP_ID: + nvvm.setmaxregister(PROD_REG_COUNT, nvvm.SetMaxRegisterAction.DECREASE) + + +@cute.jit +def _host( + problem_size: tuple, + # @@INJECT_HOST_AB_PARAMS@@ + # @@INJECT_HOST_TAP_PARAMS@@ + # @@INJECT_HOST_AUX_PARAMS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_HOST_TMA_C_PARAMS@@ + # @@TMA_STORE_ONLY:END@@ + stream: _cuda.CUstream, +) -> None: + # @@INJECT_HOST_AB_LISTS@@ + m = problem_size[0] + n = problem_size[1] + k_sym = problem_size[2] + batch = problem_size[3] + _stride_idx = 4 + _a_stride_sets = [] + for _ in cutlass.range_constexpr(num_a_operands): + _a_stride_sets.append( + ( + problem_size[_stride_idx], + problem_size[_stride_idx + 1], + problem_size[_stride_idx + 2], + ) + ) + _stride_idx += 3 + _b_stride_sets = [] + for _ in cutlass.range_constexpr(num_b_operands): + _b_stride_sets.append( + ( + problem_size[_stride_idx], + problem_size[_stride_idx + 1], + problem_size[_stride_idx + 2], + ) + ) + _stride_idx += 3 + # @@INJECT_HOST_REDUCTION_STRIDES@@ + + if cutlass.const_expr(matmul_a_batch == 1): + a_batch = 1 + else: + a_batch = batch + if cutlass.const_expr(matmul_b_batch == 1): + b_batch = 1 + else: + b_batch = batch + + # K-major A/B only (asserted at render time): TMA boxes [K_tile, cta_{m,n}]. + tma_a_desc_list = [] + for _a_idx, _a_op in enumerate(_a_operands): + a_stride_m, a_stride_k, a_stride_l = _a_stride_sets[_a_idx] + tma_a_desc_list.append( + _tma.create_tensor_map_tiled( + global_address=_a_op.iterator.toint(), + dtype=ab_tma_dtype, + global_dims=[k_sym, m, a_batch], + global_strides=[ + a_stride_m * ab_dtype.width // 128, + a_stride_l * ab_dtype.width // 128, + ], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[0], 1], + swizzle=ab_tma_swizzle, + ) + ) + tma_b_desc_list = [] + for _b_idx, _b_op in enumerate(_b_operands): + b_stride_n, b_stride_k, b_stride_l = _b_stride_sets[_b_idx] + tma_b_desc_list.append( + _tma.create_tensor_map_tiled( + global_address=_b_op.iterator.toint(), + dtype=ab_tma_dtype, + global_dims=[k_sym, n, b_batch], + global_strides=[ + b_stride_n * ab_dtype.width // 128, + b_stride_l * ab_dtype.width // 128, + ], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[1], 1], + swizzle=ab_tma_swizzle, + ) + ) + + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_HOST_TMA_C_LISTS@@ + c = _tma_c_outputs[0] + # N-major output only (asserted at render time). + tma_c_desc = _tma.create_tensor_map_tiled( + global_address=c.iterator.toint(), + dtype=cd_tma_dtype, + global_dims=[n, m, batch], + global_strides=[ + out_stride_m_0 * cd_dtype.width // 128, + out_stride_l_0 * cd_dtype.width // 128, + ], + box_dims=[epi_tile_mn[1], cta_tile_mnk[0], 1], + swizzle=(_tma.TensorMapSwizzle.s64b if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), + ) + tma_c_desc_list = [tma_c_desc] + # @@TMA_STORE_ONLY:END@@ + + # CLC persistent grid: launch the full tile grid (fort launches the same); + # CTAs that finish early cancel not-yet-launched blocks and steal their + # (m, n, l) coordinates through the response ring. No cluster launch on + # sm120 (CC 12.0 has no thread-block clusters). + cgrp_tile_m = cgrp_tile_mnk[0] + cgrp_tile_n = cgrp_tile_mnk[1] + num_tile_m_host = (m + cgrp_tile_m - 1) // cgrp_tile_m + num_tile_n_host = (n + cgrp_tile_n - 1) // cgrp_tile_n + grid_shape = (num_tile_m_host, num_tile_n_host, batch) + _kernel( + problem_size[0], + problem_size[1], + problem_size[2], + # @@INJECT_HOST_KERNEL_DESC_PASS@@ + # @@INJECT_HOST_TAP_PASS@@ + # @@INJECT_HOST_REDUCTION_STRIDE_PASS@@ + # @@INJECT_HOST_AUX_PASS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_HOST_TMA_C_PASS@@ + # @@TMA_STORE_ONLY:END@@ + ).launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + use_pdl=USE_PDL, + stream=stream, + ) + + +@lru_cache(maxsize=None) +def compile() -> Callable: + out_vec_elems = vec_bytes_epi // (cd_dtype.width // 8) + ab_stride_elems = 16 // (ab_dtype.width // 8) + sym_m = cute.sym_int64() + sym_n = cute.sym_int64(divisibility=out_vec_elems) + # K tails are supported: the K loop is ceil_div and the TMA descriptor's global K + # extent makes a partial box HW zero-filled. The only real K rule is the 16-byte + # TMA contiguous-extent one, already gated by _tma_alignment_reject. + sym_k = cute.sym_int64() + sym_l = cute.sym_int64() + if matmul_a_batch == 1: + sym_a_l = 1 + else: + sym_a_l = sym_l + if matmul_b_batch == 1: + sym_b_l = 1 + else: + sym_b_l = sym_l + + def _make_fake_a(): + return make_fake_compact_tensor( + mma_a_dtype, + (sym_m, sym_k, sym_a_l), + stride_order=(0, 1, 2) if a_is_m_major else (1, 0, 2), + assumed_align=16, + ) + + def _make_fake_b(): + return make_fake_compact_tensor( + mma_b_dtype, + (sym_n, sym_k, sym_b_l), + stride_order=(0, 1, 2) if b_is_n_major else (1, 0, 2), + assumed_align=16, + ) + + # @@TMA_STORE_ONLY:BEGIN@@ + def _make_fake_c(): + return make_fake_compact_tensor( + cd_dtype, + (sym_m, sym_n // cd_fake_n_div, sym_l), + stride_order=(0, 1, 2) if cd_out_is_m_major else (1, 0, 2), + assumed_align=16, + ) + + # @@INJECT_COMPILE_TMA_C_FAKES@@ + # @@TMA_STORE_ONLY:END@@ + def _sym_operand_strides(is_mn_major: bool) -> tuple: + # Operand is permuted to (M|N, K, L): the unit stride is mode 0 when MN-major, mode 1 when K-major, and never reaches TMA. + unit = 0 if is_mn_major else 1 + return tuple(cute.sym_int64() if i == unit else cute.sym_int64(divisibility=ab_stride_elems) for i in range(3)) + + sym_a_strides = [] + for _ in range(num_a_operands): + sym_a_strides.extend(_sym_operand_strides(a_is_m_major)) + sym_b_strides = [] + for _ in range(num_b_operands): + sym_b_strides.extend(_sym_operand_strides(b_is_n_major)) + # @@INJECT_COMPILE_REDUCTION_STRIDE_DECLS@@ + # @@INJECT_COMPILE_AB_FAKES@@ + # @@INJECT_COMPILE_TAP_FAKES@@ + problem_size = ( + sym_m, + sym_n, + sym_k, + sym_l, + *sym_a_strides, + *sym_b_strides, + # @@INJECT_COMPILE_REDUCTION_STRIDE_SYMBOLS@@ + ) + # @@INJECT_COMPILE_AUX_FAKES@@ + _fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) + return cute.compile( + _host, + problem_size, + # @@INJECT_COMPILE_AB_PASS@@ + # @@INJECT_COMPILE_TAP_PASS@@ + # @@INJECT_COMPILE_AUX_PASS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_COMPILE_TMA_C_PASS@@ + # @@TMA_STORE_ONLY:END@@ + stream=_fake_stream, + options=frost_compile_options, + ) diff --git a/python/cudnn/gemm/frost/tile_config.py b/python/cudnn/gemm/frost/tile_config.py index dfb9f2652..35b599fa8 100644 --- a/python/cudnn/gemm/frost/tile_config.py +++ b/python/cudnn/gemm/frost/tile_config.py @@ -44,7 +44,7 @@ def _sm_smem_budget_bytes(device=None) -> int: # every smem barrier, the TMEM base address and — on the MoE templates — the per-CTA TMA # tensormap scratch. The ab pipeline itself only counts the operand tensors, so these # fragments are budgeted once here instead of being modelled stage by stage. -_SMEM_FIXED_RESERVE_BY_PIPELINE = {"sm100": 2048, "sm103": 2048, "sm107": 2048} +_SMEM_FIXED_RESERVE_BY_PIPELINE = {"sm100": 2048, "sm103": 2048, "sm107": 2048, "sm120": 2048} _SMEM_FIXED_RESERVE_MOE_BY_PIPELINE = {"sm100": 4096, "sm103": 4096, "sm107": 4096} @@ -78,11 +78,11 @@ def l2_swizzle_budget_bytes(device=None) -> int: _CTA_TILE_M_MAX = 128 _CTA_TILE_N_MAX = 256 _MAX_CLUSTER_SIZE = _FROST_MAX_CLUSTER_SIZE -_CTA_TILE_K_BYTES_MAX_BY_PIPELINE = {"sm100": 128, "sm103": 384, "sm107": 128} +_CTA_TILE_K_BYTES_MAX_BY_PIPELINE = {"sm100": 128, "sm103": 384, "sm107": 128, "sm120": 128} # MMA-inst K in bytes _MMA_INST_K_BYTES = 32 -_MMA_INST_K_BYTES_BY_PIPELINE = {"sm100": 32, "sm103": 48, "sm107": 64} +_MMA_INST_K_BYTES_BY_PIPELINE = {"sm100": 32, "sm103": 48, "sm107": 64, "sm120": 32} def _pipeline_fact(table: dict, pipeline: str, what: str): @@ -371,10 +371,36 @@ class ConfigSm107(TileConfig): ``pipeline="sm107"``.""" +class ConfigSm120(TileConfig): + """sm120 geometry — warp-scoped MMA on consumer Blackwell (CC 12.x). The + cluster (none on CC 12.x, so 1x1) and the block size (a fixed 12-warp + kernel: 8 compute + TMA + CLC scheduler + 2 donors) are family-FIXED, so + ``__post_init__`` pins them rather than validating — a config crossing in + from a clustered family (``as_pipeline``) needs no sm120 knowledge. The + free axes keep ``cta_tile_n % 16 == 0`` (the 4x2 compute-warp grid owns + 8-column n-fragments per warp column) and one MMA-M block. Callers pass + ``pipeline="sm120"``.""" + + def __post_init__(self) -> None: + object.__setattr__(self, "cgrp_size_m", 1) + object.__setattr__(self, "cgrp_size_n", 1) + object.__setattr__(self, "threads_per_cta", 384) + super().__post_init__() + if self.cta_tile_n % 16 != 0: + raise NotImplementedError( + f"TileConfig {self.name!r}: sm120 needs cta_tile_n % 16 == 0 " f"(each of the 2 compute-warp columns owns 8-column n-fragments)" + ) + if self.num_mma_m != 1: + raise NotImplementedError( + f"TileConfig {self.name!r}: sm120 warp MMA has no MMA-M split; " f"mma_inst_m must equal cta_tile_m (got num_mma_m={self.num_mma_m})" + ) + + _CONFIG_CLASS_BY_PIPELINE: dict[str, type[TileConfig]] = { "sm100": ConfigSm100, "sm103": ConfigSm103, "sm107": ConfigSm107, + "sm120": ConfigSm120, } @@ -472,6 +498,21 @@ def _geom_sm107(num_mma_m: int, cta_n: int, cgrp_m: int, cgrp_n: int) -> ConfigS ) +def _geom_sm120(cta_m: int, cta_n: int, k_bytes: int) -> ConfigSm120: + """Build one sm120 config (cluster fixed at 1x1, 12-warp kernel).""" + return ConfigSm120( + cta_tile_m=cta_m, + cta_tile_n=cta_n, + cta_tile_k_bytes=k_bytes, + cgrp_size_m=1, + cgrp_size_n=1, + epi_tile_mn=(cta_m, 32), + threads_per_cta=384, + pipeline="sm120", + acc_stages=2, + ) + + def _build_catalog() -> tuple[TileConfig, ...]: cfgs: list[TileConfig] = [] for mma_m, num_mma_m in _M_AXES: @@ -492,6 +533,12 @@ def _build_catalog() -> tuple[TileConfig, ...]: for cta_n in (256, 128): for cgrp_m, cgrp_n in _CLUSTERS: cfgs.append(_geom_sm107(num_mma_m, cta_n, cgrp_m, cgrp_n)) + # sm120 warp-MMA geometries: no cluster axis (fixed 1x1); M ∈ {128, 64} + # (the MMA-inst M bounds, and the 4x2 warp grid needs cta_m % 64 == 0). + for cta_m in (128, 64): + for cta_n in range(256, 0, -32): + for k_bytes in (128, 64): + cfgs.append(_geom_sm120(cta_m, cta_n, k_bytes)) return tuple(cfgs) @@ -715,9 +762,11 @@ def select_config( def as_pipeline(cfg: TileConfig, pipeline: str) -> TileConfig: """The same geometry as a ``pipeline``-family config — only the family-fixed - MMA-inst K width moves. A family whose K axes this geometry cannot satisfy - (sm103 fixes a 384-byte K-tile) raises from the config's ``__post_init__``, - so the invariant stays in one place.""" + MMA-inst K width moves (a family that fixes MORE axes pins them in its own + ``__post_init__``, e.g. ConfigSm120's cluster and block size). A family + whose K axes this geometry cannot satisfy (sm103 fixes a 384-byte K-tile) + raises from the config's ``__post_init__``, so the invariant stays in one + place.""" if cfg.pipeline == pipeline: return cfg cls = config_class_for_pipeline(pipeline) diff --git a/test/python/gemm/frost/test_sm120_matmul.py b/test/python/gemm/frost/test_sm120_matmul.py new file mode 100644 index 000000000..862105eeb --- /dev/null +++ b/test/python/gemm/frost/test_sm120_matmul.py @@ -0,0 +1,680 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""sm120 (consumer Blackwell) matmul template gate. + +Two layers, mirroring ``test_matmul.py``: + +* Wiring tests (no sm_120 GPU needed — they run on the sm100 CI too): the + registry / tile-config / compiler plumbing for ``sm120_matmul.py`` — catalog + family, template routing, the v1 scope gates, the ≤8-element epilogue chunk + clamp with STG-only rendering, and (on any CUDA GPU) a full source render. + +* End-to-end correctness on an sm_120 GPU: (config × dtype-pair × shape) + sweeps asserting bit-tight equality vs torch-fp32 (small-integer inputs keep + the reduction exact), plus batched / batch-broadcast, epilogue fusion, + narrow output rows, and the engine's auto-select path. + +CUDNN_GEMM_TEST_FULL=1 expands the config axis to every sm120 catalog +geometry. Also runnable as a script (forwards argv to pytest). +""" + +from __future__ import annotations + +import ast +import os +import sys + +import pytest +import torch + +from gemm_test_utils import ( + Plan as _plan, + vp as _vp, + resolve as _resolve, +) + +pytestmark = [pytest.mark.L0] + + +import cudnn +import cudnn.gemm.frost # noqa: F401 — installs the cudnn.pygraph recorder hook +from cudnn.gemm.frost.graph_analyzer import analyze +from cudnn.gemm.frost.kernel_registry import PIPELINE_ARCH_RANGES +from cudnn.gemm.frost.tile_config import CATALOG, ConfigSm120, by_name + +# --- arch gate --------------------------------------------------------------- + + +def _active_sm() -> int | None: + if not torch.cuda.is_available(): + return None + major, minor = torch.cuda.get_device_capability() + return major * 10 + minor + + +_SM = _active_sm() +_SM120_RANGES = PIPELINE_ARCH_RANGES["sm120"] + +# The e2e tests JIT + LAUNCH the sm120 warp-MMA template; gate on consumer +# Blackwell so wrong-arch machines skip instead of failing in the launch. +requires_sm120 = pytest.mark.skipif( + _SM is None or not (120 <= _SM < 130), + reason="needs a consumer-Blackwell GPU (120 <= SM < 130), have " + ("none" if _SM is None else f"sm_{_SM}"), +) + +requires_any_gpu = pytest.mark.skipif( + _SM is None, + reason="rendering sizes the SMEM/L2 budgets from the active GPU", +) + + +_TORCH_DTYPE = { + "bf16": torch.bfloat16, + "fp16": torch.float16, + "fp8_e4m3": torch.float8_e4m3fn, + "fp8_e5m2": torch.float8_e5m2, +} +_CUDNN_DTYPE = { + "bf16": cudnn.data_type.BFLOAT16, + "fp16": cudnn.data_type.HALF, + "fp8_e4m3": cudnn.data_type.FP8_E4M3, + "fp8_e5m2": cudnn.data_type.FP8_E5M2, +} +_ELEM_BYTES = {"bf16": 2, "fp16": 2, "fp8_e4m3": 1, "fp8_e5m2": 1} + + +# Shape menu: tile-aligned baseline + M-OOB + N-OOB + K-OOB + combined. Every N +# is a multiple of 8, so 2-byte outputs keep 16-byte rows — one compiled anchor +# per (config, dtype) serves the whole menu (the kernel is shape-agnostic). +_WEIRD_SHAPES: tuple[tuple[int, int, int], ...] = ( + # Tile-aligned baseline. + (384, 768, 384), + (640, 384, 512), + (256, 1280, 256), + (512, 1024, 640), # K = 5×128 + # M-OOB (N, K aligned). + (255, 256, 256), # one row short of a tile + (200, 256, 256), # deep inside a partial tile + # N-OOB (predicated pair stores / TMA global-extent clip). + (256, 200, 256), + # K-OOB (bf16/fp16 only for K=200; FP8 skips via the 16B TMA stride rule). + (256, 256, 200), + (256, 256, 96), # smaller than one K_BYTES=128 BF16 tile + # M + N + K OOB. + (255, 200, 240), +) + +# (input_dtype, output_dtype) pairs — the sm120 MMA menu (fp32-accumulated). +_CORE_DTYPE_PAIRS: tuple[tuple[str, str], ...] = ( + ("bf16", "bf16"), + ("fp16", "fp16"), + ("fp8_e4m3", "fp16"), + ("fp8_e5m2", "fp16"), + ("fp8_e4m3", "bf16"), +) + +# Curated config subset — each entry covers a distinct template corner. The +# 48/144/16 N-tiles are synthesized by name (the catalog walks N in 32s). The +# full sm120 catalog sweep is opt-in via CUDNN_GEMM_TEST_FULL=1. +_QUICK_CONFIGS: tuple[str, ...] = ( + "CONFIG_sm120_128x128x128_128x128x32_cluster1x1_1ctamma", # baseline + "CONFIG_sm120_128x256x128_128x256x32_cluster1x1_1ctamma", # large N + "CONFIG_sm120_128x64x128_128x64x32_cluster1x1_1ctamma", # narrow N + "CONFIG_sm120_128x128x64_128x128x32_cluster1x1_1ctamma", # K_BYTES=64 (s64b AB swizzle) + "CONFIG_sm120_64x128x128_64x128x32_cluster1x1_1ctamma", # cta_m=64 (16-row warp tile) + "CONFIG_sm120_128x48x128_128x48x32_cluster1x1_1ctamma", # N%32 != 0, odd n-frag tail + "CONFIG_sm120_128x144x128_128x144x32_cluster1x1_1ctamma", # 9 n-frags/warp (4 pairs + tail) + "CONFIG_sm120_64x16x128_64x16x32_cluster1x1_1ctamma", # minimum tile (single-n-frag warp) + "CONFIG_sm120_64x64x64_64x64x32_cluster1x1_1ctamma", # cta_m=64 + K_BYTES=64 + "CONFIG_sm120_128x256x64_128x256x32_cluster1x1_1ctamma", # K_BYTES=64, large N +) + +_BATCHED_CONFIGS: tuple[str, ...] = ( + "CONFIG_sm120_128x128x128_128x128x32_cluster1x1_1ctamma", + "CONFIG_sm120_64x128x128_64x128x32_cluster1x1_1ctamma", +) + +_BATCHED_SHAPES: tuple[tuple[int, int, int, int], ...] = ( + (1, 384, 768, 384), + (2, 640, 384, 512), + (3, 255, 256, 256), # M-OOB + (2, 256, 200, 240), # N + K OOB +) + +_BATCH_BROADCAST_CASES = tuple((side, (2, 384, 256, 384)) for side in ("A", "B")) + + +def _sweep_config_names() -> list[str]: + """Quick subset by default; the whole sm120 catalog under CUDNN_GEMM_TEST_FULL=1.""" + if os.environ.get("CUDNN_GEMM_TEST_FULL", "0") == "1": + return [f"{c.name}_1ctamma" for c in CATALOG if c.pipeline == "sm120"] + return list(_QUICK_CONFIGS) + + +def _shape_id(s: tuple[int, int, int]) -> str: + return f"{s[0]}x{s[1]}x{s[2]}" + + +def _dtype_id(p: tuple[str, str]) -> str: + return f"{p[0]}->{p[1]}" + + +def _config_id(name: str) -> str: + return name.removeprefix("CONFIG_sm120_") + + +# --- compatibility gate ------------------------------------------------------- + + +def _compatible(cfg, M: int, N: int, K: int, in_dtype: str, out_dtype: str) -> tuple[bool, str]: + """Reject only shapes the sm120 kernel can't service. Returns (ok, reason).""" + in_eb = _ELEM_BYTES[in_dtype] + out_eb = _ELEM_BYTES[out_dtype] + if cfg.cta_tile_k_bytes % in_eb != 0: + return False, f"K_BYTES={cfg.cta_tile_k_bytes} not divisible by in_elem_bytes={in_eb}" + # K-major A and B: the TMA contiguous extent is K on both sides. + if (K * in_eb) % 16 != 0: + return False, f"K*in_eb={K * in_eb} not 16B-aligned (TMA contiguous-extent rule); " f"{in_dtype!r} needs K % {16 // in_eb} == 0" + # The pair epilogue stores 2 output elements per thread. + if (N * out_eb) % 4 != 0: + return False, f"N*out_eb={N * out_eb} not 4B-aligned — the (n, n+1) pair store needs it" + return True, "" + + +# --- graph + data + reference ------------------------------------------------ + + +def _a_stride_batched(M: int, K: int, a_major: str) -> list[int]: + return [M * K, K, 1] if a_major == "k" else [M * K, 1, M] + + +def _b_stride_batched(N: int, K: int, b_major: str) -> list[int]: + return [N * K, 1, K] if b_major == "k" else [N * K, N, 1] + + +def _build_graph( + M: int, + N: int, + K: int, + in_dtype: str = "bf16", + out_dtype: str = "bf16", + a_major: str = "k", + b_major: str = "k", + out_major: str = "n", +) -> cudnn.pygraph: + g = cudnn.pygraph( + io_data_type=_CUDNN_DTYPE[in_dtype], + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + ) + A = g.tensor(name="A", dim=[1, M, K], stride=_a_stride_batched(M, K, a_major)) + B = g.tensor(name="B", dim=[1, K, N], stride=_b_stride_batched(N, K, b_major)) + C = g.matmul(A=A, B=B, name="mm") + if out_major == "m": + C.set_stride([M * N, 1, M]) + C.set_output(True) + if out_dtype != in_dtype: + C.set_data_type(_CUDNN_DTYPE[out_dtype]) + return g + + +def _build_batched_graph(batch: int, M: int, N: int, K: int, in_dtype: str, out_dtype: str, a_batch=None, b_batch=None) -> cudnn.pygraph: + """Rank-3 batched matmul; pass ``a_batch``/``b_batch``=1 for a broadcast side.""" + g = cudnn.pygraph( + io_data_type=_CUDNN_DTYPE[in_dtype], + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + ) + A = g.tensor(name="A", dim=[a_batch or batch, M, K], stride=_a_stride_batched(M, K, "k")) + B = g.tensor(name="B", dim=[b_batch or batch, K, N], stride=_b_stride_batched(N, K, "k")) + C = g.matmul(A=A, B=B, name="mm") + C.set_output(True) + if out_dtype != in_dtype: + C.set_data_type(_CUDNN_DTYPE[out_dtype]) + return g + + +def _build_bias_relu_graph(M: int, N: int, K: int) -> cudnn.pygraph: + """matmul -> per-col bias -> relu: an aux + pointwise epilogue chain (aux + chains take the STG path — the TMA-store gate excludes aux).""" + g = cudnn.pygraph( + io_data_type=cudnn.data_type.BFLOAT16, + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + ) + A = g.tensor(name="A", dim=[1, M, K], stride=_a_stride_batched(M, K, "k")) + B = g.tensor(name="B", dim=[1, K, N], stride=_b_stride_batched(N, K, "k")) + C = g.matmul(A=A, B=B, name="mm") + bias = g.tensor(name="bias0", dim=[1, N], stride=[N, 1]) + D = g.bias(input=C, bias=bias, name="bias") + E = g.relu(input=D, name="relu") + E.set_output(True) + return g + + +def _mkdata(batch: int, M: int, N: int, K: int, in_dtype: str, out_dtype: str, seed: int = 0): + """Small-integer inputs ⇒ exact FP32 reduction ⇒ kernel and reference differ + only by the final deterministic downcast. Rank-3, K-major.""" + torch.manual_seed(seed) + rng = (-3, 3) if in_dtype.startswith("fp8") else (-2, 2) + a = torch.empty(batch, M, K, dtype=torch.int32).random_(*rng).to(dtype=_TORCH_DTYPE[in_dtype], device="cuda") + b = torch.empty(batch, N, K, dtype=torch.int32).random_(*rng).to(dtype=_TORCH_DTYPE[in_dtype], device="cuda") + c = torch.empty(batch, M, N, dtype=_TORCH_DTYPE[out_dtype], device="cuda") + return a, b, c + + +def _reference(a: torch.Tensor, b: torch.Tensor, out_dtype: str) -> torch.Tensor: + ref = torch.einsum("bmk,bnk->bmn", a.to(torch.float32), b.to(torch.float32)) + return ref.to(_TORCH_DTYPE[out_dtype]) + + +def _assert_bit_tight(c, ref, header: str) -> None: + """Both sides reduce exactly in FP32 and downcast the same way ⇒ equality.""" + diff = (c.to(torch.float32) - ref.to(torch.float32)).abs() + bad = int((diff > 0).sum().item()) + assert bad == 0, ( + f"\n {header}" + f"\n bad: {bad}/{diff.numel()} ({100 * bad / diff.numel():.2f}%)" + f"\n max|diff|: {float(diff.max().item()):.4g}" + f"\n max|ref|: {float(ref.abs().max().item()):.4g}" + f"\n hint: sample c[0,0,:8] = {c[0, 0, :8].to(torch.float32).tolist()}" + f"\n sample ref[0,0,:8] = {ref[0, 0, :8].to(torch.float32).tolist()}" + ) + + +# --- compile cache (session-scoped) -------------------------------------------- + + +@pytest.fixture(scope="session") +def _compile_cache() -> dict: + """Maps a case key → Plan | ("skip"|"fail", msg). Cases visit in (config, + dtype, shape) order, so each (config, dtype) block shares one compile.""" + return {} + + +def _cached_outcome(entry): + if isinstance(entry, tuple) and entry[0] in ("skip", "fail"): + kind, msg = entry + if kind == "skip": + pytest.skip(msg) + pytest.fail(msg, pytrace=False) + return entry + + +def _plan_or_skip(cache, key, build_graph, cfg, cta_group, **plan_kw): + """JIT the anchor graph; the engine's clean "unsupported" rejections SKIP, + any other compile error FAILS.""" + try: + compiled = _plan(build_graph(), config=cfg, cta_group=cta_group, **plan_kw) + except Exception as e: + first = str(e).splitlines()[0] if str(e) else "" + if isinstance(e, NotImplementedError) or (isinstance(e, ValueError) and "no kernel template" in str(e)): + msg = first[:300] + cache[key] = ("skip", msg) + pytest.skip(msg) + msg = f"JIT compile failed: {type(e).__name__}: {first[:200]}" + cache[key] = ("fail", msg) + pytest.fail(msg, pytrace=False) + cache[key] = compiled + return compiled + + +def _pick_anchor(cfg, in_dt: str, out_dt: str) -> tuple[int, int, int] | None: + for shape in _WEIRD_SHAPES: + ok, _ = _compatible(cfg, *shape, in_dt, out_dt) + if ok: + return shape + return None + + +def _get_compiled(cache: dict, cfg, in_dt: str, out_dt: str, cta_group: int): + key = (cfg.name, in_dt, out_dt, cta_group) + if key in cache: + return _cached_outcome(cache[key]) + anchor = _pick_anchor(cfg, in_dt, out_dt) + if anchor is None: + msg = f"no menu shape is compatible with ({cfg.name}, {in_dt}->{out_dt})" + cache[key] = ("skip", msg) + pytest.skip(msg) + return _plan_or_skip(cache, key, lambda: _build_graph(*anchor, in_dt, out_dt), cfg, cta_group) + + +# ============================================================================= +# Wiring tests — no sm_120 GPU required (registry / config / compiler gates). +# ============================================================================= + + +def test_sm120_registry_wiring() -> None: + """The sm120 template is registered with its own family, class and gates.""" + from cudnn.gemm.frost.kernel_registry import ( + MMA_TYPE_SUPPORT, + TEMPLATES, + GraphType, + Sm120KernelTemplate, + ) + + # SM 12.x is in the family's active range (whatever else the range covers). + assert any(lo <= 120 < hi for lo, hi in _SM120_RANGES) + + (tmpl,) = [t for t in TEMPLATES if t.pipeline == "sm120"] + assert tmpl.file == "sm120_matmul.py" + assert isinstance(tmpl, Sm120KernelTemplate) + # Warp-scoped MMA: 1-CTA only, no multi-GEMM (no per-GEMM operand indexing). + assert tmpl.cta_group == 1 and not tmpl.supports_multi_gemm + assert tmpl.graph_type is GraphType.MATMUL and not tmpl.mainloop + + # dtype support mirrors the sm100 matmul pipeline (incl. int8 and fp8 mixes) + assert ("bf16", "bf16", "fp32") in MMA_TYPE_SUPPORT["sm120"][GraphType.MATMUL] + assert ("int8", "int8", "int32") in MMA_TYPE_SUPPORT["sm120"][GraphType.MATMUL] + + # the template file itself ships with the package + from pathlib import Path + + import cudnn.gemm.frost.compiler as C + + assert (Path(C.__file__).parent / "kernel_templates" / "sm120_matmul.py").is_file() + + +def test_sm120_tile_config_family() -> None: + """ConfigSm120: catalog membership, name round-trip, geometry guards, and + the auto path's as_pipeline conversion (cluster + block size move).""" + from cudnn.gemm.frost.tile_config import as_pipeline + + sm120 = [c for c in CATALOG if c.pipeline == "sm120"] + assert len(sm120) == 2 * 8 * 2 # m {128,64} × n {32..256/32} × kb {128,64} + assert all(isinstance(c, ConfigSm120) for c in sm120) + assert all(c.cgrp_size_mn == (1, 1) and c.threads_per_cta == 384 for c in sm120) + + cfg = by_name("CONFIG_sm120_128x128x128_128x128x32_cluster1x1") + assert isinstance(cfg, ConfigSm120) + assert cfg.name == "CONFIG_sm120_128x128x128_128x128x32_cluster1x1" + # non-catalog sm120 geometries synthesize with the family's 12-warp block + c48 = by_name("CONFIG_sm120_128x48x128_128x48x32_cluster1x1") + assert c48.threads_per_cta == 384 and c48.cgrp_size_mn == (1, 1) + + # a clustered sm120 NAME pins to 1x1 and so fails the canonical round-trip; + # the free axes (16-col n-frag granularity, no MMA-M split) reject directly + with pytest.raises(KeyError, match="round-trips"): + by_name("CONFIG_sm120_128x128x128_128x128x32_cluster2x1") + with pytest.raises(NotImplementedError, match="cta_tile_n"): + by_name("CONFIG_sm120_128x24x128_128x24x32_cluster1x1") + with pytest.raises(NotImplementedError): + by_name("CONFIG_sm120_256x128x128_128x128x32_cluster1x1") + + # as_pipeline: an sm100 auto pick crosses over, and ConfigSm120 pins its + # family-fixed axes (cluster -> 1x1, threads -> 384) by itself + picked = by_name("CONFIG_sm100_128x256x128_128x256x32_cluster2x1") + conv = as_pipeline(picked, "sm120") + assert isinstance(conv, ConfigSm120) + assert conv.cgrp_size_mn == (1, 1) and conv.threads_per_cta == 384 + assert conv.cta_tile_mn == picked.cta_tile_mn + + # resolve() understands the legacy test names + r_cfg, r_group = _resolve("CONFIG_sm120_128x128x128_128x128x32_cluster1x1_1ctamma") + assert r_cfg is cfg and r_group == 1 + + +def test_sm120_template_routing() -> None: + """select_template: (sm120 config, cta_group=1) → the sm120 template; the + strategies sm120 doesn't have decline with the registry's clean error.""" + from cudnn.gemm.frost.kernel_registry import select_template + + chain = analyze(_build_graph(256, 256, 128)) + cfg = by_name("CONFIG_sm120_128x128x128_128x128x32_cluster1x1") + assert select_template(chain, cfg, 1).file == "sm120_matmul.py" + with pytest.raises(ValueError, match="no kernel template"): + select_template(chain, cfg, 2) + # sm100 configs keep routing to the sm100 family + cfg100 = by_name("CONFIG_sm100_128x128x128_128x128x32_cluster1x1") + assert select_template(chain, cfg100, 1).file == "sm100_matmul_1ctamma.py" + + +def test_sm120_scope_gates() -> None: + """The v1 scope contract rejects through the registry (never a template + AssertionError mid-render): K-major-only inputs, N-major output, and a + pair-storable epilogue chunk.""" + from cudnn.gemm.frost.kernel_registry import select_template + + cfg = by_name("CONFIG_sm120_128x128x128_128x128x32_cluster1x1") + tmpl = select_template(analyze(_build_graph(256, 256, 128)), cfg, 1) + + def scope_reject(**graph_kw): + return tmpl._extra_reject(analyze(_build_graph(256, 256, 128, **graph_kw)), cfg) + + assert scope_reject() is None + assert "K-major" in scope_reject(a_major="m") + assert "K-major" in scope_reject(b_major="n") + assert "N-major" in scope_reject(out_major="m") + + # the funnel never offers sm120 points for an out-of-scope chain + from cudnn.gemm.frost.kernel_registry import candidates + + nmaj = analyze(_build_graph(256, 256, 128, b_major="n")) + assert not [t for t, _c in candidates(nmaj) if t.pipeline == "sm120"] + + +def test_sm120_epi_vec_clamp_and_stg_only() -> None: + """sm120 always renders the transposed-STG epilogue, with the chunk capped + at the 8-element fragment row run; the sm100 derivation is untouched.""" + from cudnn.gemm.frost.compiler import _epi_vec_bytes, _use_tma_store_epi + + chain = analyze(_build_graph(384, 768, 384)) + for name in ( + "CONFIG_sm120_128x128x128_128x128x32_cluster1x1", + "CONFIG_sm120_64x128x128_64x128x32_cluster1x1", + "CONFIG_sm120_128x48x128_128x48x32_cluster1x1", + ): + c = by_name(name) + v = _epi_vec_bytes(chain, c, 1) + assert v == 16, name # 8-element bf16 run (template: 8 % _STG_V == 0) + assert _use_tma_store_epi(chain, c, v, 1) is False, name + + # the sm100 chunk derivation and its TMA-store gate are untouched + cfg100 = by_name("CONFIG_sm100_128x128x128_128x128x32_cluster1x1") + v100 = _epi_vec_bytes(chain, cfg100, 1) + assert v100 >= 16 and _use_tma_store_epi(chain, cfg100, v100, 1) is True + + # 8B-aligned output rows (N=100 bf16) narrow the chunk to 4 elements + narrow = analyze(_build_graph(256, 100, 256)) + assert _epi_vec_bytes(narrow, by_name("CONFIG_sm120_128x128x128_128x128x32_cluster1x1"), 1) == 8 + + +@requires_any_gpu +def test_sm120_render_smoke() -> None: + """Render the sm120 template end-to-end (real tile constants + epilogue + snippets) on whatever GPU is active — no cute.compile, so this covers the + sm100 CI too. The source must be marker-free, parseable, and carry the + STG-only sm120 contract constants.""" + from cudnn.gemm.frost.compiler import _epi_vec_bytes, _render_template + from cudnn.gemm.frost.epilogue_codegen import generate + + chain = analyze(_build_graph(384, 768, 384)) + for name in ( + "CONFIG_sm120_128x128x128_128x128x32_cluster1x1", + "CONFIG_sm120_128x48x128_128x48x32_cluster1x1", # odd n-frag tail + ): + cfg = by_name(name) + vec = _epi_vec_bytes(chain, cfg, 1) + snippets = generate(chain, vec_bytes_epi=vec, output_elem_bytes=2, use_tma_store=False) + src = _render_template(chain, snippets, cfg, 1) + assert "@@" not in src, "leftover injection markers" + ast.parse(src) + assert "cudnn_frost_sm120_matmul_" in src + assert "threads_per_cta = 384" in src + assert f"vec_bytes_epi = {vec}" in src + assert "use_tma_store_epi = False" in src + + +# ============================================================================= +# End-to-end correctness — needs an sm_120 GPU. +# ============================================================================= + + +@requires_sm120 +@pytest.mark.parametrize("shape", _WEIRD_SHAPES, ids=[_shape_id(s) for s in _WEIRD_SHAPES]) +@pytest.mark.parametrize("in_dt,out_dt", _CORE_DTYPE_PAIRS, ids=[_dtype_id(p) for p in _CORE_DTYPE_PAIRS]) +@pytest.mark.parametrize("config_name", _sweep_config_names(), ids=[_config_id(n) for n in _sweep_config_names()]) +def test_sm120_matmul( + _compile_cache, + config_name: str, + in_dt: str, + out_dt: str, + shape: tuple[int, int, int], +) -> None: + """One (config, dtype-pair, shape); incompatible combos SKIP, else bit-tight.""" + cfg, cta_group = _resolve(config_name) + ok, reason = _compatible(cfg, *shape, in_dt, out_dt) + if not ok: + pytest.skip(reason) + + compiled = _get_compiled(_compile_cache, cfg, in_dt, out_dt, cta_group) + + M, N, K = shape + a, b, c = _mkdata(1, M, N, K, in_dt, out_dt) + compiled(_vp(compiled, a, b, c)) + torch.cuda.synchronize() + + _assert_bit_tight( + c, + _reference(a, b, out_dt), + f"config: {config_name}\n dtype: {in_dt} -> {out_dt}\n shape: {M}x{N}x{K}", + ) + + +@requires_sm120 +@pytest.mark.parametrize("bshape", _BATCHED_SHAPES, ids=[f"B{s[0]}_{s[1]}x{s[2]}x{s[3]}" for s in _BATCHED_SHAPES]) +@pytest.mark.parametrize("config_name", _BATCHED_CONFIGS, ids=[_config_id(n) for n in _BATCHED_CONFIGS]) +def test_sm120_matmul_batched(_compile_cache, config_name: str, bshape) -> None: + """Rank-3 batches ride gridDim.z; the CLC scheduler steals across planes.""" + batch, M, N, K = bshape + in_dt = out_dt = "bf16" + cfg, cta_group = _resolve(config_name) + + # Keyed by batch too: batch=1 bakes the degenerate-broadcast const branch + # (matmul_a_batch == 1), so one anchor cannot serve both batch classes. + key = ("batched", cfg.name, in_dt, out_dt, batch) + if key in _compile_cache: + compiled = _cached_outcome(_compile_cache[key]) + else: + compiled = _plan_or_skip( + _compile_cache, + key, + lambda: _build_batched_graph(batch, M, N, K, in_dt, out_dt), + cfg, + cta_group, + ) + + a, b, c = _mkdata(batch, M, N, K, in_dt, out_dt) + compiled(_vp(compiled, a, b, c)) + torch.cuda.synchronize() + _assert_bit_tight(c, _reference(a, b, out_dt), f"config: {config_name} batched {batch}x{M}x{N}x{K}") + + +@requires_sm120 +@pytest.mark.parametrize("case", _BATCH_BROADCAST_CASES, ids=[f"broadcast{s}" for s, _ in _BATCH_BROADCAST_CASES]) +def test_sm120_batch_broadcast(case) -> None: + """One operand batch-broadcast (batch=1 input against a batch>1 GEMM).""" + side, (batch, M, N, K) = case + cfg, cta_group = _resolve("CONFIG_sm120_128x128x128_128x128x32_cluster1x1_1ctamma") + a_batch = 1 if side == "A" else batch + b_batch = 1 if side == "B" else batch + compiled = _plan( + _build_batched_graph(batch, M, N, K, "bf16", "bf16", a_batch=a_batch, b_batch=b_batch), + config=cfg, + cta_group=cta_group, + ) + a, _, _ = _mkdata(a_batch, M, N, K, "bf16", "bf16") + _, b, _ = _mkdata(b_batch, M, N, K, "bf16", "bf16", seed=1) + c = torch.empty(batch, M, N, dtype=torch.bfloat16, device="cuda") + compiled(_vp(compiled, a, b, c)) + torch.cuda.synchronize() + ref = torch.einsum( + "bmk,bnk->bmn", + a.to(torch.float32).expand(batch, -1, -1), + b.to(torch.float32).expand(batch, -1, -1), + ).to(torch.bfloat16) + _assert_bit_tight(c, ref, f"batch-broadcast {side}, {batch}x{M}x{N}x{K}") + + +@requires_sm120 +@pytest.mark.parametrize( + "config_name", + [ + "CONFIG_sm120_128x128x128_128x128x32_cluster1x1_1ctamma", # aligned warp N-frags + "CONFIG_sm120_128x48x128_128x48x32_cluster1x1_1ctamma", # odd n-frag tail + ], + ids=["128x128", "128x48_tail"], +) +def test_sm120_bias_relu_epilogue(config_name: str) -> None: + """Per-col bias + relu through the epilogue: the per-col vector aux load + plus the predicated pair stores.""" + M, N, K = 256, 240, 128 + cfg, cta_group = _resolve(config_name) + compiled = _plan(_build_bias_relu_graph(M, N, K), config=cfg, cta_group=cta_group) + assert compiled.aux_names == ["bias0"] + + a, b, c = _mkdata(1, M, N, K, "bf16", "bf16") + torch.manual_seed(1) + bias = torch.empty(1, N, dtype=torch.int32).random_(-2, 2).to(dtype=torch.bfloat16, device="cuda") + compiled(_vp(compiled, a, b, c, bias)) + torch.cuda.synchronize() + + ref = torch.einsum("bmk,bnk->bmn", a.to(torch.float32), b.to(torch.float32)) + ref = torch.relu(ref + bias.to(torch.float32)).to(torch.bfloat16) + _assert_bit_tight(c, ref, f"bias+relu, {config_name}, {M}x{N}x{K}") + + +@requires_sm120 +def test_sm120_narrow_output_rows() -> None: + """8B-aligned output rows (N=100 bf16) narrow the chunk to 4 elements and + still store correct predicated pairs.""" + M, N, K = 256, 100, 256 + cfg, cta_group = _resolve("CONFIG_sm120_128x128x128_128x128x32_cluster1x1_1ctamma") + compiled = _plan(_build_graph(M, N, K), config=cfg, cta_group=cta_group) + a, b, c = _mkdata(1, M, N, K, "bf16", "bf16") + compiled(_vp(compiled, a, b, c)) + torch.cuda.synchronize() + _assert_bit_tight(c, _reference(a, b, "bf16"), f"narrow output rows, {M}x{N}x{K}") + + +@requires_sm120 +def test_sm120_out_of_scope_jit_rejects() -> None: + """An out-of-scope chain raises NotImplementedError from the jit gates — + never an AssertionError from the template's render asserts.""" + from cudnn.gemm.frost.compiler import jit_from_cudnn_graph + + cfg = by_name("CONFIG_sm120_128x128x128_128x128x32_cluster1x1") + with pytest.raises(NotImplementedError, match="K-major"): + jit_from_cudnn_graph(_build_graph(256, 256, 128, b_major="n"), config=cfg, cta_group=1) + with pytest.raises(NotImplementedError, match="N-major"): + jit_from_cudnn_graph(_build_graph(256, 256, 128, out_major="m"), config=cfg, cta_group=1) + + +@requires_sm120 +def test_sm120_probe_and_auto_select() -> None: + """The engine path: probe accepts with its DEFAULT arguments (the sm100 + default geometry is re-targeted to the family the auto path builds), and + build_gemm_plan compiles + runs an sm120 kernel without a caller config.""" + from cudnn.gemm.frost.compiler import probe_supported + from cudnn.gemm.frost.graph_analyzer import build_gemm_plan + from cudnn.gemm.frost.kernel_registry import preferred_strategy + from cudnn.gemm.frost.tile_config import DEFAULT_CONFIG + + M, N, K = 384, 768, 384 + g = _build_graph(M, N, K) + cfg, grp = preferred_strategy(analyze(g), DEFAULT_CONFIG, 2) + assert cfg.pipeline == "sm120" and grp == 1 # warp MMA: cta_group clamps to 1 + probe_supported(g) # must not raise on an sm_120 GPU + + compiled = build_gemm_plan(g) + assert compiled.config.pipeline == "sm120" + a, b, c = _mkdata(1, M, N, K, "bf16", "bf16") + bd = compiled.binding + compiled({bd.a_operands[0]: a, bd.b_operands[0]: b, bd.outputs[0]: c}) + torch.cuda.synchronize() + _assert_bit_tight(c, _reference(a, b, "bf16"), f"auto-select ({compiled.config.name}), {M}x{N}x{K}") + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"] + sys.argv[1:])) From 95a7808a6a44b19324108896a4d683e0c8255244 Mon Sep 17 00:00:00 2001 From: Yihua Wei Date: Tue, 25 Aug 2026 14:56:47 -0700 Subject: [PATCH 2/4] resolve conflicts for sm120 matmul --- python/cudnn/gemm/frost/compiler.py | 14 +++-- .../frost/kernel_templates/sm120_matmul.py | 53 ++++++++++--------- test/python/gemm/frost/test_sm120_matmul.py | 9 ++-- 3 files changed, 44 insertions(+), 32 deletions(-) diff --git a/python/cudnn/gemm/frost/compiler.py b/python/cudnn/gemm/frost/compiler.py index 1957280d0..81229da1a 100644 --- a/python/cudnn/gemm/frost/compiler.py +++ b/python/cudnn/gemm/frost/compiler.py @@ -2999,6 +2999,10 @@ def _smem_d_bytes(cfg, chain, cta_group: int) -> int: def _output_store_mode(out, chain, cfg, cta_group: int) -> str: """Store mode for ONE output, from the output itself plus the geometry it is stored with -- never from what the others need.""" + # Families whose templates render the TMA-store arm at all; the sm120 warp + # kernel always takes its transposed-STG path. + if cfg.pipeline not in _TMA_STORE_EPI_PIPELINES: + return "stg" # TMA addresses its contiguous dim in 16-byte units, and truncates. The next # two rejections are that granule at a different extent. epi_n = _epi_n(cfg, cta_group, chain.output_dtype) @@ -3188,8 +3192,8 @@ def _check_executable(chain: FusionChain) -> None: def plan_config(chain: FusionChain) -> "tuple[TileConfig, int]": - from .kernel_registry import preferred_pipeline - from .tile_config import as_pipeline, select_config + from .kernel_registry import preferred_strategy + from .tile_config import select_config tile_m = chain.matmul.M if chain.moe is not None: @@ -3203,14 +3207,16 @@ def plan_config(chain: FusionChain) -> "tuple[TileConfig, int]": b_n_major=chain.matmul.b_major == "n", b_elem_bytes=DTYPE_BYTES[chain.matmul.b_dtype], ) - return as_pipeline(config, preferred_pipeline(chain)), cta_group + # Re-target at the preferred family AND clamp cta_group to a group that + # family has a template for (sm120 is warp-scoped MMA, 1-CTA only). + return preferred_strategy(chain, config, cta_group) def _precheck_plain(chain: FusionChain, config: TileConfig, cta_group: int) -> None: from .kernel_registry import select_template _check_supported(chain, config) - _arch_reason = select_template(chain, config, cta_group).active_reject(config) + _arch_reason = select_template(chain, config, cta_group).active_reject(config, chain) if _arch_reason is not None: raise NotImplementedError(_arch_reason) _check_dtype_config_compat(chain, config, cta_group) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm120_matmul.py b/python/cudnn/gemm/frost/kernel_templates/sm120_matmul.py index d5b8df84e..fc3521f25 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm120_matmul.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm120_matmul.py @@ -74,10 +74,13 @@ _EPI_N = epi_tile_mn[1] -# SMEM K-row swizzle: the TMA s{128,64,32}b pattern == cutlass.Swizzle(b, 4, 3) -# with b = log2(row_bytes / 16). ldmatrix addresses below apply the same XOR +# SMEM K-row swizzle: the K-row width IS the swizzle span (the renderer derives +# ab_tma_swizzle from cta_tile_k_bytes; cross-checked against it below). The TMA +# s{128,64,32}b pattern == cutlass.Swizzle(b, 4, 3) with b = log2(row_bytes / 16). +# ldmatrix addresses below apply the same XOR # (fort: swizzled_bank_id = bank ^ ((bank / 8) % SWIZZLE_SCALE)). -_AB_SW_BBITS = (ab_smem_swizzle_bytes // 16).bit_length() - 1 +_AB_SMEM_SWIZZLE_BYTES = _CTA_K_ELEMS * _ELEM_BYTES +_AB_SW_BBITS = (_AB_SMEM_SWIZZLE_BYTES // 16).bit_length() - 1 _AB_SWIZZLE = cutlass.Swizzle(_AB_SW_BBITS, 4, 3) # Epilogue staging tile swizzle — matches the s64b TMA-store descriptor. _EPI_SWIZZLE = cutlass.Swizzle(2, 4, 3) @@ -91,30 +94,32 @@ _STG_EPI_NGRP = (_N_FRAGS + _STG_EPI_GROUP_FRAGS - 1) // _STG_EPI_GROUP_FRAGS _STG_V = (vec_bytes_epi * 8) // cd_dtype.width -if not use_tma_store_epi: - _STG_EPI_BYTES = 4 * _STG_EPI_WARP_ELEMS * NUM_COMPUTE_WARPS - _AB_STAGE_BYTES = (cta_tile_mnk[0] + cta_tile_mnk[1]) * _CTA_K_ELEMS * _ELEM_BYTES + 16 - ab_stages = ab_stages - -(-_STG_EPI_BYTES // _AB_STAGE_BYTES) - assert ab_stages >= 1, "transposed STG epilogue: staging stream cannot be funded from the AB pipeline" +# @@STG_ONLY:BEGIN@@ +_STG_EPI_BYTES = 4 * _STG_EPI_WARP_ELEMS * NUM_COMPUTE_WARPS +_AB_STAGE_BYTES = (cta_tile_mnk[0] + cta_tile_mnk[1]) * _CTA_K_ELEMS * _ELEM_BYTES + 16 +ab_stages = ab_stages - -(-_STG_EPI_BYTES // _AB_STAGE_BYTES) +assert ab_stages >= 1, "transposed STG epilogue: staging stream cannot be funded from the AB pipeline" +# @@STG_ONLY:END@@ # ---- v1 scope guards (fail at render/import, not at runtime) --------------- assert cluster_shape_mnk == (1, 1, 1), "sm120 has no thread-block clusters (CC 12.0): cluster_shape must be (1,1,1)" assert threads_per_cta == NUM_WARPS * 32, f"sm120 template is a fixed 12-warp kernel (384 threads), got {threads_per_cta}" assert not multicast_a and not multicast_b, "sm120 has no TMA multicast (no clusters)" assert not a_is_m_major and not b_is_n_major, "sm120 template v1 supports K-major A and K-major B only (TN GEMM)" -assert not cd_out_is_m_major, "sm120 template v1 supports N-major output only" +# N-major, non-fp4 output only in v1 — enforced upstream by Sm120KernelTemplate._extra_reject. assert num_gemms == 1 and num_a_operands == 1 and num_b_operands == 1, "sm120 template v1 is single-GEMM only" -assert cd_fake_n_div == 1, "sm120 template v1 does not support fp4 output" assert cta_tile_mnk[0] % WARPS_M == 0 and _WARP_TILE_M % 16 == 0, f"cta_tile_m={cta_tile_mnk[0]} must be a multiple of {WARPS_M * 16}" assert cta_tile_mnk[1] % WARPS_N == 0 and _WARP_TILE_N % 8 == 0, f"cta_tile_n={cta_tile_mnk[1]} must be a multiple of {WARPS_N * 8}" -assert ab_smem_swizzle_bytes == _CTA_K_ELEMS * _ELEM_BYTES, "SMEM K-row width must equal the swizzle span" +_TMA_SWIZZLE_BY_BYTES = {32: _tma.TensorMapSwizzle.s32b, 64: _tma.TensorMapSwizzle.s64b, 128: _tma.TensorMapSwizzle.s128b} +assert ab_tma_swizzle == _TMA_SWIZZLE_BY_BYTES.get(_AB_SMEM_SWIZZLE_BYTES), "SMEM K-row width must equal the swizzle span" assert _NUM_K_BLOCKS * _K_BLK_ELEMS == _CTA_K_ELEMS, "cta_tile_k must be a multiple of 32 bytes" -if use_tma_store_epi: - assert vec_bytes_epi * 8 == 2 * cd_dtype.width, "sm120 TMA-store epilogue drains one (n, n+1) accumulator pair per thread" -else: - assert _STG_V >= 2 and _STG_V % 2 == 0 and 8 % _STG_V == 0, "transposed STG epilogue: the store vector must be whole pairs tiling the 8-column row run" -if use_tma_store_epi: - assert cta_tile_mnk[1] % _EPI_N == 0, "TMA-store epilogue needs cta_tile_n to be a whole number of epi subtiles" +# @@TMA_STORE_ONLY:BEGIN@@ +assert vec_bytes_epi * 8 == 2 * cd_dtype.width, "sm120 TMA-store epilogue drains one (n, n+1) accumulator pair per thread" +assert cta_tile_mnk[1] % _EPI_N == 0, "TMA-store epilogue needs cta_tile_n to be a whole number of epi subtiles" +# @@TMA_STORE_ONLY:END@@ +# @@STG_ONLY:BEGIN@@ +assert _STG_V >= 2 and _STG_V % 2 == 0 and 8 % _STG_V == 0, "transposed STG epilogue: the store vector must be whole pairs tiling the 8-column row run" +# @@STG_ONLY:END@@ # --------------------------------------------------------------------------- # The warp MMA instruction, resolved from the injected MMA dtypes. @@ -806,17 +811,17 @@ def _host( # @@TMA_STORE_ONLY:BEGIN@@ # @@INJECT_HOST_TMA_C_LISTS@@ c = _tma_c_outputs[0] - # N-major output only (asserted at render time). + # N-major output only (enforced by Sm120KernelTemplate._extra_reject). tma_c_desc = _tma.create_tensor_map_tiled( global_address=c.iterator.toint(), - dtype=cd_tma_dtype, + dtype=cd_dtype, global_dims=[n, m, batch], global_strides=[ out_stride_m_0 * cd_dtype.width // 128, out_stride_l_0 * cd_dtype.width // 128, ], box_dims=[epi_tile_mn[1], cta_tile_mnk[0], 1], - swizzle=(_tma.TensorMapSwizzle.s64b if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), + swizzle=_tma.TensorMapSwizzle.s64b, ) tma_c_desc_list = [tma_c_desc] # @@TMA_STORE_ONLY:END@@ @@ -886,11 +891,11 @@ def _make_fake_b(): ) # @@TMA_STORE_ONLY:BEGIN@@ - def _make_fake_c(): + def _make_fake_c(_dt, _div, _mm): return make_fake_compact_tensor( - cd_dtype, - (sym_m, sym_n // cd_fake_n_div, sym_l), - stride_order=(0, 1, 2) if cd_out_is_m_major else (1, 0, 2), + _dt, + (sym_m, sym_n // _div, sym_l), + stride_order=(0, 1, 2) if _mm else (1, 0, 2), assumed_align=16, ) diff --git a/test/python/gemm/frost/test_sm120_matmul.py b/test/python/gemm/frost/test_sm120_matmul.py index 862105eeb..f55531393 100644 --- a/test/python/gemm/frost/test_sm120_matmul.py +++ b/test/python/gemm/frost/test_sm120_matmul.py @@ -468,12 +468,12 @@ def test_sm120_epi_vec_clamp_and_stg_only() -> None: c = by_name(name) v = _epi_vec_bytes(chain, c, 1) assert v == 16, name # 8-element bf16 run (template: 8 % _STG_V == 0) - assert _use_tma_store_epi(chain, c, v, 1) is False, name + assert _use_tma_store_epi(chain, c, 1) is False, name # the sm100 chunk derivation and its TMA-store gate are untouched cfg100 = by_name("CONFIG_sm100_128x128x128_128x128x32_cluster1x1") v100 = _epi_vec_bytes(chain, cfg100, 1) - assert v100 >= 16 and _use_tma_store_epi(chain, cfg100, v100, 1) is True + assert v100 >= 16 and _use_tma_store_epi(chain, cfg100, 1) is True # 8B-aligned output rows (N=100 bf16) narrow the chunk to 4 elements narrow = analyze(_build_graph(256, 100, 256)) @@ -496,14 +496,15 @@ def test_sm120_render_smoke() -> None: ): cfg = by_name(name) vec = _epi_vec_bytes(chain, cfg, 1) - snippets = generate(chain, vec_bytes_epi=vec, output_elem_bytes=2, use_tma_store=False) + snippets = generate(chain, vec_bytes_epi=vec, output_elem_bytes=2) # empty tma_slots == STG everywhere src = _render_template(chain, snippets, cfg, 1) assert "@@" not in src, "leftover injection markers" ast.parse(src) assert "cudnn_frost_sm120_matmul_" in src assert "threads_per_cta = 384" in src assert f"vec_bytes_epi = {vec}" in src - assert "use_tma_store_epi = False" in src + assert "_STG_EPI_BYTES" in src, "transposed-STG arm must be rendered" + assert "tma_c_desc" not in src, "TMA-store arm must be stripped" # ============================================================================= From 7cd3805aa813e00d6215e7347a6717331e485088 Mon Sep 17 00:00:00 2001 From: Yihua Wei Date: Tue, 25 Aug 2026 15:18:30 -0700 Subject: [PATCH 3/4] add benchmark test file for sm120 matmul --- .../gemm/frost/benchmark_matmul_sm120.py | 342 ++++++++++++++++++ 1 file changed, 342 insertions(+) create mode 100644 benchmark/gemm/frost/benchmark_matmul_sm120.py diff --git a/benchmark/gemm/frost/benchmark_matmul_sm120.py b/benchmark/gemm/frost/benchmark_matmul_sm120.py new file mode 100644 index 000000000..8233f262a --- /dev/null +++ b/benchmark/gemm/frost/benchmark_matmul_sm120.py @@ -0,0 +1,342 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Multi-library GEMM comparison for the frost sm120 matmul. + +Benchmarks C = A @ B^T (TN, BF16 in / FP32 accumulate / BF16 out) across a +fixed 27-shape (M, N, K) sweep and prints a per-shape TFLOPS table comparing: + + frost the sm120_matmul_1ctamma template (this repo) + cuBLAS torch.matmul (cuBLASLt) + CUTLASS classic CUTLASS python op (if installed) + TensorRT a single-MatMul engine (if installed) + b12x local-inference-lab/b12x (if importable) + FlashInfer flashinfer (if installed and it exposes a BF16 dense GEMM) + +Libraries that are unavailable on the machine are reported as SKIP with the +reason; the table renders whatever columns actually ran. + + CUDA_VISIBLE_DEVICES=2 python benchmark/gemm/frost/benchmark_matmul_sm120.py + ... --shapes 4096x4096x4096,8192x8192x8192 # subset + ... --libs frost,cublas # subset of libraries + ... --check # bit-exactness vs cuBLAS +""" + +from __future__ import annotations + +import argparse +import os +import sys +import time + +import torch + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from benchmark_utils import time_ms_delayed, time_ms_events # noqa: E402 + +# The 27-shape sweep (M, N, K). +SHAPES: list[tuple[int, int, int]] = [(m, n, k) for m in (4096, 8192, 16384) for n in (4096, 8192, 16384) for k in (4096, 8192, 16384)] + +LIBS = ["frost", "cublas", "cutlass", "tensorrt", "b12x", "flashinfer"] + + +class Skip(Exception): + """Raised by an adapter when its library can't run on this machine.""" + + +# --------------------------------------------------------------------------- +# Adapters. Each setup_(M, N, K, data) returns a zero-arg callable that +# launches one GEMM on the (a, b, c) buffers, or raises Skip(reason). +# Layout contract: a[1,M,K] and b[1,N,K] row-major (both K-major), c[1,M,N]. +# Stream contract: the harness times torch's CURRENT (non-default) stream; a +# library that does not follow it must be handed that stream explicitly, or +# the timer measures an empty stream (torch/flashinfer follow it by default). +# --------------------------------------------------------------------------- + + +def setup_frost(M: int, N: int, K: int, data, args): + import cudnn + import cudnn.gemm.frost # noqa: F401 — installs the pygraph recorder hook + from cudnn.gemm.frost.compiler import jit_from_cudnn_graph + from cudnn.gemm.frost.tile_config import by_name + + major, minor = torch.cuda.get_device_capability() + if not (120 <= major * 10 + minor < 130): + raise Skip(f"needs a consumer-Blackwell GPU (sm120..12x), have sm_{major}{minor}") + + g = cudnn.pygraph( + io_data_type=cudnn.data_type.BFLOAT16, + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + ) + A = g.tensor(name="A", dim=[1, M, K], stride=[M * K, K, 1]) + B = g.tensor(name="B", dim=[1, K, N], stride=[N * K, 1, K]) + C = g.matmul(A=A, B=B, name="mm") + C.set_output(True) + compiled = jit_from_cudnn_graph(g, by_name(args.frost_config), cta_group=1) + a, b, c = data + bd = compiled.binding + pack = {bd.a_operands[0]: a, bd.b_operands[0]: b, bd.outputs[0]: c} + # A direct CompiledFusedGemm call launches on the DEFAULT stream when no + # stream is passed (the cuDNN-handle stream only arrives via engine + # dispatch), so hand it the harness stream. + stream = torch.cuda.current_stream().cuda_stream + return lambda: compiled(pack, stream=stream) + + +def setup_cublas(M: int, N: int, K: int, data, args): + a, b, c = data + bt = b.transpose(-1, -2) + return lambda: torch.matmul(a, bt, out=c) + + +def setup_cutlass(M: int, N: int, K: int, data, args): + # The classic CUTLASS python interface (package `nvidia-cutlass`, module + # cutlass.op). The nvidia-cutlass-dsl wheel installed alongside frost does + # NOT ship a prebuilt dense-GEMM op for sm120 (cutlass.utils.gemm only has + # sm100 tcgen05 helpers), so this probes for the classic API. + try: + import cutlass # noqa: F401 + from cutlass.op import Gemm # type: ignore[attr-defined] + except Exception: + raise Skip("no CUTLASS python GEMM op (nvidia-cutlass-dsl has no sm120 dense-GEMM entry; pip install nvidia-cutlass for cutlass.op.Gemm)") + a, b, c = data + a2, b2, c2 = a[0], b[0].transpose(0, 1), c[0] + plan = Gemm( + A=a2, + B=b2, + C=c2, + D=c2, + alpha=1.0, + beta=0.0, + element_accumulator=torch.float32, + ) + return lambda: plan.run(a2, b2, c2, c2, alpha=1.0, beta=0.0, sync=False) + + +def setup_tensorrt(M: int, N: int, K: int, data, args): + try: + import tensorrt as trt + except Exception: + raise Skip("tensorrt not installed (pip install tensorrt)") + a, b, c = data + logger = trt.Logger(trt.Logger.WARNING) + builder = trt.Builder(logger) + # Strongly typed: BF16 flows from the input dtypes to the output. The + # weakly-typed route (BuilderFlag.BF16 + the ITensor.dtype setter) was + # deprecated in TRT 10 and removed in TRT 11. + network = builder.create_network(int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + a_in = network.add_input("A", trt.DataType.BF16, (M, K)) + b_in = network.add_input("B", trt.DataType.BF16, (N, K)) + mm = network.add_matrix_multiply(a_in, trt.MatrixOperation.NONE, b_in, trt.MatrixOperation.TRANSPOSE) + out = mm.get_output(0) + out.name = "C" + network.mark_output(out) + config = builder.create_builder_config() + blob = builder.build_serialized_network(network, config) + if blob is None: + raise Skip("TensorRT engine build failed") + engine = trt.Runtime(logger).deserialize_cuda_engine(blob) + ctx = engine.create_execution_context() + ctx.set_tensor_address("A", a.data_ptr()) + ctx.set_tensor_address("B", b.data_ptr()) + ctx.set_tensor_address("C", c.data_ptr()) + stream = torch.cuda.current_stream().cuda_stream + # Keep the engine/context alive via the closure. + return lambda _refs=(engine, ctx): ctx.execute_async_v3(stream) + + +def setup_b12x(M: int, N: int, K: int, data, args): + # local-inference-lab / b12x. Not on PyPI: point PYTHONPATH (or --b12x-path) + # at the checkout, then adapt the entry-point probe below to its real API. + if args.b12x_path: + sys.path.insert(0, args.b12x_path) + mod = None + for name in ("b12x", "local_inference_lab.b12x", "local_inference_lab"): + try: + import importlib + + mod = importlib.import_module(name) + break + except ImportError: + continue + if mod is None: + raise Skip("b12x not importable (clone local-inference-lab and pass --b12x-path)") + a, b, c = data + for entry in ("matmul", "gemm", "mm"): + fn = getattr(mod, entry, None) + if callable(fn): + return lambda _fn=fn: _fn(a[0], b[0], out=c[0]) + raise Skip(f"b12x imported ({mod.__name__}) but exposes no known GEMM entry — adapt setup_b12x()") + + +def setup_flashinfer(M: int, N: int, K: int, data, args): + try: + from flashinfer.gemm import mm_bf16 + except Exception: + raise Skip("flashinfer not installed or too old for gemm.mm_bf16 (pip install flashinfer-python)") + a, b, c = data + # mm_bf16 wants A (m, k) row-major and B (k, n) column-major — exactly the + # (N, K)-row-major buffer transposed. + a2, bt, c2 = a[0], b[0].transpose(0, 1), c[0] + # Its default backend='cudnn' can be broken independently of the others + # (e.g. a cuDNN sub-library mismatch), and 'auto' dies with it rather than + # falling back — so probe explicitly, first working backend wins. + backends = [args.flashinfer_backend] if args.flashinfer_backend else ["cublaslt", "cudnn", "cutlass", "tgv", "tinygemm"] + errs = [] + for bk in backends: + run = lambda _bk=bk: mm_bf16(a2, bt, out=c2, backend=_bk) + try: + run() + torch.cuda.synchronize() + except Exception as e: + errs.append(f"{bk}: {str(e).splitlines()[0][:60]}") + continue + if not getattr(setup_flashinfer, "_noted", False): + print(f" [flashinfer: backend '{bk}']", flush=True) + setup_flashinfer._noted = True + return run + raise Skip("no working mm_bf16 backend — " + "; ".join(errs)) + + +SETUP = { + "frost": setup_frost, + "cublas": setup_cublas, + "cutlass": setup_cutlass, + "tensorrt": setup_tensorrt, + "b12x": setup_b12x, + "flashinfer": setup_flashinfer, +} + + +# --------------------------------------------------------------------------- +# Harness +# --------------------------------------------------------------------------- + + +def _mkdata(M: int, N: int, K: int, check: bool): + torch.manual_seed(0) + if check: + # Small integers => exact FP32 reduction => bit-comparable outputs. + a = torch.empty(1, M, K, dtype=torch.int32).random_(-2, 2).to(dtype=torch.bfloat16, device="cuda") + b = torch.empty(1, N, K, dtype=torch.int32).random_(-2, 2).to(dtype=torch.bfloat16, device="cuda") + else: + a = torch.randn(1, M, K, device="cuda").to(torch.bfloat16) + b = torch.randn(1, N, K, device="cuda").to(torch.bfloat16) + c = torch.empty(1, M, N, dtype=torch.bfloat16, device="cuda") + return a, b, c + + +def _iters_for(flops: int) -> int: + # ~3e13 timed FLOP per measurement, clamped to [6, 50] iterations. + return max(6, min(50, int(3e13 // flops))) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--libs", default=",".join(LIBS), help=f"comma list from {LIBS}") + parser.add_argument("--shapes", default="", help="comma list of MxNxK to run (default: the full 27-shape sweep)") + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--iters", type=int, default=0, help="timed iterations (0 = auto-scale by FLOPs)") + parser.add_argument("--timing", choices=("delayed", "events"), default="delayed", help="delayed = kernel-only (host gaps hidden behind a CUDA sleep)") + parser.add_argument("--frost-config", default="CONFIG_sm120_128x128x128_128x128x32_cluster1x1") + parser.add_argument( + "--flashinfer-backend", default="", help="pin the mm_bf16 backend (cublaslt/cudnn/cutlass/tgv/tinygemm/cutile); default: probe in order" + ) + parser.add_argument("--b12x-path", default=os.environ.get("B12X_PATH", ""), help="path to the local-inference-lab checkout for b12x") + parser.add_argument("--check", action="store_true", help="small-int inputs + compare every library bit-wise against cuBLAS") + parser.add_argument("--csv", default="", help="also write results to this CSV file") + args = parser.parse_args() + + if not torch.cuda.is_available(): + print("No CUDA device, exiting.") + return 1 + # TRT's enqueueV3 inserts extra cudaStreamSynchronize calls on the DEFAULT + # stream (it warns about exactly this); a non-default stream keeps the + # delayed timer's back-to-back pipelining honest for every library. + torch.cuda.set_stream(torch.cuda.Stream()) + libs = [x.strip() for x in args.libs.split(",") if x.strip()] + unknown = [x for x in libs if x not in SETUP] + if unknown: + sys.exit(f"unknown libs {unknown}; choose from {LIBS}") + if args.shapes: + shapes = [] + for tok in args.shapes.split(","): + m, n, k = (int(x) for x in tok.lower().split("x")) + shapes.append((m, n, k)) + else: + shapes = SHAPES + + dev = torch.cuda.get_device_name() + cap = torch.cuda.get_device_capability() + timer = time_ms_delayed if args.timing == "delayed" else time_ms_events + print(f"GPU: {dev} (sm_{cap[0]}{cap[1]}) dtype: BF16 in / FP32 accum / BF16 out layout: TN (C = A @ B^T)") + print(f"timing: {args.timing}, warmup={args.warmup}, iters={'auto' if args.iters == 0 else args.iters}") + + skip_reasons: dict[str, str] = {} + results: dict[tuple[int, int, int], dict[str, float]] = {} + t0 = time.time() + + for M, N, K in shapes: + flops = 2 * M * N * K + iters = args.iters or _iters_for(flops) + data = _mkdata(M, N, K, args.check) + row: dict[str, float] = {} + ref = None + if args.check: + a, b, _ = data + ref = torch.matmul(a.to(torch.float32), b.transpose(-1, -2).to(torch.float32)).to(torch.bfloat16) + print(f"\n--- {M}x{N}x{K} ({flops / 1e12:.1f} TFLOP, iters={iters}) ---", flush=True) + for lib in libs: + if lib in skip_reasons: + continue + try: + run = SETUP[lib](M, N, K, data, args) + run() + torch.cuda.synchronize() + if ref is not None: + data[2].zero_() + run() + torch.cuda.synchronize() + bad = (data[2] != ref).sum().item() + if bad: + print(f" {lib:10s} CHECK FAILED: {bad} mismatches vs fp32 reference", flush=True) + ms = timer(lambda i, _r=run: _r(), lambda _r=run: _r(), warmup=args.warmup, iters=iters) + row[lib] = flops / (ms * 1e-3) / 1e12 + print(f" {lib:10s} {row[lib]:8.2f} TFLOPS ({ms:.3f} ms)", flush=True) + except Skip as e: + skip_reasons[lib] = str(e) + print(f" {lib:10s} SKIP: {e}", flush=True) + except Exception as e: + msg = str(e).splitlines()[0][:70] if str(e) else type(e).__name__ + row[lib] = float("nan") + print(f" {lib:10s} ERROR: {msg}", flush=True) + results[(M, N, K)] = row + del data + torch.cuda.empty_cache() + + active = [l for l in libs if l not in skip_reasons] + width = 11 + print("\n" + "=" * (22 + width * len(active))) + print(f" {'M x N x K':20s}" + "".join(f"{l:>{width}s}" for l in active)) + print("=" * (22 + width * len(active))) + for (M, N, K), row in results.items(): + cells = "".join(f"{row.get(l, float('nan')):>{width}.2f}" for l in active) + print(f" {f'{M}x{N}x{K}':20s}" + cells) + print("=" * (22 + width * len(active))) + print(" (TFLOPS; higher is better)") + for lib, why in skip_reasons.items(): + print(f" SKIP {lib}: {why}") + + if args.csv: + with open(args.csv, "w") as f: + f.write("M,N,K," + ",".join(active) + "\n") + for (M, N, K), row in results.items(): + f.write(f"{M},{N},{K}," + ",".join(f"{row.get(l, float('nan')):.2f}" for l in active) + "\n") + print(f" CSV written to {args.csv}") + print(f"total: {time.time() - t0:.1f} s") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From 92c98adffad942aa52fba999cc25117ab28601c1 Mon Sep 17 00:00:00 2001 From: Yihua Wei Date: Tue, 25 Aug 2026 16:28:25 -0700 Subject: [PATCH 4/4] resolve issues about nvvm.elect_sync(), nvvm.griddepcontrol(wait), and drain of TMA warp for SM120 matmul --- .../frost/kernel_templates/sm120_matmul.py | 79 ++++++++++++------- 1 file changed, 50 insertions(+), 29 deletions(-) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm120_matmul.py b/python/cudnn/gemm/frost/kernel_templates/sm120_matmul.py index fc3521f25..be50941d3 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm120_matmul.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm120_matmul.py @@ -24,7 +24,7 @@ # @@INJECT_TILE_CONSTANTS@@ -CLC_SCHED_STAGES = 2 +CLC_SCHED_STAGES = 1 # Programmatic Dependent Launch (PDL, sm_90+; supported on sm_120). USE_PDL = True @@ -219,6 +219,7 @@ def _kernel( warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) + elect_one = nvvm.elect_sync() tidx = cute.arch.thread_idx()[0] bidx = cute.arch.block_idx()[0] @@ -310,12 +311,15 @@ def _kernel( # clc full: tx-count armed by the scheduler; completed by the response. # clc empty: one elected arrive per consumer warp per slot. if warp_idx == 0: - if nvvm.elect_sync(): - for i in range(ab_stages): + for i in range(ab_stages): + if elect_one: nvvm.mbarrier_init(ab_full_mbar_ptr.subview(i), 1) + if elect_one: nvvm.mbarrier_init(ab_empty_mbar_ptr.subview(i), NUM_COMPUTE_WARPS) - for i in range(CLC_SCHED_STAGES): + for i in range(CLC_SCHED_STAGES): + if elect_one: nvvm.mbarrier_init(clc_full_mbar_ptr.subview(i), 1) + if elect_one: nvvm.mbarrier_init(clc_empty_mbar_ptr.subview(i), NUM_CLC_CONSUMER_WARPS) nvvm.fence_mbarrier_init() nvvm.barrier_cta_sync(0) @@ -352,8 +356,9 @@ def _kernel( while not nvvm.mbarrier_try_wait_parity(clc_empty_mbar_ptr.subview(stage), clc_empty_phase, time_limit=10_000_000): pass - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(clc_full_mbar_ptr.subview(stage), 16) + if elect_one: cute_clc.issue_clc_query( clc_full_mbar_cute_base + stage, clc_response_ptr_base + stage, @@ -368,7 +373,7 @@ def _kernel( is_valid_sched = vld nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(clc_empty_mbar_ptr.subview(stage)) sched_iter += 1 @@ -377,8 +382,7 @@ def _kernel( if warp_idx == TMA_WARP_ID: nvvm.setmaxregister(PROD_REG_COUNT, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): - nvvm.griddepcontrol("wait") + nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) tile_m = init_tile_m @@ -408,19 +412,23 @@ def _kernel( pass coord_k = k_tile_idx * cta_tile_mnk[2] - if nvvm.elect_sync(): + # One elected lane only: the barrier's arrival count is 1, and + # the TMA copies deliver exactly num_tma_copy_bytes once. + if elect_one: nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes) - # K-major A: TMA box [K_tile, cta_m] at (k, m, l); OOB rows/cols - # are hardware zero-filled (K tails contribute 0 to the MMA). - for _ai in cutlass.range_constexpr(num_a_operands): + # K-major A: TMA box [K_tile, cta_m] at (k, m, l); OOB rows/cols + # are hardware zero-filled (K tails contribute 0 to the MMA). + for _ai in cutlass.range_constexpr(num_a_operands): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cta_global( smem_a_list[_ai].subview(sA_elems * stage), tma_a_descs[_ai].get_ptr(), (coord_k, coord_m, tile_l_a), ab_full_mbar_ptr.subview(stage), ) - # K-major B: TMA box [K_tile, cta_n] at (k, n, l). - for _bj in cutlass.range_constexpr(num_b_operands): + # K-major B: TMA box [K_tile, cta_n] at (k, n, l). + for _bj in cutlass.range_constexpr(num_b_operands): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cta_global( smem_b_list[_bj].subview(sB_elems * stage), tma_b_descs[_bj].get_ptr(), @@ -444,24 +452,37 @@ def _kernel( tile_m, tile_n = _l2_swizzle_tile(m_idx, n_idx, gridx, gridy, swizzle_w) tile_l = l_idx nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(clc_empty_mbar_ptr.subview(consumer_stage)) tile_iter += 1 - # Drain: wait until the compute warps have consumed the final stage so - # the producer never exits with a stage it would have re-armed pending. + # # Drain: wait until the compute warps have consumed the final stage so + # # the producer never exits with a stage it would have re-armed pending. + # tail_stage = ab_iter % ab_stages + # tail_phase = ab_empty_phase_bit + # if tail_stage == 0 and ab_iter != 0: + # tail_phase = tail_phase ^ 1 + # for _ in range(ab_stages - 1): + # tail_stage = tail_stage + 1 + # if tail_stage == ab_stages: + # tail_stage = cutlass.Int32(0) + # tail_phase = tail_phase ^ 1 + # if elect_one: + # while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + # pass + tail_stage = ab_iter % ab_stages tail_phase = ab_empty_phase_bit if tail_stage == 0 and ab_iter != 0: tail_phase = tail_phase ^ 1 - for _ in range(ab_stages - 1): - tail_stage = tail_stage + 1 - if tail_stage == ab_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - if nvvm.elect_sync(): - while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass + if cutlass.const_expr(cluster_shape_mnk[0] * cluster_shape_mnk[1] > 1): + for _ in range(ab_stages): + while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + tail_stage = tail_stage + 1 + if tail_stage == ab_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 # -- Compute warps: mma.sync mainloop + epilogue -------------------------- if warp_idx < NUM_COMPUTE_WARPS: @@ -574,7 +595,7 @@ def _kernel( # Stage fully consumed by this warp (ldmatrix is synchronous). nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(ab_empty_mbar_ptr.subview(stage)) ab_iter += 1 @@ -628,7 +649,7 @@ def _kernel( ) if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_global_shared_cta( tma_c_desc.get_ptr(), smem_subtile_ptr, @@ -704,7 +725,7 @@ def _kernel( tile_m, tile_n = _l2_swizzle_tile(m_idx, n_idx, gridx, gridy, swizzle_w) tile_l = l_idx nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(clc_empty_mbar_ptr.subview(consumer_stage)) tile_iter += 1 @@ -713,7 +734,7 @@ def _kernel( # (fort fires launch_dependent_grids at the same point). if cutlass.const_expr(USE_PDL): if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("launch_dependents") # @@TMA_STORE_ONLY:BEGIN@@