Merge origin/main: upstream attn forward_cached fix into Plan A/B/C work

Author: physicsrob

tests/compile/forward/test_batched_decode_with_past.py

@@ -0,0 +1,89 @@
+"""Batched decode with non-empty past — the speculative-decoding shape.
+
+When ``forward_cached`` is called with ``n_new > 1`` and a non-empty past
+(e.g. K+1 batched rows in a spec-decode step), each new row must:
+
+1. Attend unconditionally to every past key — those are verified history.
+2. Attend causally among new rows — row ``i`` sees new rows ``0..i`` only.
+
+A bug in either direction silently corrupts the rollout. This test pins
+both: drive a sequential per-row rollout to length L, then run one
+batched ``K+1``-row decode from past_len=L and check row-by-row that each
+batched-row output equals the corresponding sequential-step output.
+"""
+
+import pytest
+import torch
+
+from torchwright.compiler.forward.compile import forward_compile
+
+D = 1024
+D_HEAD = 16
+
+
+@pytest.fixture(scope="module")
+def compiled_calc():
+    from examples.calculator_v2 import create_network_parts
+
+    output_node, pos_encoding, embedding = create_network_parts(1)
+    net = forward_compile(
+        d=D,
+        d_head=D_HEAD,
+        output_node=output_node,
+        pos_encoding=pos_encoding,
+        verbose=False,
+    )
+    return net, output_node, embedding
+
+
+def _row_input(net, tokens, t):
+    """Build the t-th row of the input residual stream for ``tokens``."""
+    full_res = net.get_input_res_stream(t + 1, {"embedding_input": tokens[: t + 1]})
+    return full_res[t : t + 1].to(net.device)
+
+
+def test_batched_decode_with_past_matches_sequential(compiled_calc):
+    """K+1 batched rows from past_len=L equal L+1..L+K+1 sequential rows."""
+    net, _, _ = compiled_calc
+    tokens = ["<bos"] + list("3+5+9\n")
+    seed_len = 3  # past_len at the start of the batched step
+    batch_size = len(tokens) - seed_len  # K+1 rows in the batched step
+    assert batch_size >= 2, "test requires n_new >= 2 to exercise the new mask"
+
+    # Stage 1: sequential rollout up to seed_len, snapshot the past.
+    kvs = None
+    for t in range(seed_len):
+        kvs, _ = _step_one(net, tokens, t, kvs)
+    past_kvs_seed = [(K.clone(), V.clone()) for (K, V) in kvs]
+
+    # Stage 2a: continue sequentially to capture per-row reference outputs.
+    seq_outputs = []
+    kvs_seq = [(K.clone(), V.clone()) for (K, V) in past_kvs_seed]
+    for t in range(seed_len, len(tokens)):
+        kvs_seq, out = _step_one(net, tokens, t, kvs_seq)
+        seq_outputs.append(out.detach().clone())
+
+    # Stage 2b: run the same suffix as one batched call from the snapshot.
+    suffix_inp = torch.cat(
+        [_row_input(net, tokens, t) for t in range(seed_len, len(tokens))],
+        dim=0,
+    )
+    batched_out, _ = net.forward_cached(suffix_inp, past_kvs_seed)
+
+    # Bit-exact row-by-row equality. SDPA forces the MATH backend (see
+    # _SDPA_BACKEND in components/attn.py), and the matmul shapes for a
+    # given row are identical in the two paths once the mask is correct.
+    for i, expected in enumerate(seq_outputs):
+        actual = batched_out[i : i + 1]
+        diff = (actual - expected).abs().max().item()
+        assert diff == 0.0, (
+            f"row {i}: sequential vs batched diverged by {diff}; "
+            f"expected bit-exact equality"
+        )
+
+
+def _step_one(net, tokens, t, kvs):
+    """Run one cached single-row forward and return (new_kvs, output)."""
+    new_inp = _row_input(net, tokens, t)
+    out, new_kvs = net.forward_cached(new_inp, kvs)
+    return new_kvs, out

torchwright/compiler/components/attn.py

@@ -106,23 +106,44 @@ def forward_cached(
 
         n_new = inp.shape[0]
         n_total = K.shape[1]
-
-        # is_causal=True  → pure prefill (no past): local indices == absolute
-        #                   positions; upper-triangular mask is correct.
-        # is_causal=False → decode (has past): all K entries are strictly in
-        #                   the past relative to Q; no future positions to mask.
+        n_past = n_total - n_new
+
+        # Three regimes:
+        #   pure prefill (no past, n_new == n_total): is_causal=True gives
+        #     the standard upper-triangular mask.
+        #   single-step decode (n_new == 1, n_past > 0): every K entry is in
+        #     the past or is the lone new row, so no mask is needed.
+        #   batched decode with past (n_new > 1 and n_past > 0): every new
+        #     row sees the full past unconditionally, but among new rows the
+        #     mask is lower-triangular (row i sees new rows 0..i). This is
+        #     the speculative-decoding shape — is_causal=True would mask out
+        #     past columns for early rows; is_causal=False would let row 0
+        #     attend to future new rows.
         # scale=1.0 preserves the raw dot-product magnitudes that all attention
         # weights were compiled against (no 1/sqrt(d_head) rescaling).
-        with sdpa_kernel(_SDPA_BACKEND):
-            weighted = F.scaled_dot_product_attention(
-                Q.unsqueeze(0),
-                K.unsqueeze(0),
-                V.unsqueeze(0),
-                is_causal=(n_new == n_total),
-                scale=1.0,
-            ).squeeze(
-                0
-            )  # (n_heads, n_new, d_head)
+        if n_new > 1 and n_past > 0:
+            past_visible = torch.ones(n_new, n_past, dtype=torch.bool, device=Q.device)
+            new_block = torch.ones(
+                n_new, n_new, dtype=torch.bool, device=Q.device
+            ).tril()
+            attn_mask = torch.cat([past_visible, new_block], dim=1)
+            with sdpa_kernel(_SDPA_BACKEND):
+                weighted = F.scaled_dot_product_attention(
+                    Q.unsqueeze(0),
+                    K.unsqueeze(0),
+                    V.unsqueeze(0),
+                    attn_mask=attn_mask,
+                    scale=1.0,
+                ).squeeze(0)
+        else:
+            with sdpa_kernel(_SDPA_BACKEND):
+                weighted = F.scaled_dot_product_attention(
+                    Q.unsqueeze(0),
+                    K.unsqueeze(0),
+                    V.unsqueeze(0),
+                    is_causal=(n_new == n_total),
+                    scale=1.0,
+                ).squeeze(0)  # (n_heads, n_new, d_head)
 
         output = torch.einsum("hpk,hkd->pd", weighted, self.output_matrix)
 

torchwright/compiler/export.py

@@ -1176,6 +1176,34 @@ def step(
     ) -> tuple:
         """Cached forward step.
 
+        Supports three ``n_new`` regimes uniformly:
+
+        * Pure prefill (``past`` empty, ``n_new == seq_len``).
+        * Single-step decode (``n_new == 1`` with a non-empty ``past``).
+        * **Batched decode with past** (``n_new > 1`` with a non-empty
+          ``past``) — the speculative-decoding shape. Each new row sees
+          the entire past unconditionally, plus a lower-triangular slice
+          of the new block (row ``i`` sees new rows ``0..i``). Row ``i``
+          of the returned outputs is bit-identical to the same row from
+          a sequential rollout starting from the same past.
+
+        Speculative-decoding partial-commit pattern. After running with
+        ``n_new = K + 1`` (one current input + ``K`` drafts + bonus), a
+        caller may decide to commit only the first ``commit_count``
+        rows. Slice the returned ``new_past`` directly::
+
+            outputs, (new_K, new_V) = compiled.step(inputs_batch, past)
+            # ... compare outputs[:K] to the K drafts, decide commit_count ...
+            target_len = past[0][0].shape[1] + commit_count
+            committed = (
+                tuple(K[:, :target_len] for K in new_K),
+                tuple(V[:, :target_len] for V in new_V),
+            )
+
+        ``committed`` is the past for the next step. The discarded rows
+        (``commit_count..n_new``) carry no state outside ``new_past``,
+        so the slice is the only commit primitive needed.
+
         Args:
             inputs: ``(n_new, d_input)`` float tensor for the new rows.
             past: ``(past_K_tuple, past_V_tuple)`` from a prior step or