configs/baseline.toml

@@ -84,15 +84,30 @@ dropout      = 0.0
 # secondary-species failure. Flow (not the schema default wgan) so the
 # baseline varies only the decoder relative to the best v0.2 result.
 #
-# COST, measured (RTX 4070, bs 4096, 10 ODE steps), not estimated:
-# sample.sample_secondaries_ar loops `for k in range(k_max)` unconditionally
-# — all 15 slots regardless of predicted n_sec — so a flow AR token costs
-# k_max * steps = 150 stage-2 calls per physics step. That makes this block
-# the dominant cost on both sides:
+# COST, measured (RTX 4070, bs 4096, 10 ODE steps), not estimated — but see
+# the row-compaction note below, which changes the INFERENCE side of this:
 #   training   flow AR 29.5k samp/s  vs flow one-shot 190.7k samp/s  (6.5x)
 #   inference  flow AR  8.5k step/s  vs flow one-shot  68.7k step/s  (8.1x)
 # Accepted deliberately: one-shot is the configuration whose secondary
 # species distribution failed, and that failure is what v0.3 exists to fix.
+#
+# Row compaction (landed after the above measurement): at inference,
+# sample.sample_secondaries_ar used to loop `for k in range(k_max)`
+# unconditionally — all 15 slots regardless of predicted n_sec — so a flow
+# AR token cost k_max * steps = 150 stage-2 calls per physics step. It now
+# drops a row from the batch the moment its own secondary count is
+# exhausted, so the real inference cost is ~n_sec * steps stage-2 calls
+# (this checkpoint's own rollout measured 0.382 secondaries/step — see
+# giant-baseline-flow-ar-rollout-validation.md), not k_max * steps. A CPU
+# micro-benchmark at that multiplicity (giant/model/history.py's
+# hidden_dim=512/6-block shape, k_max=15, batch 512) measured 17.6-22.9x
+# fewer wall-clock seconds for the AR loop alone (markov/attention history
+# respectively) — directional only (CPU, synthetic n_sec distribution, not
+# an end-to-end rollout); the 8.1x inference ratio above is now stale and
+# should be re-measured on GPU via a real rollout + `eval_cost_per_step`
+# once one is run against this checkpoint. Training cost (the 6.5x/29.5k
+# figures) is untouched by this: teacher_forcing = "always" here never
+# calls the AR sampler at train time (see [stage2_model.autoregressive]).
 decoder      = "autoregressive"
 generator    = "flow"
 hidden_dim   = 512

giant/model/history.py

@@ -4,6 +4,7 @@ for the `HISTORY_REGISTRY`/`build_history` factory, which mirrors
 `giant.model.routers`'s `Router`/`ROUTER_REGISTRY` pattern (gitea #35)."""
 
 import inspect
+from typing import cast
 
 import torch
 import torch.nn as nn
@@ -34,6 +35,17 @@ class HistoryEncoder(nn.Module):
     def step(self, feat: torch.Tensor, has_prev: torch.Tensor, cache: object) -> tuple[torch.Tensor, object]:
         return self.forward(feat, has_prev), cache
 
+    def select_cache(self, cache: object, idx: torch.Tensor) -> object:
+        """Row-compacts an inference cache (`init_cache`/`step`'s state) down
+        to `idx` — used by `giant.sample.sample_secondaries_ar`'s row
+        compaction to keep a shrinking active-row set's cache aligned as rows
+        finish generating. Default here matches `init_cache`/`step`'s O(1)
+        default: `cache` is always `None`, so there's nothing to index —
+        correct for any encoder whose per-step state doesn't carry a batch
+        dimension (`MarkovHistory` has no cache at all; its running state is
+        `prev_repr`/`remaining`, compacted directly by the caller)."""
+        return cache
+
 
 HISTORY_REGISTRY: dict[str, type[HistoryEncoder]] = {}
 
@@ -215,3 +227,13 @@ class AttentionHistory(HistoryEncoder):
             x, kv_new = block.step(x, kv)
             new_cache.append(kv_new)
         return x, new_cache
+
+    def select_cache(self, cache: object, idx: torch.Tensor) -> list[torch.Tensor | None]:
+        """Row-compacts every block's `(B, T, dim)` KV cache down to `idx`
+        along its batch dimension — see `HistoryEncoder.select_cache`. `idx`
+        may be a long index tensor or a boolean mask (`giant.sample`'s AR
+        loop uses both). `None` entries (a block that has never seen a
+        `step` call yet) stay `None`."""
+        assert isinstance(cache, list)
+        cache_t = cast("list[torch.Tensor | None]", cache)
+        return [None if kv is None else kv[idx] for kv in cache_t]

giant/model/models.py

@@ -614,6 +614,14 @@ class Stage2Autoregressive(StageModel):
         slot — see `AttentionHistory.step`'s docstring."""
         return self.history_encoder.step(token_feat, has_prev, cache)
 
+    def select_history_cache(self, cache, idx: torch.Tensor):
+        """Row-compacts `cache` (from `init_history_cache`/`history_step`)
+        down to `idx` — see `HistoryEncoder.select_cache`. Used by
+        `giant.sample.sample_secondaries_ar`'s active-row compaction to keep
+        the cache aligned with a shrinking batch as rows finish generating
+        across AR slots."""
+        return self.history_encoder.select_cache(cache, idx)
+
     def forward(
         self,
         x_t: torch.Tensor,

giant/sample.py

@@ -231,6 +231,7 @@ def sample_secondaries_ar(
     stage1_out: torch.Tensor,
     n_sec_pred: torch.Tensor | None,
     steps: int = 10,
+    full_length: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """`Stage2Autoregressive` inference loop: one token at a time, in
     descending-energy slot order, up to `k_max` sequential calls. Unlike
@@ -242,19 +243,18 @@ def sample_secondaries_ar(
     expressiveness.
 
     A `{flow,ddpm}` token costs `steps` ODE substeps; `wgan` costs one pass —
-    the "K sequential forwards" cost applies per-token here, not
-    once, so a flow/ddpm AR run costs ~`k_max * steps` model calls per
-    physics step (or ~`n_sec * steps` under `n_sec_pred=None` below, once
-    every row in the batch has stopped).
+    the "K sequential forwards" cost applies per-token here, not once, so a
+    flow/ddpm AR run costs ~`n_sec * steps` model calls per physics step
+    (measured on `configs/baseline.toml`: 0.382 secondaries/step at rollout
+    time), not `k_max * steps` — see the row-compaction paragraph below.
 
-    `n_sec_pred`, if given, fixes each row's secondary count up front (as
-    resolved by `resolve_n_sec` — `n_sec.mode` in `("head", "truth")`, or a
-    stop-token decoder driven by `_assemble_stage2_ar_inputs_scheduled`'s
-    ground-truth `n_sec`, which must run the *full* `k_max`-length free-
-    running self-sample regardless of the decoder's own stop head — the
-    scheduled-sampling training contract does not truncate). This always
-    runs the full `k_max`-iteration loop, masking by the given count at the
-    end exactly as before.
+    `n_sec_pred`, if given (as resolved by `resolve_n_sec` — `n_sec.mode` in
+    `("head", "truth")`, or a stop-token decoder driven by
+    `_assemble_stage2_ar_inputs_scheduled`'s ground-truth `n_sec`, which must
+    run the *full* `k_max`-length free-running self-sample regardless of the
+    decoder's own stop head — the scheduled-sampling training contract does
+    not truncate, see `full_length` below) fixes each row's secondary count
+    up front.
 
     `n_sec_pred=None` is only valid when `sec_decoder.stop_head` is set
     (`n_sec.mode = "stop_token"`): before generating each slot's token, that
@@ -263,11 +263,35 @@ def sample_secondaries_ar(
     why this needs no extra state) decides whether generation should have
     already stopped, per `sec_decoder.n_sec_sampling` ("greedy": threshold at
     0; "sample": a Bernoulli draw at `sigmoid(logit)`). A row's own
-    `n_sec_pred` is the first slot index where this fires; once every row in
-    the batch has fired, the loop breaks before spending a model call on the
-    next slot's token — the average-case cost win the docstring above
-    describes. A row that never fires within `k_max` is capped there
-    (`K_MAX` stays a safety cap, not a modeling ceiling).
+    `n_sec_pred` is the first slot index where this fires. A row that never
+    fires within `k_max` is capped there (`K_MAX` stays a safety cap, not a
+    modeling ceiling).
+
+    **Row compaction.** A row that has already produced its `n_sec_pred`
+    tokens (or, under `stop_token`, has already fired its stop logit) has
+    nothing left to contribute — every later slot of that row is masked out
+    of `sec_valid` on return, and downstream consumers (`giant/rollout.py`,
+    `giant/cli.py`) never read it. So unless `full_length=True`, this
+    function drops such rows from the active set entirely instead of running
+    the model on them: `active_idx` starts at every row with `n_sec_pred > 0`
+    (or, under `stop_token`, every row — the first stop decision can fire at
+    slot 0) and only shrinks as rows finish, so slot `k`'s model calls cost
+    `O(active rows)` not `O(B)`. Slots a row never reaches keep their `0.0`
+    zero-init in `sec_cont`/`sec_type` — masked by `sec_valid`, identical to
+    what a full, uncompacted run would have written there before masking.
+    `AttentionHistory`'s KV cache is kept aligned to the shrinking active set
+    via `Stage2Autoregressive.select_history_cache`
+    (`giant.model.history.HistoryEncoder.select_cache`) every time the set
+    shrinks; `MarkovHistory`'s O(1) state (`prev_repr`/`remaining`, carried
+    directly rather than through a cache) is compacted the same way.
+
+    `full_length=True` disables all of the above: every row runs the full
+    `k_max`-iteration loop regardless of `n_sec_pred`/stop decisions, exactly
+    reproducing the pre-compaction behaviour. Required by
+    `_assemble_stage2_ar_inputs_scheduled`'s scheduled-sampling self-sample,
+    whose training contract needs a real prediction at every slot up to
+    `k_max` (mixed per-slot against ground truth) even past a row's own
+    `n_sec` — see that function's docstring.
 
     Under `history="attention"` the history encoding is computed once per
     slot via `Stage2Autoregressive.history_step` (a KV-cache append)
@@ -312,11 +336,6 @@ def sample_secondaries_ar(
     sec_cont = torch.zeros(B, k_max, CONT_SLOT_DIM, device=device)
     sec_type = torch.zeros(B, k_max, type_dim, device=device)
 
-    # Running per-token state, threaded from one slot to the next.
-    prev_repr = torch.zeros(B, CONT_SLOT_DIM + type_dim, device=device)
-    remaining = torch.ones(B, device=device)
-    history_cache = sec_decoder.init_history_cache()
-
     use_stop_token = n_sec_pred is None
     if use_stop_token:
         assert getattr(sec_decoder, "stop_head", None) is not None, (
@@ -324,21 +343,46 @@ def sample_secondaries_ar(
             "with no stop_head — only valid under stage2_model.n_sec.mode = "
             "'stop_token'"
         )
-        finished = torch.zeros(B, dtype=torch.bool, device=device)
         derived_n_sec = torch.full((B,), k_max, dtype=torch.long, device=device)
+        # Tracks which rows have already recorded a stop, globally by
+        # original batch index — needed even under compaction's own
+        # never-revisit guarantee, because `full_length=True` keeps every
+        # row in `active_idx` for the whole loop, so a row whose stop logit
+        # fires once but flips back below threshold at a later slot (a real
+        # possibility for an untrained/lightly-trained stop_head) must not
+        # have `derived_n_sec` overwritten by that later, spurious re-fire.
+        finished = torch.zeros(B, dtype=torch.bool, device=device)
+
+    # `active_idx`: rows still contributing tokens, indexed into the
+    # original batch. Only ever shrinks (never full_length) or stays fixed
+    # at arange(B) (full_length) — see the row-compaction docstring section.
+    active_idx = torch.arange(B, device=device)
+    if not full_length and not use_stop_token:
+        active_idx = active_idx[n_sec_pred > 0]
+
+    # Running per-token state, already compacted to `active_idx`.
+    prev_repr = torch.zeros(active_idx.numel(), CONT_SLOT_DIM + type_dim, device=device)
+    remaining = torch.ones(active_idx.numel(), device=device)
+    history_cache = sec_decoder.init_history_cache()
 
     for k in range(k_max):
-        has_prev = torch.full((B, 1), k >= 1, dtype=torch.bool, device=device)
-        history_feat = prev_repr.unsqueeze(1)  # (B, 1, CONT_SLOT_DIM + type_dim)
-        remaining_frac = remaining.unsqueeze(1)  # (B, 1)
-        slot_idx = torch.full((B, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32)
+        if active_idx.numel() == 0:
+            break
+        Bc = active_idx.numel()
+        cc = cond_cont.index_select(0, active_idx)
+        ck = cond_cat.index_select(0, active_idx)
+        s1 = stage1_out.index_select(0, active_idx)
+        has_prev = torch.full((Bc, 1), k >= 1, dtype=torch.bool, device=device)
+        history_feat = prev_repr.unsqueeze(1)  # (Bc, 1, CONT_SLOT_DIM + type_dim)
+        remaining_frac = remaining.unsqueeze(1)  # (Bc, 1)
+        slot_idx = torch.full((Bc, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32)
         hist, history_cache = sec_decoder.history_step(history_feat, has_prev, history_cache)
 
         if use_stop_token:
             stop_logit = sec_decoder.predict_stop(
-                cond_cont,
-                cond_cat,
-                stage1_out,
+                cc,
+                ck,
+                s1,
                 history_feat,
                 has_prev,
                 remaining_frac,
@@ -346,21 +390,31 @@ def sample_secondaries_ar(
                 hist=hist,
             ).squeeze(1)
             if sec_decoder.n_sec_sampling == "sample":
-                stop_now = torch.rand(B, device=device) < torch.sigmoid(stop_logit)
+                stop_now = torch.rand(Bc, device=device) < torch.sigmoid(stop_logit)
             else:
                 stop_now = stop_logit >= 0.0
-            derived_n_sec[stop_now & ~finished] = k
-            finished = finished | stop_now
-            if finished.all():
-                break
+            newly_stopped = stop_now & ~finished.index_select(0, active_idx)
+            derived_n_sec[active_idx[newly_stopped]] = k
+            finished[active_idx[stop_now]] = True
+            if not full_length:
+                keep = ~stop_now
+                active_idx = active_idx[keep]
+                cc, ck, s1 = cc[keep], ck[keep], s1[keep]
+                has_prev, remaining_frac, slot_idx = has_prev[keep], remaining_frac[keep], slot_idx[keep]
+                history_feat, hist = history_feat[keep], hist[keep]
+                history_cache = sec_decoder.select_history_cache(history_cache, keep)
+                prev_repr, remaining = prev_repr[keep], remaining[keep]
+                if active_idx.numel() == 0:
+                    break
+                Bc = active_idx.numel()
 
         if objective.is_adversarial:
-            z = torch.randn(B, 1, sec_decoder.noise_dim, device=device)
+            z = torch.randn(Bc, 1, sec_decoder.noise_dim, device=device)
             token = sec_decoder(
                 z,
-                cond_cont,
-                cond_cat,
-                stage1_out,
+                cc,
+                ck,
+                s1,
                 history_feat,
                 has_prev,
                 remaining_frac,
@@ -368,15 +422,15 @@ def sample_secondaries_ar(
                 hist=hist,
             )
         else:
-            x = torch.randn(B, 1, token_dim, device=device)
+            x = torch.randn(Bc, 1, token_dim, device=device)
             dt = 1.0 / steps
             for i in range(steps):
-                t = torch.full((B, 1), i * dt, device=device)
+                t = torch.full((Bc, 1), i * dt, device=device)
                 v = sec_decoder(
                     x,
-                    cond_cont,
-                    cond_cat,
-                    stage1_out,
+                    cc,
+                    ck,
+                    s1,
                     history_feat,
                     has_prev,
                     remaining_frac,
@@ -387,15 +441,15 @@ def sample_secondaries_ar(
                 x = x + v * dt
             token = x
 
-        token = token.squeeze(1)  # (B, token_dim)
+        token = token.squeeze(1)  # (Bc, token_dim)
         cont_k = token[:, :CONT_SLOT_DIM]
         if type_folded:
             type_k = token[:, CONT_SLOT_DIM:]
         else:
             type_k = sec_decoder.predict_type(
-                cond_cont,
-                cond_cat,
-                stage1_out,
+                cc,
+                ck,
+                s1,
                 history_feat,
                 has_prev,
                 remaining_frac,
@@ -403,8 +457,8 @@ def sample_secondaries_ar(
                 hist=hist,
             ).squeeze(1)
 
-        sec_cont[:, k] = cont_k
-        sec_type[:, k] = type_k
+        sec_cont[active_idx, k] = cont_k
+        sec_type[active_idx, k] = type_k
 
         if target == "onehot":
             type_for_history = F.one_hot(type_k.argmax(dim=-1), num_classes=type_dim).float()
@@ -415,6 +469,12 @@ def sample_secondaries_ar(
         prev_repr = torch.cat([stick_fraction.unsqueeze(-1), cont_k[:, 1:4], type_for_history], dim=-1)
         remaining = torch.clamp(remaining * (1.0 - stick_fraction), min=0.0)
 
+        if not full_length and not use_stop_token:
+            keep2 = n_sec_pred.index_select(0, active_idx) > (k + 1)
+            active_idx = active_idx[keep2]
+            prev_repr, remaining = prev_repr[keep2], remaining[keep2]
+            history_cache = sec_decoder.select_history_cache(history_cache, keep2)
+
     resolved_n_sec = derived_n_sec if use_stop_token else n_sec_pred
     sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < resolved_n_sec.unsqueeze(1)
     return sec_cont, sec_type, sec_valid

giant/training/stage2_inputs.py

@@ -291,9 +291,13 @@ def _assemble_stage2_ar_inputs_scheduled(
     skips self-sampling entirely), so callers can call this unconditionally.
 
     The free-running estimate is a REAL autoregressive self-sample —
-    `giant.sample.sample_secondaries_ar` under `torch.no_grad()` — not a
-    cheap one-step proxy, so building it costs the same `k_max` (`* steps`
-    for flow) sequential forwards `sample.py` pays at inference, EVERY batch
+    `giant.sample.sample_secondaries_ar` under `torch.no_grad()`, called here
+    with `full_length=True` — not a cheap one-step proxy, so building it
+    costs the full `k_max` (`* steps` for flow) sequential forwards for every
+    row regardless of that row's own secondary count (`full_length=True`
+    disables `sample.py`'s inference-time row compaction — see that
+    function's docstring for why: the mixing below needs a real prediction
+    at every slot up to `k_max`, not just the valid ones). Paid EVERY batch
     this is called on (paid at train time too whenever teacher_forcing !=
     "always"). Fully detached: gradient only ever flows
     through the "real" target path each stage trainer already uses
@@ -306,7 +310,7 @@ def _assemble_stage2_ar_inputs_scheduled(
 
     was_training = model.training
     sec_cont_pred, sec_type_pred, _ = sample_secondaries_ar(
-        model, cond_cont, cond_cat, stage1_ctx, n_sec, steps=sample_steps
+        model, cond_cont, cond_cat, stage1_ctx, n_sec, steps=sample_steps, full_length=True
     )
     if was_training:
         model.train()

tests/test_sample.py

@@ -341,6 +341,87 @@ def test_sample_secondaries_ar_none_n_sec_pred_without_stop_head_raises():
         sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
 
 
+# ── Row compaction: full_length=False (default, inference) must agree with
+# full_length=True (the pre-compaction behaviour, still exercised by
+# _assemble_stage2_ar_inputs_scheduled's training-time self-sample) ────────
+
+
+def _zero_randn(*size, **kwargs):
+    """Drop-in replacement for `torch.randn` that returns zeros of the same
+    shape — makes the ODE/WGAN noise deterministic so a compacted run and a
+    full_length run can be compared row-for-row regardless of how many
+    `torch.randn` calls each makes (compaction changes the batch size, and
+    therefore the RNG stream position, at every slot)."""
+    device = kwargs.get("device")
+    dtype = kwargs.get("dtype")
+    return torch.zeros(*size, device=device, dtype=dtype)
+
+
+@pytest.mark.parametrize("history", ["markov", "attention"])
+@pytest.mark.parametrize("generator", ["flow", "wgan"])
+def test_sample_secondaries_ar_compaction_matches_full_length_head_mode(generator, history, monkeypatch):
+    B, k_max, emb_dim = 4, 5, 6
+    decoder = _stage2_ar("physical", generator, emb_dim=emb_dim, k_max=k_max, history=history)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec_pred = torch.tensor([0, 1, 3, k_max])
+
+    monkeypatch.setattr(torch, "randn", _zero_randn)
+    sec_cont_c, sec_type_c, sec_valid_c = sample_secondaries_ar(
+        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2, full_length=False
+    )
+    sec_cont_f, sec_type_f, sec_valid_f = sample_secondaries_ar(
+        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2, full_length=True
+    )
+
+    assert torch.equal(sec_valid_c, sec_valid_f)
+    assert torch.equal(sec_valid_c, torch.arange(k_max).unsqueeze(0) < n_sec_pred.unsqueeze(1))
+    assert torch.allclose(sec_cont_c[sec_valid_c], sec_cont_f[sec_valid_f], atol=1e-4, rtol=1e-4)
+    assert torch.allclose(sec_type_c[sec_valid_c], sec_type_f[sec_valid_f], atol=1e-4, rtol=1e-4)
+
+
+@pytest.mark.parametrize("generator", ["flow", "wgan"])
+def test_sample_secondaries_ar_compaction_matches_full_length_stop_token(generator, monkeypatch):
+    """`n_sec_sampling="greedy"` keeps the stop decision itself deterministic
+    (no `torch.rand` draw), so only `torch.randn` needs zeroing."""
+    B, k_max = 6, 5
+    decoder = _stage2_ar_stop_token("physical", generator, n_sec_sampling="greedy", k_max=k_max)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+
+    monkeypatch.setattr(torch, "randn", _zero_randn)
+    sec_cont_c, sec_type_c, sec_valid_c = sample_secondaries_ar(
+        decoder, cond_cont, cond_cat, stage1_out, None, steps=2, full_length=False
+    )
+    sec_cont_f, sec_type_f, sec_valid_f = sample_secondaries_ar(
+        decoder, cond_cont, cond_cat, stage1_out, None, steps=2, full_length=True
+    )
+
+    assert torch.equal(sec_valid_c, sec_valid_f)
+    assert torch.allclose(sec_cont_c[sec_valid_c], sec_cont_f[sec_valid_f], atol=1e-4, rtol=1e-4)
+    assert torch.allclose(sec_type_c[sec_valid_c], sec_type_f[sec_valid_f], atol=1e-4, rtol=1e-4)
+
+
+def test_sample_secondaries_ar_full_length_ignores_n_sec_pred_zero_rows():
+    """A row with n_sec_pred == 0 would be dropped from the active set at
+    slot 0 under compaction (full_length=False) — full_length=True must
+    still run the model for it at every slot (only masked by sec_valid at
+    the end), matching _assemble_stage2_ar_inputs_scheduled's contract."""
+    B, k_max, emb_dim = 3, 4, 6
+    decoder = _stage2_ar("physical", "flow", emb_dim=emb_dim, k_max=k_max)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec_pred = torch.tensor([0, 0, 0])
+    sec_cont, sec_type, sec_valid = sample_secondaries_ar(
+        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2, full_length=True
+    )
+    assert not sec_valid.any()
+    # every slot still ran the model (not left at the zero-init default) —
+    # a real flow ODE output from randn-initialized noise is essentially
+    # never exactly zero.
+    assert not torch.allclose(sec_cont, torch.zeros_like(sec_cont))
+
+
 # ── resolve_n_sec: n_sec.mode = "head" sampling policy (gitea #86) ──────────