[autotuner][cute] add formula-based tcgen05 matmul seed heuristic

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,

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

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:

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

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,

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)