Fix CP-SAT INFEASIBLE-under-width-pressure (silent heuristic fallback)

On the DOOM forward, CP-SAT returned INFEASIBLE and forward_compile
silently fell back to the heuristic at every width where residual width
was binding (d <= 6400), so optimize>0 was secretly optimize=0. The model
proved infeasible on schedules the heuristic compiles fine, because it
over-counted two resources versus the heuristic it must replay:

1. BIRTH-dirty over-count (attention-head cumulative). The model charged
   a full-width dirty cancel to EVERY fresh allocation; the heuristic only
   clears the initial-pool dirty subset (recycled columns are already
   clean). Under width pressure columns recycle constantly, so the model
   charged hundreds of columns of phantom cancels and the pooled attention
   cumulative overflowed. Fix: compile_to_onnx / compile_headless_to_onnx
   default assume_zero_init=True -- the ONNX runtime always builds the
   residual stream from zeros, so the dirty cancels are unnecessary and
   model and replay both skip them.

2. Input-pinning over-reservation (residual cumulative). The model
   reserved every input column forever (available_residual = d -
   input_residual), but the heuristic frees consumed inputs and recycles
   their columns (the 600+-wide token Embedding is freed early to carry
   the geometry-stage intermediates). Fix: every input except pos_encoding
   is now modelled as a freeable residual interval [0, cancel); the cancel
   layers are written into the assignment and DirectedLayerScheduler frees
   the inputs at replay.

Result on the DOOM forward: d=6400 solves FEASIBLE at 32 layers (heuristic
40), d=11200 at 28 (heuristic 34). The residual floor (d <= 3040, where
the heuristic runs at ~100% residual occupancy) still falls back -- the
model cannot represent the heuristic's within-layer eager freeing -- but
the fallback is now loud, not silent.

Supporting changes:
- Extract build_cpsat_model from solve_schedule (behavior-preserving) so
  the model has one definition, plus a diagnostic constraint-family toggle.
- forward_compile(require_solver=True) raises instead of silently falling
  back; the default path warns (RuntimeWarning) and surfaces SolveStats on
  net.cpsat_solve_stats.
- Tests: input-freeing regression + require_solver guard.
- docs/cpsat_scheduler.md updated for the residual budget, the BIRTH-dirty
  conservatism, and the eager-free limitation.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: physicsrob

docs/cpsat_scheduler.md

@@ -289,6 +289,24 @@ The DEATH and BIRTH terms only run for residual-using nodes —
 chain-internal exclusive `L1` and chain-internal `ReLU` live in MLP
 hidden slots, not residual, so they have no columns to cancel.
 
+**BIRTH-dirty is over-conservative under `assume_zero_init=False`.**
+The model charges a full-width BIRTH-dirty cancel to *every* fresh
+allocation, but the heuristic only clears the *dirty subset* of the
+allocated columns — columns recycled from a previously-cancelled node
+are already clean (`ResidualStreamMap.dirty_subset`).  Under width
+pressure columns recycle constantly, so the model charges hundreds of
+columns of phantom cancels the replay never pays, and the pooled
+attention cumulative rejects schedules the heuristic compiles.  The
+exact dirty cost is allocation-order-dependent (a property of replay,
+not of the layer assignment), so a sound tight per-layer bound is not
+easily expressible.  The practical resolution: `compile_to_onnx`
+defaults `assume_zero_init=True`, because the ONNX runtime always
+builds the residual stream from zeros (`get_input_res_stream`), making
+every BIRTH-dirty cancel unnecessary — both the model and the replay
+then skip them and agree.  `assume_zero_init=False` stays the
+conservative default of `forward_compile` for callers that may pass a
+non-zero residual stream to `forward()` directly.
+
 **MLP slot budget.** Capacity `d_hidden` per layer.
 
 - For each MLP-routed standalone `Linear` `n`: an optional unit-width
@@ -305,18 +323,49 @@ hidden slots, not residual, so they have no columns to cancel.
   `layer_var[chain.relu]`, demand `len(chain.relu)` (the chain's
   hidden width).
 
-**Residual column budget.** Capacity `d − input_residual_cols`,
-where `input_residual_cols` is the sum of widths of pre-allocated
-input nodes plus `pos_encoding`.
+**Residual column budget.** Capacity `d − len(pos_encoding)`.  Only
+`pos_encoding` is reserved for the whole schedule (the attention
+sublayer reads it at nearly every layer); every other input node is
+*freeable*.
 
-- For each residual-using node `n`: a regular interval
+- For each residual-using scheduled node `n`: a regular interval
   `[layer_var[n], cancel_layer[n])`, demand `len(n)`.
+- For each freeable input node `n` (every input except `pos_encoding`):
+  a regular interval `[0, input_cancel_layer[n])`, demand `len(n)`.
+  Inputs are pre-allocated at layer 0, so their birth is fixed at 0;
+  `input_cancel_layer[n]` obeys the same consumer lower bound and
+  cancel-slack window as a scheduled node's `cancel_layer`, and is kept
+  at `max_layers` for an input feeding a terminal `Concatenate`
+  (output cone).  `solve_schedule` writes these cancels into
+  `ScheduleAssignment.node_to_cancel_layer`, and
+  `DirectedLayerScheduler._find_dead_nodes` frees the input at replay —
+  matching the heuristic, which frees consumed inputs and recycles their
+  columns (the wide token `Embedding` is freed early and its 600+
+  columns carry the geometry-stage intermediates).  Reserving every
+  input forever (the pre-fix model) starved intermediates under width
+  pressure and made the residual cumulative reject schedules the
+  heuristic compiles fine.  Each freeable input also contributes a
+  DEATH-layer cancel interval to the attention-head cumulative (its
+  columns must be zeroed before reuse).
 
 `uses_residual(n)` is `False` for chain-internal `ReLU` (lives in
 MLP hidden slots, not residual) and exclusive chain L1 (computed
 inline inside `linear1` from its input's residual columns, never
 written back to the stream). True for everything else.
 
+**Limit — within-layer (eager) freeing is not modelled.** The
+cumulative is layer-granular, so a node read by a consumer at layer
+`K` is held through `K` (`cancel ≥ K+1`).  The heuristic frees such a
+node *within* layer `K` (after the consuming attention op) and reuses
+its columns in the same layer's MLP sublayer — a density the model
+cannot represent without a sublayer-resolution residual axis.  The
+consequence shows only at the residual floor: where the heuristic runs
+at ~100% residual occupancy (e.g. the DOOM forward at `d≈2560`), the
+model needs strictly more layers than the heuristic and CP-SAT cannot
+find an incumbent in budget, so it falls back.  With any residual
+slack (the DOOM forward at `d ≥ 6400`) CP-SAT solves and beats the
+heuristic (e.g. 32 vs 40 layers).
+
 ### Objective
 
 ```

tests/compile/forward/test_cpsat_integration.py

@@ -416,13 +416,115 @@ def fake_solve(*args, **kwargs):
     monkeypatch.setattr(compile_mod, "solve_schedule", fake_solve)
 
     out, inputs = _build_branchy()
+    # The fallback now warns loudly (RuntimeWarning) so it can't masquerade
+    # as a solve; the compile is still valid.
+    with pytest.warns(RuntimeWarning, match="UNOPTIMIZED|fall(ing|s)? back"):
+        net = forward_compile(
+            d=D,
+            d_head=D_HEAD,
+            output_node=out,
+            verbose=False,
+            optimize=1,
+        )
+    actual = net.compute(2, inputs)[out].cpu()
+    expected = out.compute(2, inputs)
+    torch.testing.assert_close(actual, expected, atol=1e-4, rtol=1e-4)
+
+
+def test_cpsat_frees_wide_input_for_intermediate():
+    """A wide input consumed early must be freed so its columns carry a
+    wide downstream intermediate.
+
+    Regression for the input-freeing fix.  The token-`Embedding`-shaped
+    case in miniature: a 24-wide input ``x`` feeds one narrow ``Linear``
+    and is then dead, after which a 40-wide standalone ``Linear`` ``b``
+    must materialise.  At ``d=64`` with a 9-wide ``pos_encoding``:
+
+      * Pinning ``x`` forever (the pre-fix model) leaves only
+        ``64 - 24 - 9 = 31`` residual columns for intermediates — less
+        than ``b``'s 40, so the residual cumulative is INFEASIBLE and
+        CP-SAT falls back.  ``require_solver=True`` turns that silent
+        regression into a hard error.
+      * Freeing ``x`` once its consumer runs (the fix) frees its 24
+        columns, leaving ``64 - 9 = 55`` for ``b`` (40) plus the narrow
+        ``a`` (4) — a comfortable fit, so CP-SAT solves.
+
+    The compiled output must match the graph oracle: the directed replay
+    has to actually reclaim ``x``'s columns and reuse them for ``b``.
+    """
+    from torchwright.graph.pos_encoding import PosEncoding
+
+    torch.manual_seed(0)
+    x = create_input("x", 24)
+    a = Linear(x, torch.randn(24, 4), torch.zeros(4), name="a")
+    b = Linear(a, torch.randn(4, 40), torch.zeros(40), name="b")
+    out = Linear(b, torch.randn(40, 4), torch.zeros(4), name="out")
+
     net = forward_compile(
-        d=D,
-        d_head=D_HEAD,
+        d=64,
+        d_head=8,
         output_node=out,
+        pos_encoding=PosEncoding(9),
         verbose=False,
         optimize=1,
+        require_solver=True,
+    )
+    # require_solver=True would have raised on a fallback, so reaching here
+    # means CP-SAT produced a real assignment.
+    assert net.cpsat_solve_stats is not None
+    assert net.cpsat_solve_stats.status_name in ("OPTIMAL", "FEASIBLE")
+
+    inputs = {"x": torch.randn(2, 24)}
+    actual = net.compute(2, inputs)[out].cpu()
+    expected = out.compute(2, inputs)
+    torch.testing.assert_close(actual, expected, atol=1e-4, rtol=1e-4)
+
+
+def test_require_solver_raises_on_fallback(monkeypatch):
+    """``require_solver=True`` converts a silent CP-SAT fallback into a hard
+    error; the default (``require_solver=False``) warns loudly instead of
+    failing silently.  Regression guard for the silent-fallback footgun.
+    """
+    from torchwright.compiler.forward import compile as compile_mod
+    from torchwright.compiler.forward.cpsat_scheduler import SolveStats
+
+    fake_stats = SolveStats(
+        status_name="INFEASIBLE",
+        objective_value=-1,
+        best_objective_bound=0.0,
+        wall_time_s=0.0,
+        solver_log="",
+        total_attn_heads=-1,
+        total_mlp_bypass_slots=-1,
+        is_optimal=False,
+    )
+    monkeypatch.setattr(
+        compile_mod, "solve_schedule", lambda *a, **k: (None, fake_stats)
     )
+
+    out, inputs = _build_branchy()
+
+    # require_solver=True -> raise rather than silently fall back.
+    with pytest.raises(RuntimeError, match="no usable assignment|require_solver"):
+        forward_compile(
+            d=D,
+            d_head=D_HEAD,
+            output_node=out,
+            verbose=False,
+            optimize=1,
+            require_solver=True,
+        )
+
+    # require_solver=False (default) -> warn, fall back, still produce a
+    # correct compile.
+    with pytest.warns(RuntimeWarning, match="UNOPTIMIZED|fall(ing|s)? back"):
+        net = forward_compile(
+            d=D,
+            d_head=D_HEAD,
+            output_node=out,
+            verbose=False,
+            optimize=1,
+        )
     actual = net.compute(2, inputs)[out].cpu()
     expected = out.compute(2, inputs)
     torch.testing.assert_close(actual, expected, atol=1e-4, rtol=1e-4)

torchwright/compiler/export.py

@@ -562,6 +562,7 @@ def compile_headless_to_onnx(
     extra_metadata: Optional[dict] = None,
     d_hidden: Optional[int] = None,
     trim_heads: bool = True,
+    assume_zero_init: bool = True,
 ) -> None:
     """Compile a float-I/O graph to a KV-cached ONNX model.
 
@@ -582,6 +583,12 @@ def compile_headless_to_onnx(
     ``d_hidden`` is the per-layer MLP hidden width.  Defaults to ``d``
     when omitted; pass an explicit value to decouple the MLP intermediate
     width from the residual stream width.
+
+    ``assume_zero_init`` defaults to ``True`` (see :func:`compile_to_onnx`):
+    the ONNX runtime always builds the residual stream from zeros plus the
+    input projections, so BIRTH-layer dirty-column cancels are unnecessary and
+    skipping them keeps the CP-SAT attention-head cumulative from being
+    over-tight under width pressure.
     """
     dense_inits: list = []
     sparse_inits: list = []
@@ -608,6 +615,7 @@ def compile_headless_to_onnx(
         on_layer_compiled=on_layer_compiled,
         d_hidden=d_hidden,
         trim_heads=trim_heads,
+        assume_zero_init=assume_zero_init,
     )
     t_compile = time.perf_counter() - t0
 
@@ -759,6 +767,7 @@ def compile_to_onnx(
     verbose: bool = True,
     trim_heads: bool = True,
     optimize: int = 0,
+    assume_zero_init: bool = True,
 ) -> None:
     """Compile a token-I/O graph to a KV-cached ONNX model.
 
@@ -775,6 +784,20 @@ def compile_to_onnx(
 
     Prefill = empty past + past_len=0; decode = feed back the new_K_i /
     new_V_i from a prior run with the accumulated past_len.
+
+    ``assume_zero_init`` defaults to ``True`` here (unlike ``forward_compile``,
+    which defaults ``False``): the ONNX runtime always constructs the residual
+    stream from zeros plus the input projections (see
+    ``HeadlessTransformer.get_input_res_stream`` /
+    ``_OnnxRuntime._build_res_stream``), so the initially-free residual columns
+    are guaranteed clean on entry.  The compiler can therefore skip the
+    BIRTH-layer dirty-column cancels that defend against a non-zero residual
+    stream — those cancels are pure overhead the ONNX runtime never needs, and
+    modelling them per-allocation makes the CP-SAT attention-head cumulative
+    over-tight under width pressure (it charges every fresh allocation a
+    full-width dirty cancel, while the replay only clears the genuinely-dirty
+    initial-pool subset).  No ONNX caller can supply a non-zero stream, so
+    ``True`` is always sound for this entry point.
     """
     dense_inits: list = []
     sparse_inits: list = []
@@ -801,6 +824,7 @@ def compile_to_onnx(
         on_layer_compiled=on_layer_compiled,
         trim_heads=trim_heads,
         optimize=optimize,
+        assume_zero_init=assume_zero_init,
     )
     t_compile = time.perf_counter() - t0
 

torchwright/compiler/forward/compile.py

@@ -8,6 +8,7 @@
 import copy
 import os
 import time
+import warnings
 from typing import Callable, Optional, Set
 
 import torch
@@ -404,6 +405,7 @@ def forward_compile(
     cpsat_costs: Costs = Costs(),
     cpsat_flex_routing: bool = True,
     assume_zero_init: bool = False,
+    require_solver: bool = False,
 ) -> HeadlessTransformer:
     """Compile a computation graph into a HeadlessTransformer.
 
@@ -463,6 +465,15 @@ def forward_compile(
             directly.  Defaults to False — the conservative behaviour
             that runs BIRTH-layer dirty cancels on every fresh
             allocation regardless of the runtime contract.
+        require_solver: When True and ``optimize>0``, raise
+            ``RuntimeError`` if CP-SAT returns no usable assignment
+            instead of silently falling back to the heuristic.  Use it
+            in tests and measurement harnesses that must distinguish a
+            real optimizer solve from a fallback (a fallback returns a
+            valid-but-unoptimized compile with no error).  When False
+            (the default) the compile still falls back, but now emits a
+            ``warnings.warn`` so the fallback is never silent.  Ignored
+            when ``optimize=0``.
 
     Returns:
         A HeadlessTransformer whose compute() method reproduces
@@ -509,6 +520,9 @@ def forward_compile(
 
     # 2. Initialize
     net = HeadlessTransformer(d, d_head, pos_encoding, d_hidden=d_hidden)
+    # Solver provenance, populated only when CP-SAT runs (optimize>0); stays
+    # None for the heuristic path so callers can always query it.
+    net.cpsat_solve_stats = None
     residual_map = ResidualStreamMap(d)
     residual_map.allocate(pos_encoding)
     # pos_encoding + input_nodes are populated by get_input_res_stream at
@@ -675,22 +689,45 @@ def forward_compile(
             hint_cancel=hint_cancel if hint_cancel else None,
             log_search_progress=verbose,
         )
+        # Surface the solver provenance on the returned net so callers can
+        # distinguish a real solve from a fallback without re-deriving it
+        # (``stats.status_name``, ``is_optimal``, the LB/objective gap).
+        net.cpsat_solve_stats = _stats
         if assignment is None:
             # CP-SAT found no feasible incumbent within budget — fall
             # back to the heuristic schedule.  The warm-start was a
             # sunk cost; we already know the heuristic produces a
             # valid schedule.
+            #
+            # A fallback returns a valid-but-unoptimized compile with no
+            # error, so a layer-count-only measurement cannot tell it from
+            # a real solve.  Surface it loudly (and optionally fatally via
+            # ``require_solver``) so it can never masquerade as an
+            # optimizer result.  ``INFEASIBLE`` here is almost always a
+            # model bug (the heuristic found a feasible schedule yet
+            # CP-SAT's model proved none exists) rather than a hard
+            # instance — see ``docs/cpsat_scheduler.md``.
+            fallback_msg = (
+                f"CP-SAT returned no usable assignment "
+                f"(status={_stats.status_name}, "
+                f"{cpsat_time_budget_s:.0f}s budget); falling back to the "
+                f"heuristic schedule ({hint_n_layers} layers). The compile "
+                f"is valid but UNOPTIMIZED — optimize>0 did not take "
+                f"effect. status=INFEASIBLE on a graph the heuristic "
+                f"schedules indicates a CP-SAT model bug, not a hard "
+                f"instance."
+            )
+            if require_solver:
+                raise RuntimeError(
+                    fallback_msg
+                    + " (require_solver=True: refusing the silent fallback)"
+                )
+            warnings.warn(fallback_msg, RuntimeWarning, stacklevel=2)
             if verbose and _stats.solver_log:
                 print("--- CP-SAT solver log (last 40 lines) ---")
                 for line in _stats.solver_log.splitlines()[-40:]:
                     print(f"  {line}")
                 print("--- end CP-SAT solver log ---")
-            if verbose:
-                print(
-                    f"  CP-SAT found no feasible incumbent within "
-                    f"{cpsat_time_budget_s:.0f}s budget — falling back to "
-                    f"heuristic schedule ({hint_n_layers} layers)"
-                )
             scheduler = LayerScheduler(
                 graph,
                 d,