| """Opt-in V1 diagnostics. Imported/installed only with DSPARK_V1_DUMP_DIR. |
| |
| Selected graph replies are checked against a traced eager replay on the SAME |
| state. The reply is never replaced. A graph/eager disagreement fails the round. |
| Window KV is read from the live ring; main_x is the insert graph's actual output, |
| not a differently batched reconstruction. This mode deliberately synchronizes and |
| must be disabled for latency/engine measurements. Dumps contain private data. |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import os |
| from pathlib import Path |
| from typing import Any, Callable, cast |
|
|
| import numpy as np |
| import torch |
|
|
| from . import fused as FU |
| from .server import DraftCore |
| from .tp2 import TP2Backend |
|
|
|
|
| def chosen_rounds(value: str) -> frozenset[int]: |
| rounds = frozenset(int(x) for x in value.split(",")) |
| if not rounds or min(rounds) < 1 or len(rounds) > 1024: |
| raise ValueError("DSPARK_V1_ROUNDS requires 1..1024 positive round numbers") |
| return rounds |
|
|
|
|
| class V1Dump: |
| def __init__(self, model: Any, ident: dict[str, Any], directory: str, rounds: str = "1"): |
| self.m = model |
| self.ident = ident |
| self.directory = Path(directory).resolve() |
| source = Path(__file__).resolve().parents[3] |
| if self.directory == source or source in self.directory.parents: |
| raise ValueError("V1 dumps must be outside the repository (private capture data)") |
| self.rounds = chosen_rounds(rounds) |
| self.directory.mkdir(mode=0o700, parents=True, exist_ok=True) |
| self.main_x: dict[tuple[int, int], torch.Tensor] = {} |
| self.sources: dict[tuple[int, int], str] = {} |
| self.call = 0 |
| self.wire: dict[str, int] = {} |
| self.samples: list[dict[str, Any]] | None = None |
| self.files: list[str] = [] |
| self._install_insert_observers() |
|
|
| def _remember(self, slots: torch.Tensor, pos: torch.Tensor, mx: torch.Tensor, source: str) -> None: |
| slot_values, positions = slots.cpu().tolist(), pos.cpu().tolist() |
| saved_mx = mx.detach().cpu().clone() |
| for j, (slot, p) in enumerate(zip(slot_values, positions)): |
| if slot >= 4: |
| continue |
| self.main_x[slot, p] = saved_mx[j] |
| self.sources[slot, p] = source |
| old = [(s, q) for s, q in self.main_x if s == slot and q < p - self.m.s.window + 1] |
| for key in old: |
| self.main_x.pop(key) |
| self.sources.pop(key) |
|
|
| def _install_insert_observers(self) -> None: |
| m = self.m |
| m._dump_insert_mx = {} |
| original_mx = cast(Callable[[Any, torch.Tensor], torch.Tensor], m._main_x) |
|
|
| def traced_mx(rk: Any, taps: torch.Tensor) -> torch.Tensor: |
| y = original_mx(rk, taps) |
| m._dump_insert_mx[rk.r] = y |
| return y |
| m._main_x = traced_mx |
| original_insert = m.insert_taps_graphed |
|
|
| def insert(slots: torch.Tensor, positions: torch.Tensor, taps: torch.Tensor) -> None: |
| original_insert(slots, positions, taps) |
| bucket = min(n for n in m.insert_graphs if n >= slots.numel()) |
| self._remember(slots, positions, m.insert_graphs[bucket]["dump_main_x"][0][:slots.numel()], "taps") |
| m.insert_taps_graphed = insert |
| original_eager = m.insert_eager |
|
|
| def eager(slots: torch.Tensor, positions: torch.Tensor, taps: torch.Tensor) -> None: |
| m._dump_insert_mx.clear() |
| original_eager(slots, positions, taps) |
| self._remember(slots, positions, m._dump_insert_mx[0], "taps") |
| m.insert_eager = eager |
| original_inject = m.insert_mx_eager |
|
|
| def inject(slots: torch.Tensor, positions: torch.Tensor, mx: torch.Tensor) -> None: |
| original_inject(slots, positions, mx) |
| self._remember(slots, positions, mx, "engine_mx") |
| m.insert_mx_eager = inject |
| original_reset = m.reset_slots |
|
|
| def reset(slots: list[int]) -> None: |
| original_reset(slots) |
| for key in [key for key in self.main_x if key[0] in slots]: |
| self.main_x.pop(key) |
| self.sources.pop(key) |
| m.reset_slots = reset |
|
|
| def emit(self, host: torch.Tensor, N: int, n: int, want_margin: bool, |
| reply: tuple[np.ndarray, np.ndarray, np.ndarray | None]) -> None: |
| self.call += 1 |
| if self.call not in self.rounds: |
| return |
| if tuple(self.m.ALL) != (0, 1): |
| raise RuntimeError("V1 dumps require complete TP2; a solo shard is not a reference sample") |
| trace: dict[str, Any] = {} |
| ids, conf, margin, peer_ids = self.m.draft(host, N, n_out=n, margin=want_margin, trace=trace) |
| if not (np.array_equal(ids.cpu().numpy(), reply[0]) and |
| np.array_equal(conf.cpu().numpy(), reply[1]) and torch.equal(ids, peer_ids.to(ids.device))): |
| raise RuntimeError("V1 graph/eager dump replay differs from the actual reply") |
| if want_margin: |
| if margin is None or reply[2] is None or not np.array_equal(margin.cpu().numpy(), reply[2]): |
| raise RuntimeError("V1 graph/eager margins differ") |
| rk = self.m.ranks[0] |
| B = host.shape[1] |
| hh = trace["head_hidden"] |
| head_norm = FU.rms_norm(hh, rk.norm, self.m.eps).cpu().view(B, N, -1) |
| for j in range(B): |
| slot, anchor, P, lo = (int(t) for t in host[:, j]) |
| positions = list(range(max(lo, P - self.m.s.window + 1), P + 1)) |
| keys = [(slot, p) for p in positions] |
| if not keys or any(key not in self.main_x for key in keys): |
| raise RuntimeError("V1 dump lacks inserted main_x provenance for its window") |
| src = [self.sources[key] for key in keys] |
| sample = self.samples[j] if self.samples else {} |
| key = sample.get("key", f"r{self.call:06d}_b{j}") |
| record: dict[str, Any] = { |
| "schema": "dspark.v1.arm-dump.v1", "P": P, "anchor": anchor, "n_out": n, |
| "drafts": reply[0][j].tolist(), "conf": torch.from_numpy(reply[1][j].copy()), |
| "base_logits": trace["base_logits"][j].detach().float().cpu(), |
| "main_x_P": self.main_x[slot, P].float(), |
| "kv_window": rk.ring_kv[slot, :, [p % self.m.s.window for p in positions]].transpose(0, 1).float().cpu(), |
| "window": {"n": len(positions), "taps_rows": src.count("taps"), |
| "engine_mx_rows": src.count("engine_mx"), "P_source": src[-1], |
| "first": positions[0], "last": positions[-1]}, |
| "head_x": hh.view(B, N, -1)[j].float().cpu(), "head_hidden": head_norm[j], |
| "sample": sample, "wire": dict(self.wire), |
| "meta": {"arm": self.ident.get("arm", "k2-tp2"), "draft_head": self.ident.get("draft_head", "K6"), |
| "drafter_identity": self.ident["drafter_identity"], "weights": "official 1017 below K6 head", |
| "git_sha": os.environ.get("DSPARK_SOURCE_COMMIT", "unrecorded"), |
| "devices": os.environ.get("CUDA_VISIBLE_DEVICES"), "dtype_policy": "arm config/identity recipe", |
| "trace": "eager replay checked bit-for-bit against graph reply"}, |
| } |
| for i in range(self.m.s.n_stages): |
| tag = f"stage{i}" |
| st = rk.stages[i] |
| xin = trace[tag + ".ffn_in"].view(B, N, -1)[j] |
| biased = torch.nn.functional.softplus(xin.float() @ st.gate_w.t()).sqrt() + st.gate_bias |
| top = biased.topk(self.m.s.top_k + 1, dim=-1) |
| record[f"s{i}"] = { |
| "attn_out": trace[tag + ".attn_out"].view(B, N, -1)[j].float().cpu(), |
| "ffn_in": xin.float().cpu(), |
| "ffn_out": trace[tag + ".ffn_out"].view(B, N, -1)[j].float().cpu(), |
| "h_out": trace[tag + ".stream"].view(B, N, self.m.s.hc_mult, -1)[j].float().cpu(), |
| "route": {"idx": trace[tag + ".experts"].view(B, N, -1)[j].long().cpu(), |
| "w": trace[tag + ".route_w"].view(B, N, -1)[j].float().cpu(), |
| "biased": biased.cpu(), "top_vals": top.values.cpu(), "top_ids": top.indices.cpu(), |
| "margin": (top.values[:, -2] - top.values[:, -1]).cpu()}, |
| } |
| |
| path = self.directory / f"{key}.pt" |
| fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) |
| with os.fdopen(fd, "wb") as f: |
| torch.save(record, f) |
| self.files.append(str(path)) |
|
|
|
|
| class DumpTP2Backend(TP2Backend): |
| arm_identity: dict[str, Any] |
|
|
| def __init__(self, *args: Any, sink: V1Dump, **kwargs: Any): |
| super().__init__(*args, **kwargs) |
| self.sink = sink |
|
|
| def draft(self, reqs: Any, N: int, want_margin: bool) -> tuple[np.ndarray, np.ndarray, np.ndarray | None]: |
| reply = super().draft(reqs, N, want_margin) |
| self.sink.emit(self.inp_buf[len(reqs)], N, max(k for *_r, k in reqs), want_margin, reply) |
| return cast(tuple[np.ndarray, np.ndarray, np.ndarray | None], reply) |
|
|
|
|
| class DumpDraftCore(DraftCore): |
| def _step(self, h: Any, body: Any, t_in: float) -> bytes | None: |
| self.backend.sink.wire = {"session": h.session, "seq": h.seq} |
| return super()._step(h, body, t_in) |
|
|
