Redesign PosEncoding: counter in its own column, exact 1.0 query constant

Move the raw position counter out of the slowest-sin slot into its own
last column (counter_col = d_pos-1), leaving a complete sin/cos block
(trig_width = d_pos-1) and requiring odd d_pos. This decouples the
encoding width from d_head and removes the slow-cosine "almost one"
that several attention primitives leaned on.

- pos_encoding: add trig_width/trig_slice/counter_col; reject even d_pos;
  rename get_pos_delta_matrix -> trig_shift_matrix (row convention, key
  matrix numerically identical to before); drop slow_sin_freq and the
  dead get_position_scalar max_pos arg.
- get_prev_value: rebuild on an exact LiteralValue([1.0]) query plus
  raw-counter recency, documenting the TF32 max_pos=256 budget. The old
  query "constant" was the sum of every sinusoid at the query position,
  which is negative at ~19% of the first 2048 positions and inverted the
  recency ordering there.
- attention_ops: switch all query-side constants from the slow cosine
  (now overwritten by the counter) to LiteralValue([1.0]); shrink the
  selection heads' d_qk to 1-2 so it no longer scales with d_pos, and
  auto-split value pass-through (drop the len(value) <= d_pos cap).
- compiler: auto-create PosEncoding(17); validate d_head >= trig_width
  early with an actionable error; copy heads zero the counter column
  explicitly (trig-only self-match); affine rule reads counter_col.

create_pos_encoding now returns PosEncoding(17) (16 trig columns + 1
counter), whose trig_width matches the default d_head=16.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: physicsrob

tests/compile/forward/test_compiler_assertions.py

@@ -125,7 +125,7 @@ def test_C_literal_write_rejects_extra_target_cols():
 def test_B_attn_rejects_v_source_cols_wrong_length():
     """_write_compute_attn raises when V source_cols length doesn't match
     value_in width."""
-    pos = PosEncoding(d_pos=D_HEAD)
+    pos = PosEncoding(17)
     v_in = InputNode("v", 4, value_range=(-100.0, 100.0))
     node = Attn(
         query_in=pos,
@@ -155,7 +155,7 @@ def test_B_attn_rejects_v_source_cols_wrong_length():
 
 
 def test_B_attn_rejects_q_source_cols_wrong_length():
-    pos = PosEncoding(d_pos=D_HEAD)
+    pos = PosEncoding(17)
     v_in = InputNode("v", 4, value_range=(-100.0, 100.0))
     node = Attn(
         query_in=pos,
@@ -198,7 +198,7 @@ def test_A_end_of_layer_catches_freed_too_early(monkeypatch):
     weight = torch.eye(4, 4)
     bias = torch.zeros(4)
     l = Linear(x, weight, bias)
-    pos = PosEncoding(d_pos=D_HEAD)
+    pos = PosEncoding(17)
 
     graph = GraphAnalyzer(l)
     rmap = ResidualStreamMap(D)
@@ -225,7 +225,7 @@ def test_A_require_live_raises_for_unallocated_input():
     op context before the downstream KeyError."""
     from torchwright.compiler.forward.scheduler import LayerScheduler
 
-    pos = PosEncoding(d_pos=D_HEAD)
+    pos = PosEncoding(17)
     x = InputNode("x", 4, value_range=(-100.0, 100.0))
     weight = torch.eye(4, 4)
     bias = torch.zeros(4)
@@ -275,7 +275,7 @@ def test_scheduler_pins_never_marks_pinned_dead():
     even when its effective consumers have all been computed."""
     from torchwright.compiler.forward.scheduler import LayerScheduler
 
-    pos = PosEncoding(d_pos=D_HEAD)
+    pos = PosEncoding(17)
     x = InputNode("x", 4, value_range=(-100.0, 100.0))
     weight = torch.eye(4, 4)
     bias = torch.zeros(4)
@@ -306,22 +306,25 @@ def test_delta_transfer_guard_catches_reallocation():
         _verify_overlay_target_protection,
     )
 
-    pos = PosEncoding(d_pos=D_HEAD)
+    pos = PosEncoding(17)
     x = InputNode("x", 4, value_range=(-1.0, 1.0))
     bait = InputNode("bait", 1, value_range=(-1.0, 1.0))
     weight = torch.eye(4, 4)
     bias = torch.zeros(4)
     y = Linear(x, weight, bias)
 
     rmap = ResidualStreamMap(D)
-    rmap.allocate(pos)  # [0..15]
-    rmap.allocate(x)  # [16..19]
-    rmap.allocate(bait)  # [20]
+    rmap.allocate(pos)
+    rmap.allocate(x)
+    rmap.allocate(bait)
+    bait_col = rmap.get_indices(bait)[0]
 
-    # Poison: overlay's target col 20 is owned by `bait`, which is neither
+    # Poison: overlay's target col is owned by `bait`, which is neither
     # pos_encoding nor a pinned input.  The guard must fire.
-    overlays = {y: (x, [20])}
-    with pytest.raises(AssertionError, match=r"Overlay target column 20 is owned by"):
+    overlays = {y: (x, [bait_col])}
+    with pytest.raises(
+        AssertionError, match=rf"Overlay target column {bait_col} is owned by"
+    ):
         _verify_overlay_target_protection(
             overlays, rmap, pos_encoding=pos, overlay_pinned_inputs=set()
         )

tests/compile/forward/test_cpsat_chain_overlap.py

@@ -69,7 +69,7 @@ def test_chain_detector_no_overlap_after_fusion():
         f"{L_b.inputs[0]!r}"
     )
 
-    pos = PosEncoding(8)
+    pos = PosEncoding(9)
     gm = build_graph_model(L_c, pos)
 
     # The structural assertion in build_graph_model would have raised
@@ -98,7 +98,7 @@ def test_chain_detector_picks_upstream_chain():
     x, L_pre, R, L_a, L_b, R_prime, L_c = _build_post_fusion_repro()
     fuse_consecutive_linears({L_c})
 
-    pos = PosEncoding(8)
+    pos = PosEncoding(9)
     gm = build_graph_model(L_c, pos)
 
     # Exactly one chain — the upstream one.  The downstream chain
@@ -119,7 +119,7 @@ def test_cpsat_solves_post_fusion_repro():
     x, L_pre, R, L_a, L_b, R_prime, L_c = _build_post_fusion_repro()
     fuse_consecutive_linears({L_c})
 
-    pos = PosEncoding(8)
+    pos = PosEncoding(9)
     assignment, stats = solve_schedule(
         L_c, pos,
         d=64, d_head=8, d_hidden=128,
@@ -148,7 +148,7 @@ def test_heuristic_chain_detector_no_overlap_after_fusion():
     x, L_pre, R, L_a, L_b, R_prime, L_c = _build_post_fusion_repro()
     fuse_consecutive_linears({L_c})
 
-    pos = PosEncoding(8)
+    pos = PosEncoding(9)
     graph = GraphAnalyzer(L_c)
     scheduler = LayerScheduler(graph, d=64, d_head=8, pos_encoding=pos)
     # Build a ready set that includes both L_pre and L_b — the case

tests/compile/forward/test_d_hidden_decoupling.py

@@ -11,14 +11,16 @@
 from torchwright.compiler.export import compile_headless
 from torchwright.compiler.forward.compile import forward_compile
 from torchwright.compiler.groups.mlp_sublayer import MLPSubLayer
-from torchwright.ops.inout_nodes import create_input, create_pos_encoding
+from torchwright.graph import PosEncoding
+from torchwright.ops.inout_nodes import create_input
 from torchwright.ops.linear_relu_linear import linear_relu_linear
 
 
 def _build_relu_chain_graph(d_input: int, d_hidden_chain: int, d_output: int):
     """Build a tiny graph: input -> Linear -> ReLU -> Linear with random weights."""
     inp = create_input("x", d_input)
-    pos = create_pos_encoding()
+    # These tests compile at d_head=8, so the encoding's trig_width must be 8.
+    pos = PosEncoding(9)
     torch.manual_seed(0)
     input_proj = torch.randn(d_hidden_chain, d_input)
     input_bias = torch.randn(d_hidden_chain)

tests/compile/forward/test_forward_compile.py

@@ -602,7 +602,8 @@ def test_compile_split_vo(d_qk, d_v):
 
 def test_compile_split_vo_large_ratio():
     """d_v=64, d_head=8 — 8 V/O heads with tiny d_qk=2."""
-    pos = create_pos_encoding()
+    # d_head=8 requires trig_width <= 8, i.e. PosEncoding(9).
+    pos = PosEncoding(9)
     x = create_input("x", 8)
     out = _build_attn(x, x, x, d_qk=2, d_v=64, d_out=8)
 
@@ -622,6 +623,23 @@ def test_compile_split_vo_large_ratio():
     ), f"Max diff: {(actual.cpu() - expected).abs().max().item():.6f}"
 
 
+def test_compile_rejects_d_head_below_trig_width():
+    """forward_compile rejects d_head < pos_encoding.trig_width early, before
+    the copy heads would silently match on only part of the trig grid."""
+    pos = PosEncoding(17)  # trig_width = 16
+    x = create_input("x", 4)
+    out = pos.attend_to_offset(x, delta_pos=-1)
+    with pytest.raises(ValueError, match="trig_width"):
+        forward_compile(
+            d=256,
+            d_head=8,
+            output_node=out,
+            pos_encoding=pos,
+            verbose=False,
+            max_layers=10,
+        )
+
+
 def test_compile_dqk_equals_dv_unchanged():
     """d_qk == d_v == 8 — backward compat, single head, padded to d_head=16."""
     x = create_input("x", 8)

tests/compile/forward/test_scheduler.py

@@ -31,7 +31,7 @@
 
 
 def _make_pos_encoding():
-    return PosEncoding(d_pos=D_HEAD)
+    return PosEncoding(17)  # trig_width 16 == D_HEAD
 
 
 def _make_linear(inp, d_out, name=""):
@@ -476,14 +476,14 @@ def test_mlp_slot_exhaustion():
 def test_schedule_under_column_pressure():
     """Stream full with dead nodes: scheduler cancels to make room for new computes.
 
-    Setup: D=64, pos=16, filler=40, x=4, a=4 → 0 free.
+    Setup: D=64, pos=17, filler=39, x=4, a=4 → 0 free.
     x is dead (consumer a is computed). Relu chain needs 3 output cols.
     Scheduler must cancel x to free space, then schedule the chain.
     """
     pos = _make_pos_encoding()
     filler = InputNode(
-        "filler", D - D_HEAD - 8, value_range=(-100.0, 100.0)
-    )  # 40 cols, not in graph
+        "filler", D - len(pos) - 8, value_range=(-100.0, 100.0)
+    )  # fills the stream to 0 free alongside pos + x + a
     x = InputNode("x", 4, value_range=(-100.0, 100.0))
     a = _make_linear(x, 4, "a")
     l2, r, l1 = _make_relu_chain(a, 8, 3, "out")
@@ -882,9 +882,9 @@ def test_no_progress_raises_error():
 
     graph = GraphAnalyzer(out)
     # Tiny d: just enough for pos + a + b, nothing spare
-    small_d = D_HEAD + 4 + 4  # 24
+    small_d = len(pos) + 4 + 4  # 17 + 4 + 4
     rmap = ResidualStreamMap(small_d)
-    rmap.allocate(pos)  # 16 cols
+    rmap.allocate(pos)  # len(pos) cols
     rmap.allocate(a)  # 4 cols
     rmap.allocate(b)  # 4 cols
     assert rmap.get_free_count() == 0

tests/compile/forward/test_weight_writer.py

@@ -36,7 +36,7 @@
 
 
 def _make_pos_encoding():
-    return PosEncoding(d_pos=D_HEAD)
+    return PosEncoding(17)  # trig_width 16 == D_HEAD
 
 
 def _build_residual_stream(
@@ -264,11 +264,13 @@ def test_attn_compute_multiposition():
     cond_in = InputNode("c", 1, value_range=(-100.0, 100.0))
 
     attn_node = pos.get_prev_value(value_in, cond_in)
+    query_one = attn_node.inputs[0]  # LiteralValue([1.0]) query constant
 
     rmap = ResidualStreamMap(D)
     rmap.allocate(pos)
     rmap.allocate(value_in)
     rmap.allocate(cond_in)
+    rmap.allocate(query_one)
     out_cols = rmap.allocate(attn_node)
 
     layer = TransformerLayer(D, D_HEAD, pos)
@@ -287,15 +289,15 @@ def test_attn_compute_multiposition():
     c_values = torch.tensor([[1.0], [0.0], [0.0], [1.0]])
     pe_values = pos.compute(N_POS, {})
 
-    # get_prev_value needs Concatenate([pos, cond]) as key_in
-    # We need to also place the concatenated node's children
-    concat_node = attn_node.inputs[1]  # key_in is Concatenate([pos, cond])
+    # get_prev_value's key_in is Concatenate([cond, pos]); the query is an
+    # exact 1.0 literal.  Place every leaf/constant the head reads.
     res = _build_residual_stream(
         rmap,
         {
             pos: pe_values,
             value_in: v_values,
             cond_in: c_values,
+            query_one: query_one.compute(N_POS, {}),
         },
     )
 

tests/graph/test_affine_bounds.py

@@ -724,16 +724,16 @@ def test_pos_encoding_identity_bound(self):
         """PosEncoding produces an identity A-matrix with [-1, 1] per column."""
         from torchwright.graph import PosEncoding
 
-        pe = PosEncoding(d_pos=8)
+        pe = PosEncoding(d_pos=9)
         ab = pe.affine_bound
         assert ab.n_cols > 0, "PosEncoding must have non-degenerate bound"
-        assert ab.d_output == 8
+        assert ab.d_output == 9
         assert pe.node_id in ab.columns
         assert pe.node_id in ab.input_ranges
-        assert torch.equal(ab.A_lo, torch.eye(8, dtype=torch.float64))
+        assert torch.equal(ab.A_lo, torch.eye(9, dtype=torch.float64))
         intervals = ab.to_interval()
         for i, iv in enumerate(intervals):
-            if i == 6:  # d_pos - 2: raw position counter
+            if i == pe.counter_col:  # last column: raw position counter
                 assert iv.lo == pytest.approx(0.0)
                 assert iv.hi == pytest.approx(100000.0)
             else:
@@ -746,14 +746,14 @@ def test_attn_propagates_value_range(self):
         with fresh_graph_session():
             from torchwright.graph import PosEncoding, Attn
 
-            pe = PosEncoding(d_pos=8)
+            pe = PosEncoding(d_pos=9)
             value = LiteralValue(torch.tensor([2.0, 3.0]))
             attn = Attn(
                 query_in=pe,
                 key_in=pe,
                 value_in=value,
-                query_matrix=torch.eye(8, 2),
-                key_matrix=torch.eye(8, 2),
+                query_matrix=torch.eye(9, 2),
+                key_matrix=torch.eye(9, 2),
                 value_matrix=torch.eye(2),
                 output_matrix=torch.eye(2),
             )

tests/graph/test_pos_encoding.py

@@ -1,24 +1,42 @@
 from torchwright.graph import PosEncoding, InputNode
-from torchwright.graph.pos_encoding import get_pos_delta_matrix
+from torchwright.graph.pos_encoding import trig_shift_matrix
+import pytest
 import torch
 
 atol = 1.0e-4  # Absolute tolerance for comparing tensors
 
-d_pos = 32
+d_pos = 33
 N_pos = 16
 
 pos_encoding_node = PosEncoding(d_pos)
 pos_encoding = pos_encoding_node.compute(n_pos=N_pos, input_values={})
 
 
+def test_layout_properties():
+    pe = PosEncoding(17)
+    assert pe.trig_width == 16
+    assert pe.counter_col == 16
+    assert pe.trig_slice == slice(0, 16)
+
+
+def test_even_d_pos_rejected():
+    with pytest.raises(ValueError):
+        PosEncoding(16)
+
+
+def test_counter_column_is_exact_position():
+    counter = pos_encoding[:, pos_encoding_node.counter_col]
+    assert torch.allclose(counter, torch.arange(N_pos, dtype=counter.dtype))
+
+
 def test_get_prev_value():
     input_values = torch.tensor([[1.0], [2.0], [3.0], [4.0], [5.0]])
     cond_values = torch.tensor([[1.0], [0.0], [0.0], [1.0], [0.0]])
     expected_prev_values = torch.tensor([[1.0], [1.0], [1.0], [4.0], [4.0]])
 
     value_input = InputNode("value", 1, value_range=(-100.0, 100.0))
     cond_input = InputNode("cond", 1, value_range=(-100.0, 100.0))
-    pos_encoding = PosEncoding(16)
+    pos_encoding = PosEncoding(17)
     last_input = pos_encoding.get_prev_value(value_input, cond_input)
     output = last_input.compute(
         n_pos=5, input_values={"value": input_values, "cond": cond_values}
@@ -29,7 +47,7 @@ def test_get_prev_value():
 def test_attend_to_offset():
     input_values = torch.tensor([[1.0], [2.0], [3.0], [4.0], [5.0]])
     value_input = InputNode("value", 1, value_range=(-100.0, 100.0))
-    pos_encoding = PosEncoding(16)
+    pos_encoding = PosEncoding(17)
     last_input = pos_encoding.attend_to_offset(value_input, delta_pos=-1)
     output = last_input.compute(n_pos=5, input_values={"value": input_values})
     assert torch.allclose(output[0], input_values[0])
@@ -39,35 +57,31 @@ def test_attend_to_offset():
 def test_attend_to_offset_value_can_be_wider_than_position_encoding():
     input_values = torch.arange(5 * 20, dtype=torch.float32).reshape(5, 20)
     value_input = InputNode("value", 20, value_range=(-100.0, 100.0))
-    pos_encoding = PosEncoding(16)
+    pos_encoding = PosEncoding(17)
     last_input = pos_encoding.attend_to_offset(value_input, delta_pos=-1)
     output = last_input.compute(n_pos=5, input_values={"value": input_values})
     assert torch.allclose(output[0], input_values[0])
     assert torch.allclose(output[1:], input_values[:-1])
 
 
-def test_delta_matrix():
-    # The last pair (d_pos-2, d_pos-1) contains the counter — exclude it.
-    s = d_pos - 2
-    delta_matrix = get_pos_delta_matrix(delta_pos=1, d=d_pos)[:s, :s]
-    sin_part = pos_encoding[:, :s]
-    delta_pos_encoding = sin_part @ delta_matrix
-    assert torch.allclose(delta_pos_encoding[1:, :], sin_part[:-1, :], atol=atol)
+def _check_shift(k: int):
+    # Row-convention shift over the trig block: trig(pos) @ S == trig(pos+k).
+    tw = pos_encoding_node.trig_width
+    trig = pos_encoding[:, :tw]
+    shifted = trig @ trig_shift_matrix(k, tw)
+    if k >= 0:
+        assert torch.allclose(shifted[: N_pos - k], trig[k:], atol=atol)
+    else:
+        assert torch.allclose(shifted[-k:], trig[: N_pos + k], atol=atol)
+
 
+def test_trig_shift_matrix_pos1():
+    _check_shift(1)
 
-def test_delta_matrix2():
-    s = d_pos - 2
-    delta_matrix = get_pos_delta_matrix(delta_pos=2, d=d_pos)[:s, :s]
-    sin_part = pos_encoding[:, :s]
-    delta_pos_encoding = sin_part @ delta_matrix
-    assert torch.allclose(delta_pos_encoding[2:, :], sin_part[:-2, :], atol=atol)
 
+def test_trig_shift_matrix_pos2():
+    _check_shift(2)
 
-def test_delta_matrix_neg():
-    s = d_pos - 2
-    delta_matrix = get_pos_delta_matrix(delta_pos=-1, d=d_pos)[:s, :s]
-    sin_part = pos_encoding[:, :s]
-    delta_pos_encoding = sin_part @ delta_matrix
-    assert torch.allclose(delta_pos_encoding[:-1, :], sin_part[1:, :], atol=atol)
 
-    assert torch.allclose(delta_pos_encoding[1], sin_part[2], atol=atol)
+def test_trig_shift_matrix_neg1():
+    _check_shift(-1)