fix: route2 Arm A flags per-rollout not per-token (external review)

The hook gate is necessarily per-token ([G*s, r], nn.Linear flattens the
batch). _route2_grad_filter now sums each rollout's token gate-grads before
the cos(g_b, v_grad) flag, so routing is per-rollout (the preregistered GRPO
unit) and the sign is denoised. Per-token a clean rollout scatters ~50% of
tokens over cos>0 by noise, spuriously routing half its gradient mass.

Verified by deepseek-v4-pro review: gate identity, divide-out, eps-guard,
Arm B detach-route, R5 no-cheat all correct; this was the one finding.

Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>

src/projected_grpo/antipasto.py

@@ -111,9 +111,11 @@ def _delta_hook(layer: nn.Linear, args: tuple, y: Tensor) -> Tensor:
 
     if layer._antipasto_mask_mode == "grad":
         if torch.is_grad_enabled():
-            # per-rollout gate [b, 1.., r], identity at c=1 so the forward value is
-            # unchanged. After backward c.grad = delta_S * (per-rollout delta_S grad);
-            # train.py divides out delta_S to recover g_b and routes post-backward.
+            # gate c, one entry per (token, axis) since nn.Linear flattens the batch
+            # ([G*s, r]); identity at c=1 so the forward value is unchanged. After
+            # backward c.grad = delta_S * g_b (per-token). train.py reshapes to
+            # [G, s, r], sums each rollout's tokens, divides out delta_S to recover
+            # the per-rollout g_b, and routes post-backward.
             c = torch.ones(a.shape[0], *([1] * (a.dim() - 2)), a.shape[-1],
                            device=a.device, dtype=a.dtype, requires_grad=True)
             layer._antipasto_gate = c

src/projected_grpo/train.py

@@ -1135,16 +1135,24 @@ def main(cfg: Config) -> int:
         GATE_EPS = 1e-6
         step_flagged: list[float] = []
 
-        def _route2_grad_filter(info) -> torch.Tensor:
-            g = info["delta_S"].grad                          # [r] summed over rollouts
-            cg = info["layer"]._antipasto_gate.grad.reshape(-1, g.shape[0])      # [b, r]
+        def _route2_grad_filter(info, n_rollouts: int) -> torch.Tensor:
+            g = info["delta_S"].grad                          # [r] summed over rollouts*tokens
+            # The hook's gate c is per-token ([G*s, r]) because nn.Linear sees a
+            # flattened batch. Sum each rollout's token gate-grads -> per-rollout
+            # delta_S*g_b: reshape [G*s, r] -> [G, s, r] -> sum tokens -> [G, r].
+            # Pad tokens carry ~0 grad (masked in the loss), so summing every
+            # position is safe. Per-rollout (not per-token) is the preregistered
+            # unit: GRPO advantage is per-rollout, and summing first denoises the
+            # cos(g_b, v_grad) sign (a clean rollout's individual tokens scatter
+            # ~50% over cos>0; its token-sum points reliably clean-ward).
+            cg = info["layer"]._antipasto_gate.grad.reshape(n_rollouts, -1, g.shape[0]).sum(1)  # [G, r]
             dS = info["delta_S"].detach()                     # [r]
             reliable = dS.abs() > GATE_EPS                    # [r]
             dS_safe = torch.where(reliable, dS, torch.ones_like(dS))
-            g_b = torch.where(reliable, cg / dS_safe, torch.zeros_like(cg))  # [b, r]
+            g_b = torch.where(reliable, cg / dS_safe, torch.zeros_like(cg))  # [G, r] per-rollout
             vg = v_grad[name]                                 # [r] unit, hack-ward
-            cos_b = (g_b @ vg) / g_b.norm(dim=1).clamp_min(1e-12)            # [b]
-            flagged = (cos_b > 0).float()                     # [b]
+            cos_b = (g_b @ vg) / g_b.norm(dim=1).clamp_min(1e-12)            # [G]
+            flagged = (cos_b > 0).float()                     # [G]
             step_flagged.append(flagged.mean().item())
             sub = torch.where(reliable, (cg * flagged.unsqueeze(1)).sum(0) / dS_safe,
                               torch.zeros_like(g))             # flagged rollouts' contribution
@@ -1451,7 +1459,7 @@ def main(cfg: Config) -> int:
                     # grad-mask routes here: strip flagged rollouts from delta_S.grad
                     # (quarantine still learns them via its always-on forward path).
                     if is_route2_grad:
-                        g = _route2_grad_filter(info)
+                        g = _route2_grad_filter(info, merged.shape[0])
                     step_grad_s[name] = (step_grad_s[name] + g.detach().clone()
                                          if name in step_grad_s
                                          else g.detach().clone())