SIGN IN SIGN UP

[Refactor] Give reducers a first-class PartialFragment layout solved by layout inference (#3093)

* Steer reducer destinations during layout inference

* Remove extra reducer steering attempt filtering

* handle reducer problem3

* bug fix

* [Refactor] Give reducers a first-class PartialFragment layout solved by layout inference

The reducer's physical plan (narrow vs wide partials) used to be decided
after inference by ReducerPlanAndMaterialize, with dst-steering emulating
the missing propagation through side channels (wake tables, hint plumbing,
ownership gates); the same induced layout was computed twice and kept
consistent by convention only, and the reducer split the use-def graph
into disconnected components.

This makes the reducer's layout a first-class solver value that fully
determines the physical plan:

* New layout kind `tl.PartialFragment` (PartialFragmentNode : FragmentNode):
  Fragment algebra whose replication coordinate carries an explicit
  combine decomposition — under a low-bits convention enforced at
  construction, `_rep % combine_size` enumerates addend lanes the finalize
  collective must reduce and `_rep / combine_size` enumerates equal-value
  copy groups. (storage algebra, combine width) is the complete plan: the
  collective steps are derivable from the node alone (CombineSteps), the
  FullParticipant object (combine == replicate == participants, thread ==
  _rep) is the wide fallback, and combine width is part of layout identity
  (same-storage different-decomposition proposals widen at the commit
  point instead of unifying silently). A boundary-preserving condensation
  (CondenseReplicateVarKeepingBoundary) keeps the convention true through
  projection; kind-aware IsEqual in both dispatch directions prevents
  silent unification with plain Fragments.

* Update nests (ParallelOp) record their reducer_update sites and, once
  their loop layout solves, propose the induced partial (per-site
  AnalyzeReducerUpdateSite proofs) or the participant-wide plan. The
  commit point resolves multi-site disagreement with a monotone
  widen-on-conflict rule; zero-update epochs are seeded wide before any
  inference level runs; tl.reducer_force_baseline pins the proposal.

* finalize_reducer reads the solved partial and steers an unconstrained
  destination to its post-collective reading. The wake table, the
  LayoutInferArgs.reducer_update_sites hint plumbing and the
  reducer_use_list_ component glue are deleted: reducers travel ordinary
  use_list_ buffer edges.

* ReducerPlanAndMaterialize reads the plan off the node: storage =
  AsPostCollective(), steps = CombineSteps(), wide lowering iff the node
  is the FullParticipant object. TryNarrowPlan is retired; what remains
  (DeriveNarrowLowering) derives only the lowering details that live
  outside the layout — the packed-lane heuristic and the
  destination-containment gate (which still demotes gracefully to the
  wide plan, preserving the unsatisfiable-steering fallback contract).
  The destination override chains are kept as a verified-no-op fallback
  with retirement telemetry.

* T.annotate_layout accepts a PartialFragment on reducer buffers, pinning
  plans the automatic solution cannot reach (e.g. a middle plan between
  narrow and wide). Annotated partials are strict: a disagreeing update
  site raises with the narrow-rejection reason, and
  ComputeReducerPinnedCandidate back-propagates the annotation through
  the update loop into its input fragments.

* Layout cost models price register-array spills: a fragment access whose
  physical slot index depends on the thread variable (detected by the
  shared FragmentThreadIndexProbe, extracted from
  HasThreadDependentFragmentIndex) demotes the whole array to local
  memory, charged as byte-denominated traffic in both models' `mem`
  channel so it competes with — rather than vetoes — global-bandwidth
  estimates. ChooseBestCandidate's anti-spill preference now applies
  regardless of where the loop stores (a fragment->fragment publication
  copy spills exactly like a fragment->global one). With generation and
  pricing combined, inference produces the wide publication chain's
  replicated layouts on its own: the reducer_v2 suite and the
  mhc_post_bwd production kernels pass with the destination override
  fully disabled.

Testing (all with the kernel cache disabled): reducer_v2 suite 70 passed
(new coverage: pinned middle plans, annotation steering input fragments,
conflict diagnostics, combine-step roundtrip, same-storage
different-combine widening at commit); language suite 1261 passed;
layout-inference cost-model tests and maint/layout_inference goldens
unchanged; generated CUDA for existing reducer kernels is byte-identical;
mhc_post_bwd correctness passes and its six benchmark sizes show
no latency change with zero STL/LDL in SASS.

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
L
Lei Wang committed
6279e735f991a78d860fc242915c702d00ec8434
Parent: 905c75a
Committed by GitHub <noreply@github.com> on 8/27/2026, 6:20:55 AM