[cute] Leaner masked store for single padded-M tcgen05 tile

The padded-M tcgen05 fp8 path (block_m > static_m, added in the parent
commit) ran the epilogue store through the SIMT R2G hybrid's scalar masked
loop, leaving the memory pipe under-fed. For a single padded tile this is
pure overhead: the kernel issued a Python-unrolled per-element loop that
recomputed a 2-D `cute.elem_less(coord, (m_size, n_size))` predicate and a
scalar `if`-guarded store for every element of every subtile.

Root of the waste: when `block_m > m_size` (e.g. M=16 on a bm=64 tile),
`m_size % bm != 0` marks the output tile non-static-full, so the store
dispatches `if tcgen05_full_tile: <vector copy> else: <scalar masked loop>`.
The full-tile predicate `m_offset + bm <= m_size` (0+64 <= 16) is statically
false, so the vector branch is dead and every store takes the scalar path.
NCU showed this capping DRAM throughput well below the equivalent full tile.

Fix (single padded tile only, gated by
`tcgen05_flat_m_edge_single_tile = tcgen05_flat_m_edge_tma and m_size <= bm`,
threaded via a new `CuteTcgen05StoreValue.flat_m_edge` field):
- Drop the dead `if full_tile` branch.
- Emit a single vectorized `cute.copy(..., pred=mask)` (the same
  `logical_divide` predicated-copy path already used by `simt_edge_only`)
  instead of the scalar per-element loop.
- Use an M-only predicate `_coord[0] < m_size` instead of the 2-D
  `elem_less`, since this fast path already requires N % bn == 0 so every
  in-bounds column is valid — fewer scalar ops on the epilogue warps.

The gate is deliberately narrowed to `m_size <= bm` (a single M tile). For
multi-M-tile edges (m_size > bm, e.g. M=80 with bm=64) the first tile is a
genuine full tile and keeps its unpredicated vector store; only the original
load-side TMA relaxation (`tcgen05_flat_m_edge_tma`) applies there.

Measured (B200, M=16 K=4096 N=14336, fp8 e4m3, do_bench_wrapper, cudagraph,
A/B against the parent commit in one session):
- warm-L2:     18.36 us / 3.23 TB/s  ->  16.07 us / 3.69 TB/s  (1.14x)
- cache-clear: 26.67 us / 2.22 TB/s  ->  24.45 us / 2.42 TB/s
- DRAM %peak (NCU): 38.4% -> 42.9% (closes ~57% of the gap to the 46.6%
  full-tile ceiling).

Correctness: output shape is the true (M, N); relerr vs a float32 matmul is
at the fp8-rounding floor (~1.7e-3) for the real rows; generated code still
contains `cute.nvgpu.tcgen05`. Verified across static_m in {16,32,48,64,80}
x block_m in {64,128} (the M=80 multi-tile case confirms the full-tile fast
branch is retained). Regression test
`test_matmul_mma_fp8_small_static_m_padded_tile` passes; full
`test_cute_backend.py` passes (114 + 6 subtests).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

stack-info: PR: #2844, branch: yushangdi/stack/61

Author: yushangdi

helion/_compiler/cute/cute_mma.py

@@ -2040,6 +2040,14 @@ def _tcgen05_tma_2d_major(t: torch.Tensor) -> str | None:
         and n_size % bn == 0
         and k_total_size % bk == 0
     )
+    # The leaner SIMT-store optimization (drop the full-tile fast path, emit a
+    # single M-only predicated copy) is only valid when there is exactly ONE
+    # padded M tile -- i.e. the real M does not even fill one ``bm`` tile, so the
+    # full-tile predicate ``m_offset + bm <= m_size`` is statically unsatisfiable.
+    # When ``m_size > bm`` (multiple M tiles, the last one partial) the first M
+    # tile IS a genuine full tile and must keep its unpredicated vectorized store,
+    # so only the load-side TMA relaxation (``tcgen05_flat_m_edge_tma``) applies.
+    tcgen05_flat_m_edge_single_tile = tcgen05_flat_m_edge_tma and m_size <= bm
     tcgen05_double_edge_tma = (
         mma_impl == "tcgen05"
         and tcgen05_use_tma_pipeline
@@ -5117,6 +5125,7 @@ def _tcgen05_tma_tile_predicate(
                 use_tma_store_epilogue=tcgen05_use_tma_store_epilogue,
                 tma_store_full_tiles_only=tcgen05_tma_store_full_tiles_only,
                 partial_output_tma_store=tcgen05_partial_output_tma_store,
+                flat_m_edge=tcgen05_flat_m_edge_single_tile,
                 # Mirror the value passed to `_make_tcgen05_layout_plan_setup`
                 # above; the store path compares this against `target_dtype`
                 # to enforce the kernel/store equality contract on

helion/_compiler/cute/device_state.py

@@ -45,6 +45,21 @@ class CuteTcgen05StoreValue:
     use_tma_store_epilogue: bool = False
     tma_store_full_tiles_only: bool = False
     partial_output_tma_store: bool = False
+    # Single padded-M tcgen05 tile: the real problem M does not even fill one
+    # ``bm`` tile (``m_size <= bm``, e.g. M=16 on a 64/128-row tile) while N
+    # divides ``bn`` and K divides ``bk`` cleanly, so the ONLY masked output
+    # boundary is the M rows AND there is exactly one M tile. The full-tile
+    # predicate is then statically false (``m_offset + bm <= m_size`` can never
+    # hold), so the SIMT edge store always runs; the edge branch can drop the
+    # dead full-tile fast path and use a single vectorized ``logical_divide``
+    # predicated copy with an M-only row predicate (one boolean mask + one
+    # ``cute.copy``) instead of the per-element scalar ``elem_less``/``if`` loop,
+    # cutting epilogue issue overhead on the streaming warps. Set from
+    # ``tcgen05_flat_m_edge_single_tile`` in ``_emit_mma_pipeline`` (note this is
+    # narrower than the load-side ``tcgen05_flat_m_edge_tma``, which also relaxes
+    # the AB-load predicate for the multi-M-tile ``m_size > bm`` edge case where
+    # the first tile is a genuine full tile and must keep its plain store).
+    flat_m_edge: bool = False
     # Output element dtype (cutlass type string, e.g. "cutlass.BFloat16")
     # used when computing the tcgen05 epilogue tile.
     epi_elem_dtype_str: str = ""

helion/language/memory_ops.py

@@ -3297,15 +3297,27 @@ def _simt_edge_logical_divide_copy_source(
         include_coord_setup: bool = True,
         var_prefix: str = "tcgen05_edge",
         copy_atom: str | None = None,
+        m_only_pred: bool = False,
     ) -> str:
         # Shared edge-only vector copy emitter. The make_layout(1) retile gives
         # cute.copy a per-element predicate, while var_prefix/copy_atom let the
         # same shape drive D stores or exact-aux G2R register loads.
+        #
+        # ``m_only_pred`` is the padded-M tcgen05 fast path: when N divides ``bn``
+        # cleanly the only out-of-bounds output coordinate is the M row, so the
+        # predicate collapses from the 2-D ``cute.elem_less(coord, (m, n))`` to a
+        # single ``coord[0] < m_size`` row test. Fewer live registers / scalar
+        # ops on the epilogue warps for the tiny-M streaming case.
         copy_atom = copy_atom or simt_atom
         edge_src = df.new_var(f"{var_prefix}_src")
         edge_dst = df.new_var(f"{var_prefix}_dst")
         edge_coord = df.new_var(f"{var_prefix}_coord")
         edge_pred = df.new_var(f"{var_prefix}_pred")
+        pred_rhs = (
+            f"_coord[0] < cutlass.Int32({m_size})"
+            if m_only_pred
+            else f"cute.elem_less(_coord, ({m_size}, {n_size}))"
+        )
         return (
             (_simt_edge_coord_subtile_source(indent) if include_coord_setup else "")
             + f"{indent}{edge_src} = cute.logical_divide({src}, cute.make_layout(1))\n"
@@ -3314,7 +3326,7 @@ def _simt_edge_logical_divide_copy_source(
             f"{indent}{edge_pred} = cute.make_rmem_tensor((1, {edge_src}.shape[1]), cutlass.Boolean)\n"
             f"{indent}for _edge_i in range(cute.size({edge_src}.shape[1])):\n"
             f"{indent}    _coord = {edge_coord}[0, _edge_i]\n"
-            f"{indent}    {edge_pred}[0, _edge_i] = cute.elem_less(_coord, ({m_size}, {n_size}))\n"
+            f"{indent}    {edge_pred}[0, _edge_i] = {pred_rhs}\n"
             f"{indent}cute.copy({copy_atom}, {edge_src}, {edge_dst}, pred={edge_pred})\n"
         )
 
@@ -4083,6 +4095,21 @@ def store_common_setup(
             ttr_gc_subtile,
             include_coord_setup=not simt_store_edge_coord_preloaded,
         )
+    elif tcgen05_value.flat_m_edge:
+        # Padded-M tcgen05 tile (block_m > real M, N divides bn). The full-tile
+        # predicate ``m_offset + bm <= m_size`` is statically unsatisfiable here
+        # (there is a single padded M tile), so the ``if full_tile`` fast path is
+        # dead code and the scalar per-element ``elem_less``/``if`` edge loop runs
+        # every subtile -- inflating epilogue register pressure / issue overhead
+        # on the streaming warps. Emit ONLY the vectorized predicated copy with an
+        # M-only row predicate (N is in-bounds because it divides bn): one boolean
+        # mask + one ``cute.copy`` per subtile, no scalar branch per element.
+        simt_store_copy_source = _simt_edge_logical_divide_copy_source(
+            "        ",
+            ttr_rd,
+            ttr_gc_subtile,
+            m_only_pred=True,
+        )
     else:
         simt_store_copy_source = (
             f"        if {full_tile}:\n"