fix(eval): structured-mode greedy decode capped at in-range layout bound (369) — probe-window traces reach 633 tokens and were silently truncated; cap now keys to positional table (1024) so OOD traces decode fully. Regression-verified vs stub model.

1 files changed, 9 insertions, 1 deletions

diff --git a/src/model_api.py b/src/model_api.py
index 07d5d8a..c82baf4 100644
--- a/src/model_api.py
+++ b/src/model_api.py
@@ -41,7 +41,15 @@ def greedy_decode(model: PrimeModel, x: torch.Tensor, cfg: Config, max_len: int
     x = pad_inputs(x, cfg)
     from src.data import _global_lengths
     _, out_max = _global_lengths(cfg)
-    max_len = max_len or max(getattr(cfg, "max_out_len", 6), out_max)
+    if max_len is None:
+        max_len = max(getattr(cfg, "max_out_len", 6), out_max)
+        if cfg.scratch_mode == "structured":
+            # Probe-window traces ([range_end+1, +1000]) exceed the in-range layout
+            # bound (measured: 633-token trace at n=1327 vs out_max=369). Cap at the
+            # model's positional table instead so OOD traces decode fully.
+            pos = getattr(model, "pos", None)
+            if pos is not None:
+                max_len = max(max_len, pos.num_embeddings - x.shape[1] - 1)
     B = x.shape[0]
     y_in = torch.full((B, 1), cfg.eos_id, dtype=torch.long, device=device)
     for _ in range(max_len):