Add long-span recency regression

Author: physicsrob

modal_test.py

@@ -25,7 +25,9 @@
 # own container; everything else is batched together.
 # New test files are caught by the catch-all shard automatically.
 
-_HEAVY_FILES: list[str] = []
+_HEAVY_FILES: list[str] = [
+    "tests/compile/forward/test_most_recent_long_span.py",
+]
 
 # Each inner list becomes one container.  Splitting medium files across
 # multiple groups keeps any single shard from dominating wall time.

tests/compile/forward/test_most_recent_long_span.py

@@ -0,0 +1,93 @@
+"""Long-span regression for ``attend_most_recent_matching``.
+
+Plan C guards the DOOM renderer's long recency reads.  The important
+behavior is the compiled attention backend, not the small oracle
+``Attn.compute`` path: at 32,768 positions the oracle would build a
+large CPU attention matrix, while the compiled path uses the same SDPA
+backend the real graph uses.
+"""
+
+import pytest
+import torch
+
+from torchwright.compiler.device import get_device
+from torchwright.compiler.forward.compile import forward_compile
+from torchwright.ops.attention_ops import attend_most_recent_matching
+from torchwright.ops.inout_nodes import create_input, create_pos_encoding
+
+MATCH_GAIN_LONG = 300_000.0
+SPANS = (8_500, 32_768)
+D = 64
+D_HEAD = 16
+
+
+@pytest.fixture(scope="module")
+def compiled_long_span_pick():
+    device = get_device(verbose=False)
+    if device.type != "cuda":
+        pytest.skip("long-span recency regression requires the CUDA backend")
+
+    assert torch.get_default_dtype() == torch.float32
+    assert torch.get_float32_matmul_precision() == "highest"
+    assert not torch.backends.cuda.matmul.allow_tf32
+
+    pos = create_pos_encoding()
+    query = create_input("query", 1, value_range=(0.0, 1.0))
+    key = create_input("key", 1, value_range=(0.0, 1.0))
+    value = create_input("value", 1, value_range=(0.0, float(max(SPANS))))
+    out = attend_most_recent_matching(
+        pos,
+        query,
+        key,
+        value,
+        match_gain=MATCH_GAIN_LONG,
+    )
+
+    net = forward_compile(
+        d=D,
+        d_head=D_HEAD,
+        output_node=out,
+        pos_encoding=pos,
+        verbose=False,
+        max_layers=10,
+    )
+    assert net.device.type == "cuda"
+    return net, out
+
+
+def _run_case(compiled_long_span_pick, n_pos: int, key: torch.Tensor) -> torch.Tensor:
+    net, out = compiled_long_span_pick
+    inputs = {
+        "query": torch.ones(n_pos, 1),
+        "key": key,
+        "value": torch.arange(float(n_pos)).unsqueeze(1),
+    }
+    return net.compute(n_pos, inputs)[out].detach().cpu().squeeze(1)
+
+
+@pytest.mark.parametrize("n_pos", SPANS)
+def test_long_span_dense_adjacent_matches_pick_most_recent(
+    compiled_long_span_pick,
+    n_pos: int,
+) -> None:
+    """Dense adjacent matches still resolve the one-position recency gap."""
+    result = _run_case(compiled_long_span_pick, n_pos, torch.ones(n_pos, 1))
+
+    for probe in (n_pos // 2, n_pos - 1):
+        assert result[probe].item() == pytest.approx(float(probe), abs=0.05)
+
+
+@pytest.mark.parametrize("n_pos", SPANS)
+def test_long_span_sparse_match_beats_recent_non_matches(
+    compiled_long_span_pick,
+    n_pos: int,
+) -> None:
+    """At ``300_000`` gain, one old unit match beats newer non-matches."""
+    match_pos = n_pos // 3
+    key = torch.zeros(n_pos, 1)
+    key[match_pos, 0] = 1.0
+
+    result = _run_case(compiled_long_span_pick, n_pos, key)
+
+    for probe in (match_pos, match_pos + 1, n_pos // 2, n_pos - 1):
+        assert result[probe].item() == pytest.approx(float(match_pos), abs=0.05)

torchwright/ops/attention_ops.py

@@ -1224,8 +1224,11 @@ def attend_most_recent_matching(
     Quick sizing guide:
 
     - Unit dot products (one-hot match, ``match_dot=1`` on match,
-      ``0`` off) at ``max_n_pos ≈ 2000``: pick
-      ``match_gain ≥ 20000 + margin`` (e.g. ``30000``).
+      ``0`` off): the safe span is roughly
+      ``match_gain / _QUERY_GAIN``.  At ``match_gain = 300_000`` this
+      covers spans below ``37,500`` positions, including the DOOM
+      renderer's measured ~8,500-position rollout and the 32,768-position
+      regression test, but not a 65,536-position hard cap.
     - 10×-scaled E8 codes (``match_dot=1600``, worst off-diagonal
       ``~800``) at ``max_n_pos ≈ 2000``: the ``800``-unit dot gap gives
       plenty of headroom even at the default ``match_gain = 200``
@@ -1239,20 +1242,16 @@ def attend_most_recent_matching(
     even dense matches resolve cleanly.  Set ``assert_hardness_gt`` if
     you need runtime enforcement.
 
-    **TF32 caveat.**  On Ampere GPUs (A100), PyTorch's default matmul
-    path uses TF32 (~10-bit mantissa, ~1e-3 relative precision).  When
-    ``match_gain · match_dot`` is large and ``_QUERY_GAIN · max_n_pos``
-    is also large, the sum can exceed the regime where TF32 resolves
-    the recency-tiebreak gap cleanly.  For example, at match_gain=20
-    with 10×-scaled E8 codes (match_dot=1600) and max_n_pos≈1000, the
-    peak logit is ~30000, where TF32 absolute precision ~30 eats the
-    `_QUERY_GAIN = 8` unit-position gap.  If your callsite hits that
-    regime, either (a) use smaller match vectors so the match
-    contribution is ≲ a few thousand, or (b) switch to
-    :func:`attend_argmax_where` with an explicit integer score if one
-    is available at the callsite — its logits stay around
-    ``_VALIDITY_DIRECT ≈ 1000`` where integer-score gaps resolve
-    cleanly on TF32.  (M3 in the DOOM renderer chose option (b).)
+    **Precision policy.**  Torchwright's current oracle path uses fp32
+    ``torch.matmul`` / ``torch.softmax`` and the compiled path routes
+    SDPA through the MATH backend with fp32 tensors.  PyTorch's default
+    CUDA TF32 matmul flag is also off in the supported environment.
+    Under that policy, ``match_gain = 300_000`` leaves ample precision
+    headroom for the 8-logit adjacent-position recency gap at the spans
+    above.  If a future global precision setting enables TF32, re-run
+    the long-span regression before relying on this gain; TF32's
+    coarser mantissa can erase the 8-logit tiebreak when the total logit
+    is in the hundreds of thousands.
 
     **When no position matches.**  If no causal-window position has
     ``query_vector · key_vector`` above the unmatched baseline, the