diff --git a/helion/_compiler/autotuner_heuristics/__init__.py b/helion/_compiler/autotuner_heuristics/__init__.py index 745e5fd48..0b45cd790 100644 --- a/helion/_compiler/autotuner_heuristics/__init__.py +++ b/helion/_compiler/autotuner_heuristics/__init__.py @@ -14,6 +14,8 @@ from .cute import CuteTileVecHeuristic from .cute import CuteTileVecWarpPerRowHeuristic from .cute import CuteTileVecWarpReduceHeuristic +from .cute_matmul_formula import CuteTcgen05FormulaFfiAltHeuristic +from .cute_matmul_formula import CuteTcgen05FormulaMatmulHeuristic from .pallas import PallasMatmulF32NoTilingSeedHeuristic from .pallas import PallasMatmulNoTilingSeedHeuristic from .triton import TritonB200MatmulHeuristic @@ -40,6 +42,12 @@ CuteFlashAttentionCausalLptHeuristic, CuteTcgen05ClusterM2FfiHeuristic, CuteTcgen05ClusterM2Heuristic, + # The formula heuristic subsumes the 3 cluster_m=2 producers above and is + # registered AFTER them so its promote_seed_to_default wins (last-promote-wins). + # The FFI alt-seed is a second ranked (non-promoting) config, benchmarked beside + # the promoted Bucket-A default for 16-bit compute. + CuteTcgen05FormulaMatmulHeuristic, + CuteTcgen05FormulaFfiAltHeuristic, CuteReductionTileHeuristic, CuteReductionWideChunkHeuristic, CuteTileVecHeuristic, diff --git a/helion/_compiler/autotuner_heuristics/cute_matmul_formula.py b/helion/_compiler/autotuner_heuristics/cute_matmul_formula.py new file mode 100644 index 000000000..09c8eb574 --- /dev/null +++ b/helion/_compiler/autotuner_heuristics/cute_matmul_formula.py @@ -0,0 +1,757 @@ +"""Formulaic tcgen05 matmul autotuner-seeding heuristic — the CuTe analog of #3007. + +The three shipped CuTe matmul seed producers +(``CuteTcgen05ClusterM2Heuristic``, ``CuteTcgen05ClusterM2FfiHeuristic``, +``CuteFp8GemmSkinnyMHeuristic``) are special-case, ``cluster_m=2``-only templates +that (a) structurally cannot even *propose* whole regimes (M=64 decode / cluster_m=1, +medium-M single-wave rectangular) and (b) couple their ``is_eligible`` to +``enforce_dot_requirements``' *search-restriction* gate — so a shape the wave-quant +gate suppressed from cm2 *search* also got NO seed (the "A1 gap"). + +``CuteTcgen05FormulaMatmulHeuristic`` replaces that with ONE shape-aware formula: +``f(M, K, N, dtype, epilogue, num_sm, smem_budget) -> Config``. Its genuine surface is +the ~5 knobs the codegen defaults get wrong for a given regime (see +``cute-matmul-heuristic-plan.md`` §4.3.5): + + 1. ``tcgen05_cluster_m`` — the regime selector (1 decode / 2 compute+medium-M) + 2. ``block_sizes`` — [bm, bn, bk]: collective tile, wave-fill-shrunk + 3. ``tcgen05_ab_stages`` — depth-fill to the ~196 KB AB-SMEM isobar (dtype-capped) + 4. ``l2_groupings`` — wave-count-aware ([1] many-wave / [4] single-wave) + 5. ``pid_type`` / persistence — persistent_blocked (decode) / _interleaved (compute) + +Everything else is INHERITED from the ``bn``/``num_stages``-keyed codegen defaults +(``acc_stages``, ``c_stages``, ``num_epi_warps``, role/warp-spec) — emitted for +completeness because a seed is a full Config, but written to the default value. + +Two design invariants proven by the run-3 hill-climbs + the design-validation smoke: + + * **The depth-fill fills to a fixed ~196 608-byte AB-SMEM isobar.** ``bk`` and ``ab`` + trade off along it (``TCGEN05_DIRECT_ENTRY_STAGE_TUPLES_BY_BK`` encodes exactly + this). So ``_pick_bk_ab`` picks the ``(bk, ab)`` that MAXIMIZES AB-SMEM bytes used + within the per-CTA budget and the dtype cap (tie -> deepest ab). This reproduces + the pretuned decode bn=32/bk=256/ab=8 (196 608) over bk=128/ab=12 (147 456) AND the + fp8 compute bk=128/ab=6 AND the fp8 decode bn=64/bk=128/ab=12. + * **Wave/occupancy counts must be in CTAs, not output tiles** — ``cluster_m=2`` spends + 2 CTAs per output tile, so ``_wave_eff`` / ``l2_groupings`` multiply by ``cluster_m``. + +The seed is ORTHOGONAL to ``enforce_dot_requirements``' search restrictions +(``cute-seed-orthogonal-to-search`` memory): ``is_eligible`` reads the ``MatmulFact`` +directly and NEVER asks the search-restriction gate for permission. The formula emits +the best config that passes genuine (SMEM-budget / physical) validation; gate-3 +"artificial-but-errors" caps (bf16 ab>3) are handled by the bundled prerequisite edits. + +Bucket-A fills 16-bit ``ab`` to the dtype cap (6) on the DEFAULT path — the bundled +bf16-deep-AB prerequisite made deep 16-bit AB (e.g. compute ``bk=64/ab=6``) run on the plain +path, so it is NOT restricted to the FFI topology (measured +1–18% over the old ab≤3 cap). +The FFI ``explicit_epi_tile`` config still ships as a SECOND ranked seed +(``CuteTcgen05FormulaFfiAltHeuristic``, Bucket B); the autotuner keeps whichever wins. +""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING +from typing import Any + +import torch + +from ...runtime.config import Config +from ..cute.strategies import TCGEN05_L2_SWIZZLE_SIZE_CONFIG_KEY +from ..cute.strategies import TCGEN05_LAYOUT_OVERRIDES_D_STORE_BOX_N_KEY +from ..cute.strategies import TCGEN05_LAYOUT_OVERRIDES_EPI_TILE_M_KEY +from ..cute.strategies import TCGEN05_LAYOUT_OVERRIDES_EPI_TILE_N_KEY +from ..cute.strategies import TCGEN05_LAYOUT_STRATEGY_CONFIG_KEY +from ..cute.strategies import TCGEN05_PERSISTENCE_MODEL_CONFIG_KEY +from ..cute.strategies import Tcgen05LayoutStrategy +from ..cute.strategies import Tcgen05PersistenceModel +from ..cute.tcgen05_config import CuteTcgen05Config +from ..cute.tcgen05_constants import TCGEN05_FLAT_ROLE_COORDINATES_CONFIG_KEY +from ..cute.tcgen05_constants import TCGEN05_TVM_FFI_LAUNCH_CONFIG_KEY +from ..cute.tcgen05_constants import TCGEN05_TWO_CTA_BLOCK_M +from ..cute.tcgen05_constants import TCGEN05_TWO_CTA_BLOCK_N +from ..cute.tcgen05_constants import TCGEN05_TWO_CTA_SEED_L2_GROUPING +from ..cute.tcgen05_constants import tcgen05_ab_smem_bytes_per_cta +from .registry import AutotunerHeuristic + +if TYPE_CHECKING: + from ...autotuner.config_spec import ConfigSpec + from ...autotuner.config_spec import MatmulFact + from ..compile_environment import CompileEnvironment + from ..device_ir import DeviceIR + +# --- device / geometry constants (mirrors tcgen05_constants.py; sm100 / B200) --- +_STATIC_PERSISTENT = Tcgen05PersistenceModel.STATIC_PERSISTENT.value +# fp8=1 -> ab cap 12; 16-bit (bf16/fp16)=2 -> cap 6 (Bucket-A clamps to 3, see below). +_DTYPE_AB_CAP = {1: 12, 2: 6, 4: 3} +# Candidate block-K per dtype, largest first (the isobar fill tries all, keeps the best). +# 16-bit includes 256 because the narrow decode tile (bm=64/bn=32) fills the isobar at +# bk=256/ab=4 (196 608 = the R3 #8 bf16-decode key); for the 256² compute tile bk=256 +# only reaches ab=1 so bk=128 still wins there — adding 256 is safe. +_DTYPE_BK_CHOICES = {1: (256, 128), 2: (256, 128, 64), 4: (64,)} +_WAVE_FULL = 0.8 # a tile "fills a wave" above this CTA occupancy (compute/medm classifier + FFI) +# Decode (memory-bound) accepts a WIDER bn at LOWER occupancy than the compute classifier does: +# the loop returns the widest bn clearing this bar, and a bandwidth-bound decode prefers the wider +# (fatter/more-contiguous TMA B-read) tile down to ~half-device fill. Measured: dropping the decode +# bar 0.8->0.5 improves 64x3584x3584 (bn64@0.38w 100 -> bn32@0.76w 115, +14%) and 64x1536x1536 +# (bn64@0.16w 25 -> bn32@0.32w 30, +20%), and changes NO other curriculum shape (0/40 regime flips) +# — but only because it is DECODE-LOCAL: sharing 0.8 with the compute/medm classifier would risk +# re-routing a medium-M shape into the 256² compute tile (the regression the Phase-2 medm fix removed). +_DECODE_WAVE_FULL = 0.5 +_MANY_WAVE = 4 # waves >= this -> l2_groupings=[1] +_DECODE_M_MAX = 128 # M <= this is the cluster_m=1 decode regime (ONE_CTA bm cap) +# M < this is below the tcgen05 decode admission floor: matmul_ops.py +# enforce_dot_requirements gates the whole tcgen05 block on static_m>=64, so deep-AB is +# never admitted below it and the decode ab-lever is unavailable. +_DECODE_M_MIN = 64 +# M <= this is the skinny-M SIMT regime (owned by CuteFp8GemmSkinnyMHeuristic). +_SKINNY_M_MAX = 16 +_DECODE_BM = 64 # decode M-tile (bm=64 in every decode answer key / pretuned row) +_DECODE_BN_CHOICES = (128, 64, 32, 16) # decode N-tile menu (from the pretuned table) +_FP8_SMALL_GRID = 128 # fp8 small-grid cluster_m=2 tile (bm=bn=128, per-CTA 64x128) + +_FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2) +_SIXTEEN_BIT_DTYPES = (torch.bfloat16, torch.float16) + + +def _itemsize(dtype: torch.dtype) -> int: + if dtype in _FP8_DTYPES: + return 1 + if dtype in _SIXTEEN_BIT_DTYPES: + return 2 + return dtype.itemsize + + +def _n_ctas(bm: int, bn: int, m: int, n: int, cluster_m: int) -> int: + """CTA count = output tiles x cluster_m (cluster_m=2 spends 2 CTAs / tile).""" + tiles = max(1, math.ceil(m / bm)) * max(1, math.ceil(n / bn)) + return tiles * cluster_m + + +# ``num_sm <= 0`` is the "SM count UNKNOWN" sentinel from ``_num_sm`` (get_num_sm itself is +# always >= 1; the wrapper returns 0 only when it raises for a non-CUDA/unimplemented device). +# In the current formula this is unreachable — ``_formula_eligible`` gates on +# ``per_cta_ab_smem_budget_bytes(device) > 0``, which is 0 for non-CUDA, so these helpers only run +# once the device is CUDA with a real count. The guard is kept as a defensive floor (these are pure +# helpers a future caller could invoke off the eligibility-gated path) that also avoids a +# divide-by-zero; unknown SM count => "assume the tile fills a wave" (don't shrink for occupancy). +def _wave_eff(bm: int, bn: int, m: int, n: int, cluster_m: int, num_sm: int) -> float: + """Occupancy in CTAs (NOT tiles). Getting this wrong (counting tiles) picked + bn=64 instead of the correct bn=128 for fp8 medium-M.""" + if num_sm <= 0: + return 1.0 + ctas = _n_ctas(bm, bn, m, n, cluster_m) + waves = max(1, math.ceil(ctas / num_sm)) + return ctas / (waves * num_sm) + + +def _num_waves(bm: int, bn: int, m: int, n: int, cluster_m: int, num_sm: int) -> int: + if num_sm <= 0: # unknown SM count sentinel — see note above _wave_eff + return 1 + return max(1, math.ceil(_n_ctas(bm, bn, m, n, cluster_m) / num_sm)) + + +def _pick_bk_ab( + itemsize: int, cluster_m: int, bm: int, bn: int, budget: int, k: int +) -> tuple[int, int]: + """Fill the AB pipeline to the ~196 KB SMEM isobar; pick ``(bk, ab)`` on it. + + Every climbed / pretuned answer key lands at ~196 608 AB-SMEM bytes (verified vs + ``tcgen05_ab_smem_bytes_per_cta`` + ``TCGEN05_DIRECT_ENTRY_STAGE_TUPLES_BY_BK``). + ``bk`` and ``ab`` trade off along that isobar. Rule: for each candidate bk (dtype + menu), the deepest ab that fits is ``budget // per_stage`` capped by the dtype cap; + choose the ``(bk, ab)`` that MAXIMIZES AB-SMEM bytes actually used (fills the isobar), + tie -> deepest ab, tie -> larger bk. Reproduces: + fp8 decode bn=32 : bk=256/ab=8 (196 608) > bk=128/ab=12 (147 456) [pretuned] + fp8 decode bn=64 : bk=128/ab=12 (196 608, hits the fp8 cap) [R1 M1] + fp8 medium-M : bk=128/ab=8 (196 608) [climbed] + fp8 compute 256² : bk=128/ab=6 (196 608) [R2 #3] + bf16 decode bn=32: bk=256/ab=4 (196 608, cluster_m=1) [R3 #8 key] + bf16 compute 256²: bk=64/ab=6 (196 608) [R2 #11; deep AB on the DEFAULT path] + + **16-bit fills the isobar to the dtype cap (6) on the DEFAULT path — no cluster_m + special-case.** The bundled bf16-deep-AB prerequisite admits 16-bit ab>3 within the + SMEM budget on the plain (non-FFI) path, for cluster_m=1 AND cluster_m=2. Measured + (bf16, cluster_m=2, default path, all COMPILE+ACC PASS): the deep pipeline is a WIN, not + just legal — compute 256² bk64/ab6 = 1443 vs the old bk128/ab3 = 1428 TFLOP/s (4096³: + +1%; 2048³: +4%), and rectangular medm [256,128] bk128/ab4 = 382 vs the old ab-3-capped + bk256/ab2 = 323 (+18%). An earlier draft capped 16-bit cm2 at 3 believing deep bk64/ab6 + "needs the FFI explicit_epi_tile topology" — the prerequisite made that stale (deep AB + now runs on the default path), so the cap only left perf on the table. The Bucket-B FFI + alt-seed still ships as a ranked alternative; the autotuner keeps whichever wins. + """ + cap = _DTYPE_AB_CAP.get(itemsize, 3) + best: tuple[int, int, int] | None = None # (bytes_used, ab, bk) + for bk in _DTYPE_BK_CHOICES.get(itemsize, (128, 64)): + # bk must tile K: bk<=K and K a multiple of bk (a partial K-tile is not a valid + # static-full tcgen05 tile). Skips the degenerate small-K case (e.g. K=64/bk=256). + if bk > k or k % bk != 0: + continue + per = tcgen05_ab_smem_bytes_per_cta( + bm=bm, bn=bn, bk=bk, dtype_bytes=itemsize, ab_stages=1, cluster_m=cluster_m + ) + if per <= 0: + continue + ab = min(cap, max(1, budget // per)) + bytes_used = per * ab + cand = (bytes_used, ab, bk) + if best is None or cand > best: + best = cand + if best is None: + # No dtype-menu bk divides K (e.g. K=48). Fall back to the largest power-of-2 + # that tiles K, so the seed still emits a valid tile. + bk = 1 + while bk * 2 <= k and k % (bk * 2) == 0: + bk *= 2 + per = tcgen05_ab_smem_bytes_per_cta( + bm=bm, bn=bn, bk=bk, dtype_bytes=itemsize, ab_stages=1, cluster_m=cluster_m + ) + ab = min(cap, max(1, budget // per)) if per > 0 else 1 + return bk, ab + _, ab, bk = best + return bk, ab + + +def _single_wave_rect_tile(m: int, n: int, num_sm: int) -> tuple[int, int]: + """Medium-M: keep the tall bm=256 M-tile (needs cluster_m=2) but shrink bn until the + CTA count reaches ~one wave. e.g. M512/N4096 -> [256,128] = 128 CTAs = 0.86 waves. + + bn candidates are capped at N (never emit bn>N — wasted padding) so a narrow-N tall shape + routed here from the compute branch's ``n>=256`` guard still gets a sane bn<=N.""" + bm = TCGEN05_TWO_CTA_BLOCK_M + for bn in (256, 128, 64): + if bn <= n and _wave_eff(bm, bn, m, n, 2, num_sm) >= _WAVE_FULL: + return bm, bn + return bm, min(64, n) + + +def _decode_bm(m: int) -> int | None: + """Pick the cluster_m=1 decode M-tile, or None if M is outside the decode envelope. + + The tcgen05 cluster_m=1 static-full-tile runtime path validated on GPU0 only for a + **power-of-2 M** (M=32->bm=32, M=64->bm=64, M=128->bm=64 all PASS; M=96 = 3*32 FAILS + at runtime — "validated single-root static full tiles" — even with bm=32). So the + decode branch requires M to be a power of two; bm = min(M, 64) (also a power of 2).""" + if not _is_pow2(m): + return None + return min(m, _DECODE_BM) + + +def _pick_decode_bn(m: int, bm: int, n: int, num_sm: int) -> int: + """Decode (cluster_m=1): pick the decode N-tile — the widest bn in {128,64,32,16} whose CTA + grid reaches the decode occupancy bar ``_DECODE_WAVE_FULL`` (0.5, NOT the compute 0.8), else + the bn=64 default. + + Decode is MEMORY-bound (M<=64: trivial compute, the kernel just streams B), so the objective + is per-CTA B-read WIDTH balanced against keeping ~half the device busy — NOT compute-style + last-wave occupancy. Measured trade-off (fp8, cluster_m=1, bm=64, cold-L2 cudagraph): + * N=5120: bn=64 (0.54w) = 205 TFLOP/s BEATS bn=32 (1.08w) = 183 and bn=16 (2.16w) = 165 — + once ~half the SMs are busy, WIDER bn wins (fatter/more-contiguous TMA B-reads); extra CTAs + from a narrower bn are just fragmentation on an already bandwidth-bound kernel. + * N=3584: bn=64 gives only 0.38w (56 CTAs — badly under-fills), so bn=32 (0.76w) = 115 BEATS + bn=64 (0.38w) = 100 (+14%). N=1536: bn=64 = 0.16w, bn=32 (0.32w)/bn=16 (0.65w) both +20%. + So the objective is a SINGLE consistent rule: **"the WIDEST bn reaching ~>=half the device + (>=0.5 waves); if NONE does (even the narrowest), the shape is device-starved so grab MAX + occupancy = the NARROWEST bn."** ``_DECODE_WAVE_FULL=0.5`` (vs the old 0.8) captures the + 3584/1536 wins with ZERO change to any other curriculum shape (0/40 regime flips), decode-local + so it can't perturb the compute/medium-M classifier (which stays at 0.8). + + The fallback returns the narrowest (``_DECODE_BN_CHOICES[-1]``, bn=16), NOT bn=64 — measured: + in the fallback regime (even bn=16 < 0.5 waves, i.e. tiny N<=~1168 at M=64) bn=16 ties-or-beats + bn=64 on every probed shape (64x4096x{256,512}: tie; x1024: bn16=43.7 > bn64=37.5, +14%; + 128x4096x512: tie), and bn=128 is worst. The old ``return 64`` was an unvalidated mid-point + inconsistent with the loop's own shrink-for-occupancy premise (raised in review). Still an + approximation — the fully-grounded version keys bn off DRAM read-width / the pretuned (K,N)->bn + table; see run-log.md "decode-bn".""" + for bn in _DECODE_BN_CHOICES: + if _wave_eff(bm, bn, m, n, 1, num_sm) >= _DECODE_WAVE_FULL: + return bn + return _DECODE_BN_CHOICES[-1] + + +def _is_pow2(x: int) -> bool: + return x > 0 and (x & (x - 1)) == 0 + + +def _epilogue_aux_rank(env: CompileEnvironment, device_ir: DeviceIR) -> int: + """0 = pure matmul / transparent (unary act, rank-1 rowvec bias, rowwise scale); + 2 = source-C / rank-2 exact-shape ``[M,N]`` residual (residual_add, bias_residual_gelu). + + Only a rank-2 EXACT-SHAPE ``[M,N]`` aux (a "source-C" residual) feeds back into the + config: it needs its own SMEM ring that competes with the AB pipeline, forcing the + aux-TMA path capped at ab=2 (the C2/C5 ceiling). A rank-1 rowvec/colvec broadcast + (``bias[N]``, an fp8 rowwise ``scale_a[M,1]``/``scale_b[1,N]``) is transparent — it + reuses the compute config unchanged. + + The authoritative discriminator is the aux descriptor's ``broadcast_axis`` (``None`` = + exact ``[M,N]``; ``1`` = rowvec), which is NOT the ``memory_op_facts`` ``ndim`` (that + reflects the rank-2 subscript ``aux[tile_m, tile_n]``, so an ``[M,1]`` scale and an + ``[M,N]`` residual look identical there — an ``ndim``-based rule miscategorizes both a + 16-bit ``[M,1]`` scale AND an 8-bit ``[M,N]`` residual). So use the backend's own + pre-codegen graph detector ``host_function_has_tcgen05_exact_shape_aux_kernel_pattern``, + which reuses the epilogue-chain analyzer to key on the real ``[M,N]``-vs-broadcast graph + shape (dtype-agnostic). It runs off ``device_ir.host_function.device_ir.graphs`` (the + ``HostFunction`` back-ref on the DeviceIR — available at seed time), and requires the + ``CompileEnvironment`` to be active, hence ``with env``. The post-seed + ``cute_tcgen05_exact_shape_aux_kernel_detected`` flag is the SAME detector, just not yet + populated when ``compiler_seed_configs`` runs — this is the identical call the framework + makes right after (kernel.py, ``host_function_has_tcgen05_exact_shape_aux_kernel_pattern``). + """ + from ..cute.aux_tensor import ( + host_function_has_tcgen05_exact_shape_aux_kernel_pattern, + ) + + host_function = getattr(device_ir, "host_function", None) + if host_function is None: + return 0 + try: + with env: + is_exact = host_function_has_tcgen05_exact_shape_aux_kernel_pattern( + host_function + ) + except Exception: + # The detector walks the FX graphs; on any unexpected graph shape fall back to the + # transparent (aux_rank=0) assumption rather than wrongly capping ab. + return 0 + return 2 if is_exact else 0 + + +def _single_matmul_fact(spec: ConfigSpec) -> MatmulFact | None: + facts = spec.matmul_facts + if len(facts) != 1: + return None + fact = facts[0] + if fact.static_m is None or fact.static_n is None or fact.static_k is None: + return None + if fact.lhs_ndim != 2 or fact.rhs_ndim != 2: + return None + if fact.lhs_dtype is not fact.rhs_dtype: + return None + return fact + + +def _num_sm(env: CompileEnvironment) -> int: + from ...runtime import get_num_sm + + try: + return get_num_sm(env.device) + except (AssertionError, NotImplementedError): + # get_num_sm is always >= 1 for CUDA but RAISES for non-CUDA / unimplemented devices. + # Return 0 as the "SM count unknown" sentinel the wave-fill helpers key off (see the + # note above _wave_eff). Unreachable in practice — the eligibility SMEM-budget gate + # already excludes non-CUDA — but keeps the helpers total if called off that path. + return 0 + + +def _fp8_small_grid_fits_one_wave(m: int, n: int, num_sm: int) -> bool: + """True when the fp8 small-grid 128x128 cluster grid fits within ~one wave. + + Each 128x128 cluster spans 2 CTAs; the small-grid tile is the right seed only at/below + ``clusters <= num_sm // 2`` (the existing producer's `_small_grid_within_one_wave` + threshold). Measured head-to-head (cold-L2 cudagraph): where it holds, small-grid + [128,128,128]/ab12 BEATS the rectangular [256,bn] tile on fp8 medium-M — + 512x2048x2048 347 vs 298 (+16%), 512x8192x2048 932 vs 762 (+22%), 256x4096x4096 526 vs + 467 (+13%); above it (512x2048x4096 = 128 clusters > 74) they tie and rect is kept.""" + if num_sm <= 0 or n < _FP8_SMALL_GRID: + return False + clusters = (m // _FP8_SMALL_GRID) * (n // _FP8_SMALL_GRID) + return 0 < clusters <= num_sm // 2 + + +def _regime_tile( + m: int, n: int, num_sm: int, itemsize: int +) -> tuple[int, int, int, str] | None: + """Classify the regime and return ``(cluster_m, bm, bn, pid_type)``, or None to decline. + + Declines shapes outside the tcgen05 static-full-tile envelope (M not tileable by the + regime's bm): the cluster_m=1 decode path needs a power-of-2 M; the cluster_m=2 tiles + need ``M % bm == 0`` (bm=256, or bm=128 for the fp8 small-grid). Declined shapes fall + through to the default fragment (exactly as #3007's Triton formula declines jagged).""" + # decode: 64 <= M <= 128, cluster_m=1 + if m <= _DECODE_M_MAX: + if m < _DECODE_M_MIN: + return None # below the tcgen05 decode admission floor (static_m>=64) + bm = _decode_bm(m) + if bm is None: + return ( + None # non-pow2 M (e.g. 96) — outside the cluster_m=1 static envelope + ) + return 1, bm, _pick_decode_bn(m, bm, n, num_sm), "persistent_blocked" + # cluster_m=2 tiles carry bm=256 -> M must tile it cleanly (single-root static tile). + if m % TCGEN05_TWO_CTA_BLOCK_M != 0: + return None + tile2_waves = _wave_eff( + TCGEN05_TWO_CTA_BLOCK_M, TCGEN05_TWO_CTA_BLOCK_N, m, n, 2, num_sm + ) + if ( + m >= TCGEN05_TWO_CTA_BLOCK_M + and n >= TCGEN05_TWO_CTA_BLOCK_N + and tile2_waves >= _WAVE_FULL + ): + # compute: the 256² TMEM square fills a wave. + # The ``n >= 256`` guard is a STRUCTURAL safety floor, not just occupancy: without it a + # tall+narrow shape (huge M, N in [64,256)) can clear tile2_waves>=0.8 on M-tiles alone + # (e.g. M=16384/N=64 -> 0.86 waves) yet the emitted bn=256 tile is mostly wasted padding + # over an N=64 output. Requiring n>=256 routes those to the medium-M rectangular tile + # below, which shrinks bn to <=128 so it never emits bn>N. (Curriculum-rare — extreme + # aspect ratios — but a real correctness/efficiency hole the occupancy check alone misses.) + return ( + 2, + TCGEN05_TWO_CTA_BLOCK_M, + TCGEN05_TWO_CTA_BLOCK_N, + "persistent_interleaved", + ) + # medium-M: the 256² tile underfills. The 128x128 SMALL-GRID collective (per-CTA 64x128) shrinks + # BOTH dims -> more tiles -> more CTAs, beating the rectangular [256,bn] tile via OCCUPANCY where + # it fits ~one wave (measured fp8 512x2048x2048, all cm2: small-grid 351 / rect256x128 264 / + # full256² 232 TFLOP/s — tracks CTA count 128/64/32, the reverse of arithmetic intensity). + # + # fp8-ONLY (`itemsize == 1`) is a backend CORRECTNESS gate, NOT a perf/SMEM/hardware limit: the + # (bm=128, cm2) coordinate is dtype-multiplexed (`cute_mma.py:_tcgen05_use_2cta_instrs`), routing + # fp8 to the validated CtaGroup.TWO path but bf16/fp16 to a legacy CtaGroup.ONE family whose + # multi-tile runtime ownership isn't yet validated (guarded off). Experiment (GPU1) confirmed the + # collective is dtype-general — bf16 128² cm2 is bit-exact + 8–18% faster — so 16-bit is blocked + # only on that backend gate landing. Kept fp8-gated so the formula never emits a rejected config; + # follow-up to lift it (backend validation + a 16-bit mirror of this branch) is in SUMMARY.md. + # Needs M,N multiples of 128 (single-root static tile) + the small-grid one-wave fit test. + if ( + itemsize == 1 + and m % _FP8_SMALL_GRID == 0 + and n % _FP8_SMALL_GRID == 0 + and _fp8_small_grid_fits_one_wave(m, n, num_sm) + ): + return 2, _FP8_SMALL_GRID, _FP8_SMALL_GRID, "persistent_interleaved" + # else single-wave rectangular cluster_m=2 tile (bn shrunk) — wins at wider N. + bm, bn = _single_wave_rect_tile(m, n, num_sm) + return 2, bm, bn, "persistent_interleaved" + + +def _formula_seed( + fact: MatmulFact, + spec: ConfigSpec, + num_sm: int, + budget: int, + aux_rank: int, +) -> dict[str, Any]: + """The Bucket-A core: regime-classify -> pick collective -> depth-fill the pipeline. + + Pure function of ``(fact, spec, num_sm, budget, aux_rank)`` — the env-dependent + aux-rank detection is done by the caller (``_epilogue_aux_rank(env, device_ir)``) and + passed in, so this stays unit-testable without a live CompileEnvironment. ``budget`` is + the per-CTA AB-SMEM budget in bytes; ``aux_rank`` is 0 (transparent) or 2 (source-C). + """ + m, n, k = fact.static_m, fact.static_n, fact.static_k + assert m is not None and n is not None and k is not None + itemsize = _itemsize(fact.lhs_dtype) + + # ---- 1+2. classify regime + select collective/tile (None => decline) ---- + tile = _regime_tile(m, n, num_sm, itemsize) + assert tile is not None # is_eligible already checked _regime_tile is not None + cluster_m, bm, bn, pid = tile + + # ---- 3+4. bk + ab: fill the AB pipeline to the ~196 KB SMEM isobar (dtype-aware) ---- + # A rank-2 source-C residual forces the aux-TMA path, hard-capped at ab=2 (the + # C2/C5 ceiling — its C-ring can't coexist with a deeper AB pipeline). That cap + # alone determines the depth; bk is still picked on the full budget (measured: + # for the residual rect [256,128] tile, bk128 = 200 TFLOP/s beats bk64 = 161). + bk, ab = _pick_bk_ab(itemsize, cluster_m, bm, bn, budget, k) + if aux_rank == 2: + ab = min(ab, 2) + + # ---- 5. l2_groupings: wave-count-aware (waves in CTAs). Grouping (G×G tile-walk swizzle) aims + # to reuse a shared operand panel across the tiles of a block WHILE it's L2-resident. + # few-wave: block finishes fast, short reuse distance -> panel still cached -> fewer DRAM + # bytes -> [4] wins. + # many-wave: each CTA does ~waves tiles, so a block's reuse is spread over many wave-steps; + # the panel's required residency exceeds L2 against the churn -> evicted before reuse -> + # REFETCHED. Measured (ncu, 4096x8192x8192 bf16 ~7 waves): [4] gets a higher L2 hit RATIO + # (69.6 vs 65.5%) but reads +18% DRAM BYTES (715 vs 607 MB) — short-range reuse works, + # the cross-block reuse it reached for gets evicted. At ~98% tensor pipe (no slack) that + # extra traffic hits the critical path (−3–8%) -> [1] wins. + # So [1] many-wave / [4] few-wave; the sign flips because reuse-distance-vs-L2 scales with + # waves. (Fact is ncu-solid; exact cause = reuse-distance eviction vs a bad G=4 grid map is + # not separable from these counters. _MANY_WAVE=4 is the empirical crossover, least-swept.) + waves = _num_waves(bm, bn, m, n, cluster_m, num_sm) + l2 = [1] if waves >= _MANY_WAVE else [TCGEN05_TWO_CTA_SEED_L2_GROUPING] + + # ---- 6. inherit the codegen defaults for the bn-keyed knobs (§4.3.5) ---- + acc = 2 if bn <= 256 else 1 + c = 4 if bn <= 16 else 2 + + seed: dict[str, Any] = { + "block_sizes": [bm, bn, bk], + "l2_groupings": l2, + "num_warps": 8, + "num_stages": 4, + "pid_type": pid, + "tcgen05_cluster_m": cluster_m, + "tcgen05_cluster_n": 1, + "tcgen05_ab_stages": ab, + "tcgen05_acc_stages": acc, + "tcgen05_c_stages": c, + "tcgen05_num_epi_warps": 4, + TCGEN05_PERSISTENCE_MODEL_CONFIG_KEY: _STATIC_PERSISTENT, + # Bucket-A is a DEFAULT-layout, non-FFI seed. Pin these EXPLICITLY: on a + # direct-entry-eligible shape the fragment defaults are FFI-biased + # (``_base_default_config`` -> layout=explicit_epi_tile, tvm_ffi_launch=True, + # flat_role=True, because PR2's search fragment keeps True as choices[0]). When + # ``default_config()`` layers this promoted seed over ``_base_default_config()`` + # via dict.update, any key this seed OMITS keeps the base's FFI value — grafting + # explicit_epi_tile + ffi=True onto our deep-AB DEFAULT config (e.g. bf16 + # [256,256,64] ab6 c2) yields an invalid hybrid (the FFI direct-entry path needs + # the (bk=64,ab=6,c=4) tuple, but we emit c=2) that raises InvalidConfig in the + # no-autotune / baseline compile. Emitting them keeps the promoted default + # self-consistent on the DEFAULT layout. (The FFI topology ships separately as + # the Bucket-B alt-seed, which sets its own layout/ffi keys.) + TCGEN05_LAYOUT_STRATEGY_CONFIG_KEY: Tcgen05LayoutStrategy.DEFAULT.value, + TCGEN05_FLAT_ROLE_COORDINATES_CONFIG_KEY: False, + TCGEN05_TVM_FFI_LAUNCH_CONFIG_KEY: False, + # Clear the explicit-epi-tile layout overrides too: ``_base_default_config`` + # carries epi_tile_m=128 / epi_tile_n=32 / d_store_box_n=32 (the FFI topology), + # which under the DEFAULT layout trip the "tcgen05 strategy invariants violated" + # check. None = let codegen derive the DEFAULT-layout epi tile. + TCGEN05_LAYOUT_OVERRIDES_EPI_TILE_M_KEY: None, + TCGEN05_LAYOUT_OVERRIDES_EPI_TILE_N_KEY: None, + TCGEN05_LAYOUT_OVERRIDES_D_STORE_BOX_N_KEY: None, + } + # Pure matmul has exactly the A/B/C indexing slots; only emit the explicit indexing + # list then (a fused epilogue adds memory ops -> leave those to the spec default). + if spec.indexing.length == 3: + seed["indexing"] = ["tensor_descriptor"] * 3 + return seed + + +def _formula_eligible(env: CompileEnvironment) -> bool: + """Shared eligibility for both formula heuristics: a single static 2-D tcgen05-native + matmul whose shape lands inside the static-full-tile envelope. + + Gates on ``cute_tcgen05_search_enabled`` — the authoritative "tcgen05 applies to this + kernel" signal (matmul_ops sets it only when the MMA-support probe reports tcgen05 for + the dtype AND static_m>=64/static_n>=8/static_k>=mma_k). This is a genuine backend + CAPABILITY gate (gate-2), NOT the cluster_m2 SEARCH-restriction gate the seed must stay + orthogonal to: when tcgen05 is unavailable (e.g. the warp/universal-only path) the + formula declines so the non-tcgen05 default is preserved.""" + spec = env.config_spec + if not spec.cute_tcgen05_search_enabled: + return False + fact = _single_matmul_fact(spec) + if fact is None: + return False + # tcgen05 covers fp8 e4m3 + 16-bit (bf16/fp16). e5m2 / fp32 are out of the tcgen05 + # search envelope (universal-atom only) -> decline, fall through. + if fact.lhs_dtype not in (torch.float8_e4m3fn, *_SIXTEEN_BIT_DTYPES): + return False + # Skinny-M (M<=16) is the SIMT-vec regime owned by CuteFp8GemmSkinnyMHeuristic (a bm=1 + # row-per-block kernel, not a tcgen05 tile) — decline so it stays default. + if fact.static_m is not None and fact.static_m <= _SKINNY_M_MAX: + return False + if len(spec.block_sizes) != 3: + return False + # Genuine SMEM floor (gate-2): a device that can't hold the AB ring at all. + if CuteTcgen05Config.per_cta_ab_smem_budget_bytes(env.device) <= 0: + return False + # Decline shapes outside the tcgen05 static-full-tile envelope (M not tileable by the + # regime bm — e.g. non-pow2 decode M, or M not a multiple of 256 for cluster_m=2). + return ( + _regime_tile( + fact.static_m, fact.static_n, _num_sm(env), _itemsize(fact.lhs_dtype) + ) + is not None + ) + + +class CuteTcgen05FormulaMatmulHeuristic(AutotunerHeuristic): + """Bucket-A formula seed for tcgen05 matmul — the #3007 analog (promote-to-default). + + Reads the ``MatmulFact`` directly (NOT the search-restriction gate) so it covers the + regimes the 3 cluster_m=2 producers can't: cluster_m=1 decode, single-wave medium-M. + Registered AFTER the demoted 3 producers so ``compiler_default_config`` (last-promote- + wins) is this formula's config.""" + + name = "cute_tcgen05_formula_matmul" + backend = "cute" + promote_seed_to_default = True + + @classmethod + def is_eligible(cls, env: CompileEnvironment, device_ir: DeviceIR) -> bool: + return _formula_eligible(env) + + @classmethod + def get_seed_config( + cls, env: CompileEnvironment, device_ir: DeviceIR + ) -> Config | None: + spec = env.config_spec + fact = _single_matmul_fact(spec) + if fact is None: + return None + budget = CuteTcgen05Config.per_cta_ab_smem_budget_bytes(env.device) + if budget <= 0: + return None + num_sm = _num_sm(env) + if ( + _regime_tile( + fact.static_m, fact.static_n, num_sm, _itemsize(fact.lhs_dtype) + ) + is None + ): + return None + aux_rank = _epilogue_aux_rank(env, device_ir) + seed = _formula_seed(fact, spec, num_sm, budget, aux_rank) + return Config(**seed) + + +class CuteTcgen05FormulaFfiAltHeuristic(AutotunerHeuristic): + """Bucket-B FFI ``explicit_epi_tile`` alt-seed for 16-bit full-tile compute. + + A SECOND ranked seed (there is no ``get_seed_configs`` plural — the alt-seed is its + own heuristic, like ``CuteTcgen05ClusterM2FfiHeuristic`` beside the DEFAULT producer). + The ``explicit_epi_tile`` / flat-role / tvm_ffi direct-entry config won 3 climbed bf16 + keys (R1 M2, R1 M3, R2 #11). Since the bf16-deep-AB prerequisite now lets the Bucket-A + DEFAULT path reach the same deep ``bk=64/ab=6`` pipeline, this FFI seed is no longer the + *only* way to get deep 16-bit AB — but it remains a distinct topology (flat-role + + TMA-store epilogue) that can still win on launch/epilogue-bound shapes, so it ships as a + ranked alternative. It is NOT promoted-to-default: the DEFAULT seed is Bucket-A, the FFI + seed is benchmarked beside it and kept only where it wins / compiles (the autotuner drops + it on any InvalidConfig / accuracy failure). + + Only fires for 16-bit (bf16/fp16) full-tile compute (the FFI ``explicit_epi_tile`` / + flat-role path is validated only there); fp8 runs deep ab on the DEFAULT path so it + needs no alt-seed.""" + + name = "cute_tcgen05_formula_ffi_alt" + backend = "cute" + promote_seed_to_default = False + + @classmethod + def is_eligible(cls, env: CompileEnvironment, device_ir: DeviceIR) -> bool: + # Share the capability gate (tcgen05 enabled, single static tcgen05-native fact). + if not _formula_eligible(env): + return False + spec = env.config_spec + fact = _single_matmul_fact(spec) + if fact is None: + return False + # FFI explicit_epi_tile is validated only for 16-bit (bf16/fp16) operands. + if fact.lhs_dtype not in _SIXTEEN_BIT_DTYPES: + return False + # Only the full-tile 256² compute regime (the FFI path is full-tile-only) — and + # only when the 256² tile actually fills a wave (else Bucket-A picks a rectangular + # / decode tile the FFI topology can't express). + m, n = fact.static_m, fact.static_n + assert m is not None and n is not None + if m < TCGEN05_TWO_CTA_BLOCK_M: + return False + num_sm = _num_sm(env) + tile2_waves = _wave_eff( + TCGEN05_TWO_CTA_BLOCK_M, TCGEN05_TWO_CTA_BLOCK_N, m, n, 2, num_sm + ) + if tile2_waves < _WAVE_FULL: + return False + # bk must be in the direct-entry stage-tuple table (64 admits the deep (6,4) tuple). + return _ffi_bk_ab(fact, spec, env, device_ir) is not None + + @classmethod + def get_seed_config( + cls, env: CompileEnvironment, device_ir: DeviceIR + ) -> Config | None: + spec = env.config_spec + fact = _single_matmul_fact(spec) + if fact is None: + return None + bk_ab = _ffi_bk_ab(fact, spec, env, device_ir) + if bk_ab is None: + return None + bk, ab, c = bk_ab + waves = _num_waves( + TCGEN05_TWO_CTA_BLOCK_M, + TCGEN05_TWO_CTA_BLOCK_N, + fact.static_m, + fact.static_n, + 2, + _num_sm(env), + ) + l2 = [1] if waves >= _MANY_WAVE else [TCGEN05_TWO_CTA_SEED_L2_GROUPING] + seed: dict[str, Any] = { + "block_sizes": [TCGEN05_TWO_CTA_BLOCK_M, TCGEN05_TWO_CTA_BLOCK_N, bk], + "l2_groupings": l2, + "num_warps": 8, + "num_stages": 4, + "pid_type": "persistent_interleaved", + "tcgen05_cluster_m": 2, + "tcgen05_cluster_n": 1, + "tcgen05_ab_stages": ab, + "tcgen05_acc_stages": 2, + "tcgen05_c_stages": c, + TCGEN05_L2_SWIZZLE_SIZE_CONFIG_KEY: 1, + "tcgen05_num_epi_warps": 4, + TCGEN05_PERSISTENCE_MODEL_CONFIG_KEY: _STATIC_PERSISTENT, + TCGEN05_LAYOUT_STRATEGY_CONFIG_KEY: Tcgen05LayoutStrategy.EXPLICIT_EPI_TILE.value, + TCGEN05_LAYOUT_OVERRIDES_EPI_TILE_M_KEY: 128, + TCGEN05_LAYOUT_OVERRIDES_EPI_TILE_N_KEY: 32, + TCGEN05_LAYOUT_OVERRIDES_D_STORE_BOX_N_KEY: 32, + TCGEN05_FLAT_ROLE_COORDINATES_CONFIG_KEY: True, + TCGEN05_TVM_FFI_LAUNCH_CONFIG_KEY: True, + } + if spec.indexing.length in (3, 4): + seed["indexing"] = ["tensor_descriptor"] * spec.indexing.length + return Config(**seed) + + +# The direct-entry stage tuples that the FFI codegen accepts, tried in order +# (mirrors TCGEN05_DIRECT_ENTRY_STAGE_TUPLES_BY_BK): bk=64 admits (ab=6,c=4) or (3,2); +# bk=128 admits only (3,2). The deep bk=64/ab=6 pipeline won the bf16 compute keys, so it +# is tried first. When the deep tuple is skipped (a source-C residual — deep AB loses to +# the residual C-ring), the SHALLOW bk=128/ab=3 tuple is preferred over bk=64/ab=3 +# (measured: residual_add 4096³ bk128/ab3 = 1017 > bk64/ab3 = 945). +_FFI_STAGE_TUPLES_BY_BK: tuple[tuple[int, int, int], ...] = ( + (64, 6, 4), + (128, 3, 2), + (64, 3, 2), +) + + +def _ffi_bk_ab( + fact: MatmulFact, spec: ConfigSpec, env: CompileEnvironment, device_ir: DeviceIR +) -> tuple[int, int, int] | None: + """Pick the deepest FFI (bk, ab, c) direct-entry tuple that (a) tiles K evenly and + (b) fits the per-CTA AB-SMEM budget for the 256² 16-bit tile. + + For a rank-2 source-C residual, the residual C-ring competes with the AB pipeline, so + the deep (64,6,4) tuple loses to the shallow (bk,3,2): measured on residual_add + 4096³ old_ffi ab3 = 1017 > formula_ffi ab6 = 972.8. So skip the deep tuple when a + source-C residual is present.""" + k = fact.static_k + assert k is not None + itemsize = _itemsize(fact.lhs_dtype) + aux_rank = _epilogue_aux_rank(env, device_ir) + budget = CuteTcgen05Config.per_cta_ab_smem_budget_bytes(env.device) + # Respect the bk fragment range if present (bk must be reachable in the spec). + bk_low, bk_high = _bk_fragment_bounds(spec) + for bk, ab, c in _FFI_STAGE_TUPLES_BY_BK: + if aux_rank == 2 and ab > 3: + continue # deep AB loses to the shallow tuple under a source-C residual ring + if k % bk != 0: + continue + if bk_low is not None and not (bk_low <= bk <= bk_high): + continue + per = tcgen05_ab_smem_bytes_per_cta( + bm=TCGEN05_TWO_CTA_BLOCK_M, + bn=TCGEN05_TWO_CTA_BLOCK_N, + bk=bk, + dtype_bytes=itemsize, + ab_stages=ab, + cluster_m=2, + ) + if 0 < per <= budget: + return bk, ab, c + return None + + +def _bk_fragment_bounds(spec: ConfigSpec) -> tuple[int | None, int]: + if len(spec.block_sizes) != 3: + return None, 0 + frag = spec.block_sizes[2]._fragment(spec) + low = getattr(frag, "low", None) + high = getattr(frag, "high", None) + if isinstance(low, int) and isinstance(high, int): + return low, high + return None, 0 diff --git a/helion/_compiler/cute/cute_mma.py b/helion/_compiler/cute/cute_mma.py index 1c3a1347e..397b27714 100644 --- a/helion/_compiler/cute/cute_mma.py +++ b/helion/_compiler/cute/cute_mma.py @@ -3283,10 +3283,24 @@ def _emit_tcgen05_tmem_setup() -> None: # bias on the SIMT load path (no TMA), so the explicit # epilogue-tile family stays validated for the T2 envelope: the # store side still uses the same TMA-store + epi-tile shape as - # T1/T3/T4/T5. Exact-shape rank-2 aux tensors (broadcast_axis= - # None) and any other broadcast shape remain rejected here. + # T1/T3/T4/T5. + # + # Shape-5 (bias_residual_gelu) widens this to also admit a rank-2 + # exact-shape residual aux (``broadcast_axis is None``, e.g. + # ``residual[tile_m, tile_n]``). Under the explicit-epi-tile / + # flat-role envelope ``c_input_warps == 0`` is enforced below, so the + # aux-TMA *productive* body never fires (its gate + # ``aux_tma_productive_body_gate_open`` requires ``c_input_warps > 0``): + # the residual is read by the epi warps through the direct SIMT + # exact-shape GMEM gather in ``_codegen_cute_store_tcgen05_tile`` (the + # same read path the DEFAULT-layout residual_add family already uses). + # That input-side gather is independent of the D-output TMA-store box + # the explicit epi-tile shape governs, so the store side is unchanged. + # Colvec (axis 2) and leading-axis (axis 0) broadcast aux remain + # rejected here -- only the rank-1 rowvec and rank-2 exact-shape forms + # are validated on this path. aux_descriptors_compatible_with_explicit_epi_tile = all( - d.broadcast_axis == 1 for d in aux_tensor_descriptors_value + d.broadcast_axis in (1, None) for d in aux_tensor_descriptors_value ) # The explicit-epi-tile / flat-role store path is dtype-general for any # 16-bit operand: bf16 and fp16 produce the same epilogue tile diff --git a/helion/_compiler/cute/tcgen05_config.py b/helion/_compiler/cute/tcgen05_config.py index dc7efc5c8..8a8eb6845 100644 --- a/helion/_compiler/cute/tcgen05_config.py +++ b/helion/_compiler/cute/tcgen05_config.py @@ -1547,28 +1547,38 @@ def _validate_direct_entry_ab_stage_envelope( # the 16-bit hard cap (6) AND the SMEM budget, so the real ab6/bk128 # overflow (294912 B) is still rejected -> clamped to what fits. The # canonical 256x256 cm2 tile only fits ab<=3, so a snapped ab4 there - # is clamped back to 3 here. cluster_m=1 keeps the strict ab<=3 cap - # (its reachable bf16 tiles overflow beyond ab3), and EXPLICIT_EPI_TILE - # / FFI configs are handled by the direct-entry tuple branch above. + # is clamped back to 3 here. EXPLICIT_EPI_TILE / FFI configs are + # handled by the direct-entry tuple branch above. # The batched leading-passthrough family is admitted too: per-CTA AB # SMEM is batch-invariant (``tcgen05_ab_smem_bytes_per_cta`` takes only # bm/bn/bk/dtype/stages/cluster_m, and the leading axis is squeezed to # block size 1 in codegen), so a batched ``[*,256,128,128]`` cm2 tile # fits ab=4 identically and ``max_ab_stages_that_fit`` still enforces # the real per-CTA cap. + # + # PR-5 (formula seed) extension: 16-bit ``cluster_m=1`` DEFAULT layout + # gets the SAME SMEM-clamped admission. The formula's decode regime + # emits a narrow cm1 tile ([64,32,256] bf16 ab4 = 196608 B, the R3 #8 + # decode answer key) that fits the budget and is dtype-general in the + # role_local_monolithic codegen (fp8 cm1 already runs deep ab>3 through + # it). The earlier cm1 ab<=3 restriction was about SMEM OVERFLOW, which + # ``max_ab_stages_that_fit`` already enforces: a cm1 256^2 bf16 tile + # per-stage is 65536 B (bk64) / 131072 B (bk128), so fit_max clamps it + # to 3 / 1 respectively -- the overflow is still rejected. Only the + # fitting decode tile is admitted at ab4. constraints = self.ab_stages_search_constraints is_fp8 = constraints is not None and constraints.dtype_bytes == 1 layout = config.get( TCGEN05_LAYOUT_STRATEGY_CONFIG_KEY, Tcgen05LayoutStrategy.DEFAULT.value, ) - is_16bit_default_cm2 = ( + is_16bit_default = ( constraints is not None and constraints.dtype_bytes == 2 - and config.get("tcgen05_cluster_m") == 2 + and config.get("tcgen05_cluster_m") in (1, 2) and layout == Tcgen05LayoutStrategy.DEFAULT.value ) - if is_fp8 or is_16bit_default_cm2: + if is_fp8 or is_16bit_default: config_view = self._matmul_config_view(config) cluster_m = cast("int", config.get("tcgen05_cluster_m", 1)) if config_view is not None: diff --git a/helion/language/_gelu_tanh_approx.py b/helion/language/_gelu_tanh_approx.py index 1ff0bad0d..01b34d17a 100644 --- a/helion/language/_gelu_tanh_approx.py +++ b/helion/language/_gelu_tanh_approx.py @@ -84,10 +84,24 @@ # inputs to fp32 around ``cute.math.tanh`` automatically, so the # absence of an explicit cast is intentional and safe for both call # sites. +# +# ``fastmath=True`` lowers ``cute.math.tanh`` to the single hardware +# ``tanh.approx.f32`` MUFU instruction instead of the accurate, +# multi-instruction software polynomial the default (``fastmath=False``) +# emits. In a fused tcgen05 GEMM+GELU epilogue the accurate tanh is a +# throughput bottleneck: its extra MUFU/ALU ops do not overlap the UMMA +# and expose ~14% of runtime on a 2048x4096x4096 bf16 GEMM (a plain-matmul +# and a fused-ReLU epilogue both tie the reference at ~1.0x, but the +# accurate-tanh GELU stalled at ~0.87x). The approximation matches the +# reference CuTe kernel (quack ``gemm_act`` uses the same +# ``tanh.approx.f32`` for its GELU) and stays within bf16 rounding of the +# exact ``F.gelu`` oracle (max abs diff ~0.016 on that shape), so it is +# the correct lowering for the tanh-approximation GELU the user opted into +# via ``approximate="tanh"``. _GELU_TANH_APPROX_EXPR_CUTE = ( f"(0.5 * ({{inner}}) * (1.0 + cute.math.tanh(({{inner}}) *" f" ({GELU_TANH_APPROX_KAPPA!r} + {GELU_TANH_APPROX_LAMBDA!r}" - f" * ({{inner}}) * ({{inner}})))))" + f" * ({{inner}}) * ({{inner}})), fastmath=True)))" ) # Exact erf GELU uses a helper so fp32 TensorSSA carriers, including # tcgen05 epilogue fragments, can use packed f32x2 mul/fma around the diff --git a/test/test_autotuner_heuristics.py b/test/test_autotuner_heuristics.py index 37ff1ad21..782b6f8cc 100644 --- a/test/test_autotuner_heuristics.py +++ b/test/test_autotuner_heuristics.py @@ -1325,26 +1325,34 @@ def _assert_cute_tcgen05_cluster_m2_seeded( for config in configs if config.config["tcgen05_cluster_m"] == 2 ] - # FFI-eligible shapes have both DEFAULT-layout and direct-entry seeds. - # Callers decide whether both are expected in the supplied population; - # every cluster_m=2 seed must still match the common tile envelope. + # FFI-eligible shapes have both DEFAULT-layout and direct-entry seeds, and + # the promote-to-default formula heuristic additionally emits a deep-AB + # compute seed on a different bk (e.g. [256,256,64] ab=6 alongside the + # canonical bk=128 tile). Assert the expected-envelope seed is PRESENT + # among the cluster_m=2 seeds (the property under test -- that tile is + # seeded rather than mutation-discovered) rather than requiring every + # cluster_m=2 seed to be it. self.assertGreaterEqual(len(seeded), 1) - for seed in seeded: - self.assertEqual( - seed["block_sizes"][:3], - [ - TCGEN05_TWO_CTA_BLOCK_M, - TCGEN05_TWO_CTA_BLOCK_N, - expected_block_k, - ], - ) - self.assertEqual( - seed["indexing"], - ["tensor_descriptor"] * expected_indexing_length, - ) - self.assertEqual(seed["pid_type"], "persistent_interleaved") - self.assertEqual(seed["tcgen05_num_epi_warps"], 4) - return seeded[0] + matching = [ + seed + for seed in seeded + if seed["block_sizes"][:3] + == [ + TCGEN05_TWO_CTA_BLOCK_M, + TCGEN05_TWO_CTA_BLOCK_N, + expected_block_k, + ] + and seed["indexing"] == ["tensor_descriptor"] * expected_indexing_length + and seed["pid_type"] == "persistent_interleaved" + and seed["tcgen05_num_epi_warps"] == 4 + ] + self.assertGreaterEqual( + len(matching), + 1, + f"expected cluster_m=2 seed [256,256,{expected_block_k}] not found among " + f"{[s['block_sizes'] for s in seeded]}", + ) + return matching[0] def _assert_cute_tcgen05_edge_k_tail_seed_overrides( self, diff --git a/test/test_cute_lowerings.py b/test/test_cute_lowerings.py index 23b2dc6d8..7e6374e96 100644 --- a/test/test_cute_lowerings.py +++ b/test/test_cute_lowerings.py @@ -1024,12 +1024,22 @@ def cute_matmul_mma_codegen_only( with patch_cute_mma_support(): bound = cute_matmul_mma_codegen_only.bind(args) - # Keep the narrowed cluster_m=1 search. Explicit flat - # cluster_m=2 configs are rejected until G3 runtime ownership is - # validated, and this auto-path test only needs to pin tcgen05. bound.env.config_spec.cute_tcgen05_search_enabled = True bound.env.config_spec.restrict_tcgen05_cluster_m_search((1,)) - config = bound.config_spec.default_config() + # Pin the cluster_m=1 flat auto-path config explicitly (matching the + # sibling test_tcgen05_default_store_arrives_with_exec_warp). This test + # validates the cluster_m=1 tcgen05 MMA codegen markers, so it must not + # depend on ``default_config()`` — the promote-to-default formula heuristic + # (CuteTcgen05FormulaMatmulHeuristic) now owns the default and legitimately + # emits a cluster_m=2 config for this shape. + config = helion.Config( + block_sizes=[128, 32, 16], + l2_groupings=[4], + loop_orders=[[0, 1]], + num_stages=2, + num_warps=4, + pid_type="flat", + ) code = bound.to_triton_code(config) self.assertEqual(config.config["block_sizes"][2], 16) diff --git a/test/test_cute_matmul_formula_heuristic.py b/test/test_cute_matmul_formula_heuristic.py new file mode 100644 index 000000000..2836fd059 --- /dev/null +++ b/test/test_cute_matmul_formula_heuristic.py @@ -0,0 +1,233 @@ +"""Emission tests for the CuTe tcgen05 formula matmul seed heuristic. + +These exercise the pure formula logic (regime classification -> collective -> depth-fill) +against the hill-climbed / pretuned answer keys, with a stubbed MatmulFact/ConfigSpec so +no GPU is required. B200 sm100 geometry: num_sm=148, per-CTA AB-SMEM budget 203776 bytes. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import torch + +from helion._compiler.autotuner_heuristics import cute_matmul_formula as F + +_NUM_SM = 148 +_BUDGET = 232448 - 28 * 1024 # 203776 (B200 per-CTA AB-SMEM budget) + +_FP8 = torch.float8_e4m3fn +_BF16 = torch.bfloat16 +_FP16 = torch.float16 + + +def _fact(m, k, n, dtype): + return SimpleNamespace( + static_m=m, + static_k=k, + static_n=n, + lhs_dtype=dtype, + rhs_dtype=dtype, + lhs_ndim=2, + rhs_ndim=2, + m_block_id=0, + n_block_id=1, + k_block_id=2, + ) + + +def _spec(indexing_length=3): + return SimpleNamespace(indexing=SimpleNamespace(length=indexing_length)) + + +def _seed(m, k, n, dtype, indexing_length=3, aux_rank=0): + # aux_rank (0 transparent / 2 source-C residual) is computed from the live env by + # _epilogue_aux_rank in the heuristic; the pure formula takes it as a parameter, so the + # tests pass it directly (aux-rank detection itself is covered by the compile/run tests). + return F._formula_seed( + _fact(m, k, n, dtype), _spec(indexing_length), _NUM_SM, _BUDGET, aux_rank + ) + + +def _knobs(seed): + return ( + tuple(seed["block_sizes"]), + seed["tcgen05_cluster_m"], + seed["tcgen05_ab_stages"], + seed["pid_type"], + ) + + +def test_fp8_decode_key(): + # R1 M1 answer key: [64,64,128] cluster_m=1 ab=12 persistent_blocked. + assert _knobs(_seed(64, 8192, 8192, _FP8)) == ( + (64, 64, 128), + 1, + 12, + "persistent_blocked", + ) + + +def test_fp8_medium_m_key(): + # Climbed key (§1.2): [256,128,128] cluster_m=2 ab=8 persistent_interleaved. + assert _knobs(_seed(512, 2048, 4096, _FP8)) == ( + (256, 128, 128), + 2, + 8, + "persistent_interleaved", + ) + + +def test_bf16_compute_key(): + # R2 #11: bf16 compute fills the isobar to the deep bk64/ab6 pipeline on the DEFAULT path + # (the bf16-deep-AB prerequisite removed the old ab<=3 cap; measured +1-4% over ab3). + seed = _seed(4096, 8192, 8192, _BF16) + assert _knobs(seed) == ((256, 256, 64), 2, 6, "persistent_interleaved") + assert seed["l2_groupings"] == [1] # many-wave + + +def test_fp8_compute_key(): + # R2 #3 / #4: fp8 compute square/flip-point -> [256,256,128] cluster_m=2 ab=6. + assert _knobs(_seed(4096, 4096, 4096, _FP8)) == ( + (256, 256, 128), + 2, + 6, + "persistent_interleaved", + ) + assert _knobs(_seed(2048, 4096, 4096, _FP8)) == ( + (256, 256, 128), + 2, + 6, + "persistent_interleaved", + ) + + +def test_fp8_decode_pretuned_rows(): + # Pretuned AOT table decode rows: bn=32/bk=256/ab=8 (medium K/N) and bn=64/bk=128/ab=12. + assert _knobs(_seed(64, 2048, 4096, _FP8)) == ( + (64, 32, 256), + 1, + 8, + "persistent_blocked", + ) + assert _knobs(_seed(64, 4096, 4096, _FP8)) == ( + (64, 32, 256), + 1, + 8, + "persistent_blocked", + ) + assert _knobs(_seed(64, 5120, 5120, _FP8)) == ( + (64, 64, 128), + 1, + 12, + "persistent_blocked", + ) + + +def test_bf16_decode_deep_ab_key(): + # R3 #8 key: cluster_m=1 bm=64 bn=32 bk=256 ab=4 (deep-AB is the lever; needs the + # bf16-deep-AB prerequisite so 16-bit cluster_m=1 fills the isobar like fp8). + assert _knobs(_seed(64, 4096, 4096, _BF16)) == ( + (64, 32, 256), + 1, + 4, + "persistent_blocked", + ) + + +def test_isobar_invariant(): + # Every Bucket-A key lands at the ~196608-byte AB-SMEM isobar. + for m, k, n, dt in [ + (64, 8192, 8192, _FP8), + (512, 2048, 4096, _FP8), + (4096, 8192, 8192, _BF16), + (4096, 4096, 4096, _FP8), + (64, 4096, 4096, _BF16), + ]: + seed = _seed(m, k, n, dt) + bm, bn, bk = seed["block_sizes"] + b = F.tcgen05_ab_smem_bytes_per_cta( + bm=bm, + bn=bn, + bk=bk, + dtype_bytes=F._itemsize(dt), + ab_stages=seed["tcgen05_ab_stages"], + cluster_m=seed["tcgen05_cluster_m"], + ) + assert b == 196608, (m, k, n, dt, b) + + +def test_transparent_epilogue_keeps_deep_ab(): + # A transparent epilogue (aux_rank=0: unary act, rank-1 rowvec bias, OR an fp8/16-bit + # rowwise [M,1]/[1,N] scale) must NOT clamp the pipeline: fp8 decode ab stays 12. + # (Distinguishing an [M,1] scale from an [M,N] residual is done dtype-AGNOSTICALLY by + # the graph detector in _epilogue_aux_rank — covered end-to-end by the compile/run test.) + assert _seed(64, 8192, 8192, _FP8, aux_rank=0)["tcgen05_ab_stages"] == 12 + + +def test_rank2_residual_caps_ab_at_2(): + # A rank-2 exact-shape [M,N] source-C residual (residual_add / bias_residual_gelu) caps + # ab at 2 — the C2/C5 aux-TMA ceiling. Dtype-agnostic (the cap is physical, not fp8-vs-16bit). + assert _seed(8192, 8192, 8192, _BF16, aux_rank=2)["tcgen05_ab_stages"] == 2 + assert _seed(4096, 4096, 4096, _FP8, aux_rank=2)["tcgen05_ab_stages"] == 2 + + +def test_fp16_shares_16bit_path(): + # fp16 shares the bf16 16-bit path (compute cluster_m=2, deep bk64/ab6 default path). + assert _knobs(_seed(4096, 8192, 8192, _FP16)) == ( + (256, 256, 64), + 2, + 6, + "persistent_interleaved", + ) + + +def test_decode_regime_boundary(): + # M<=128 is decode (cluster_m=1); M>=256 that fills a wave is compute (cluster_m=2). + assert _seed(128, 4096, 4096, _FP8)["tcgen05_cluster_m"] == 1 + assert _seed(4096, 4096, 4096, _FP8)["tcgen05_cluster_m"] == 2 + + +def test_fp8_medium_m_small_grid_vs_rect(): + # fp8 medium-M: the small-grid [128,128,128]/ab12 wins when it fills ~one wave + # (clusters <= num_sm//2 = 74); else the rectangular tile is kept. Measured head-to-head. + # 512x2048x2048: 4*16=64 clusters <= 74 -> small-grid + assert _knobs(_seed(512, 2048, 2048, _FP8))[:3] == ((128, 128, 128), 2, 12) + # 256x4096x4096: 2*32=64 clusters <= 74 -> small-grid + assert _knobs(_seed(256, 4096, 4096, _FP8))[:3] == ((128, 128, 128), 2, 12) + # 512x8192x2048: 4*16=64 clusters <= 74 -> small-grid + assert _knobs(_seed(512, 8192, 2048, _FP8))[:3] == ((128, 128, 128), 2, 12) + # 512x2048x4096: 4*32=128 clusters > 74 -> rectangular [256,128,128]/ab8 (the climbed key) + assert _knobs(_seed(512, 2048, 4096, _FP8))[:3] == ((256, 128, 128), 2, 8) + + +def test_regime_declines_outside_static_tile_envelope(): + # Non-pow2 decode M (96) and below-floor M (32) decline (return None tile). + assert F._regime_tile(96, 4096, _NUM_SM, 1) is None + assert F._regime_tile(32, 4096, _NUM_SM, 1) is None + # M not a multiple of 256 for cluster_m=2 (e.g. 384) declines. + assert F._regime_tile(384, 4096, _NUM_SM, 1) is None + # Valid decode M and compute M return a tile. + assert F._regime_tile(64, 4096, _NUM_SM, 1) is not None + assert F._regime_tile(4096, 4096, _NUM_SM, 1) is not None + + +def test_compute_branch_n_guard_never_emits_bn_gt_n(): + # A tall+narrow shape (huge M, N<256) can clear the 256² occupancy bar on M-tiles alone, but + # must NOT emit a bn=256 tile over a narrow-N output — the n>=256 guard routes it to the + # rectangular medium-M tile, which caps bn<=N. (Regression for the review-found safety hole.) + for m, n in [(16384, 64), (16384, 128), (32768, 192)]: + cluster_m, bm, bn, _pid = F._regime_tile(m, n, _NUM_SM, 2) # bf16 + assert bn <= n, (m, n, bn) + assert (cluster_m, bm) == (2, 256) + + +def test_decode_fallback_picks_narrowest_bn(): + # In the fallback regime (tiny N: even bn=16 < 0.5 waves at M=64) the decode bn picker must + # return the NARROWEST bn (max occupancy), not the old bn=64 mid-point. Measured: bn=16 + # ties-or-beats bn=64 on every probed tiny-N shape. Consistent with the loop's shrink premise. + for n in (256, 512, 1024): + assert _knobs(_seed(64, 4096, n, _FP8))[0][1] == 16, n + # And the answer-key decode shapes are unchanged by the fallback edit. + assert _knobs(_seed(64, 8192, 8192, _FP8))[0][1] == 64 + assert _knobs(_seed(64, 4096, 4096, _FP8))[0][1] == 32 diff --git a/test/test_dot_requirements.py b/test/test_dot_requirements.py index 8e058b8de..46f620b89 100644 --- a/test/test_dot_requirements.py +++ b/test/test_dot_requirements.py @@ -310,10 +310,19 @@ def cute_matmul_mma(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: self.assertLessEqual(default_block_sizes[0], 256) self.assertGreaterEqual(default_block_sizes[1], 8) self.assertLessEqual(default_block_sizes[1], 128) - self.assertEqual(spec.default_config().config["l2_groupings"], [1]) - # This small-N problem cannot form the validated 256x256 CtaGroup.TWO - # tile, so the autotuner keeps cluster_m narrowed to 1. - self.assertEqual(spec.default_config().config["tcgen05_cluster_m"], 1) + # The promote-to-default formula heuristic emits a wave-count-aware + # l2_grouping: this tiny single-wave 256x64x128 problem gets the few-wave + # grouping [4] (was [1] under the old fixed-grouping default). + self.assertEqual(spec.default_config().config["l2_groupings"], [4]) + # The small-N shape cannot form the validated 256x256 CtaGroup.TWO tile, so + # the SEARCH keeps cluster_m narrowed to 1. The formula seed is orthogonal + # to that search restriction (cute-seed-orthogonal-to-search): it promotes + # the best genuinely-valid config, which here is the rectangular cluster_m=2 + # tile [256,64,64] (bn shrunk to N=64) -- GPU-verified to compile and match + # x@y exactly. So the promoted default is cluster_m=2 even though the search + # arm stays cluster_m=1. + self.assertEqual(spec.default_config().config["tcgen05_cluster_m"], 2) + self.assertEqual(spec.default_config().config["block_sizes"][:2], [256, 64]) self.assertEqual(spec._tcgen05_cluster_m_search_choices, (1,)) self.assertIn("persistent_blocked", spec.allowed_pid_types) self.assertIn("persistent_interleaved", spec.allowed_pid_types) @@ -350,12 +359,14 @@ def cute_matmul_mma(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: self.assertLessEqual(default_block_sizes[0], 256) self.assertGreaterEqual(default_block_sizes[1], 8) self.assertLessEqual(default_block_sizes[1], 256) - # 16-bit 8192^3 is FFI-eligible (fp16 has full bf16 parity), so the - # default is now the validated TVM-FFI full-tile envelope: the FFI - # direct-entry seed's L2 grouping and the CtaGroup.TWO 256x256x128 tile. - self.assertEqual(spec.default_config().config["l2_groupings"], [2]) + # 16-bit 8192^3 is a full-wave compute shape; the promote-to-default + # formula heuristic emits the DEFAULT-layout deep-AB CtaGroup.TWO tile + # ([256,256,64] ab=6) with the wave-count-aware many-wave grouping [1] + # (this many-CTA shape exceeds the _MANY_WAVE crossover; was [2] under the + # old FFI-envelope default). + self.assertEqual(spec.default_config().config["l2_groupings"], [1]) # K=8192 can form validated CtaGroup.TWO products at bk >= 32 even - # though bk=16 is over the K-tile cap. The FFI full-tile default lands + # though bk=16 is over the K-tile cap. The compute full-tile default lands # on cluster_m=2, and the search exposes both arms. self.assertEqual(spec.default_config().config["tcgen05_cluster_m"], 2) self.assertEqual(spec._tcgen05_cluster_m_search_choices, (1, 2)) @@ -405,10 +416,12 @@ def cute_matmul_mma(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: bound = cute_matmul_mma.bind(args) config = bound.config_spec.default_config() code = bound.to_triton_code(config) - # 16-bit 8192^3 is FFI-eligible (fp16 == bf16 parity); the default is the - # validated TVM-FFI full-tile envelope whose bk is 128, not the old - # non-persistent bk=16 default. It still codegens on the tcgen05 path. - self.assertEqual(config.config["block_sizes"][2], 128) + # 16-bit 8192^3 is a full-wave compute shape; the promote-to-default + # formula heuristic emits a DEFAULT-layout deep-AB tile ([256,256,64] ab=6, + # bk=64) rather than the old bk=128 envelope -- still a validated tcgen05 + # full tile (bk in the 32..128 range), not the old non-persistent bk=16 + # default. It still codegens on the tcgen05 path. + self.assertIn(config.config["block_sizes"][2], (64, 128)) self.assertGreaterEqual(config.config["block_sizes"][0], 128) self.assertLessEqual(config.config["block_sizes"][0], 256) self.assertGreaterEqual(config.config["block_sizes"][1], 8) @@ -1118,23 +1131,35 @@ def test_cute_tcgen05_ab_stages_three_seeded_in_initial_population( bound = _bind_cute_4096_matmul_kernel_with_mocked_smem_budget(b200_budget_bytes) spec = bound.config_spec - # 16-bit 4096^3 is FFI-eligible (fp16 == bf16 parity), so the initial - # population now carries TWO cluster_m=2 seeds: the DEFAULT-layout - # cluster_m=2 ab=3 seed and the generalized TVM-FFI direct-entry seed. - # Both must carry the canonical ab=3 fast-config envelope (the point of - # this test — that ab=3 is seeded rather than discovered by mutation). + # 16-bit 4096^3 is FFI-eligible (fp16 == bf16 parity). The initial + # population carries the DEFAULT-layout cluster_m=2 ab=3 seed and the + # generalized TVM-FFI direct-entry seed, both on the canonical ab=3 + # fast-config envelope. The formula matmul heuristic additionally emits + # a deep-AB compute seed for this shape ([256,256,64] ab=6, which fills + # the AB-SMEM isobar and runs faster than the ab=3 tile); that extra seed + # is legitimate, so this test asserts the canonical ab=3 envelope is + # PRESENT among the cluster_m=2 seeds (the point of the test — ab=3 is + # seeded rather than discovered by mutation) rather than requiring every + # cluster_m=2 seed to be it. cluster_m2_seeds = [ config.config for config in spec.compiler_seed_configs if config.config.get("tcgen05_cluster_m") == 2 ] self.assertGreaterEqual(len(cluster_m2_seeds), 1) - for seed in cluster_m2_seeds: - self.assertEqual( - seed["block_sizes"][:3], - [TCGEN05_TWO_CTA_BLOCK_M, TCGEN05_TWO_CTA_BLOCK_N, 128], - ) - self.assertEqual(seed["tcgen05_ab_stages"], 3) + canonical_ab3_seeds = [ + seed + for seed in cluster_m2_seeds + if seed["block_sizes"][:3] + == [TCGEN05_TWO_CTA_BLOCK_M, TCGEN05_TWO_CTA_BLOCK_N, 128] + and seed["tcgen05_ab_stages"] == 3 + ] + self.assertGreaterEqual( + len(canonical_ab3_seeds), + 1, + f"canonical [256,256,128] ab=3 seed missing from cluster_m=2 seeds: " + f"{[s['block_sizes'] for s in cluster_m2_seeds]}", + ) @onlyBackends(["cute"]) def test_cute_universal_matmul_lane_loop_correctness(self) -> None: @@ -2157,6 +2182,12 @@ def cute_matmul_mma(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: # rejected by ``normalize``'s validation pass). spec._tcgen05_num_epi_warps_validation_choices = None spec.restrict_tcgen05_num_epi_warps_search((2,)) + # The promote-to-default formula heuristic pins num_epi_warps=4 explicitly + # in ``compiler_default_config``, which would shadow the search-view default + # in ``default_config()``. This assertion exercises the search-view + # fill-missing routing, so clear the promoted seed to expose the raw + # search-view fragment default (the property under test). + spec.compiler_default_config = None new_default = spec.default_config() self.assertEqual(new_default.config["tcgen05_num_epi_warps"], 2) winning_2 = helion.Config(**new_default.config) @@ -2179,13 +2210,14 @@ def test_cute_tcgen05_strategy_data_model_round_trip(self) -> None: spec = _bind_cute_strategy_kernel().config_spec - # The 256^2 16-bit shape is FFI-eligible (fp16 == bf16 parity), so the - # default is the validated TVM-FFI full-tile envelope: the - # ROLE_LOCAL_MONOLITHIC strategy is still the pin, but the FFI seed pins - # persistent_interleaved / static_persistent / explicit_epi_tile (vs the - # old non-eligible flat / non_persistent / default). The persistence - # model agrees with the persistent pid_type so the serialized config is - # still internally consistent. + # The 256^2 16-bit shape is a full-wave compute shape; the promote-to- + # default formula heuristic emits the DEFAULT-layout CtaGroup.TWO compute + # tile. The ROLE_LOCAL_MONOLITHIC strategy is still the pin, and the seed + # pins persistent_interleaved / static_persistent (vs the old non-eligible + # flat / non_persistent), but on the DEFAULT layout rather than the FFI + # explicit_epi_tile envelope. The persistence model agrees with the + # persistent pid_type so the serialized config is still internally + # consistent. default_cfg = spec.default_config() self.assertEqual( default_cfg.config["tcgen05_strategy"], "role_local_monolithic" @@ -2194,9 +2226,7 @@ def test_cute_tcgen05_strategy_data_model_round_trip(self) -> None: self.assertEqual( default_cfg.config["tcgen05_persistence_model"], "static_persistent" ) - self.assertEqual( - default_cfg.config["tcgen05_layout_strategy"], "explicit_epi_tile" - ) + self.assertEqual(default_cfg.config["tcgen05_layout_strategy"], "default") self.assertEqual(default_cfg.config["tcgen05_warp_spec_ab_load_warps"], 1) self.assertEqual(default_cfg.config["tcgen05_warp_spec_mma_warps"], 1) # ``epi_warps`` is the existing tcgen05_num_epi_warps knob. @@ -2214,15 +2244,14 @@ def test_cute_tcgen05_strategy_data_model_round_trip(self) -> None: self.assertEqual(default_cfg.config["tcgen05_warp_spec_c_input_warps"], 0) self.assertEqual(default_cfg.config["tcgen05_warp_spec_register_decrease"], 120) self.assertEqual(default_cfg.config["tcgen05_warp_spec_register_increase"], 256) - # The FFI explicit_epi_tile default pins the epilogue-tile / D-store-box - # layout overrides (the validated 128/32/32 envelope); the SMEM swizzle - # overrides remain unset so the layout helper picks them. - self.assertEqual(default_cfg.config["tcgen05_layout_overrides_epi_tile_m"], 128) - self.assertEqual(default_cfg.config["tcgen05_layout_overrides_epi_tile_n"], 32) - self.assertEqual( - default_cfg.config["tcgen05_layout_overrides_d_store_box_n"], 32 - ) + # The DEFAULT-layout compute default leaves every layout override unset so + # the layout helper derives the epilogue tile / D-store box / SMEM swizzle + # (the FFI explicit_epi_tile 128/32/32 envelope ships only on the Bucket-B + # FFI alt-seed, not the promoted DEFAULT-layout default). for key in ( + "tcgen05_layout_overrides_epi_tile_m", + "tcgen05_layout_overrides_epi_tile_n", + "tcgen05_layout_overrides_d_store_box_n", "tcgen05_layout_overrides_smem_swizzle_a", "tcgen05_layout_overrides_smem_swizzle_b", ): @@ -3150,6 +3179,17 @@ def test_cute_tcgen05_strategy_flat_round_trip_with_force_persistent( spec.allowed_pid_types, ("persistent_blocked", "persistent_interleaved"), ) + # This test guards the persistence-model derivation round-trip on the + # SEARCH representation. The promote-to-default formula heuristic pins a + # cluster_m=2 [256,256,*] compute config in ``compiler_default_config``, + # which ``default_flat()`` would flatten as the baseline; that promoted + # config is not flat-round-trip-identity in this force-persistent narrowed + # spec (its block_m=256 projects back to the flat block_m default of 128), + # which is a general promoted-seed property, not the persistence-model + # invariant under test. Clear the promoted seed so ``default_flat()`` uses + # the search-view fragment default (verified idempotent: fragment-default + # default_flat DOES round-trip to identity). + spec.compiler_default_config = None cg = ConfigGeneration(spec) default_flat = cg.default_flat() round_tripped = cg.flatten(cg.unflatten(default_flat))