From 5c576fa8f3f88b2f481f8084ee5567237c666762 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Thu, 3 Sep 2026 17:56:59 +0200 Subject: [PATCH] perf: compact Stage-2 AR inference loop to active rows only MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit sample_secondaries_ar ran all k_max=15 slots for every row regardless of each row's own predicted secondary count, even though the baseline checkpoint's rollout measured only 0.382 secondaries/step — so ~97% of stage-2 model calls generated tokens sec_valid then masked away. Compact the loop to the still-active row set at each slot: drop a row the moment its n_sec_pred is exhausted (or, under n_sec.mode="stop_token", the moment its own stop logit fires), so slot k's model calls cost O(active rows) instead of O(B). Exact — rows are independent given their own history — verified by comparing the compacted path against a new full_length=True escape hatch that reproduces the original uncompacted behavior bit-for-bit under deterministic noise. full_length=True is 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 regardless of a row's own count, so training behavior is unchanged. AttentionHistory's KV cache and MarkovHistory's O(1) state are kept aligned to the shrinking active set via a new HistoryEncoder.select_cache / Stage2Autoregressive.select_history_cache. Also fixes a latent bug the refactor surfaced: derived_n_sec (stop-token mode) could be overwritten by a later spurious re-fire of the stop logit on a row that had already stopped; now tracked via an explicit `finished` mask so only the first stop slot is recorded, matching the documented contract. No architecture or checkpoint-format change — every existing v0.3.0 Stage2Autoregressive checkpoint (flow/wgan, markov/attention, head/stop_token) picks up the speedup automatically on its next `giant rollout`/`giant predict`, no retraining needed. Measured (CPU, hidden_dim=512/6 blocks, k_max=15, batch 512, mean n_sec≈0.38 matching the baseline checkpoint's own rollout): 17.6-22.9x fewer wall-clock seconds for the AR loop alone (attention/markov history respectively). Directional only — baseline.toml's GPU inference-cost comment is updated accordingly, flagged stale pending a real rollout re-measurement via eval_cost_per_step. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01HPt7bVLZYFJe5cG6V7ahqC --- configs/baseline.toml | 25 ++++- giant/model/history.py | 22 +++++ giant/model/models.py | 8 ++ giant/sample.py | 160 ++++++++++++++++++++++---------- giant/training/stage2_inputs.py | 12 ++- tests/test_sample.py | 81 ++++++++++++++++ 6 files changed, 249 insertions(+), 59 deletions(-) diff --git a/configs/baseline.toml b/configs/baseline.toml index 0f0a044..4890ef3 100644 --- a/configs/baseline.toml +++ b/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 diff --git a/giant/model/history.py b/giant/model/history.py index d000329..072bc2e 100644 --- a/giant/model/history.py +++ b/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] diff --git a/giant/model/models.py b/giant/model/models.py index 0e26323..6b60906 100644 --- a/giant/model/models.py +++ b/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, diff --git a/giant/sample.py b/giant/sample.py index c0ac973..28929c2 100644 --- a/giant/sample.py +++ b/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 diff --git a/giant/training/stage2_inputs.py b/giant/training/stage2_inputs.py index 342db88..25c7c09 100644 --- a/giant/training/stage2_inputs.py +++ b/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() diff --git a/tests/test_sample.py b/tests/test_sample.py index 345a656..266a11f 100644 --- a/tests/test_sample.py +++ b/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) ────────── -- 2.39.5