gemma-byteshark/shared_resources/accept_hist_byteshark/accept_hist_splitkv.patch

*** Begin Patch
*** Update File: sitecustomize.py
@@
 DIXIE_FUSED_ACCEPT_PREP = os.environ.get("DIXIE_FUSED_ACCEPT_PREP") == "1"
 DIXIE_FUSED_ACCEPT_PREP_REQUIRE = (
     os.environ.get("DIXIE_FUSED_ACCEPT_PREP_REQUIRE") == "1"
 )
+SPEC_ACCEPT_HISTOGRAM = os.environ.get("SPEC_ACCEPT_HISTOGRAM") == "1"
+SPEC_ACCEPT_HISTOGRAM_EVERY = int(
+    os.environ.get("SPEC_ACCEPT_HISTOGRAM_EVERY", "2048")
+)
 _FUSED_SPARSE_ARGMAX_KERNELS: Any | None = None
 _FUSED_ACCEPT_PREP_KERNEL: Any | None = None
 _FUSED_ACCEPT_PREP_CACHE: dict[int, tuple[Any, Any]] = {}
 _LOOPGRAPH_SLOT_EVENTS_BY_PTR: dict[int, Any] = {}
 _LOOPGRAPH_SLOT_EVENT_RECORDED_BY_PTR: dict[int, bool] = {}
+_ACCEPT_HIST_STATE: dict[str, Any] = {}
+
+
+def _record_accept_histogram(valid_counts: Any, max_spec_len: int) -> None:
+    if not SPEC_ACCEPT_HISTOGRAM:
+        return
+    try:
+        import torch
+
+        device = valid_counts.device
+        hist = _ACCEPT_HIST_STATE.get("hist")
+        if hist is None or hist.device != device:
+            hist = torch.zeros(64, dtype=torch.int64, device=device)
+            _ACCEPT_HIST_STATE["hist"] = hist
+            _ACCEPT_HIST_STATE["steps"] = 0
+
+        idx = valid_counts.to(torch.long).clamp(0, hist.shape[0] - 1)
+        ones = torch.ones_like(idx, dtype=torch.int64, device=device)
+        hist.index_add_(0, idx, ones)
+        _ACCEPT_HIST_STATE["steps"] = int(_ACCEPT_HIST_STATE.get("steps", 0)) + int(
+            valid_counts.numel()
+        )
+
+        steps = int(_ACCEPT_HIST_STATE["steps"])
+        if steps in (256, 1024) or (
+            SPEC_ACCEPT_HISTOGRAM_EVERY > 0 and steps % SPEC_ACCEPT_HISTOGRAM_EVERY == 0
+        ):
+            counts = hist.tolist()
+            top = max(
+                min(len(counts) - 1, max_spec_len + 1),
+                max((i for i, value in enumerate(counts) if value), default=0),
+            )
+            total = max(sum(counts), 1)
+            weighted = sum(i * value for i, value in enumerate(counts))
+            full = counts[max_spec_len + 1] if max_spec_len + 1 < len(counts) else 0
+            print(
+                "[accept-hist] "
+                f"steps={steps} max_spec_len={max_spec_len} "
+                f"mean_emit={weighted / total:.4f} "
+                f"full={full} ({full / total:.4f}) "
+                f"valid_counts_hist={counts[: top + 1]}",
+                file=sys.stderr,
+                flush=True,
+            )
+    except Exception as exc:
+        if not _ACCEPT_HIST_STATE.get("warned"):
+            _ACCEPT_HIST_STATE["warned"] = True
+            print(
+                f"[accept-hist] disabled after error: {exc!r}",
+                file=sys.stderr,
+                flush=True,
+            )
@@
         _FUSED_ACCEPT_PREP_CACHE[output_token_ids.data_ptr()] = (
             next_token_ids,
             valid_counts,
         )
+        _record_accept_histogram(valid_counts, max_spec_len)
         if not getattr(_dixie_fused_accept_prep, "_active_logged", False):
             _dixie_fused_accept_prep._active_logged = True
             print(
*** End Patch