add padding-trigger ablation for the 4096-cap parity bug

diag/diagnose_padding_trigger.py

@@ -102,17 +102,6 @@ def main() -> None:
     n_r = int(lengths[r])
     lone_ids = ids[r:r + 1, :n_r].clone()
 
-    with torch.inference_mode():
-        ref = model(input_ids=lone_ids.cuda(), use_cache=False).logits[0]
-
-    def forward_row(x, a, positions):
-        with torch.inference_mode():
-            out = model(input_ids=x.cuda(),
-                        attention_mask=(a.cuda() if a is not None else None),
-                        position_ids=(positions.cuda() if positions is not None else None),
-                        use_cache=False).logits
-        return out[r] if out.dim() == 3 else out[0]
-
     def positions_for(a):
         p = a.long().cumsum(-1) - 1
         p.masked_fill_(a == 0, 1)
@@ -168,26 +157,66 @@ def main() -> None:
     pad_ids[0, :n_r] = ids[r, :n_r]
     pad_attn[0, :n_r] = 1
 
-    plans: list[tuple[str, Any, Any, Any]] = [
-        ("base", ids, attn, pos_stored),
-        ("base_nopos", ids, attn, None),
-        ("pad_zero", torch.where(attn.bool(), ids, torch.zeros_like(ids)), attn, pos_stored),
-        ("attn_prefix", ids, prefix, None),
-        ("attn_all_ones", ids, torch.ones_like(attn), None),
-        ("long_only_pad", pad_ids, pad_attn, None),
-        ("long_only", lone_ids, torch.ones_like(lone_ids), None),
+    n_rows = int(ids.shape[0])
+    # One reference per row: that row forwarded alone, no mask and no position_ids.
+    # This is the semantics SGLang serves, and it is what the fix must reproduce.
+    refs: list[Any] = []
+    for i in range(n_rows):
+        n_i = int(lengths[i])
+        with torch.inference_mode():
+            refs.append(model(input_ids=ids[i:i + 1, :n_i].cuda(),
+                              use_cache=False).logits[0])
+        torch.cuda.empty_cache()
+
+    plans: list[tuple[str, Any, Any, Any, list[int]]] = [
+        ("base", ids, attn, pos_stored, list(range(n_rows))),
+        ("base_nopos", ids, attn, None, list(range(n_rows))),
+        ("pad_zero", torch.where(attn.bool(), ids, torch.zeros_like(ids)), attn,
+         pos_stored, list(range(n_rows))),
+        ("attn_prefix", ids, prefix, None, list(range(n_rows))),
+        ("attn_all_ones", ids, torch.ones_like(attn), None, list(range(n_rows))),
+        ("long_only_pad", pad_ids, pad_attn, None, [r]),
+        ("long_only", lone_ids, torch.ones_like(lone_ids), None, [r]),
     ]
     variants: dict[str, dict[str, Any]] = {}
-    for name, x, a, p in plans:
-        got = forward_row(x, a, p)
-        if name == "long_only":
-            variants[name] = max_logit_delta(got, ref)
-        else:
-            variants[name] = max_logit_delta(got, ref)
-        del got
+    per_variant_row: dict[str, dict[str, Any]] = {}
+    for name, x, a, p, row_map in plans:
+        with torch.inference_mode():
+            out = model(input_ids=x.cuda(),
+                        attention_mask=(a.cuda() if a is not None else None),
+                        position_ids=(p.cuda() if p is not None else None),
+                        use_cache=False).logits
+        out3 = out if out.dim() == 3 else out.unsqueeze(0)
+        del out
         torch.cuda.empty_cache()
+        per_variant_row[name] = {}
+        for j, capture_row in enumerate(row_map):
+            n_i = int(lengths[capture_row])
+            d = max_logit_delta(out3[j, :n_i], refs[capture_row])
+            per_variant_row[name][str(capture_row)] = d
+            print("TRIGGER_VARIANT_ROW=" + json.dumps(
+                {"variant": name, "row": capture_row, "length": n_i,
+                 "max_abs_dlogit": d["max_abs_dlogit"],
+                 "mean_abs_dlogit": round(d["mean_abs_dlogit"], 6)}, sort_keys=True),
+                flush=True)
+        del out3
+        torch.cuda.empty_cache()
+        # Headline = the longest row, which is the one the parity guard trips on.
+        variants[name] = per_variant_row[name][str(r)]
         print("TRIGGER_VARIANT=" + json.dumps({"variant": name, **variants[name]},
                                              sort_keys=True), flush=True)
+    del refs
+    torch.cuda.empty_cache()
+    report["per_variant_row"] = per_variant_row
+    # Does dropping position_ids reproduce the per-row reference for EVERY row?
+    fixed = {k: round(v["max_abs_dlogit"], 6)
+             for k, v in per_variant_row.get("base_nopos", {}).items()}
+    broken = {k: round(v["max_abs_dlogit"], 6)
+              for k, v in per_variant_row.get("base", {}).items()}
+    report["fix_check"] = {"base_max_per_row": broken, "base_nopos_max_per_row": fixed,
+                           "fix_removes_the_effect_on_all_rows":
+                               all(v == 0.0 for v in fixed.values()) if fixed else None}
+    print("TRIGGER_FIX_CHECK=" + json.dumps(report["fix_check"], sort_keys=True), flush=True)
 
     report["variants"] = variants
     ranked = sorted((v["max_abs_dlogit"], k) for k, v in variants.items())