From fb8a1391f868c546ec160c626ca86c18cc70c5d5 Mon Sep 17 00:00:00 2001 From: Caleb Kim Date: Mon, 27 Jul 2026 21:43:14 +0000 Subject: [PATCH] [autotuner][cute] repair partial FFI candidates away from the seed, not onto it The parent commit made the search draw tcgen05_tvm_ffi_launch False-only, but that flag is just ONE of the six clauses in _target1_tvm_ffi_promotion_requested. The other five -- tcgen05_flat_role_coordinates, tcgen05_layout_strategy='explicit_epi_tile', and the three tcgen05_layout_overrides_{epi_tile_m,epi_tile_n,d_store_box_n} -- are independent search draws, so ~96% of candidates still tripped the gate and were overwritten wholesale with the FFI seed (measured, S3 2048x4096x4096 bf16: 293/300 candidates -> 1 distinct config, 0 of 30 keys varying). That is not a search. The projection was never needed for FFI coverage: the envelope is already in the population as an eligibility-gated seed (CuteTcgen05ClusterM2FfiHeuristic plus the formula FFI-alt), which _generate_best_available_population_flat additionally records as a pinned finalist, so it is benchmarked and protected from pruning regardless. Verified on 4 FFI-eligible shapes with the projection bypassed: the seed still reaches the population every time, canonical-hash identical to the normalized seed. _fix_target1_tvm_ffi_search_config becomes a three-way gate: 1. A config already carrying the COMPLETE envelope (_is_coherent_target1_tvm_ffi_config over the 9 keys in _TARGET1_TVM_FFI_ENVELOPE_KEYS) is left untouched -- it is a valid direct-entry kernel, and projecting would clobber out-of-envelope knobs it legitimately chose. Not reachable from a random draw (0/2000 pure draws, since tvm_ffi_launch is drawn False-only and is an envelope key); this branch exists for the seed and for user-supplied / cache-restored configs. 2. An EXPLICIT but partial tvm_ffi_launch=True (user config or cache entry; the search never draws True) is still projected onto the validated envelope -- unchanged behavior. An explicit request must not be silently downgraded, and a partial envelope cannot compile on its own. 3. Everything else -- a search candidate that tripped the gate via one of the five other keys -- is repaired AWAY from FFI via _strip_target1_tvm_ffi_promotion_surface, preserving the ~20 other keys it sampled. The four layout keys move as one unit because validate_tcgen05_strategy_invariants couples them. fix_search_config now runs the FFI stage FIRST, ahead of _fix_cluster_m2_search_config. Its bn-snapping guard only fires when the layout is already DEFAULT, so under the old order a repaired candidate had already been forced to bn=256 and could never reach the un-seeded [256,128,*] tile. Measured contribution, same seeded population: [256,128,*] 71 -> 139 (S3) and 72 -> 140 (medM). The original position (added by #2518) predates any such dependency -- the old projection wiped every tcgen05_* key, so ordering was moot. Effect (single B200, 300-draw populations, interleaved old/new control in one process): shape distinct ffi=True cm2 [256,128,*] S3 2048x4096x4096 bf16 18 -> 296 283 -> 1 7 -> 132 medM 512x4096x4096 bf16 18 -> 296 283 -> 1 8 -> 133 S3resid 2048x4096x4096 resid_add 19 -> 300 282 -> 1 9 -> 144 8192x1024x1024 bf16 11 -> 299 290 -> 1 4 -> 149 Correctness: 36 distinct post-repair configs compiled, run, and compared against torch.matmul (atol=rtol=2e-2) -- 36/36 correct, zero wrong-number results. Repaired candidates skip codegen slightly more often than naturally-non-FFI ones (20% vs ~13%, all the same pre-existing non-whitelisted-fused-epilogue gate), which the autotuner handles by skipping; net usable configs still rise by more than an order of magnitude. Branch 1 verified directly: a coherent config with a mutated out-of-envelope knob (l2_swizzle_size, l2_groupings, acc_stages, num_stages) keeps that knob and compiles correct, while breaking one envelope key sends it back through branch 2. Known adjacent issue, pre-existing and NOT addressed here: on residual_add shapes the normalized FFI seed does not compile (flatten/unflatten rewrites ab_stages 3->2 and c_stages 2->4 via _fix_aux_tma_full_tile_search_config, after which the flat-role codegen guard rejects it). Confirmed identical on the unpatched parent. Co-Authored-By: Claude Opus 5 (1M context) stack-info: PR: https://github.com/pytorch/helion/pull/3167, branch: calebmkim/stack/37 --- helion/_compiler/cute/tcgen05_config.py | 139 +++++++++++++++++++----- 1 file changed, 109 insertions(+), 30 deletions(-) diff --git a/helion/_compiler/cute/tcgen05_config.py b/helion/_compiler/cute/tcgen05_config.py index 17d518fc1..866e73c41 100644 --- a/helion/_compiler/cute/tcgen05_config.py +++ b/helion/_compiler/cute/tcgen05_config.py @@ -2078,39 +2078,110 @@ def _target1_tvm_ffi_promotion_requested(config: dict[str, object]) -> bool: or config.get(TCGEN05_LAYOUT_OVERRIDES_D_STORE_BOX_N_KEY) is not None ) - @staticmethod - def _clear_target1_tvm_ffi_promotion_surface(config: dict[str, object]) -> None: - for key in list(config): - if key.startswith("tcgen05_") or key == "epilogue_subtile": - config.pop(key, None) + def _strip_target1_tvm_ffi_promotion_surface( + self, config: dict[str, object] + ) -> None: + """Repair a config AWAY from the FFI envelope, in place. + + Turns off the FFI-only controls and returns the layout to DEFAULT, + leaving every other sampled key (tile, cluster, stage counts, swizzle, + pid_type, warp spec, indexing) untouched. The four layout keys move as + one unit because ``validate_tcgen05_strategy_invariants`` couples them: + ``explicit_epi_tile`` requires all three overrides set (and + ``d_store_box_n == epi_tile_n``), while any override set requires + ``explicit_epi_tile``. The result is an ordinary DEFAULT-layout + candidate whose epilogue tile is derived from bm/bn. + """ + config[TCGEN05_TVM_FFI_LAUNCH_CONFIG_KEY] = False + config[TCGEN05_FLAT_ROLE_COORDINATES_CONFIG_KEY] = False + if ( + config.get(TCGEN05_LAYOUT_STRATEGY_CONFIG_KEY) + == Tcgen05LayoutStrategy.EXPLICIT_EPI_TILE.value + ): + config[TCGEN05_LAYOUT_STRATEGY_CONFIG_KEY] = ( + Tcgen05LayoutStrategy.DEFAULT.value + ) + for key in ( + TCGEN05_LAYOUT_OVERRIDES_EPI_TILE_M_KEY, + TCGEN05_LAYOUT_OVERRIDES_EPI_TILE_N_KEY, + TCGEN05_LAYOUT_OVERRIDES_D_STORE_BOX_N_KEY, + ): + config[key] = None + + # Keys that define the FFI direct-entry envelope: the six promotion-request + # controls plus the tile/cluster shape the ``flat_role`` codegen guard pins. + # Deliberately NOT every ``tcgen05_`` key in the seed -- derived fields such + # as ``tcgen05_persistence_model`` are still unset at this point in + # normalization (it is derived from ``pid_type`` later), so comparing them + # would make even the seed itself look incoherent. + _TARGET1_TVM_FFI_ENVELOPE_KEYS: tuple[str, ...] = ( + TCGEN05_TVM_FFI_LAUNCH_CONFIG_KEY, + TCGEN05_FLAT_ROLE_COORDINATES_CONFIG_KEY, + TCGEN05_LAYOUT_STRATEGY_CONFIG_KEY, + TCGEN05_LAYOUT_OVERRIDES_EPI_TILE_M_KEY, + TCGEN05_LAYOUT_OVERRIDES_EPI_TILE_N_KEY, + TCGEN05_LAYOUT_OVERRIDES_D_STORE_BOX_N_KEY, + "block_sizes", + "tcgen05_cluster_m", + "tcgen05_cluster_n", + ) + + def _is_coherent_target1_tvm_ffi_config( + self, config: dict[str, object], seed: Config + ) -> bool: + """True when *config* already carries the seed's full FFI envelope. + + Such a config is already a valid direct-entry kernel, so there is + nothing to project: it is left exactly as sampled. + """ + return all( + config.get(key) == seed.config[key] + for key in self._TARGET1_TVM_FFI_ENVELOPE_KEYS + if key in seed.config + ) def _fix_target1_tvm_ffi_search_config(self, config: dict[str, object]) -> None: - # The generalized direct-entry seed projects FFI-requesting search - # candidates onto the validated CtaGroup.TWO envelope for ANY eligible - # shape (returns None for ineligible shapes, in which case the - # promotion surface is stripped back to the DEFAULT layout below). - seed = self.full_tile_direct_entry_seed_config() + # Three populations reach this gate, and they want different handling. + # + # (1) A config that ALREADY carries the complete, valid FFI envelope -- + # however it got here (sampled, seeded, user-written). It is a valid + # direct-entry kernel as-is, so keep it untouched; projecting would + # only overwrite keys outside the envelope that it legitimately + # chose. + # + # (2) An EXPLICIT but PARTIAL ``tvm_ffi_launch=True`` -- a user config or + # a cache entry that asks for FFI without the rest of the envelope. + # The search never produces this: the flag is drawn False-only + # (``EnumFragment((True, False), search_choices=(False,))`` in + # ``optional_fragments``), so True is always deliberate. Project it + # onto the validated envelope, exactly as before -- an explicit + # request must not be silently downgraded, and a partial one cannot + # compile on its own. + # + # (3) A SEARCH candidate that trips the gate through one of the five + # OTHER keys in ``_target1_tvm_ffi_promotion_requested`` (flat-role, + # ``explicit_epi_tile``, or the three epi-tile overrides). Those are + # independent draws, so ~95% of candidates land here carrying a + # partial envelope. Projecting each of them onto the seed collapsed + # the population to a single config (measured: 293/300 candidates -> + # 1 distinct config, 0 of 30 keys varying), which is not a search -- + # and it was never needed for coverage, because the FFI envelope is + # already in the population as an eligibility-gated seed + # (``CuteTcgen05ClusterM2FfiHeuristic`` + the formula FFI-alt), which + # ``_generate_best_available_population_flat`` also pins as a + # finalist. These are repaired AWAY from FFI, keeping the ~20 other + # keys the search chose (measured: 299/300 distinct). if not self._target1_tvm_ffi_promotion_requested(config): return - if seed is None: - config[TCGEN05_TVM_FFI_LAUNCH_CONFIG_KEY] = False - config[TCGEN05_FLAT_ROLE_COORDINATES_CONFIG_KEY] = False - if ( - config.get(TCGEN05_LAYOUT_STRATEGY_CONFIG_KEY) - == Tcgen05LayoutStrategy.EXPLICIT_EPI_TILE.value - ): - config[TCGEN05_LAYOUT_STRATEGY_CONFIG_KEY] = ( - Tcgen05LayoutStrategy.DEFAULT.value - ) - for key in ( - TCGEN05_LAYOUT_OVERRIDES_EPI_TILE_M_KEY, - TCGEN05_LAYOUT_OVERRIDES_EPI_TILE_N_KEY, - TCGEN05_LAYOUT_OVERRIDES_D_STORE_BOX_N_KEY, - ): - config[key] = None - return - self._clear_target1_tvm_ffi_promotion_surface(config) - config.update(seed.config) + seed = self.full_tile_direct_entry_seed_config() + if seed is not None: + if self._is_coherent_target1_tvm_ffi_config(config, seed): + return + if config.get(TCGEN05_TVM_FFI_LAUNCH_CONFIG_KEY) is True: + self._strip_target1_tvm_ffi_promotion_surface(config) + config.update(seed.config) + return + self._strip_target1_tvm_ffi_promotion_surface(config) def aux_load_mode_autotune_fragments(self) -> dict[str, ConfigSpecFragment]: if not self._aux_tma_search_enabled(): @@ -2675,11 +2746,19 @@ def _validate_diagnostic_mode( ) def fix_search_config(self, config: dict[str, object]) -> None: + # The FFI repair runs FIRST, before the cluster_m=2 block shaping. It + # only rewrites the FFI-only controls + layout keys, so the shaping below + # then sees a settled layout: a repaired candidate is DEFAULT-layout and + # so its sampled bn<=128 can snap to the un-seeded [256,128,*] tile, + # instead of being forced to 256 by the ``layout == DEFAULT`` guard in + # ``_fix_cluster_m2_search_config`` reading a pre-repair + # ``explicit_epi_tile``. (When this ran after the shaping, every repaired + # candidate landed on the already-seeded 256x256 tile.) + self._fix_target1_tvm_ffi_search_config(config) self._fix_aux_edge_search_config(config) self._fix_cluster_m2_search_config(config) self._fix_cluster_m1_persistent_search_config(config) self._fix_ab_stages_three_search_config(config) - self._fix_target1_tvm_ffi_search_config(config) self._fix_aux_tma_full_tile_search_config(config) self._fix_with_scheduler_search_config(config) self._fix_aux_tma_search_config(config)