| """Identity-checked loader for the DSpark drafter tensors that lane S2 copied to x86host. |
| |
| Source of truth: ``TENSORS.json`` written by S2 next to the extracted safetensors files. Every entry records the |
| tensor's name, dtype, shape, the destination file and byte range, and the sha256 of the *source* byte range on the |
| Spark (``/srv/dsv41/legal-v2a``, hardlinked from HeNryous EXL3 2.75 bpw @ 86f85253). Loading re-hashes each |
| destination byte range and refuses the checkpoint on any difference, on a missing or unexpected tensor, or on a |
| header that disagrees with the manifest. Nothing is converted here: tensors come back as zero-copy CPU views of the |
| mapped files, in the checkpoint's own dtypes. |
| |
| The tensor set is derived from config.json (three ``mtp.{0,1,2}`` stages, 128 drafter experts, the target's embedding |
| and EXL3 head), not from the manifest, so a manifest that lost or gained tensors is caught too. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| import mmap |
| import os |
| import struct |
| from concurrent.futures import ThreadPoolExecutor |
| from dataclasses import dataclass |
|
|
| import torch |
|
|
| DTYPES = { |
| "BF16": torch.bfloat16, |
| "F16": torch.float16, |
| "F32": torch.float32, |
| "I16": torch.int16, |
| "I32": torch.int32, |
| "F8_E4M3": torch.float8_e4m3fn, |
| "F8_E8M0": torch.float8_e8m0fnu, |
| } |
|
|
| EXL3_PARTS = ("trellis", "suh", "svh", "mul1") |
|
|
|
|
| class IdentityError(RuntimeError): |
| """The files on disk are not the tensors the manifest promises.""" |
|
|
|
|
| @dataclass(frozen=True) |
| class DrafterShape: |
| """The numbers of config.json the worker depends on (checked, never assumed).""" |
| hidden: int |
| vocab: int |
| n_heads: int |
| head_dim: int |
| rope_dim: int |
| q_lora: int |
| o_lora: int |
| o_groups: int |
| inter: int |
| n_experts: int |
| top_k: int |
| n_stages: int |
| n_target_layers: int |
| window: int |
| block: int |
| noise_token: int |
| markov_rank: int |
| hc_mult: int |
| hc_iters: int |
| hc_eps: float |
| rms_eps: float |
| swiglu_limit: float |
| route_scale: float |
| rope_theta: float |
| taps: tuple[int, ...] |
|
|
| @classmethod |
| def from_config(cls, cfg: dict) -> "DrafterShape": |
| t = cfg["text_config"] |
| n_layers = int(t["num_hidden_layers"]) |
| ratios = t["compress_ratios"] |
| n_stages = int(t["num_nextn_predict_layers"]) |
| if [int(r) for r in ratios[n_layers:n_layers + n_stages]] != [0] * n_stages: |
| raise IdentityError(f"drafter stages are not pure sliding-window layers: {ratios[n_layers:]}") |
| if t.get("scoring_func") != "sqrtsoftplus" or t.get("topk_method") != "noaux_tc" or not t.get("norm_topk_prob"): |
| raise IdentityError("routing is not sqrtsoftplus / noaux_tc / normalized top-k") |
| if t.get("hidden_act") != "silu": |
| raise IdentityError(f"hidden_act {t.get('hidden_act')!r}") |
| return cls( |
| hidden=int(t["hidden_size"]), vocab=int(t["vocab_size"]), n_heads=int(t["num_attention_heads"]), |
| head_dim=int(t["head_dim"]), rope_dim=int(t["qk_rope_head_dim"]), q_lora=int(t["q_lora_rank"]), |
| o_lora=int(t["o_lora_rank"]), o_groups=int(t["o_groups"]), inter=int(t["moe_intermediate_size"]), |
| n_experts=int(t["dspark_n_routed_experts"]), top_k=int(t["dspark_num_experts_per_tok"]), |
| n_stages=n_stages, n_target_layers=n_layers, window=int(t["sliding_window"]), |
| block=int(t["dspark_block_size"]), noise_token=int(t["dspark_noise_token_id"]), |
| markov_rank=int(t["dspark_markov_rank"]), hc_mult=int(t["hc_mult"]), hc_iters=int(t["hc_sinkhorn_iters"]), |
| hc_eps=float(t["hc_eps"]), rms_eps=float(t["rms_norm_eps"]), swiglu_limit=float(t["swiglu_limit"]), |
| route_scale=float(t["routed_scaling_factor"]), rope_theta=float(t["rope_theta"]), |
| taps=tuple(int(x) for x in t["dspark_target_layer_ids"]), |
| ) |
|
|
|
|
| def required_names(s: DrafterShape) -> list[str]: |
| """Every tensor the worker reads (plus ffn.gate.bias_vl, checked though text drafts never use it).""" |
| names = ["embed.weight"] + [f"head.{p}" for p in EXL3_PARTS] |
| for i in range(s.n_stages): |
| pre = f"mtp.{i}." |
| for lin in ("attn.wq_a", "attn.wkv", "attn.wq_b", "attn.wo_b", |
| "ffn.shared_experts.w1", "ffn.shared_experts.w2", "ffn.shared_experts.w3"): |
| names += [f"{pre}{lin}.{p}" for p in EXL3_PARTS] |
| names += [f"{pre}attn.wo_a.weight", f"{pre}attn.wo_a.scale", f"{pre}attn.attn_sink", |
| f"{pre}attn.q_norm.weight", f"{pre}attn.kv_norm.weight", f"{pre}attn_norm.weight", |
| f"{pre}ffn_norm.weight", f"{pre}ffn.gate.weight", f"{pre}ffn.gate.bias", f"{pre}ffn.gate.bias_vl"] |
| names += [f"{pre}hc_{s_}_{k}" for s_ in ("attn", "ffn") for k in ("fn", "base", "scale")] |
| for e in range(s.n_experts): |
| for w in ("w1", "w2", "w3"): |
| names += [f"{pre}ffn.experts.{e}.{w}.{p}" for p in EXL3_PARTS] |
| names += [f"mtp.0.main_proj.{p}" for p in EXL3_PARTS] + ["mtp.0.main_norm.weight"] |
| last = f"mtp.{s.n_stages - 1}." |
| names += [f"{last}norm.weight", f"{last}markov_head.embed.weight", f"{last}markov_head.head.weight", |
| f"{last}confidence_head.proj.weight"] |
| return names |
|
|
|
|
| def _read_header(path: str) -> tuple[int, dict]: |
| with open(path, "rb") as f: |
| n = struct.unpack("<Q", f.read(8))[0] |
| if n > 64 << 20: |
| raise IdentityError(f"{path}: implausible safetensors header length {n}") |
| return 8 + n, json.loads(f.read(n)) |
|
|
|
|
| class Checkpoint: |
| """The verified drafter tensors of one S2 delivery directory. ``get(name)`` returns a CPU view (no copy).""" |
|
|
| def __init__(self, root: str, verify: bool = True, threads: int = 16): |
| self.root = root |
| with open(os.path.join(root, "config.json")) as f: |
| self.config = json.load(f) |
| with open(os.path.join(root, "TENSORS.json")) as f: |
| self.manifest = json.load(f) |
| with open(os.path.join(root, "exl3_moe_plan.json")) as f: |
| self.plan_json = json.load(f) |
| self.shape = DrafterShape.from_config(self.config) |
| self._check_quant_config() |
| entries = {e["name"]: e for e in self.manifest["tensors"]} |
| need = required_names(self.shape) |
| missing = [n for n in need if n not in entries] |
| if missing: |
| raise IdentityError(f"{len(missing)} drafter tensors missing from TENSORS.json, e.g. {missing[:5]}") |
| |
| extra = sorted(set(entries) - set(need) - {"norm.weight"}) |
| if extra: |
| raise IdentityError(f"unexpected tensors in TENSORS.json: {extra[:5]} ({len(extra)})") |
| self.entries = {n: entries[n] for n in need} |
| self._maps: dict[str, mmap.mmap] = {} |
| self._headers: dict[str, tuple[int, dict]] = {} |
| from .identity import IdentityMismatch, resolve_inside |
| for e in self.entries.values(): |
| fn = e["destination_file"] |
| if fn not in self._maps: |
| try: |
| path = resolve_inside(root, fn) |
| except IdentityMismatch as err: |
| raise IdentityError(f"{e['name']}: {err}") from err |
| self._headers[fn] = _read_header(path) |
| with open(path, "rb") as f: |
| self._maps[fn] = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) |
| for name, e in self.entries.items(): |
| self._check_header(name, e) |
| self.verified = False |
| if verify: |
| self.verify(threads) |
|
|
| def _check_quant_config(self) -> None: |
| q = self.config.get("quantization_config", {}) |
| x = q.get("exl3_moe", {}) |
| if q.get("quant_method") != "deepseek_v41_exl3" or x.get("exllamav3_version") != "1.6.0": |
| raise IdentityError(f"quantization_config is not deepseek_v41_exl3 / exllamav3 1.6.0: {q.get('quant_method')}") |
| if x.get("head") != "exl3": |
| raise IdentityError(f"head is {x.get('head')!r}, the worker implements the EXL3 head only") |
| if self.manifest.get("codebook") != "mul1" or self.plan_json.get("codebook") != "mul1": |
| raise IdentityError("codebook is not mul1") |
| self.head_k = int(x["head_k"]) |
| self.drafter_k = {k: int(v) for k, v in x["dense"]["drafter"].items()} |
|
|
| def _check_header(self, name: str, e: dict) -> None: |
| fn = e["destination_file"] |
| hlen, hdr = self._headers[fn] |
| h = hdr.get(name) |
| if h is None: |
| raise IdentityError(f"{name}: not in {fn}'s header") |
| if h["dtype"] != e["dtype"] or list(h["shape"]) != list(e["shape"]): |
| raise IdentityError(f"{name}: header {h['dtype']} {h['shape']} vs manifest {e['dtype']} {e['shape']}") |
| a, b = h["data_offsets"] |
| if [hlen + a, hlen + b] != list(e["destination_byte_range"]): |
| raise IdentityError(f"{name}: header range {[hlen + a, hlen + b]} vs manifest {e['destination_byte_range']}") |
| if e["dtype"] not in DTYPES: |
| raise IdentityError(f"{name}: dtype {e['dtype']} not supported") |
| numel = 1 |
| for d in e["shape"]: |
| numel *= int(d) |
| if numel * DTYPES[e["dtype"]].itemsize != b - a: |
| raise IdentityError(f"{name}: {b - a} bytes for {e['shape']} {e['dtype']}") |
| sa, sb = e["source_byte_range"] |
| if sb - sa != b - a: |
| raise IdentityError(f"{name}: source range {sb - sa} B vs destination {b - a} B") |
|
|
| def _bytes(self, name: str) -> memoryview: |
| e = self.entries[name] |
| a, b = e["destination_byte_range"] |
| return memoryview(self._maps[e["destination_file"]])[a:b] |
|
|
| def verify(self, threads: int = 16) -> dict: |
| """sha256 of every destination byte range == the manifest's source-range sha256. Raises on any mismatch.""" |
| def one(name: str) -> tuple[str, bool]: |
| return name, hashlib.sha256(self._bytes(name)).hexdigest() == self.entries[name]["source_byte_range_sha256"] |
| |
| order = sorted(self.entries, key=lambda n: -int(self.entries[n]["resident_packed_bytes"])) |
| with ThreadPoolExecutor(threads) as ex: |
| bad = [n for n, ok in ex.map(one, order) if not ok] |
| if bad: |
| raise IdentityError(f"sha256 mismatch for {len(bad)} tensors, e.g. {bad[:5]}") |
| self.verified = True |
| return {"tensors": len(order), "bytes": sum(int(self.entries[n]["resident_packed_bytes"]) for n in order)} |
|
|
| def get(self, name: str) -> torch.Tensor: |
| """Zero-copy CPU view of the tensor bytes in the mapped file (read-only memory; copy before writing).""" |
| if not self.verified: |
| raise IdentityError("checkpoint not verified") |
| e = self.entries[name] |
| buf = self._bytes(name) |
| dt = DTYPES[e["dtype"]] |
| if len(buf) == 0: |
| return torch.empty(e["shape"], dtype=dt) |
| import warnings |
| with warnings.catch_warnings(): |
| warnings.simplefilter("ignore") |
| t = torch.frombuffer(buf, dtype=torch.uint8) |
| return t.view(dt).view(*e["shape"]) if e["shape"] else t.view(dt).reshape(()) |
|
|
| def exl3(self, prefix: str) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: |
| """(trellis int16 [in/16, out/16, 16K], suh fp16 [in], svh fp16 [out], K) of an EXL3 linear; mul1 checked.""" |
| trellis, suh, svh, mul1 = (self.get(f"{prefix}.{p}") for p in EXL3_PARTS) |
| words = trellis.shape[-1] |
| if trellis.dim() != 3 or words % 16: |
| raise IdentityError(f"{prefix}.trellis shape {tuple(trellis.shape)}") |
| if int(mul1.reshape(-1)[0].item()) & 0xFFFFFFFF != 0x83DCD12D: |
| raise IdentityError(f"{prefix}.mul1 marker {int(mul1.reshape(-1)[0].item()) & 0xFFFFFFFF:#x} is not mul1") |
| if suh.shape[0] != trellis.shape[0] * 16 or svh.shape[0] != trellis.shape[1] * 16: |
| raise IdentityError(f"{prefix}: suh/svh {tuple(suh.shape)}/{tuple(svh.shape)} vs trellis {tuple(trellis.shape)}") |
| return trellis, suh, svh, words // 16 |
|
|
| def close(self) -> None: |
| for m in self._maps.values(): |
| m.close() |
| self._maps.clear() |
|
|
