From ab0fdc2025de179801a47d47bf165a41a00facc4 Mon Sep 17 00:00:00 2001 From: davidtai Date: Sun, 9 Aug 2026 19:52:21 -0500 Subject: [PATCH 01/24] Add DeepSeek V4 0731 DSpark model support --- mtplx/deepseek_v4_attention_island.py | 5 + mtplx/models/deepseek_v4.py | 641 ++++++++++++++++++++- tests/test_deepseek_v4_attention_island.py | 20 + tests/test_deepseek_v4_dspark.py | 633 ++++++++++++++++++++ 4 files changed, 1291 insertions(+), 8 deletions(-) create mode 100644 tests/test_deepseek_v4_dspark.py diff --git a/mtplx/deepseek_v4_attention_island.py b/mtplx/deepseek_v4_attention_island.py index 93b28a24..b0e54e74 100644 --- a/mtplx/deepseek_v4_attention_island.py +++ b/mtplx/deepseek_v4_attention_island.py @@ -627,6 +627,11 @@ def install_deepseek_v4_attention_island( ) -> dict[str, Any]: """Validate the canonical checkpoint, prebind nine tapes, and install.""" + if getattr(model, "_dspark", None) is not None: + raise AttentionIslandError( + "DSpark requires tap-aware attention-island support; refusing to " + "install a target route that its tap collector would bypass" + ) _validate_model(model, config) stock = getattr(model, "_target_hc_hidden_route", model.model.hc_hidden) bound_by_width: dict[int, _BoundWidthBody] = {} diff --git a/mtplx/models/deepseek_v4.py b/mtplx/models/deepseek_v4.py index c9e390f4..abd59f1a 100644 --- a/mtplx/models/deepseek_v4.py +++ b/mtplx/models/deepseek_v4.py @@ -273,7 +273,7 @@ import os from dataclasses import dataclass, field, replace from functools import lru_cache -from typing import List, Optional +from typing import List, Optional, Tuple import mlx.core as mx import mlx.nn as nn @@ -1381,6 +1381,14 @@ class _DirectDenseMTPOLora(_DirectDenseOLora): __slots__ = () +_DSPARK_MANIFEST_KEYS = ( + "dspark_block_size", + "dspark_noise_token_id", + "dspark_target_layer_ids", + "dspark_markov_rank", +) + + @dataclass class ModelArgs(BaseModelArgs): model_type: str = "deepseek_v4" @@ -1434,6 +1442,26 @@ class ModelArgs(BaseModelArgs): # upstream as ``mtp.0.*``; a conversion that drops it leaves this field at 1 # while shipping no weights, which :meth:`Model.sanitize` detects and honours. num_nextn_predict_layers: int = 0 + temperature: float = 1.0 + # DeepSeek-V4-Flash-0731's DSpark draft is not the legacy one-layer MTP + # block above. These fields are deliberately separate: the manifest selects + # one installed implementation at construction, never in the decode path. + # ``None`` preserves whether the artifact actually carried a DSpark field. + # This lets construction distinguish an absent legacy field from an explicit + # corrupt value such as ``dspark_block_size: 0``. + dspark_block_size: Optional[int] = None + dspark_noise_token_id: Optional[int] = None + dspark_target_layer_ids: Optional[List[int]] = None + dspark_markov_rank: Optional[int] = None + _dspark_signature_present: bool = field(default=False, init=False, repr=False) + + @classmethod + def from_dict(cls, params): + args = super().from_dict(params) + args._dspark_signature_present = any( + key in params for key in _DSPARK_MANIFEST_KEYS + ) + return args def __post_init__(self): # Accept the HF rope_scaling block and mirror it into the flat YaRN fields @@ -3451,6 +3479,518 @@ def __call__( return (logits, x) if return_hidden else logits +# DSpark arithmetic below is transcribed from the official dedicated repository, +# not inferred from the earlier preview-MTP implementation: +# https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-DSpark/blob/aa22cb07426656189b2573b8e77a9b7333b8ae0f/inference/model.py +# The cited line numbers refer to that exact immutable source revision. +def get_dspark_topk_idxs( + window_size: int, batch_size: int, block_size: int, start_pos: int +) -> mx.array: + """Exact 0731 DSpark visibility matrix (official model.py L744-747).""" + if int(start_pos) <= 0: + raise ValueError("DSpark decode visibility requires start_pos > 0") + main = mx.arange(min(int(window_size), int(start_pos) + 1), dtype=mx.int32) + draft = int(window_size) + mx.arange(int(block_size), dtype=mx.int32) + row = mx.concatenate([main, draft]) + return mx.broadcast_to(row[None, None, :], (int(batch_size), int(block_size), row.shape[0])) + + +class DeepseekV4DSparkCache: + """Stage-owned fixed ring used by official ``DSparkAttention``.""" + + def __init__(self, window_size: int, head_dim: int): + self.window_size = int(window_size) + self.head_dim = int(head_dim) + self.ring: Optional[mx.array] = None + self.prefill_length = 0 + + def prefill(self, main_kv: mx.array) -> None: + b, seqlen, d = main_kv.shape + if d != self.head_dim: + raise ValueError("DSpark cache head dimension mismatch") + win = self.window_size + if seqlen <= win: + pad = mx.zeros((b, win - seqlen, d), dtype=main_kv.dtype) + self.ring = mx.concatenate([main_kv, pad], axis=1) + else: + last = main_kv[:, -win:] + cutoff = seqlen % win + self.ring = last if cutoff == 0 else mx.concatenate( + [last[:, win - cutoff:], last[:, : win - cutoff]], axis=1 + ) + self.prefill_length = int(seqlen) + + def commit_main(self, start_pos: int, main_kv: mx.array) -> None: + """Commit consecutive authoritative target rows into the fixed ring.""" + if self.ring is None: + raise RuntimeError("DSpark decode requires attention-only prefill first") + if main_kv.ndim != 3 or main_kv.shape[0] != self.ring.shape[0]: + raise ValueError("DSpark committed main KV must match the ring batch") + rows = int(main_kv.shape[1]) + if rows <= 0 or rows > self.window_size: + raise ValueError("DSpark committed main KV width is outside its ring") + index = int(start_pos) % self.window_size + first = min(rows, self.window_size - index) + ring = mx.concatenate( + [ + self.ring[:, :index], + main_kv[:, :first], + self.ring[:, index + first :], + ], + axis=1, + ) + remaining = rows - first + if remaining: + ring = mx.concatenate( + [main_kv[:, first:], ring[:, remaining:]], axis=1 + ) + self.ring = ring + + def replace_main(self, start_pos: int, main_kv: mx.array) -> None: + """Compatibility name for the one-row official proposal update.""" + if int(main_kv.shape[1]) != 1: + raise ValueError("DSpark decode replaces exactly one current-main KV") + self.commit_main(start_pos, main_kv) + + +class DeepseekV4DSparkAttention(DeepseekV4Attention): + """Official 0731 DSpark attention, distinct from trunk CSA attention.""" + + def __init__(self, args: ModelArgs, layer_id: int): + super().__init__(args, layer_id) + if self.compress_ratio != 0: + raise ValueError("DSpark attention requires compress_ratio=0") + + def _kv(self, x: mx.array, positions: mx.array) -> mx.array: + rd = self.rope_head_dim + cos, sin = self._rope_tables(positions) + kv = self.kv_norm(self.wkv(x)) + return mx.concatenate( + [kv[..., :-rd], _apply_interleaved_rope(kv[..., -rd:], cos[None], sin[None])], + axis=-1, + ) + + def __call__( + self, + x: mx.array, + *, + start_pos: int, + main_x: mx.array, + cache: DeepseekV4DSparkCache, + ) -> mx.array: + b, main_len, _ = main_x.shape + main_pos = mx.arange(int(start_pos), int(start_pos) + main_len) + main_kv = self._kv(main_x, main_pos) + if int(start_pos) == 0: + cache.prefill(main_kv) + return x + + if int(x.shape[1]) != _DSPARK_BLOCK_SIZE: + raise ValueError("DSpark decode requires one complete five-token block") + cache.replace_main(start_pos, main_kv) + block = int(x.shape[1]) + positions = mx.arange(int(start_pos) + main_len, int(start_pos) + main_len + block) + cos, sin = self._rope_tables(positions) + rd = self.rope_head_dim + + qr = self.q_norm(self.wq_a(x)) + q = self.wq_b(qr).reshape(b, block, self.n_heads, self.head_dim) + q = q * mx.rsqrt( + mx.mean(mx.square(q.astype(mx.float32)), axis=-1, keepdims=True) + self.eps + ) + q = q.astype(x.dtype) + q = mx.concatenate( + [q[..., :-rd], _apply_interleaved_rope(q[..., -rd:], cos[None, :, None], sin[None, :, None])], + axis=-1, + ) + draft_kv = self._kv(x, positions) + full_kv = mx.concatenate([cache.ring, draft_kv], axis=1) + topk = get_dspark_topk_idxs(self.window_size, b, block, start_pos) + # Every row has the same official index vector. Slice once; the query + # dimension is still fully retained in q. + visible_kv = full_kv[:, topk[0, 0]] + o = self._attend(q.transpose(0, 2, 1, 3), visible_kv, None) + o = o.transpose(0, 2, 1, 3) + o = mx.concatenate( + [o[..., :-rd], _apply_interleaved_rope(o[..., -rd:], cos[None, :, None], -sin[None, :, None])], + axis=-1, + ) + return self._o_lora(o.reshape(b, block, self.n_heads * self.head_dim)) + + +class DSparkMarkovHead(nn.Module): + """The 0731 sequential token-id bias, kept separate from the lm head.""" + + def __init__(self, vocab_size: int, rank: int): + super().__init__() + self.markov_w1 = nn.Embedding(vocab_size, rank) + self.markov_w2 = nn.Linear(rank, vocab_size, bias=False) + + def __call__(self, token_ids: mx.array) -> Tuple[mx.array, mx.array]: + embed = self.markov_w1(token_ids) + return self.markov_w2(embed), embed + + +class DSparkConfidenceHead(nn.Module): + """DSpark's fp32 confidence projection (not a vocabulary-logit head).""" + + def __init__(self, hidden_size: int, markov_rank: int): + super().__init__() + self.proj = nn.Linear(hidden_size + markov_rank, 1, bias=False) + + def __call__(self, hidden: mx.array, markov_embed: mx.array) -> mx.array: + x = mx.concatenate([hidden, markov_embed], axis=-1).astype(mx.float32) + # MLX stores Linear's parameters at the module dtype. Cast both here so + # the confidence contract remains fp32 even when the model is bf16. + return (x @ self.proj.weight.astype(mx.float32).T).squeeze(-1) + + +class DeepseekV4DSparkStage(DeepseekV4DecoderLayer): + """One of the three native 0731 DSpark stages. + + Prefill writes this stage's attention cache only, as the upstream + ``DSparkBlock.forward`` does at ``start_pos == 0``. Decode takes the normal + HC-attention-MoE block path. The cache is supplied by its owning stage; no + stage ever borrows a trunk or sibling cache. + """ + + def __init__(self, args: ModelArgs, stage_id: int): + layer_id = args.num_hidden_layers + stage_id + ratios = list(args.compress_ratios) + if len(ratios) <= layer_id: + ratios.extend([0] * (layer_id + 1 - len(ratios))) + args = replace(args, compress_ratios=ratios) + super().__init__(args, layer_id) + self.attn = DeepseekV4DSparkAttention(args, layer_id) + self.stage_id = int(stage_id) + self.block_size = int(args.dspark_block_size) + self.noise_token_id = int(args.dspark_noise_token_id) + self.main_proj = None + self.main_norm = None + self.norm = None + self.hc_head = None + self.markov_head = None + self.confidence_head = None + if stage_id == 0: + self.main_proj = nn.Linear( + args.hidden_size * len(args.dspark_target_layer_ids), args.hidden_size, + bias=False, + ) + self.main_norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + if stage_id == _DSPARK_STAGE_COUNT - 1: + self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.hc_head = HeadHC(args.hidden_size, args.hc_mult, args.hc_eps) + self.markov_head = DSparkMarkovHead(args.vocab_size, args.dspark_markov_rank) + self.confidence_head = DSparkConfidenceHead( + args.hidden_size, args.dspark_markov_rank + ) + + def fuse_main(self, main_hidden: mx.array) -> mx.array: + if self.main_proj is None or self.main_norm is None: + raise RuntimeError("DSpark main fusion belongs exclusively to stage 0") + return self.main_norm(self.main_proj(main_hidden)) + + def prefill(self, h: mx.array, cache, main_x: mx.array) -> mx.array: + """Populate only this stage's attention cache; do not run its MoE.""" + # The stage has a pure sliding-window cache in the 0731 manifest. The + # cache is deliberately built from stage 0's projected target state on + # every stage, matching DSparkAttention's ``main_x`` prefill operand. + # The attention result is discarded: upstream prefill exists to seed KV + # state, and DSpark's draft output is produced on decode. + self.attn(h, start_pos=0, main_x=main_x, cache=cache) + return h + + def __call__( + self, h: mx.array, *, start_pos: int, cache=None, input_ids=None, + main_x=None, + ) -> mx.array: + if int(start_pos) == 0: + if main_x is None: + raise ValueError("DSpark prefill requires stage-0 main_x") + return self.prefill(h, cache, main_x) + residual = h + x, post, comb = self.attn_hc.pre(h) + x = self.attn_norm(x) + x = self.attn(x, start_pos=start_pos, main_x=main_x, cache=cache) + h = self.attn_hc.post(x, residual, post, comb) + residual = h + x, post, comb = self.ffn_hc.pre(h) + x = self.ffn_norm(x) + x = self.ffn(x, input_ids=input_ids) + return self.ffn_hc.post(x, residual, post, comb) + + +_DSPARK_STAGE_COUNT = 3 +_DSPARK_BLOCK_SIZE = 5 +_DSPARK_NOISE_TOKEN_ID = 128799 +_DSPARK_TARGET_LAYER_IDS = (40, 41, 42) +_DSPARK_MARKOV_RANK = 256 + + +def _has_dspark_signature(args: ModelArgs) -> bool: + """Whether any 0731-only manifest value is present, complete or corrupt.""" + return bool(args._dspark_signature_present) or any( + value is not None + for value in ( + args.dspark_block_size, + args.dspark_noise_token_id, + args.dspark_target_layer_ids, + args.dspark_markov_rank, + ) + ) + + +def _config_has_dspark_signature(config: dict) -> bool: + config = config or {} + return any(key in config for key in _DSPARK_MANIFEST_KEYS) + + +def _validate_dspark_manifest(args: ModelArgs) -> None: + """Fail before installation if this is not the exact 0731 DSpark artifact.""" + if int(args.dspark_block_size or 0) != _DSPARK_BLOCK_SIZE: + raise ValueError("DSpark-0731 requires dspark_block_size=5") + if int(args.num_nextn_predict_layers) != 1: + raise ValueError("DSpark-0731 requires num_nextn_predict_layers=1") + if int(args.dspark_noise_token_id or 0) != _DSPARK_NOISE_TOKEN_ID: + raise ValueError("DSpark-0731 requires dspark_noise_token_id=128799") + if tuple(int(x) for x in (args.dspark_target_layer_ids or ())) != _DSPARK_TARGET_LAYER_IDS: + raise ValueError("DSpark-0731 requires target taps (40, 41, 42)") + if args.num_hidden_layers <= _DSPARK_TARGET_LAYER_IDS[-1]: + raise ValueError("DSpark-0731 target taps are absent from this trunk") + if args.vocab_size <= _DSPARK_NOISE_TOKEN_ID: + raise ValueError("DSpark-0731 vocabulary omits its noise token") + if int(args.dspark_markov_rank or 0) != _DSPARK_MARKOV_RANK: + raise ValueError("DSpark-0731 requires dspark_markov_rank=256") + ratios = list(args.compress_ratios) + for layer_id in range(args.num_hidden_layers, args.num_hidden_layers + _DSPARK_STAGE_COUNT): + if layer_id < len(ratios) and int(ratios[layer_id]) != 0: + raise ValueError("DSpark-0731 stages require uncompressed attention") + + +def _sample_dspark_token( + logits: mx.array, temperature: float, *, greedy: bool = False, key=None +) -> mx.array: + """Official Gumbel-max sampler plus an explicit canonical greedy control.""" + temperature = float(temperature) + if greedy or temperature == 0.0: + return mx.argmax(logits, axis=-1) + scaled = logits / max(temperature, 1e-5) + uniform = mx.random.uniform(shape=scaled.shape, key=key) + uniform = mx.clip(uniform, 1e-30, 1.0 - mx.finfo(mx.float32).eps) + gumbel = -mx.log(-mx.log(uniform)) + return mx.argmax(scaled.astype(mx.float32) + gumbel, axis=-1) + + +class DeepseekV4DSpark: + """Installed 0731 DSpark layer set; intentionally not generation routing.""" + + def __init__(self, args: ModelArgs): + _validate_dspark_manifest(args) + self.args = args + self.block_size = _DSPARK_BLOCK_SIZE + self.noise_token_id = _DSPARK_NOISE_TOKEN_ID + self.target_layer_ids = _DSPARK_TARGET_LAYER_IDS + self.stages = [DeepseekV4DSparkStage(args, i) for i in range(_DSPARK_STAGE_COUNT)] + + def draft_input_ids(self, target_ids: mx.array) -> mx.array: + if target_ids.ndim != 1: + raise ValueError("DSpark target ids must be a [batch] tensor") + noise = mx.full((target_ids.shape[0], self.block_size), self.noise_token_id, + dtype=target_ids.dtype) + return mx.concatenate([target_ids[:, None], noise[:, 1:]], axis=1) + + def make_cache(self) -> list: + return [ + DeepseekV4DSparkCache( + window_size=stage.attn.window_size, + head_dim=stage.attn.head_dim, + ) + for stage in self.stages + ] + + def prefill(self, main_hidden: mx.array, caches) -> None: + """Seed all three stage rings from the authoritative prompt taps.""" + if len(caches) != _DSPARK_STAGE_COUNT: + raise ValueError("DSpark requires one cache owned by each stage") + main_x = self.stages[0].fuse_main(main_hidden) + # At start_pos=0 DSparkAttention reads only main_x. Passing a narrow + # view avoids constructing the five noise-token embeddings discarded by + # the official attention-only prefill branch. + ignored = main_x[:, :1] + for stage, cache in zip(self.stages, caches): + stage.attn( + ignored, + start_pos=0, + main_x=main_x, + cache=cache, + ) + + def commit_main(self, main_hidden: mx.array, caches, *, start_pos: int) -> None: + """Commit only the target-verified proposal prefix to every stage ring.""" + if len(caches) != _DSPARK_STAGE_COUNT: + raise ValueError("DSpark requires one cache owned by each stage") + if int(main_hidden.shape[1]) <= 0: + return + main_x = self.stages[0].fuse_main(main_hidden) + positions = mx.arange(int(start_pos), int(start_pos) + int(main_x.shape[1])) + for stage, cache in zip(self.stages, caches): + cache.commit_main(start_pos, stage.attn._kv(main_x, positions)) + + def finish( + self, logits: mx.array, hidden: mx.array, target_ids: mx.array, + *, greedy: bool = False, key=None, + ) -> Tuple[mx.array, mx.array, mx.array]: + """Apply the sequential Markov recurrence and return fp32 confidence.""" + final = self.stages[-1] + if final.markov_head is None or final.confidence_head is None: + raise RuntimeError("DSpark final stage is missing its output heads") + if logits.shape[1] != self.block_size or hidden.shape[1] != self.block_size: + raise ValueError("DSpark finish requires exactly one five-token block") + output_ids = [target_ids] + biased_rows = [] + markov_embeds = [] + previous = target_ids + keys = [None] * self.block_size if key is None else list(mx.random.split(key, self.block_size)) + for i in range(self.block_size): + bias, markov_embed = final.markov_head(previous) + row = logits[:, i] + bias + biased_rows.append(row) + markov_embeds.append(markov_embed) + previous = _sample_dspark_token( + row, self.args.temperature, greedy=greedy, key=keys[i] + ).astype(target_ids.dtype) + output_ids.append(previous) + confidence = final.confidence_head(hidden, mx.stack(markov_embeds, axis=1)) + return (mx.stack(output_ids, axis=1), mx.stack(biased_rows, axis=1), confidence) + + def finish_ids( + self, + logits: mx.array, + target_ids: mx.array, + *, + width: int, + forced_first_token_ids: mx.array | None = None, + ) -> mx.array: + """Return a greedy proposal prefix without unused heads or rows. + + ``forced_first_token_ids`` installs the target-owned primary at row zero + and uses it to seed the sequential Markov bias for the genuinely future + rows. The neural DSpark rows remain the same fixed parallel block; only + the token-id recurrence stops asking the drafter to overrule a token the + target has already sampled. + """ + width = int(width) + if width < 1 or width > self.block_size: + raise ValueError("DSpark ids-only width must be between one and five") + if int(logits.shape[1]) != width: + raise ValueError("DSpark ids-only logits must match proposal width") + final = self.stages[-1] + if final.markov_head is None: + raise RuntimeError("DSpark final stage is missing its Markov head") + output_ids = [target_ids] + if forced_first_token_ids is None: + previous = target_ids + first_row = 0 + else: + if forced_first_token_ids.shape != target_ids.shape: + raise ValueError("forced DSpark primary must match target id shape") + previous = forced_first_token_ids.astype(target_ids.dtype) + output_ids.append(previous) + first_row = 1 + for index in range(first_row, width): + bias, _markov_embed = final.markov_head(previous) + previous = mx.argmax(logits[:, index] + bias, axis=-1).astype( + target_ids.dtype + ) + output_ids.append(previous) + return mx.stack(output_ids, axis=1) + + def forward( + self, + main_hidden: mx.array, + target_ids: mx.array, + embed_tokens: nn.Module, + lm_head: nn.Module, + caches=None, + *, + start_pos: int, + greedy: bool = False, + key=None, + ids_only_width: int | None = None, + forced_first_token_ids: mx.array | None = None, + ): + """Execute the three-stage 0731 layer without generation integration. + + ``main_hidden`` is the target route's already-concatenated HC means. + ``start_pos == 0`` is the sole prefill signal: all three stages only write + their attention caches and return no draft output. Positive positions run + all three full HC-attention-MoE stages. + """ + if caches is None: + caches = self.make_cache() + if len(caches) != _DSPARK_STAGE_COUNT: + raise ValueError("DSpark requires one cache owned by each of its three stages") + main_x = self.stages[0].fuse_main(main_hidden) + ids = self.draft_input_ids(target_ids) + h = embed_tokens(ids) + h = mx.broadcast_to(h[:, :, None, :], (*h.shape[:2], self.args.hc_mult, h.shape[-1])) + for stage, cache in zip(self.stages, caches): + h = stage( + h, start_pos=start_pos, cache=cache, input_ids=ids, + main_x=main_x, + ) + if int(start_pos) == 0: + return None + final = self.stages[-1] + if final.hc_head is None or final.norm is None: + raise RuntimeError("DSpark final stage is missing its shared-head route") + collapsed = final.hc_head(h) + if ids_only_width is not None: + width = int(ids_only_width) + if width < 1 or width > self.block_size: + raise ValueError("DSpark ids-only width must be between one and five") + logits = lm_head(final.norm(collapsed[:, :width])) + if forced_first_token_ids is None: + return self.finish_ids(logits, target_ids, width=width) + return self.finish_ids( + logits, + target_ids, + width=width, + forced_first_token_ids=forced_first_token_ids, + ) + logits = lm_head(final.norm(collapsed)) + return self.finish(logits, collapsed, target_ids, greedy=greedy, key=key) + + +class _LegacyTargetRoute: + """Installed target route for pre-0731 checkpoints.""" + + def __call__(self, owner, inputs: mx.array, cache): + h = owner._target_hc_hidden_route(inputs, cache) + return h, h + + +class _DSparkTargetRoute: + """Installed target route that captures the three HC-collapsed tap means.""" + + def __call__(self, owner, inputs: mx.array, cache): + h = owner.model.embed_tokens(inputs) + h = mx.broadcast_to(h[:, :, None, :], (*h.shape[:2], owner.args.hc_mult, h.shape[-1])) + if cache is None: + cache = [None] * len(owner.model.layers) + taps = [] + wanted = owner._dspark.target_layer_ids + for layer_id, (layer, layer_cache) in enumerate(zip(owner.model.layers, cache)): + h = layer(h, mask=None, cache=layer_cache, input_ids=inputs) + if layer_id in wanted: + # This is intentionally inside the layer loop: DSpark consumes + # the HC mean from the exact post-layer state, not a later state. + taps.append(mx.mean(h, axis=2)) + if len(taps) != _DSPARK_STAGE_COUNT: + raise RuntimeError("DSpark target route did not observe every required tap") + return h, mx.concatenate(taps, axis=-1) + + class DeepseekV4Model(nn.Module): def __init__(self, args: ModelArgs): super().__init__() @@ -3503,10 +4043,22 @@ def __init__(self, args: ModelArgs): # Reference ``Transformer.mtp`` (model.py L789-793): a top-level list, so # the parameter paths are ``mtp.{i}.*`` — exactly the upstream checkpoint's # names. Dropped again by :meth:`sanitize` if the weights are not there. - self.mtp = [ - DeepseekV4MTP(args, args.num_hidden_layers + i) - for i in range(max(int(args.num_nextn_predict_layers or 0), 0)) - ] + # A DSpark manifest installs a different, typed target route and exactly + # three owned stage objects. Do not even construct the legacy preview-MTP + # type for that artifact: the manifest selects one representation once. + if _has_dspark_signature(args): + self._dspark = DeepseekV4DSpark(args) + # Preserve the checkpoint's upstream ``mtp.{stage}.*`` namespace. + # `_dspark` is the installed type/control surface, while this list is + # the only registered parameter owner. + self.mtp = self._dspark.stages + else: + self._dspark = None + self.mtp = [ + DeepseekV4MTP(args, args.num_hidden_layers + i) + for i in range(max(int(args.num_nextn_predict_layers or 0), 0)) + ] + self._target_hidden_route = _DSparkTargetRoute() if self._dspark else _LegacyTargetRoute() # Construction-time performance installers may replace this with a # typed phase/width router. The stock callable is explicit and direct; # decoder layers never probe candidate eligibility or fall back. @@ -3550,7 +4102,7 @@ def __call__( "the DeepSeek-V4 backend does not support input_embeddings " "(no vision splice path)" ) - h = self._target_hc_hidden_route(inputs, cache) + h, exposed_hidden = self._target_hidden_route(self, inputs, cache) logits = None if emit_logits: source = h @@ -3559,7 +4111,7 @@ def __call__( logits = self.logits_from_hc_hidden(source) if not return_hidden: return logits - return logits, h + return logits, exposed_hidden @property def layers(self): @@ -3583,12 +4135,44 @@ def mtp_blocks(self) -> list: @property def has_mtp(self) -> bool: - return bool(self.mtp_blocks) + # DSpark has its own five-token output protocol and has deliberately not + # been connected to the generic preview-MTP generation path yet. + return self._dspark is None and bool(self.mtp_blocks) def hc_hidden(self, inputs: mx.array, cache=None) -> mx.array: """Trunk forward stopping at the pre-head state the MTP block consumes.""" return self.model.hc_hidden(inputs, cache) + def _collect_dspark_taps( + self, h: mx.array, *, start_layer: int = 0, cache=None, input_ids=None + ) -> mx.array: + """Collect DSpark's post-layer HC means, primarily for exactness gates. + + The installed target route above uses the same operation during a real + forward. Keeping this small helper makes the boundary observable without + creating a second model-forward implementation for tests or loaders. + """ + if self._dspark is None: + raise RuntimeError("DSpark taps requested from a legacy V4 model") + if cache is None: + cache = [None] * len(self.model.layers) + taps = [] + wanted = self._dspark.target_layer_ids + for layer_id in range(int(start_layer), len(self.model.layers)): + h = self.model.layers[layer_id]( + h, mask=None, cache=cache[layer_id], input_ids=input_ids + ) + if layer_id in wanted: + taps.append(mx.mean(h, axis=2)) + if len(taps) != _DSPARK_STAGE_COUNT: + raise RuntimeError("DSpark target tap collection was incomplete") + return mx.concatenate(taps, axis=-1) + + def make_dspark_cache(self): + if self._dspark is None: + raise RuntimeError("this checkpoint does not install DSpark") + return self._dspark.make_cache() + def logits_from_hc_hidden(self, h: mx.array) -> mx.array: """``[b, s, hc, dim]`` -> target logits; the other half of :meth:`hc_hidden`. @@ -3635,6 +4219,8 @@ def mtp_forward( RoPE at the wrong absolute position instead of failing. """ blocks = self.mtp_blocks + if self._dspark is not None: + raise RuntimeError("DSpark-0731 generation routing is not installed") if not blocks: raise RuntimeError("this checkpoint ships no MTP block") if isinstance(cache, (list, tuple)): @@ -3741,6 +4327,40 @@ def sanitize(self, weights: dict) -> dict: ``load_weights(strict=True)`` still sees an exact match instead of 58 spurious "missing" keys. """ + # Official PyTorch HC tensors are flat fields; the MLX modules group the + # same three arrays under their installed HC objects. Translate once at + # the load boundary for both trunk and DSpark blocks. + hc_suffixes = { + ".hc_attn_fn": ".attn_hc.fn", + ".hc_attn_base": ".attn_hc.base", + ".hc_attn_scale": ".attn_hc.scale", + ".hc_ffn_fn": ".ffn_hc.fn", + ".hc_ffn_base": ".ffn_hc.base", + ".hc_ffn_scale": ".ffn_hc.scale", + ".hc_head_fn": ".hc_head.fn", + ".hc_head_base": ".hc_head.base", + ".hc_head_scale": ".hc_head.scale", + } + translated = {} + for key, value in weights.items(): + target = str(key) + for source_suffix, target_suffix in hc_suffixes.items(): + if target.endswith(source_suffix): + target = target[: -len(source_suffix)] + target_suffix + break + translated[target] = value + weights = translated + if self._dspark is not None: + missing = [ + stage_id for stage_id in range(_DSPARK_STAGE_COUNT) + if not any(str(k).startswith(f"mtp.{stage_id}.") for k in weights) + ] + if missing: + raise ValueError( + "DSpark-0731 checkpoint is missing required stage tensors: " + + ", ".join(f"mtp.{stage_id}.*" for stage_id in missing) + ) + return weights if self.mtp_blocks and not any(str(k).startswith("mtp.") for k in weights): self.mtp = [] return weights @@ -4272,6 +4892,11 @@ def is_deepseek_v4_mtp_config(config: dict) -> bool: mlx-community conversions declare the layer and ship no tensors, which is what the runtime's degrade-to-autoregressive branch exists for). """ + if _config_has_dspark_signature(config or {}): + # 0731 uses the same num_nextn_predict_layers=1 marker as preview MTP but + # has a different three-stage protocol. It must wait for its dedicated + # runtime route instead of being injected into the legacy adapter. + return False model_type = str((config or {}).get("model_type") or "").lower() architectures = [str(a) for a in (config or {}).get("architectures") or []] if model_type != "deepseek_v4" and not any( diff --git a/tests/test_deepseek_v4_attention_island.py b/tests/test_deepseek_v4_attention_island.py index 74dbb558..3c714e85 100644 --- a/tests/test_deepseek_v4_attention_island.py +++ b/tests/test_deepseek_v4_attention_island.py @@ -394,6 +394,26 @@ def test_runtime_installs_island_after_loaded_o_lora_routes(): assert "deepseek_v4_attention_island_report" in runtime_source +def test_install_rejects_dspark_before_publishing_a_bypassed_route(monkeypatch): + stock = object() + model = SimpleNamespace( + _dspark=object(), + _target_hc_hidden_route=stock, + model=SimpleNamespace(hc_hidden=stock), + ) + + def forbidden_validation(*_args, **_kwargs): + raise AssertionError("DSpark must fail before attention-island binding") + + monkeypatch.setattr(AI, "_validate_model", forbidden_validation) + + with pytest.raises(AI.AttentionIslandError, match="DSpark.*tap-aware"): + AI.install_deepseek_v4_attention_island(model, {}) + + assert model._target_hc_hidden_route is stock + assert not hasattr(model, "_mtplx_dsv4_attention_island_selector") + + def test_arm_selector_switches_prebound_model_route_without_hot_checks(): model = SimpleNamespace(_target_hc_hidden_route=None) stock = object() diff --git a/tests/test_deepseek_v4_dspark.py b/tests/test_deepseek_v4_dspark.py new file mode 100644 index 00000000..0933ae3a --- /dev/null +++ b/tests/test_deepseek_v4_dspark.py @@ -0,0 +1,633 @@ +"""Synthetic contract gates for the native 0731 DSpark model layer. + +These deliberately exercise structure and state transitions only. They do not +need a checkpoint or a GPU. +""" + +import importlib.util +import os +import sys +from dataclasses import fields + +import numpy as np +import pytest + +pytest.importorskip("mlx.core") +import mlx.core as mx # noqa: E402 +from mlx.utils import tree_flatten # noqa: E402 + + +@pytest.fixture(autouse=True) +def _cpu_default_device(): + previous = mx.default_device() + mx.set_default_device(mx.cpu) + try: + yield + finally: + mx.set_default_device(previous) + + +_HERE = os.path.dirname(os.path.abspath(__file__)) +_MODEL = os.path.join(_HERE, "..", "mtplx", "models", "deepseek_v4.py") +_spec = importlib.util.spec_from_file_location("dsv4_dspark_undertest", _MODEL) +D = importlib.util.module_from_spec(_spec) +sys.modules["dsv4_dspark_undertest"] = D +_spec.loader.exec_module(D) + + +def _args(**over): + cfg = dict( + vocab_size=128800, + hidden_size=8, + num_hidden_layers=43, + num_hash_layers=0, + num_attention_heads=1, + head_dim=8, + qk_rope_head_dim=4, + q_lora_rank=8, + o_lora_rank=8, + o_groups=1, + moe_intermediate_size=4, + n_routed_experts=2, + num_experts_per_tok=1, + index_n_heads=1, + index_head_dim=8, + index_topk=2, + sliding_window=8, + compress_ratios=[0] * 46, + dspark_block_size=5, + dspark_noise_token_id=128799, + dspark_target_layer_ids=[40, 41, 42], + dspark_markov_rank=256, + num_nextn_predict_layers=1, + temperature=1.0, + ) + cfg.update(over) + return D.ModelArgs(**cfg) + + +@pytest.fixture +def tiny_model(monkeypatch): + """Keep synthetic memory small without weakening the production manifest.""" + real_markov = D.DSparkMarkovHead + real_confidence = D.DSparkConfidenceHead + + class TinyMarkov(real_markov): + def __init__(self, vocab_size, _rank): + real_markov.__init__(self, vocab_size, 3) + + class TinyConfidence(real_confidence): + def __init__(self, hidden_size, _rank): + real_confidence.__init__(self, hidden_size, 3) + + monkeypatch.setattr(D, "DSparkMarkovHead", TinyMarkov) + monkeypatch.setattr(D, "DSparkConfidenceHead", TinyConfidence) + return D.Model(_args()) + + +_OFFICIAL_FILTERED = { + "model_type": "deepseek_v4", + "vocab_size": 129280, + "hidden_size": 4096, + "num_hidden_layers": 43, + "num_hash_layers": 3, + "num_attention_heads": 64, + "num_key_value_heads": 1, + "head_dim": 512, + "qk_rope_head_dim": 64, + "q_lora_rank": 1024, + "o_lora_rank": 1024, + "o_groups": 8, + "moe_intermediate_size": 2048, + "n_routed_experts": 256, + "n_shared_experts": 1, + "num_experts_per_tok": 6, + "index_n_heads": 64, + "index_head_dim": 128, + "index_topk": 512, + "sliding_window": 128, + "compress_ratios": [0, 0] + [4, 128] * 20 + [4, 0, 0, 0], + "num_nextn_predict_layers": 1, + "dspark_block_size": 5, + "dspark_noise_token_id": 128799, + "dspark_target_layer_ids": [40, 41, 42], + "dspark_markov_rank": 256, +} + + +def test_real_filtered_config_derives_three_stages_without_n_mtp_layers(): + assert "n_mtp_layers" not in _OFFICIAL_FILTERED + allowed = {f.name for f in fields(D.ModelArgs)} + args = D.ModelArgs(**{k: v for k, v in _OFFICIAL_FILTERED.items() if k in allowed}) + D._validate_dspark_manifest(args) + assert D.is_deepseek_v4_mtp_config(_OFFICIAL_FILTERED) is False + assert ( + D.inject_deepseek_v4_mtp_support(object(), config=_OFFICIAL_FILTERED) is False + ) + + +def test_legacy_model_call_uses_prebound_target_route_exactly_once(): + inputs = mx.array([[7]], dtype=mx.int32) + cache = object() + hidden = object() + calls = [] + + class InnerModel: + def hc_hidden(self, *_args, **_kwargs): + raise AssertionError("legacy dispatch bypassed the installed target route") + + class Owner: + model = InnerModel() + _target_hidden_route = D._LegacyTargetRoute() + + owner = Owner() + + def installed_route(got_inputs, got_cache): + calls.append((got_inputs, got_cache)) + return hidden + + owner._target_hc_hidden_route = installed_route + logits, got_hidden = D.Model.__call__( + owner, + inputs, + cache=cache, + return_hidden=True, + emit_logits=False, + ) + + assert logits is None + assert got_hidden is hidden + assert calls == [(inputs, cache)] + + +def test_dspark_manifest_installs_exactly_three_stages(tiny_model): + model = tiny_model + assert isinstance(model._dspark, D.DeepseekV4DSpark) + assert len(model._dspark.stages) == 3 + assert model._dspark.block_size == 5 + assert model._dspark.noise_token_id == 128799 + assert model._dspark.target_layer_ids == (40, 41, 42) + assert all(type(stage) is D.DeepseekV4DSparkStage for stage in model.mtp) + assert not model.has_mtp # generic preview-MTP routing is intentionally absent + keys = {k for k, _ in tree_flatten(model.parameters())} + assert "mtp.0.main_proj.weight" in keys + assert "mtp.2.markov_head.markov_w1.weight" in keys + assert "mtp.2.confidence_head.proj.weight" in keys + assert "mtp.2.confidence_head.proj.bias" not in keys + assert type(model._dspark.stages[0]) is D.DeepseekV4DSparkStage + assert model._dspark.stages[0].main_proj is not None + assert model._dspark.stages[1].main_proj is None + assert model._dspark.stages[2].markov_head is not None + + +@pytest.mark.parametrize( + "mut", + [ + {"dspark_block_size": 4}, + {"num_nextn_predict_layers": 0}, + {"dspark_noise_token_id": 0}, + {"dspark_target_layer_ids": [39, 41, 42]}, + {"dspark_markov_rank": 255}, + ], +) +def test_dspark_manifest_fails_loudly_on_missing_or_wrong_invariants(mut): + with pytest.raises(ValueError, match="DSpark"): + D._validate_dspark_manifest(_args(**mut)) + + +def test_partial_dspark_signature_never_becomes_legacy_mtp(): + corrupt = _args(dspark_block_size=0) + with pytest.raises(ValueError, match="dspark_block_size=5"): + D.Model(corrupt) + partial_config = dict(_OFFICIAL_FILTERED, dspark_block_size=0) + assert D.is_deepseek_v4_mtp_config(partial_config) is False + assert D._has_dspark_signature(D.ModelArgs(dspark_block_size=0)) is True + assert ( + D.is_deepseek_v4_mtp_config( + { + "model_type": "deepseek_v4", + "num_nextn_predict_layers": 1, + "dspark_block_size": 0, + } + ) + is False + ) + + +@pytest.mark.parametrize( + "null_key", + [ + "dspark_block_size", + "dspark_noise_token_id", + "dspark_target_layer_ids", + "dspark_markov_rank", + ], +) +def test_explicit_null_dspark_key_selects_validation_not_legacy_mtp(null_key): + config = vars(_args()).copy() + for key in ( + "dspark_block_size", + "dspark_noise_token_id", + "dspark_target_layer_ids", + "dspark_markov_rank", + ): + config.pop(key) + config[null_key] = None + + args = D.ModelArgs.from_dict(config) + + assert D._has_dspark_signature(args) is True + with pytest.raises(ValueError, match="DSpark-0731"): + D.Model(args) + + +def test_start_zero_runs_attention_only_on_all_three_stages(tiny_model): + calls = [] + h = mx.zeros((1, 5, 4, 8), dtype=mx.float32) + main_x = mx.zeros((1, 3, 8), dtype=mx.float32) + + class AttentionSpy: + def __init__(self, stage_id): + self.stage_id = stage_id + + def __call__(self, x, *, start_pos, main_x, cache): + calls.append((self.stage_id, start_pos, cache)) + return x + + def forbidden(*_args, **_kwargs): + raise AssertionError("HC/FFN path ran during DSpark prefill") + + for stage in tiny_model._dspark.stages: + stage.attn = AttentionSpy(stage.stage_id) + stage.attn_hc.pre = forbidden + stage.ffn = forbidden + cache = object() + got = stage(h, start_pos=0, main_x=main_x, cache=cache) + assert got is h + assert [(stage_id, start_pos) for stage_id, start_pos, _ in calls] == [ + (0, 0), + (1, 0), + (2, 0), + ] + + +def test_positive_start_runs_all_three_full_stages(tiny_model, monkeypatch): + calls = [] + + def full_stage(self, h, *, start_pos, cache=None, input_ids=None, main_x=None): + calls.append((self.stage_id, start_pos, cache)) + return h + + monkeypatch.setattr(D.DeepseekV4DSparkStage, "__call__", full_stage) + tiny_model._dspark.finish = lambda *_args, **_kwargs: "draft-output" + result = tiny_model._dspark.forward( + mx.zeros((1, 1, 24), dtype=mx.float32), + mx.array([7], dtype=mx.int32), + tiny_model.model.embed_tokens, + tiny_model.lm_head, + tiny_model.make_dspark_cache(), + start_pos=9, + greedy=True, + ) + assert result == "draft-output" + assert [(stage_id, start_pos) for stage_id, start_pos, _ in calls] == [ + (0, 9), + (1, 9), + (2, 9), + ] + + +def test_ids_only_forward_projects_only_requested_m3_rows(tiny_model, monkeypatch): + projected = [] + forced_primary = mx.array([19], dtype=mx.int32) + forced_seen = [] + + def full_stage(self, h, **_kwargs): + return h + + def lm_head(rows): + projected.append(tuple(rows.shape)) + return mx.zeros((rows.shape[0], rows.shape[1], 128800)) + + monkeypatch.setattr(D.DeepseekV4DSparkStage, "__call__", full_stage) + tiny_model._dspark.finish = lambda *_args, **_kwargs: (_ for _ in ()).throw( + AssertionError("ids-only forward used the confidence/logit-stack route") + ) + + def finish_ids(logits, _target_ids, *, width, forced_first_token_ids=None): + forced_seen.append(forced_first_token_ids) + return "ids-only", tuple(logits.shape), width + + tiny_model._dspark.finish_ids = finish_ids + + result = tiny_model._dspark.forward( + mx.zeros((1, 1, 24), dtype=mx.float32), + mx.array([7], dtype=mx.int32), + tiny_model.model.embed_tokens, + lm_head, + tiny_model.make_dspark_cache(), + start_pos=9, + greedy=True, + ids_only_width=3, + forced_first_token_ids=forced_primary, + ) + + assert result == ("ids-only", (1, 3, 128800), 3) + assert projected == [(1, 3, 8)] + assert forced_seen == [forced_primary] + + +def test_zero_start_model_prefill_returns_no_draft_output(tiny_model, monkeypatch): + calls = [] + + def attention_only(self, h, *, start_pos, cache=None, input_ids=None, main_x=None): + calls.append((self.stage_id, start_pos)) + return h + + monkeypatch.setattr(D.DeepseekV4DSparkStage, "__call__", attention_only) + result = tiny_model._dspark.forward( + mx.zeros((1, 3, 24), dtype=mx.float32), + mx.array([7], dtype=mx.int32), + tiny_model.model.embed_tokens, + tiny_model.lm_head, + tiny_model.make_dspark_cache(), + start_pos=0, + ) + assert result is None + assert calls == [(0, 0), (1, 0), (2, 0)] + + +def test_dspark_prefill_seeds_all_three_stage_owned_caches(tiny_model): + caches = tiny_model.make_dspark_cache() + hidden = mx.arange(3 * 24, dtype=mx.float32).reshape(1, 3, 24) + + tiny_model._dspark.prefill(hidden, caches) + + mx.eval(*(cache.ring for cache in caches)) + assert [cache.prefill_length for cache in caches] == [3, 3, 3] + assert all(cache.ring.shape == (1, 8, 8) for cache in caches) + assert len({id(cache.ring) for cache in caches}) == 3 + + +def test_dspark_commit_main_updates_every_stage_ring_with_accepted_prefix(tiny_model): + dspark = tiny_model._dspark + caches = tiny_model.make_dspark_cache() + dspark.prefill(mx.zeros((1, 3, 24), dtype=mx.float32), caches) + accepted = mx.arange(3 * 24, dtype=mx.float32).reshape(1, 3, 24) + main_x = dspark.stages[0].fuse_main(accepted) + positions = mx.arange(7, 10) + expected = [stage.attn._kv(main_x, positions) for stage in dspark.stages] + + dspark.commit_main(accepted, caches, start_pos=7) + + mx.eval(*(cache.ring for cache in caches), *expected) + for cache, stage_expected in zip(caches, expected): + got = cache.ring[:, [7, 0, 1]] + assert mx.allclose(got, stage_expected, rtol=1e-6, atol=1e-6) + + +def test_target_taps_are_hc_collapsed_in_exact_order(tiny_model): + args = tiny_model.args + model = tiny_model + # Avoid a 43-layer numerical run: each selected layer emits a distinctive HC + # tensor, and the target collector must collapse it at the layer boundary. + for i in (40, 41, 42): + value = float(i) + model.model.layers[i] = lambda h, *a, _v=value, **k: mx.full(h.shape, _v) + h = mx.zeros((1, 1, args.hc_mult, args.hidden_size)) + taps = model._collect_dspark_taps(h, start_layer=40) + got = np.array(taps) + assert got.shape == (1, 1, 3 * args.hidden_size) + assert np.all(got[..., :8] == 40) + assert np.all(got[..., 8:16] == 41) + assert np.all(got[..., 16:] == 42) + + +def test_draft_ids_are_target_then_four_noise_tokens(tiny_model): + dspark = tiny_model._dspark + ids = dspark.draft_input_ids(mx.array([7, 9], dtype=mx.int32)) + assert np.array(ids).tolist() == [ + [7, 128799, 128799, 128799, 128799], + [9, 128799, 128799, 128799, 128799], + ] + + +def test_stage_caches_are_distinct_and_stage_owned(tiny_model): + model = tiny_model + caches = model.make_dspark_cache() + assert len(caches) == 3 + assert len({id(c) for c in caches}) == 3 + assert all(type(c) is D.DeepseekV4DSparkCache for c in caches) + assert [c.window_size for c in caches] == [8, 8, 8] + + +def test_dspark_missing_stage_weights_fail_at_load_boundary(tiny_model): + with pytest.raises(ValueError, match=r"mtp\.1\.\*"): + tiny_model.sanitize({"mtp.0.main_proj.weight": mx.zeros((8, 24))}) + + +def test_markov_is_sequential_and_confidence_is_fp32(tiny_model): + dspark = tiny_model._dspark + # Make the Markov head depend only on the previous sampled token; this avoids + # accidental dependence on the target logits while proving the recurrence. + stage = dspark.stages[-1] + stage.markov_head.markov_w1.weight = mx.arange( + 128800 * 3, dtype=mx.float32 + ).reshape(128800, 3) + stage.markov_head.markov_w2.weight = mx.ones((128800, 3), dtype=mx.float32) + logits = np.zeros((1, 5, 128800), dtype=np.float32) + hidden = mx.ones((1, 5, 8), dtype=mx.float32) + # Greedy target row selects 1, while row 1's Markov-biased output selects a + # different token once the previous id is changed. + logits[:, :, 1] = 1.0 + logits = mx.array(logits) + ids_a, out_a, conf_a = dspark.finish( + logits, hidden, mx.array([2], dtype=mx.int32), greedy=True + ) + ids_b, out_b, conf_b = dspark.finish( + logits, hidden, mx.array([3], dtype=mx.int32), greedy=True + ) + assert np.array(ids_a)[:, 0].tolist() == [2] + assert np.array(ids_b)[:, 0].tolist() == [3] + assert not np.array_equal(np.array(out_a)[:, 0], np.array(out_b)[:, 0]) + assert conf_a.dtype == mx.float32 and conf_b.dtype == mx.float32 + + +def test_greedy_ids_only_finish_seeds_future_recurrence_from_target_primary( + tiny_model, +): + dspark = tiny_model._dspark + stage = dspark.stages[-1] + + class NextTokenMarkov: + def __call__(self, previous): + previous = np.asarray(previous, dtype=np.int32) + bias = np.full((previous.shape[0], 128800), -1000.0, dtype=np.float32) + bias[np.arange(previous.shape[0]), previous + 1] = 1000.0 + return mx.array(bias), mx.zeros((previous.shape[0], 1)) + + stage.markov_head = NextTokenMarkov() + logits = mx.zeros((1, 3, 128800), dtype=mx.float32) + target_ids = mx.array([7], dtype=mx.int32) + primary_ids = mx.array([19], dtype=mx.int32) + + ids = dspark.finish_ids( + logits, + target_ids, + width=3, + forced_first_token_ids=primary_ids, + ) + + mx.eval(ids) + assert np.asarray(ids).tolist() == [[7, 19, 20, 21]] + + +def test_official_hc_names_map_to_exact_installed_parameter_keys(tiny_model): + raw = {} + expected = set() + for stage in range(3): + raw[f"mtp.{stage}.hc_attn_fn"] = mx.zeros((1,)) + raw[f"mtp.{stage}.hc_attn_base"] = mx.zeros((1,)) + raw[f"mtp.{stage}.hc_attn_scale"] = mx.zeros((1,)) + raw[f"mtp.{stage}.hc_ffn_fn"] = mx.zeros((1,)) + raw[f"mtp.{stage}.hc_ffn_base"] = mx.zeros((1,)) + raw[f"mtp.{stage}.hc_ffn_scale"] = mx.zeros((1,)) + expected |= { + f"mtp.{stage}.attn_hc.fn", + f"mtp.{stage}.attn_hc.base", + f"mtp.{stage}.attn_hc.scale", + f"mtp.{stage}.ffn_hc.fn", + f"mtp.{stage}.ffn_hc.base", + f"mtp.{stage}.ffn_hc.scale", + } + raw |= { + "mtp.2.hc_head_fn": mx.zeros((1,)), + "mtp.2.hc_head_base": mx.zeros((1,)), + "mtp.2.hc_head_scale": mx.zeros((1,)), + } + expected |= {"mtp.2.hc_head.fn", "mtp.2.hc_head.base", "mtp.2.hc_head.scale"} + mapped = tiny_model.sanitize(raw) + assert set(mapped) == expected + installed = {k for k, _ in tree_flatten(tiny_model.parameters())} + assert expected <= installed + + +def test_dspark_visibility_is_exact_and_includes_all_five_draft_rows(): + got = np.array(D.get_dspark_topk_idxs(8, 2, 5, 3)) + expected = [0, 1, 2, 3, 8, 9, 10, 11, 12] + assert got.shape == (2, 5, 9) + assert got.tolist() == [[expected] * 5] * 2 + wrapped = np.array(D.get_dspark_topk_idxs(8, 1, 5, 10)) + assert wrapped.tolist() == [[list(range(8)) + [8, 9, 10, 11, 12]] * 5] + + +def _rms(x, eps=1e-6): + return x / np.sqrt(np.mean(x * x, axis=-1, keepdims=True) + eps) + + +def _rope(x, positions, inverse=False): + inv = 1.0 / (10000.0 ** (np.arange(0, 4, 2, dtype=np.float64) / 4)) + ang = np.asarray(positions, dtype=np.float64)[:, None] * inv[None] + cos, sin = np.cos(ang), np.sin(ang) + if inverse: + sin = -sin + out = x.copy() + tail = x[..., -4:].reshape(*x.shape[:-1], 2, 2) + a, b = tail[..., 0], tail[..., 1] + trig_shape = (1, len(positions)) + (1,) * (a.ndim - 3) + (2,) + rot = np.stack( + [ + a * cos.reshape(trig_shape) - b * sin.reshape(trig_shape), + a * sin.reshape(trig_shape) + b * cos.reshape(trig_shape), + ], + axis=-1, + ) + out[..., -4:] = rot.reshape(*x.shape[:-1], 4) + return out + + +def _identity_dspark_attention(): + attn = D.DeepseekV4DSparkAttention(_args(), 43) + eye = mx.eye(8, dtype=mx.float32) + attn.wq_a.weight = eye + attn.wq_b.weight = eye + attn.wkv.weight = eye + attn.wo_a.weight = eye + attn.wo_b.weight = eye + attn.q_norm.weight = mx.ones((8,), dtype=mx.float32) + attn.kv_norm.weight = mx.ones((8,), dtype=mx.float32) + attn.attn_sink = mx.array([0.2], dtype=mx.float32) + return attn + + +def _np_dspark_oracle(prefill, current, draft, start_pos, win=8): + prekv = _rope(_rms(prefill), np.arange(prefill.shape[1])) + if prefill.shape[1] <= win: + ring = np.concatenate([prekv, np.zeros((1, win - prefill.shape[1], 8))], axis=1) + else: + last = prekv[:, -win:] + cut = prefill.shape[1] % win + ring = ( + last + if cut == 0 + else np.concatenate([last[:, win - cut :], last[:, : win - cut]], axis=1) + ) + mainkv = _rope(_rms(current), [start_pos]) + ring[:, start_pos % win : start_pos % win + 1] = mainkv + positions = np.arange(start_pos + 1, start_pos + 6) + q = _rms(_rms(draft))[:, :, None, :] + q = _rope(q, positions) + dkv = _rope(_rms(draft), positions) + full = np.concatenate([ring, dkv], axis=1) + idx = list(range(min(win, start_pos + 1))) + list(range(win, win + 5)) + visible = full[:, idx] + scores = np.einsum("bshd,btd->bhst", q, visible) * (8**-0.5) + sink = np.array(0.2).reshape(1, 1, 1, 1) + maximum = np.maximum(scores.max(-1, keepdims=True), sink) + exp = np.exp(scores - maximum) + probs = exp / (exp.sum(-1, keepdims=True) + np.exp(sink - maximum)) + out = np.einsum("bhst,btd->bshd", probs, visible) + return _rope(out, positions, inverse=True).reshape(1, 5, 8), ring + + +@pytest.mark.parametrize("prefill_len", [3, 10]) +def test_dspark_attention_matches_prefill_decode_oracle_and_ring_wrap(prefill_len): + rng = np.random.default_rng(13 + prefill_len) + prefill = rng.normal(size=(1, prefill_len, 8)).astype(np.float32) + current = rng.normal(size=(1, 1, 8)).astype(np.float32) + draft = rng.normal(size=(1, 5, 8)).astype(np.float32) + attn = _identity_dspark_attention() + cache = D.DeepseekV4DSparkCache(8, 8) + dummy = mx.zeros((1, 5, 8), dtype=mx.float32) + assert tuple( + attn(dummy, start_pos=0, main_x=mx.array(prefill), cache=cache).shape + ) == (1, 5, 8) + got = attn( + mx.array(draft), start_pos=prefill_len, main_x=mx.array(current), cache=cache + ) + ref, ref_ring = _np_dspark_oracle(prefill, current, draft, prefill_len) + assert np.allclose(np.array(got), ref, rtol=3e-5, atol=3e-5) + assert np.allclose(np.array(cache.ring), ref_ring, rtol=2e-6, atol=2e-6) + + +def test_dspark_cache_commits_an_accepted_prefix_across_ring_wrap(): + cache = D.DeepseekV4DSparkCache(4, 2) + cache.prefill(mx.array([[[0.0, 0.5], [1.0, 1.5], [2.0, 2.5]]])) + accepted = mx.array([[[30.0, 30.5], [40.0, 40.5], [50.0, 50.5]]]) + + cache.commit_main(3, accepted) + + assert np.array(cache.ring).tolist() == [ + [[40.0, 40.5], [50.0, 50.5], [2.0, 2.5], [30.0, 30.5]] + ] + + +def test_dspark_seeded_gumbel_matches_reference_and_greedy_is_explicit(): + logits = mx.array([[0.1, 1.2, -0.7, 0.8]], dtype=mx.float32) + key = mx.random.key(90210) + got = D._sample_dspark_token(logits, 0.7, key=key) + uniform = mx.random.uniform(shape=logits.shape, key=key) + uniform = mx.clip(uniform, 1e-30, 1.0 - mx.finfo(mx.float32).eps) + ref = mx.argmax(logits / 0.7 - mx.log(-mx.log(uniform)), axis=-1) + assert np.array_equal(np.array(got), np.array(ref)) + assert np.array(D._sample_dspark_token(logits, 0.7, greedy=True)).tolist() == [1] From 8c847d4585cfa44cd6b13b505fda81cd088652ef Mon Sep 17 00:00:00 2001 From: davidtai Date: Wed, 12 Aug 2026 19:08:25 -0500 Subject: [PATCH 02/24] Add construction-bound DeepSeek V4 0731 DSpark runtime --- mtplx/benchmarks/runners/mtp_depth_grid.py | 58 +- mtplx/benchmarks/runners/mtp_depth_sweep.py | 35 +- mtplx/cli.py | 26 +- mtplx/commands/public.py | 199 ++- mtplx/deepseek_v4_0731_dspark_ffn.py | 353 +++++ mtplx/deepseek_v4_0731_full_install.py | 484 +++++++ mtplx/deepseek_v4_0731_m3_target.py | 191 +++ mtplx/deepseek_v4_0731_m3_wob.py | 416 ++++++ mtplx/deepseek_v4_0731_m3_wqb_qnorm_rope.py | 514 ++++++++ mtplx/deepseek_v4_0731_moe.py | 217 ++++ mtplx/deepseek_v4_attention_island.py | 187 ++- mtplx/deepseek_v4_dspark_generation.py | 176 +++ mtplx/generation.py | 607 +++++---- mtplx/models/deepseek_v4.py | 484 ++++--- mtplx/native_block_speculation.py | 536 ++++++++ mtplx/runtime.py | 357 ++++-- mtplx/server/openai.py | 179 ++- scripts/deepseek_v4_0731_k2_bench.py | 1114 ++++++++++++++++ scripts/deepseek_v4_guard_window.py | 13 +- tests/test_deepseek_v4_0731_dspark_ffn.py | 496 ++++++++ tests/test_deepseek_v4_0731_full_install.py | 663 ++++++++++ tests/test_deepseek_v4_0731_k2_bench.py | 762 +++++++++++ tests/test_deepseek_v4_0731_m3_target.py | 229 ++++ tests/test_deepseek_v4_0731_m3_wob.py | 260 ++++ ...test_deepseek_v4_0731_m3_wqb_qnorm_rope.py | 297 +++++ tests/test_deepseek_v4_0731_moe.py | 165 +++ tests/test_deepseek_v4_attention_island.py | 181 ++- tests/test_deepseek_v4_dspark.py | 12 + tests/test_deepseek_v4_dspark_generation.py | 1131 +++++++++++++++++ tests/test_deepseek_v4_loader.py | 185 ++- ...neration_deepseek_v4_dspark_integration.py | 108 ++ tests/test_mtp_depth_grid.py | 68 + tests/test_mtp_depth_sweep.py | 95 +- tests/test_public_cli.py | 219 +++- tests/test_runtime_deepseek_v4_dspark.py | 605 +++++++++ tests/test_server_openai.py | 237 ++++ 36 files changed, 11159 insertions(+), 700 deletions(-) create mode 100644 mtplx/deepseek_v4_0731_dspark_ffn.py create mode 100644 mtplx/deepseek_v4_0731_full_install.py create mode 100644 mtplx/deepseek_v4_0731_m3_target.py create mode 100644 mtplx/deepseek_v4_0731_m3_wob.py create mode 100644 mtplx/deepseek_v4_0731_m3_wqb_qnorm_rope.py create mode 100644 mtplx/deepseek_v4_0731_moe.py create mode 100644 mtplx/deepseek_v4_dspark_generation.py create mode 100644 mtplx/native_block_speculation.py create mode 100644 scripts/deepseek_v4_0731_k2_bench.py create mode 100644 tests/test_deepseek_v4_0731_dspark_ffn.py create mode 100644 tests/test_deepseek_v4_0731_full_install.py create mode 100644 tests/test_deepseek_v4_0731_k2_bench.py create mode 100644 tests/test_deepseek_v4_0731_m3_target.py create mode 100644 tests/test_deepseek_v4_0731_m3_wob.py create mode 100644 tests/test_deepseek_v4_0731_m3_wqb_qnorm_rope.py create mode 100644 tests/test_deepseek_v4_0731_moe.py create mode 100644 tests/test_deepseek_v4_dspark_generation.py create mode 100644 tests/test_generation_deepseek_v4_dspark_integration.py create mode 100644 tests/test_mtp_depth_grid.py create mode 100644 tests/test_runtime_deepseek_v4_dspark.py diff --git a/mtplx/benchmarks/runners/mtp_depth_grid.py b/mtplx/benchmarks/runners/mtp_depth_grid.py index 2abcd751..571e2fe8 100644 --- a/mtplx/benchmarks/runners/mtp_depth_grid.py +++ b/mtplx/benchmarks/runners/mtp_depth_grid.py @@ -52,10 +52,7 @@ def _mean_present(values: list[float | None]) -> float | None: def _rate_by_depth(accepted: list[int], drafted: list[int]) -> list[float | None]: - return [ - (a / d if d else None) - for a, d in zip(accepted, drafted) - ] + return [(a / d if d else None) for a, d in zip(accepted, drafted)] def _sum_lists(values: list[list[int]], length: int) -> list[int]: @@ -66,6 +63,10 @@ def _sum_lists(values: list[list[int]], length: int) -> list[int]: return totals +def _cycle_count(events: list[dict], verify_calls: int) -> int: + return len(events) or max(0, int(verify_calls)) + + def run_mtp_depth_policy_grid( model_path: Path | str, prompt_suite: Path | str, @@ -178,7 +179,7 @@ def run_mtp_depth_policy_grid( validations = [asdict(validate_no_degenerate_loop(out.text))] if case.category == "json_tool": validations.append(asdict(validate_json_text(out.text.strip()))) - cycles = len(out.stats.events) + cycles = _cycle_count(out.stats.events, out.stats.verify_calls) ar_row = ar_rows[index] if compare_ar else None row: dict[str, Any] = { "prompt_id": case.id, @@ -249,18 +250,30 @@ def run_mtp_depth_policy_grid( "rows": rows, "summary": { "prompts": len(rows), - "generated_tokens": sum(row["generated_tokens"] for row in rows), - "mean_tok_s": statistics.mean([row["tok_s"] for row in rows]) if rows else 0.0, - "mean_ar_tok_s": _mean_present([row["ar_tok_s"] for row in rows]), - "mean_speedup_vs_ar": _mean_present([row["speedup_vs_ar"] for row in rows]), - "mean_model_path_tok_s": _mean_present([row["model_path_tok_s"] for row in rows]), + "generated_tokens": sum( + row["generated_tokens"] for row in rows + ), + "mean_tok_s": statistics.mean([row["tok_s"] for row in rows]) + if rows + else 0.0, + "mean_ar_tok_s": _mean_present( + [row["ar_tok_s"] for row in rows] + ), + "mean_speedup_vs_ar": _mean_present( + [row["speedup_vs_ar"] for row in rows] + ), + "mean_model_path_tok_s": _mean_present( + [row["model_path_tok_s"] for row in rows] + ), "cycles": sum(row["cycles"] for row in rows), "accepted_drafts": sum(row["accepted_drafts"] for row in rows), "rejected_drafts": sum(row["rejected_drafts"] for row in rows), "drafted_tokens": sum(row["drafted_tokens"] for row in rows), "accepted_by_depth": accepted_by_depth, "drafted_by_depth": drafted_by_depth, - "acceptance_by_depth": _rate_by_depth(accepted_by_depth, drafted_by_depth), + "acceptance_by_depth": _rate_by_depth( + accepted_by_depth, drafted_by_depth + ), "accepted_drafts_per_cycle": ( sum(row["accepted_drafts"] for row in rows) / max(1, sum(row["cycles"] for row in rows)) @@ -277,17 +290,24 @@ def run_mtp_depth_policy_grid( ), "verify_time_s": sum(row["verify_time_s"] for row in rows), "draft_time_s": sum(row["draft_time_s"] for row in rows), - "target_forward_time_s": sum(row["target_forward_time_s"] for row in rows), - "validations_passed": sum(1 for v in validations if v["passed"]), + "target_forward_time_s": sum( + row["target_forward_time_s"] for row in rows + ), + "validations_passed": sum( + 1 for v in validations if v["passed"] + ), "validations_total": len(validations), - "peak_memory_bytes": max([row["peak_memory_bytes"] for row in rows] or [0]), + "peak_memory_bytes": max( + [row["peak_memory_bytes"] for row in rows] or [0] + ), }, } ) results.sort( key=lambda item: ( - item["summary"]["validations_passed"] == item["summary"]["validations_total"], + item["summary"]["validations_passed"] + == item["summary"]["validations_total"], item["summary"]["mean_tok_s"], ), reverse=True, @@ -306,9 +326,13 @@ def run_mtp_depth_policy_grid( "mtp_cache_policy": mtp_cache_policy, "mtp_history_policy": mtp_history_policy, "verify_strategy": verify_strategy, - "mtp_corrector_path": str(mtp_corrector_path) if mtp_corrector_path is not None else None, + "mtp_corrector_path": str(mtp_corrector_path) + if mtp_corrector_path is not None + else None, "mtp_corrector_blend": mtp_corrector_blend, - "mtp_corrector_kind": getattr(mtp_corrector, "kind", None) if mtp_corrector is not None else None, + "mtp_corrector_kind": getattr(mtp_corrector, "kind", None) + if mtp_corrector is not None + else None, "thresholds": threshold_values, "min_depths": min_depth_values, "ar_rows": ar_rows, diff --git a/mtplx/benchmarks/runners/mtp_depth_sweep.py b/mtplx/benchmarks/runners/mtp_depth_sweep.py index 43e4bc58..3b707e83 100644 --- a/mtplx/benchmarks/runners/mtp_depth_sweep.py +++ b/mtplx/benchmarks/runners/mtp_depth_sweep.py @@ -33,11 +33,17 @@ def _rate_by_depth(accepted: list[int], drafted: list[int]) -> list[float | None return [(a / d if d else None) for a, d in zip(accepted, drafted)] +def _cycle_count(events: list[dict], verify_calls: int) -> int: + return len(events) or max(0, int(verify_calls)) + + def _token_budget(max_tokens: int, case_max_tokens: int) -> int: return min(int(max_tokens), int(case_max_tokens)) -def _hit_token_budget(generated_tokens: int, token_budget: int, finish_reason: str | None) -> bool: +def _hit_token_budget( + generated_tokens: int, token_budget: int, finish_reason: str | None +) -> bool: if finish_reason == "length": return True return int(generated_tokens) >= int(token_budget) @@ -133,8 +139,12 @@ def run_mtp_depth_sweep( ) contract = getattr(rt, "contract", None) is_gemma4_assistant = getattr(rt, "backend_id", None) == "gemma4_assistant" - resolved_base_hidden_variant = str(getattr(contract, "base_hidden_variant", "gemma4_assistant")) - resolved_mtp_hidden_variant = str(getattr(contract, "hidden_variant", "gemma4_assistant")) + resolved_base_hidden_variant = str( + getattr(contract, "base_hidden_variant", "gemma4_assistant") + ) + resolved_mtp_hidden_variant = str( + getattr(contract, "hidden_variant", "gemma4_assistant") + ) resolved_concat_order = str(getattr(contract, "concat_order", "assistant_pair")) draft_lm_head_report: dict[str, Any] | None = None if draft_lm_head_bits is not None and not is_gemma4_assistant: @@ -287,6 +297,7 @@ def run_mtp_depth_sweep( ) ] ar_row = ar_rows[index] if compare_ar else None + cycles = _cycle_count(out.stats.events, out.stats.verify_calls) rows.append( { "prompt_id": case.id, @@ -344,7 +355,7 @@ def run_mtp_depth_sweep( out.stats.drafted_by_depth, ), "mean_accepted_drafts_per_cycle": ( - out.stats.accepted_drafts / max(1, len(out.stats.events)) + out.stats.accepted_drafts / max(1, cycles) ), "acceptance_rate": ( out.stats.accepted_drafts / out.stats.drafted_tokens @@ -392,7 +403,9 @@ def run_mtp_depth_sweep( validations = [v for row in rows for v in row["validations"]] finish_reasons = _finish_reason_counts(rows) - accepted_by_depth = _sum_lists([row["accepted_by_depth"] for row in rows], depth) + accepted_by_depth = _sum_lists( + [row["accepted_by_depth"] for row in rows], depth + ) drafted_by_depth = _sum_lists([row["drafted_by_depth"] for row in rows], depth) accept_probability_sum_by_depth = _sum_float_lists( [row["accept_probability_sum_by_depth"] for row in rows], @@ -408,12 +421,8 @@ def run_mtp_depth_sweep( target_distribution_windows = sum( row["target_distribution_materialized_windows"] for row in rows ) - lazy_bonus_verify_calls = sum( - row["lazy_bonus_verify_calls"] for row in rows - ) - lazy_bonus_commit_time_s = sum( - row["lazy_bonus_commit_time_s"] for row in rows - ) + lazy_bonus_verify_calls = sum(row["lazy_bonus_verify_calls"] for row in rows) + lazy_bonus_commit_time_s = sum(row["lazy_bonus_commit_time_s"] for row in rows) depth_results.append( { "depth": depth, @@ -512,9 +521,7 @@ def run_mtp_depth_sweep( "verify_target_distribution_time_s": ( verify_target_distribution_time_s ), - "target_distribution_materialized_rows": ( - target_distribution_rows - ), + "target_distribution_materialized_rows": (target_distribution_rows), "target_distribution_materialized_windows": ( target_distribution_windows ), diff --git a/mtplx/cli.py b/mtplx/cli.py index 40bf633c..95a90c95 100644 --- a/mtplx/cli.py +++ b/mtplx/cli.py @@ -652,6 +652,14 @@ def _add_mtp_toggle_args(parser: argparse.ArgumentParser) -> None: "where live switching is supported." ), ) + parser.add_argument( + "--deepseek-v4-0731-k2", + action="store_true", + help=( + "Select the exact construction-bound DeepSeek-V4-Flash-0731 " + "DSpark K2 stack. Requires explicit --depth 2 and MTP." + ), + ) SCHEDULER_MODE_CHOICES = ( @@ -2376,7 +2384,11 @@ def build_parser() -> argparse.ArgumentParser: ) quickstart_server_p.add_argument("--host", default="127.0.0.1") quickstart_server_p.add_argument("--port", type=int, default=8000) - quickstart_server_p.add_argument("--model-id", default=DEFAULT_PUBLIC_MODEL_ID, help="Served OpenAI model id; defaults to the loaded artifact identity") + quickstart_server_p.add_argument( + "--model-id", + default=DEFAULT_PUBLIC_MODEL_ID, + help="Served OpenAI model id; defaults to the loaded artifact identity", + ) quickstart_server_p.add_argument( "--embedding-model", action="append", @@ -2417,8 +2429,16 @@ def build_parser() -> argparse.ArgumentParser: "(jina-style model.py/rerank.py) to execute it; off by default" ), ) - quickstart_server_p.add_argument("--dry-run", action="store_true", help="Preview the server launch command without loading MLX") - quickstart_server_p.add_argument("--json", action="store_true", help="Emit machine-readable JSON for --dry-run and errors") + quickstart_server_p.add_argument( + "--dry-run", + action="store_true", + help="Preview the server launch command without loading MLX", + ) + quickstart_server_p.add_argument( + "--json", + action="store_true", + help="Emit machine-readable JSON for --dry-run and errors", + ) quickstart_server_p.add_argument("--depth", type=int, default=3) _add_mtp_toggle_args(quickstart_server_p) quickstart_server_p.add_argument( diff --git a/mtplx/commands/public.py b/mtplx/commands/public.py index 5b983e32..3876d890 100644 --- a/mtplx/commands/public.py +++ b/mtplx/commands/public.py @@ -28,7 +28,7 @@ from types import SimpleNamespace from typing import Any, Callable -from mtplx.artifacts import inspect_model +from mtplx.artifacts import inspect_model, load_config from mtplx.benchmarks.validators.basic import ( summarize_benchmark_quality, validate_balanced_delimiters, @@ -657,6 +657,38 @@ def _validate_public_depth(args: Any, *, printer=print) -> int | None: return None +def _deepseek_v4_0731_k2_entrypoint_error(args: Any) -> dict[str, str] | None: + """Reject an explicit K2 selection before any public model construction.""" + + if not bool(getattr(args, "deepseek_v4_0731_k2", False)): + return None + cli_flags = getattr(args, "_cli_flags", set()) or set() + if "depth" not in cli_flags or int(getattr(args, "depth", 3)) != 2: + return { + "error": "DeepSeek-V4-0731 K2 requires explicit --depth 2", + "detail": "The selected construction owns exactly two future drafts.", + } + if ( + _generation_mode_from_args(args) != GENERATION_MODE_MTP + or getattr(args, "load_mtp", True) is False + or bool(getattr(args, "no_mtp", False)) + ): + return { + "error": "DeepSeek-V4-0731 K2 requires MTP generation", + "detail": "Remove target-only AR or --no-load-mtp/--no-mtp options.", + } + return None + + +def _runtime_load_kwargs(args: Any) -> dict[str, Any]: + kwargs: dict[str, Any] = { + "mtp": getattr(args, "load_mtp", True) is not False, + } + if bool(getattr(args, "deepseek_v4_0731_k2", False)): + kwargs["deepseek_v4_0731_k2"] = True + return kwargs + + def _normalize_generation_mode(value: Any) -> str: text = str(value or GENERATION_MODE_MTP).strip().lower() if text == "auto": @@ -8320,6 +8352,14 @@ def cmd_serve_public(args: Any) -> int: if depth_error is not None: return depth_error generation_mode = _generation_mode_from_args(args) + k2_error = _deepseek_v4_0731_k2_entrypoint_error(args) + if k2_error is not None: + _print_command_error( + k2_error, + command=_server_command_name(args), + json_output=bool(getattr(args, "json", False)), + ) + return 2 fan_mode = _fan_mode_from_args(args) if ( generation_mode == GENERATION_MODE_MTP @@ -8592,6 +8632,14 @@ def cmd_serve_public(args: Any) -> int: _apply_backend_serve_defaults(args, inspection) _apply_qwen36_35b_optimized_speed_defaults(args, model_id) backend_descriptor = descriptor_from_inspection(inspection) + from mtplx.models.deepseek_v4 import _config_has_dspark_signature + + try: + dspark_request_defaults = _config_has_dspark_signature( + load_config(runtime_model) + ) + except (FileNotFoundError, OSError, ValueError): + dspark_request_defaults = False draft_lm_head = _model_draft_lm_head_spec(inspection, profile) or { "bits": 4, "group_size": 64, @@ -8601,8 +8649,48 @@ def cmd_serve_public(args: Any) -> int: draft_sampler_override = _explicit_draft_sampler_override(args, draft_sampler) if draft_sampler_override is not None: draft_sampler = draft_sampler_override + cli_flags = getattr(args, "_cli_flags", set()) or set() + if dspark_request_defaults and ("depth" not in cli_flags or int(args.depth) != 2): + _print_command_error( + { + "error": "DeepSeek-V4-0731 DSpark requires explicit --depth 2", + "detail": ( + "The installed proposer supports exactly two future drafts; " + "other depths are not silently coerced." + ), + }, + command=_server_command_name(args), + json_output=bool(getattr(args, "json", False)), + ) + return 2 if not quiet_json: _print_serve_handoff(args, runtime_model, profile.name) + depth_args = ( + ["--depth", str(args.depth)] + if not dspark_request_defaults or "depth" in cli_flags or int(args.depth) != 3 + else [] + ) + verify_strategy = str( + getattr(args, "verify_strategy", "capture_commit") or "capture_commit" + ) + verify_strategy_args = ( + ["--verify-strategy", verify_strategy] + if not dspark_request_defaults + or "verify-strategy" in cli_flags + or verify_strategy != "capture_commit" + else [] + ) + verify_core = str( + getattr(args, "verify_core", "linear-gdn-from-conv-tape") + or "linear-gdn-from-conv-tape" + ) + verify_core_args = ( + ["--verify-core", verify_core] + if not dspark_request_defaults + or "verify-core" in cli_flags + or verify_core != "linear-gdn-from-conv-tape" + else [] + ) cmd = [ sys.executable, "-m", @@ -8615,8 +8703,7 @@ def cmd_serve_public(args: Any) -> int: args.host, "--port", str(args.port), - "--depth", - str(args.depth), + *depth_args, "--generation-mode", generation_mode, "--profile", @@ -8627,13 +8714,8 @@ def cmd_serve_public(args: Any) -> int: _pi_preserve_thinking_policy(args) if bool(getattr(args, "quickstart_pi", False)) else _preserve_thinking_policy(args), - "--verify-strategy", - str(getattr(args, "verify_strategy", "capture_commit") or "capture_commit"), - "--verify-core", - str( - getattr(args, "verify_core", "linear-gdn-from-conv-tape") - or "linear-gdn-from-conv-tape" - ), + *verify_strategy_args, + *verify_core_args, "--draft-lm-head-bits", str(draft_lm_head["bits"]), "--draft-lm-head-group-size", @@ -8659,6 +8741,8 @@ def cmd_serve_public(args: Any) -> int: "--fan-mode", fan_mode, ] + if bool(getattr(args, "deepseek_v4_0731_k2", False)): + cmd.append("--deepseek-v4-0731-k2") for attr, flag in ( ("max_active_requests", "--max-active-requests"), ("decode_batch_max", "--decode-batch-max"), @@ -8671,7 +8755,10 @@ def cmd_serve_public(args: Any) -> int: # Retrieval models. The server runs as a subprocess with an explicitly # rebuilt argv, so anything not forwarded here never reaches it — the # endpoints would stay unconfigured on every path but a bare `mtplx serve`. - for flag, attr in (("--embedding-model", "embedding_model"), ("--reranker-model", "reranker_model")): + for flag, attr in ( + ("--embedding-model", "embedding_model"), + ("--reranker-model", "reranker_model"), + ): for reference in getattr(args, attr, None) or []: if str(reference).strip(): cmd.extend([flag, str(reference)]) @@ -9273,6 +9360,9 @@ def _generate_one_shot_public( profile = get_profile(getattr(args, "profile", None) or DEFAULT_PROFILE_NAME) apply_profile_env(profile.name) generation_mode = _generation_mode_from_args(args) + k2_error = _deepseek_v4_0731_k2_entrypoint_error(args) + if k2_error is not None: + return 2, k2_error, [] draft_lm_head = ( _model_draft_lm_head_spec(inspection, profile) if generation_mode == GENERATION_MODE_MTP @@ -9311,7 +9401,7 @@ def _emit(line: str) -> None: from mtplx.sampling import SamplerConfig try: - rt = load(runtime_model, mtp=getattr(args, "load_mtp", True) is not False) + rt = load(runtime_model, **_runtime_load_kwargs(args)) draft_report = None if ( draft_lm_head is not None @@ -9380,20 +9470,17 @@ def _emit_smart(line: str) -> None: seed=args.seed, ) else: - out = generate_mtpk( + request_kwargs = _public_mtpk_request_kwargs( rt, - prompt_ids, + args, max_tokens=max_tokens_value, sampler=sampler, draft_sampler=_draft_sampler_from_spec(draft_sampler), speculative_depth=args.depth, seed=args.seed, - mtp_hidden_variant="post_norm", mtp_cache_policy="persistent", - mtp_history_policy="committed", - verify_strategy="capture_commit", - verify_core="linear-gdn-from-conv-tape", ) + out = generate_mtpk(rt, prompt_ids, **request_kwargs) finally: if smart_fans is not None and smart_request_id is not None: smart_fans.end_request(smart_request_id, wait_for_restore=True) @@ -10016,6 +10103,62 @@ def _print_label(self) -> None: print(f"{self._label}:") +def _public_mtpk_request_kwargs( + rt: Any, + args: Any, + **common: Any, +) -> dict[str, Any]: + from mtplx.runtime import build_mtpk_request_kwargs + + cli_flags = getattr(args, "_cli_flags", set()) or set() + explicit_legacy = {} + for flag, key, value, baseline in ( + ( + "mtp-hidden-variant", + "mtp_hidden_variant", + getattr(args, "mtp_hidden_variant", "post_norm"), + "post_norm", + ), + ( + "mtp-history-policy", + "mtp_history_policy", + getattr(args, "mtp_history_policy", "committed"), + "committed", + ), + ( + "verify-strategy", + "verify_strategy", + getattr(args, "verify_strategy", "capture_commit"), + "capture_commit", + ), + ( + "verify-core", + "verify_core", + getattr(args, "verify_core", "linear-gdn-from-conv-tape"), + "linear-gdn-from-conv-tape", + ), + ( + "draft-core", + "draft_core", + getattr(args, "draft_core", "stock"), + "stock", + ), + ): + if flag in cli_flags or value != baseline: + explicit_legacy[key] = value + return build_mtpk_request_kwargs( + rt, + common=common, + legacy_defaults={ + "mtp_hidden_variant": "post_norm", + "mtp_history_policy": "committed", + "verify_strategy": "capture_commit", + "verify_core": "linear-gdn-from-conv-tape", + }, + explicit_legacy=explicit_legacy, + ) + + def _quickstart_stats_line(payload: dict[str, Any]) -> str: stats = payload.get("stats") or {} generated_tokens = int(stats.get("generated_tokens") or 0) @@ -10167,21 +10310,18 @@ def record_tokens(new_tokens: list[int]) -> None: token_callback=record_tokens, ) else: - out = generate_mtpk( + request_kwargs = _public_mtpk_request_kwargs( rt, - prompt_ids, + args, max_tokens=max_tokens_value, sampler=sampler, draft_sampler=_draft_sampler_from_spec(draft_sampler), speculative_depth=int(getattr(args, "depth", 3)), seed=seed, - mtp_hidden_variant="post_norm", mtp_cache_policy="persistent", - mtp_history_policy="committed", - verify_strategy="capture_commit", - verify_core="linear-gdn-from-conv-tape", token_callback=record_tokens, ) + out = generate_mtpk(rt, prompt_ids, **request_kwargs) finally: if smart_fans is not None and smart_request_id is not None: smart_fans.end_request(smart_request_id, wait_for_restore=False) @@ -11608,6 +11748,7 @@ def _with_server_policy_args(target: Any, source: Any) -> Any: ("retrieval_trust_remote_code", False), ("api_key_file", None), ("api_key_source", "none"), + ("deepseek_v4_0731_k2", False), ("default_presence_penalty", 0.0), ("default_frequency_penalty", 0.0), ("paged_kv_quantization", "off"), @@ -12267,6 +12408,14 @@ def _quickstart_run_terminal_chat_body( profile = get_profile(getattr(args, "profile", None) or DEFAULT_PROFILE_NAME) apply_profile_env(profile.name) generation_mode = _generation_mode_from_args(args) + k2_error = _deepseek_v4_0731_k2_entrypoint_error(args) + if k2_error is not None: + _print_command_error( + k2_error, + command="start", + json_output=bool(getattr(args, "json", False)), + ) + return 2 draft_lm_head = ( _model_draft_lm_head_spec(inspection, profile) if getattr(args, "load_mtp", True) is not False @@ -12313,7 +12462,7 @@ def _quickstart_run_terminal_chat_body( quiet_progress = not sys.stdout.isatty() with ModelLoadProgress("Loading model", quiet=quiet_progress) as progress: progress.set_subtitle(f"profile {profile.name}") - rt = load(runtime_model, mtp=getattr(args, "load_mtp", True) is not False) + rt = load(runtime_model, **_runtime_load_kwargs(args)) progress.set_subtitle("ready") _quickstart_line(f"Model ready in {time.perf_counter() - started:.1f}s") _quickstart_line(f"Generation mode: {_generation_mode_label(generation_mode)}") diff --git a/mtplx/deepseek_v4_0731_dspark_ffn.py b/mtplx/deepseek_v4_0731_dspark_ffn.py new file mode 100644 index 00000000..20af7e71 --- /dev/null +++ b/mtplx/deepseek_v4_0731_dspark_ffn.py @@ -0,0 +1,353 @@ +"""Retained native M=5 DSpark Q3 gate/up packing for Flash-0731. + +The three installed DSpark stages use affine Q3/group-128 routed projections +with the fixed physical layout recorded by the 0731 receipt. This module +validates loaded storage once, then replaces each gate/up pair with the native +MLX packed projection. Activation and the Q3 down projection stay unchanged. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import mlx.core as mx +import mlx.nn as nn +from mlx_lm.models.switch_layers import QuantizedSwitchLinear + +from .models.deepseek_v4 import ( + DeepseekV4DSpark, + DeepseekV4DSparkStage, + DeepseekV4MoE, + Model, + MoEGate, +) +from .moe_packed_projections import PackedSwitchGLU, _pack_pair + + +DSPARK_Q3_GATE_UP_GEOMETRY = { + "rows": 5, + "hidden_size": 4096, + "width": 2048, + "experts": 256, + "top_k": 6, + "bits": 3, + "group_size": 128, + "weight_shape": (256, 2048, 384), + "metadata_shape": (256, 2048, 32), +} + + +@dataclass(frozen=True) +class DSparkQ3GateUpContract: + rows: int + hidden_size: int + width: int + experts: int + top_k: int + bits: int + group_size: int + weight_shape: tuple[int, int, int] + metadata_shape: tuple[int, int, int] + + +class DeepseekV40731DSparkM5PackedSwitchGLU(PackedSwitchGLU): + """Construction-qualified fixed-M5 route with one unsorted packed gather.""" + + def __call__(self, x: mx.array, indices: mx.array) -> mx.array: + expanded = mx.expand_dims(x, (-2, -3)) + packed = self.gate_up_proj.gather(expanded, indices, False) + gate, up = mx.split(packed, [self._split_at], axis=-1) + routed = self.down_proj( + self.activation(up, gate), + indices, + sorted_indices=False, + ) + return routed.squeeze(-2) + + +def _shape(value: Any) -> tuple[int, ...]: + shape = getattr(value, "shape", None) + if shape is None: + getter = getattr(value, "get_shape", None) + shape = () if getter is None else getter() + return tuple(int(item) for item in shape) + + +def _q3_contract( + *, hidden_size: int, width: int, experts: int, top_k: int, rows: int +) -> DSparkQ3GateUpContract: + hidden_size = int(hidden_size) + width = int(width) + experts = int(experts) + top_k = int(top_k) + rows = int(rows) + if ( + hidden_size <= 0 + or hidden_size % 128 + or hidden_size * 3 % 32 + or width <= 0 + or experts <= 0 + or top_k != 6 + or rows != 5 + ): + raise ValueError("DSpark Q3 M=5 geometry is invalid") + return DSparkQ3GateUpContract( + rows=rows, + hidden_size=hidden_size, + width=width, + experts=experts, + top_k=top_k, + bits=3, + group_size=128, + weight_shape=(experts, width, hidden_size * 3 // 32), + metadata_shape=(experts, width, hidden_size // 128), + ) + + +def validate_dspark_q3_gate_up( + gate: nn.Module, + up: nn.Module, + *, + hidden_size: int, + width: int, + experts: int, + top_k: int, + rows: int, +) -> DSparkQ3GateUpContract: + """Validate physical Q3 storage once before binding the candidate route.""" + + contract = _q3_contract( + hidden_size=hidden_size, + width=width, + experts=experts, + top_k=top_k, + rows=rows, + ) + for label, projection in (("gate", gate), ("up", up)): + if not isinstance(projection, QuantizedSwitchLinear): + raise ValueError( + f"DSpark Q3 {label} projection is not QuantizedSwitchLinear" + ) + if int(getattr(projection, "bits", 0)) != contract.bits: + raise ValueError(f"DSpark Q3 {label} projection must use Q3") + if int(getattr(projection, "group_size", 0)) != contract.group_size: + raise ValueError(f"DSpark Q3 {label} projection must use group-128") + if str(getattr(projection, "mode", "")) != "affine": + raise ValueError( + f"DSpark Q3 {label} projection must use affine quantization" + ) + if getattr(projection.weight, "dtype", None) != mx.uint32: + raise ValueError(f"DSpark Q3 {label} packed weight must be U32") + if _shape(projection.weight) != contract.weight_shape: + raise ValueError(f"DSpark Q3 {label} packed weight shape is invalid") + if _shape(projection.scales) != contract.metadata_shape: + raise ValueError(f"DSpark Q3 {label} scales shape is invalid") + if _shape(getattr(projection, "biases", None)) != contract.metadata_shape: + raise ValueError(f"DSpark Q3 {label} biases shape is invalid") + if ( + getattr(projection.scales, "dtype", None) != mx.bfloat16 + or getattr(projection.biases, "dtype", None) != mx.bfloat16 + ): + raise ValueError(f"DSpark Q3 {label} affine metadata must be BF16") + if "bias" in projection: + raise ValueError(f"DSpark Q3 {label} projection must not have output bias") + return contract + + +def _validate_dspark_q3_down( + projection: nn.Module, contract: DSparkQ3GateUpContract +) -> None: + """Prove the unchanged down projection is the paired native Q3 bank.""" + + expected_weight = ( + contract.experts, + contract.hidden_size, + contract.width * contract.bits // 32, + ) + expected_metadata = ( + contract.experts, + contract.hidden_size, + contract.width // contract.group_size, + ) + if not isinstance(projection, QuantizedSwitchLinear): + raise ValueError("DSpark Q3 down projection is not QuantizedSwitchLinear") + if ( + int(getattr(projection, "bits", 0)) != contract.bits + or int(getattr(projection, "group_size", 0)) != contract.group_size + or str(getattr(projection, "mode", "")) != "affine" + or getattr(projection.weight, "dtype", None) != mx.uint32 + or _shape(projection.weight) != expected_weight + or _shape(projection.scales) != expected_metadata + or _shape(getattr(projection, "biases", None)) != expected_metadata + or getattr(projection.scales, "dtype", None) != mx.bfloat16 + or getattr(projection.biases, "dtype", None) != mx.bfloat16 + or "bias" in projection + ): + raise ValueError("DSpark Q3 down projection storage is invalid") + + +def _validate_dspark_q3_switch( + switch: nn.Module, + *, + hidden_size: int, + width: int, + experts: int, + top_k: int, + rows: int, +) -> DSparkQ3GateUpContract: + contract = validate_dspark_q3_gate_up( + switch.gate_proj, + switch.up_proj, + hidden_size=hidden_size, + width=width, + experts=experts, + top_k=top_k, + rows=rows, + ) + _validate_dspark_q3_down(switch.down_proj, contract) + if not callable(getattr(switch, "activation", None)): + raise ValueError("DSpark Q3 switch activation is absent") + return contract + + +def build_dspark_q3_packed_gate_up( + switch: nn.Module, + *, + hidden_size: int, + width: int, + experts: int, + top_k: int, + rows: int, +) -> DeepseekV40731DSparkM5PackedSwitchGLU: + """Pack native Q3 gate/up rows without changing activation or down math.""" + + _validate_dspark_q3_switch( + switch, + hidden_size=hidden_size, + width=width, + experts=experts, + top_k=top_k, + rows=rows, + ) + result = _pack_pair(switch.gate_proj, switch.up_proj, axis=1) + if isinstance(result, str): + raise ValueError(f"DSpark native Q3 gate/up packing failed: {result}") + packed, split_at = result + if int(split_at) != int(width): + raise ValueError("DSpark native Q3 gate/up split is invalid") + return DeepseekV40731DSparkM5PackedSwitchGLU( + packed, + switch.down_proj, + switch.activation, + split_at, + ) + + +def _validate_dspark_m5_owner(model: Any) -> tuple[DeepseekV4DSparkStage, ...]: + if type(model) is not Model: + raise ValueError("DSpark native Q3 packing requires the exact model owner") + dspark = getattr(model, "_dspark", None) + if type(dspark) is not DeepseekV4DSpark: + raise ValueError("DSpark native Q3 packing requires the exact model owner") + stages = tuple(getattr(dspark, "stages", ())) + if len(stages) != 3: + raise ValueError("DSpark native Q3 packing requires exactly three stages") + if any(type(stage) is not DeepseekV4DSparkStage for stage in stages): + raise ValueError("DSpark native Q3 packing stage identity is invalid") + if tuple(int(stage.stage_id) for stage in stages) != (0, 1, 2): + raise ValueError("DSpark native Q3 packing stage order must be 0, 1, 2") + if any(type(stage.ffn) is not DeepseekV4MoE for stage in stages): + raise ValueError("DSpark native Q3 packing FFN identity is invalid") + if any(type(stage.ffn.gate) is not MoEGate for stage in stages): + raise ValueError("DSpark native Q3 packing router identity is invalid") + published = tuple(getattr(model.mtp, "layers", model.mtp)) + if len(published) != 3 or any( + registered is not owned + for registered, owned in zip(published, stages, strict=True) + ): + raise ValueError("DSpark native Q3 packing model.mtp ownership is invalid") + if int(getattr(dspark, "block_size", 0)) != 5 or any( + int(getattr(stage, "block_size", 0)) != 5 for stage in stages + ): + raise ValueError("DSpark native Q3 packing requires exact M=5 ownership") + if any(int(getattr(stage.ffn.gate, "topk", 0)) != 6 for stage in stages): + raise ValueError("DSpark native Q3 packing requires router top-k=6") + return stages + + +def _receipt() -> dict[str, Any]: + return { + "candidate": "dspark-native-packed-q3-gate-up-m5", + "stages": 3, + "geometry": dict(DSPARK_Q3_GATE_UP_GEOMETRY), + "gate_up_dispatches_per_stage": 1, + "stock_gate_up_dispatches_per_stage": 2, + "explicit_dequantize": False, + "resident_weight_bytes_added": 0, + } + + +@dataclass(frozen=True, slots=True) +class PreparedDSparkQ3PackedGateUpM5: + """Fully constructed three-stage FFN replacements awaiting publication.""" + + stages: tuple[DeepseekV4DSparkStage, ...] + originals: tuple[Any, ...] + replacements: tuple[DeepseekV40731DSparkM5PackedSwitchGLU, ...] + receipt: dict[str, Any] + + def publish(self) -> None: + try: + for stage, replacement in zip(self.stages, self.replacements, strict=True): + stage.ffn.switch_mlp = replacement + except Exception as publication_error: + try: + self.restore() + except ExceptionGroup as restoration_errors: + publication_error.add_note( + f"DSpark FFN publication rollback also failed: {restoration_errors}" + ) + raise + + def restore(self) -> None: + errors = [] + for stage, original in zip(self.stages, self.originals, strict=True): + try: + stage.ffn.switch_mlp = original + except Exception as exc: + errors.append(exc) + if errors: + raise ExceptionGroup("DSpark FFN restoration failed", errors) + + +def prepare_dspark_q3_packed_gate_up_m5( + model: Any, +) -> PreparedDSparkQ3PackedGateUpM5: + """Validate and build the exact model-owned M5 stages without publishing.""" + + stages = _validate_dspark_m5_owner(model) + geometry = { + key: DSPARK_Q3_GATE_UP_GEOMETRY[key] + for key in ("hidden_size", "width", "experts", "top_k", "rows") + } + switches = tuple(stage.ffn.switch_mlp for stage in stages) + for switch in switches: + _validate_dspark_q3_switch(switch, **geometry) + replacements = tuple( + build_dspark_q3_packed_gate_up(switch, **geometry) for switch in switches + ) + return PreparedDSparkQ3PackedGateUpM5( + stages=stages, + originals=switches, + replacements=replacements, + receipt=_receipt(), + ) + + +def install_dspark_q3_packed_gate_up_m5(model: Any) -> dict[str, Any]: + """Convenience wrapper that prepares, then publishes all three stages.""" + + prepared = prepare_dspark_q3_packed_gate_up_m5(model) + prepared.publish() + return prepared.receipt diff --git a/mtplx/deepseek_v4_0731_full_install.py b/mtplx/deepseek_v4_0731_full_install.py new file mode 100644 index 00000000..86feb76b --- /dev/null +++ b/mtplx/deepseek_v4_0731_full_install.py @@ -0,0 +1,484 @@ +"""Atomic construction installer for the measured Flash-0731 target stack. + +There is one supported configuration: packed affine-Q2/group-128 routed +gate/up, row-owned M1 reduction, compiled packed-Q2 physical M3 tails, fused +WQB-qhead, and fused WOB. Pinned config/index metadata, loaded ownership, +storage, and every 43-layer self-check complete before any route is published. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import hashlib +from pathlib import Path +from typing import Any, Callable + +import mlx.core as mx + +from . import deepseek_v4_attention_island as AI +from .deepseek_v4_0731_m3_target import ( + build_0731_m3_target_route, + build_m3_compiled_tail_layer, +) +from .deepseek_v4_0731_moe import ( + build_routed_q2_pair, + build_row_owned_combine_m1, + exact_selfcheck_row_owned_combine_m1, + validate_routed_q2_pair, +) + + +EXPECTED_FULL_CONFIG_SHA256 = ( + "44735712733fcf8f299bdf1faa1d87fac88f1917efe1d3876d6d4c582f79a68f" +) +EXPECTED_FULL_INDEX_SHA256 = ( + "f1332b2b209769c2db335954c2651652a8048e7d7dbf60296c2f2c0198715861" +) +RECORDED_ARTIFACT_LABEL = "mlx-community/DeepSeek-V4-Flash-0731-2.4bit-mixed" +EXPECTED_SOURCE_REVISION = "10001e0065f8394e03e968e652cbbe7cd2ca122c" + +_EXPECTED_CONFIG = { + "model_type": "deepseek_v4", + "hidden_size": 4096, + "num_hidden_layers": 43, + "num_attention_heads": 64, + "num_key_value_heads": 1, + "head_dim": 512, + "n_routed_experts": 256, + "num_experts_per_tok": 6, + "moe_intermediate_size": 2048, + "n_shared_experts": 1, + "swiglu_limit": 10.0, + "num_nextn_predict_layers": 1, + "dspark_block_size": 5, + "dspark_noise_token_id": 128799, + "dspark_target_layer_ids": [40, 41, 42], + "dspark_markov_rank": 256, +} +_TARGET_LAYER_IDS = (40, 41, 42) +_LAYERS = 43 +_STAGE_COUNT = 3 +_M3_WQB_SHAPE = (1, 3, 1024) +_M3_WOB_SHAPE = (1, 3, 8192) + + +@dataclass(frozen=True, slots=True) +class Full0731DSparkContract: + layers: int + target_layer_ids: tuple[int, int, int] + stage_count: int + source_revision: str + config_sha256: str + index_sha256: str + + +def _contract_error(detail: str) -> ValueError: + return ValueError(f"DeepSeek-V4 0731 full DSpark contract failed: {detail}") + + +def _metadata_revision(path: Path) -> str: + try: + return path.read_text(encoding="utf-8").splitlines()[0] + except (OSError, IndexError) as exc: + raise _contract_error(f"artifact metadata is unreadable: {exc}") from exc + + +def validate_full_0731_dspark_artifact( + model_path: str | Path, + config: dict[str, Any], +) -> Full0731DSparkContract: + """Require the exact measured config/index and recorded metadata revision.""" + + root = Path(model_path) + try: + config_bytes = (root / "config.json").read_bytes() + index_bytes = (root / "model.safetensors.index.json").read_bytes() + except OSError as exc: + raise _contract_error(f"artifact identity is unreadable: {exc}") from exc + config_sha = hashlib.sha256(config_bytes).hexdigest() + index_sha = hashlib.sha256(index_bytes).hexdigest() + if config_sha != EXPECTED_FULL_CONFIG_SHA256: + raise _contract_error("config SHA-256 does not match the pinned artifact") + if index_sha != EXPECTED_FULL_INDEX_SHA256: + raise _contract_error("index SHA-256 does not match the pinned artifact") + for name, expected in _EXPECTED_CONFIG.items(): + if config.get(name) != expected: + raise _contract_error( + f"config {name}={config.get(name)!r}, expected {expected!r}" + ) + metadata_root = root / ".cache/huggingface/download" + revisions = ( + _metadata_revision(metadata_root / "config.json.metadata"), + _metadata_revision(metadata_root / "model.safetensors.index.json.metadata"), + ) + if revisions != (EXPECTED_SOURCE_REVISION, EXPECTED_SOURCE_REVISION): + raise _contract_error("recorded metadata revision does not match") + return Full0731DSparkContract( + layers=_LAYERS, + target_layer_ids=_TARGET_LAYER_IDS, + stage_count=_STAGE_COUNT, + source_revision=EXPECTED_SOURCE_REVISION, + config_sha256=config_sha, + index_sha256=index_sha, + ) + + +def _validate_loaded_dspark_owner( + model: Any, + contract: Full0731DSparkContract, +) -> tuple[tuple[Any, ...], Callable]: + try: + layers = tuple(model.model.layers) + dspark = model._dspark + stages = tuple(dspark.stages) + published_stages = tuple(getattr(model.mtp, "layers", model.mtp)) + native_route = model._target_hidden_route + except AttributeError as exc: + raise _contract_error( + f"loaded model lacks native DSpark ownership: {exc}" + ) from exc + if len(layers) != contract.layers: + raise _contract_error(f"loaded layer count is {len(layers)}, expected 43") + if dspark is None: + raise _contract_error("loaded model does not own native DSpark stages") + if tuple(getattr(dspark, "target_layer_ids", ())) != contract.target_layer_ids: + raise _contract_error("loaded DSpark target taps are not (40, 41, 42)") + if len(stages) != contract.stage_count: + raise _contract_error("loaded DSpark does not own exactly three stages") + if len(published_stages) != len(stages) or any( + published is not owned + for published, owned in zip(published_stages, stages, strict=True) + ): + raise _contract_error("DSpark stage ownership is not preserved by model.mtp") + if not callable(native_route): + raise _contract_error("native DSpark target route is absent") + return layers, native_route + + +class _BoundDSparkM1Body: + """Prebound M1 trunk retaining post-layer taps 40, 41, and 42.""" + + __slots__ = ("_body", "_layers", "_target_layer_ids") + + def __init__(self, body: Any, layers, target_layer_ids: tuple[int, int, int]): + self._body = body + self._layers = tuple(layers) + self._target_layer_ids = target_layer_ids + + def __call__(self, input_ids: mx.array, cache=None): + hidden = self._body.embed_tokens(input_ids) + hidden = mx.broadcast_to( + hidden[:, :, None, :], + (*hidden.shape[:2], self._body.hc_mult, hidden.shape[-1]), + ) + entries = (None,) * len(self._layers) if cache is None else cache + taps = [] + for layer_id, ((layer, tail), entry) in enumerate(zip(self._layers, entries)): + residual = hidden + attention_in, post, comb = layer.attn_hc.pre(hidden) + attention_in = layer.attn_norm(attention_in) + attention_out = layer.attn(attention_in, mask=None, cache=entry) + hidden = tail( + attention_out, + residual, + post, + comb, + input_ids, + ) + if layer_id in self._target_layer_ids: + taps.append(mx.mean(hidden, axis=2)) + return hidden, mx.concatenate(taps, axis=-1) + + +@dataclass(frozen=True, slots=True) +class _FullDSparkTargetRoute: + native: Callable + m1: Callable + + def __call__(self, owner: Any, input_ids: mx.array, cache=None): + if int(input_ids.shape[1]) == 1: + return self.m1(input_ids, cache) + return self.native(owner, input_ids, cache) + + +def _require_wqb_receipt(receipt: Any) -> dict[str, Any]: + published = tuple(getattr(receipt, "published_routes", ())) + if ( + len(published) != _LAYERS + or int(getattr(receipt, "q6_count", -1)) != _LAYERS + or int(getattr(receipt, "exact_selfchecked", -1)) != _LAYERS + or not callable(getattr(receipt, "publish", None)) + or not callable(getattr(receipt, "restore", None)) + ): + raise _contract_error("fused WQB-qhead preparation is not 43/43 exact") + return { + "candidate": "official-wheel-custom-fixed-m3-wqb-qhead-fused", + "layers_installed": _LAYERS, + "q6_g128_layers": _LAYERS, + "exact_selfchecked_layers": _LAYERS, + "shape": [1, 3, 1024], + "output_shape": [1, 3, 64, 512], + } + + +def _require_wob_receipt(receipt: Any) -> dict[str, Any]: + published = tuple(getattr(receipt, "published_routes", ())) + if ( + len(published) != _LAYERS + or int(getattr(receipt, "q6_count", -1)) != _LAYERS + or int(getattr(receipt, "exact_selfchecked", -1)) != _LAYERS + or int(getattr(receipt, "o_lora_sink_count", -1)) != _LAYERS + or not callable(getattr(receipt, "publish", None)) + or not callable(getattr(receipt, "restore", None)) + ): + raise _contract_error("fused WOB preparation is not 43/43 exact") + return { + "candidate": "official-wheel-custom-fixed-m3-affine-qmv", + "layers_installed": _LAYERS, + "q6_g128_layers": _LAYERS, + "exact_selfchecked_layers": _LAYERS, + "active_o_lora_sinks_installed": _LAYERS, + "shape": [1, 3, 8192], + "output_size": 4096, + } + + +def _m3_wqb_qhead_exact_selfcheck(): + """Build the construction-only fused M3-versus-three-M1 q-head gate.""" + + key = mx.random.key(73_103_1025) + qr_key, phase_key = mx.random.split(key) + qr = mx.random.normal(_M3_WQB_SHAPE, key=qr_key).astype(mx.bfloat16) + phase = mx.random.normal((3, 32), key=phase_key) + cos, sin = mx.cos(phase), mx.sin(phase) + + def check(stock: Callable, candidate: Callable, _layer_index: int) -> bool: + actual = candidate(qr, cos, sin) + expected = mx.concatenate( + tuple( + stock(qr[:, row : row + 1], cos[row : row + 1], sin[row : row + 1]) + for row in range(3) + ), + axis=1, + ) + mx.eval(actual, expected) + return bool(mx.array_equal(actual, expected)) + + return check + + +def _m3_wob_exact_selfcheck(): + """Build the construction-only real-weight M3-versus-three-M1 gate.""" + + key = mx.random.key(73_103_8192) + value = mx.random.normal(_M3_WOB_SHAPE, key=key).astype(mx.bfloat16) + + def check(stock: Callable, candidate: Callable, _layer_index: int) -> bool: + actual = candidate(value) + expected = mx.concatenate( + tuple(stock(value[:, row : row + 1]) for row in range(3)), + axis=1, + ) + mx.eval(actual, expected) + return bool(mx.array_equal(actual, expected)) + + return check + + +@dataclass(frozen=True, slots=True) +class PreparedFull0731DSparkTargetStack: + """Fully checked target stack whose publication is one reversible step.""" + + model: Any + native_route: Callable + target_route: Callable + layers: tuple[Any, ...] + stock_switches: tuple[Any, ...] + replacement_switches: tuple[Any, ...] + wqb: Any + wob: Any + receipt: dict[str, Any] + + def publish(self) -> None: + """Publish every prepared route, rolling the whole option back on error.""" + + try: + for layer, replacement in zip(self.layers, self.replacement_switches): + layer.ffn.switch_mlp = replacement + self.wqb.publish() + self.wob.publish() + self.model._target_hidden_route = self.target_route + except Exception as publication_error: + try: + self.restore() + except Exception as restoration_error: + publication_error.add_note( + f"target-stack rollback also raised: {restoration_error}" + ) + raise + + def restore(self) -> None: + """Restore in reverse publication order.""" + + errors = [] + try: + self.model._target_hidden_route = self.native_route + except Exception as exc: + errors.append(exc) + try: + self.wob.restore() + except Exception as exc: + errors.append(exc) + try: + self.wqb.restore() + except Exception as exc: + errors.append(exc) + for layer, switch in zip(self.layers, self.stock_switches): + try: + layer.ffn.switch_mlp = switch + except Exception as exc: + errors.append(exc) + if errors: + raise ExceptionGroup("target-stack route restoration failed", errors) + + +def prepare_full_0731_dspark_compiled_tail_q2_pair( + model: Any, + config: dict[str, Any], + model_path: str | Path, + *, + prepare_wqb_qhead: Callable[[tuple[Any, ...]], Any], + prepare_wob: Callable[[tuple[Any, ...]], Any], +) -> PreparedFull0731DSparkTargetStack: + """Build and check the measured target stack without publishing any route.""" + + if not callable(prepare_wqb_qhead) or not callable(prepare_wob): + raise _contract_error("both fused projection preparers are required") + contract = validate_full_0731_dspark_artifact(model_path, config) + layers, native_route = _validate_loaded_dspark_owner(model, contract) + try: + for layer in layers: + switch = layer.ffn.switch_mlp + validate_routed_q2_pair( + switch.gate_proj, + switch.up_proj, + hidden_size=4096, + width=2048, + experts=256, + ) + except AttributeError as exc: + raise _contract_error(f"trunk MoE topology is incomplete: {exc}") from exc + + row_owned = build_row_owned_combine_m1(hidden_size=4096, top_k=6) + exact_selfcheck_row_owned_combine_m1(row_owned) + staged = [] + m3_layer_routes = [] + for layer in layers: + original = layer.ffn.switch_mlp + replacement = build_routed_q2_pair( + original, + hidden_size=4096, + width=2048, + experts=256, + ) + m1_tail = AI._bind_attention_island_layer( + layer, + width=1, + allowed_widths=(1,), + shared_bits=8, + routed_pair=True, + routed_combine=row_owned, + routed_switch=replacement, + ) + m3_tail = AI._bind_attention_island_layer( + layer, + width=3, + allowed_widths=(3,), + shared_bits=8, + routed_pair=True, + routed_switch=replacement, + ) + m3_layer = build_m3_compiled_tail_layer(layer, m3_tail) + staged.append((layer, original, replacement, m1_tail)) + m3_layer_routes.append(m3_layer) + + m1_body = _BoundDSparkM1Body( + model.model, + tuple((layer, m1_tail) for layer, _, _, m1_tail in staged), + contract.target_layer_ids, + ) + m1_route = _FullDSparkTargetRoute(native=native_route, m1=m1_body) + target_route = build_0731_m3_target_route( + model, + full_layer_routes=tuple(m3_layer_routes), + base_route=m1_route, + ) + + # Projection preparation is also staging-only. Until both prepared objects + # and their exact receipts are accepted, every model route remains stock. + wqb_prepared = prepare_wqb_qhead( + layers, + exact_selfcheck=_m3_wqb_qhead_exact_selfcheck(), + ) + wqb_receipt = _require_wqb_receipt(wqb_prepared) + wob_prepared = prepare_wob( + layers, + exact_selfcheck=_m3_wob_exact_selfcheck(), + ) + wob_receipt = _require_wob_receipt(wob_prepared) + + receipt = { + "candidate": "mtplx-full-dspark-compiled-tail-packed-q2-pair-m1-m3", + "artifact_label": RECORDED_ARTIFACT_LABEL, + "validated_config_sha256": contract.config_sha256, + "validated_index_sha256": contract.index_sha256, + "validated_metadata_revision": contract.source_revision, + "layers_installed": len(staged), + "decode_m": 1, + "fixed_k": 2, + "physical_target_rows": 3, + "m3_tail": "fixed-width3-compiled-tail", + "m3_wqb": wqb_receipt, + "m3_wob": wob_receipt, + "row_owned_combine": True, + "non_m1_m3_route": "native-dspark", + "routed_bits": 2, + "routed_group_size": 128, + "routed_gate_up_paired": True, + "shared_bits": 8, + "target_taps": contract.target_layer_ids, + "dspark_stages": contract.stage_count, + "stage_ownership": "native", + } + return PreparedFull0731DSparkTargetStack( + model=model, + native_route=native_route, + target_route=target_route, + layers=layers, + stock_switches=tuple(original for _, original, _, _ in staged), + replacement_switches=tuple(replacement for _, _, replacement, _ in staged), + wqb=wqb_prepared, + wob=wob_prepared, + receipt=receipt, + ) + + +def install_full_0731_dspark_compiled_tail_q2_pair( + model: Any, + config: dict[str, Any], + model_path: str | Path, + *, + prepare_wqb_qhead: Callable[[tuple[Any, ...]], Any], + prepare_wob: Callable[[tuple[Any, ...]], Any], +) -> dict[str, Any]: + """Prepare and atomically publish the sole measured target stack.""" + + prepared = prepare_full_0731_dspark_compiled_tail_q2_pair( + model, + config, + model_path, + prepare_wqb_qhead=prepare_wqb_qhead, + prepare_wob=prepare_wob, + ) + prepared.publish() + return prepared.receipt diff --git a/mtplx/deepseek_v4_0731_m3_target.py b/mtplx/deepseek_v4_0731_m3_target.py new file mode 100644 index 00000000..91baff08 --- /dev/null +++ b/mtplx/deepseek_v4_0731_m3_target.py @@ -0,0 +1,191 @@ +"""Construction-bound physical-M3 target traversal for Flash-0731 DSpark.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Callable, Sequence + +import mlx.core as mx + + +_ROWS = 3 +_LAYERS = 43 +_HIDDEN = 4096 +_HEADS = 64 +_HEAD_DIM = 512 +_HC = 4 +_TAP_LAYERS = (40, 41, 42) + + +class M3TargetContractError(ValueError): + """The loaded owner cannot bind the fixed physical-M3 route.""" + + +def _contract_error(detail: str) -> M3TargetContractError: + return M3TargetContractError( + f"DeepSeek-V4 0731 physical-M3 contract failed: {detail}" + ) + + +@dataclass(frozen=True, slots=True) +class M3TargetContract: + layers: int + hidden_size: int + hc_mult: int + target_layer_ids: tuple[int, int, int] + + +def validate_0731_m3_target(model: Any) -> M3TargetContract: + """Validate full-artifact ownership once before route construction.""" + + args = getattr(model, "args", None) + body = getattr(model, "model", None) + dspark = getattr(model, "_dspark", None) + if args is None or body is None or dspark is None: + raise _contract_error("requires a loaded full 0731 DSpark owner") + expected = { + "hidden_size": _HIDDEN, + "num_hidden_layers": _LAYERS, + "num_attention_heads": _HEADS, + "num_key_value_heads": 1, + "head_dim": _HEAD_DIM, + } + for name, expected_value in expected.items(): + if int(getattr(args, name, -1)) != expected_value: + raise _contract_error(f"{name} is not {expected_value}") + layers = tuple(getattr(body, "layers", ())) + if len(layers) != _LAYERS: + raise _contract_error(f"body owns {len(layers)} layers, expected {_LAYERS}") + if int(getattr(body, "hc_mult", -1)) != _HC: + raise _contract_error(f"body hc_mult is not {_HC}") + taps = tuple(int(value) for value in getattr(dspark, "target_layer_ids", ())) + if taps != _TAP_LAYERS: + raise _contract_error("DSpark target taps are not (40, 41, 42)") + if len(tuple(getattr(dspark, "stages", ()))) != 3: + raise _contract_error("DSpark does not own exactly three stages") + if not callable(getattr(body, "embed_tokens", None)): + raise _contract_error("body embedding route is absent") + if not callable(getattr(model, "logits_from_hc_hidden", None)): + raise _contract_error("target logit route is absent") + return M3TargetContract( + layers=_LAYERS, + hidden_size=_HIDDEN, + hc_mult=_HC, + target_layer_ids=_TAP_LAYERS, + ) + + +class _M3TargetBody: + """Exact three-row trunk traversal through 43 prebound layer routes.""" + + __slots__ = ("_body", "_layer_routes", "_tap_layers") + + def __init__( + self, + body: Any, + layer_routes: Sequence[Callable], + tap_layers: tuple[int, int, int], + ) -> None: + self._body = body + self._layer_routes = tuple(layer_routes) + self._tap_layers = tap_layers + + def __call__(self, input_ids: mx.array, cache=None) -> tuple[mx.array, mx.array]: + hidden = self._body.embed_tokens(input_ids) + hidden = mx.broadcast_to( + hidden[:, :, None, :], + (*hidden.shape[:2], self._body.hc_mult, hidden.shape[-1]), + ) + entries = (None,) * len(self._layer_routes) if cache is None else cache + taps = [] + for layer_id, (route, entry) in enumerate(zip(self._layer_routes, entries)): + hidden = route(hidden, input_ids, entry) + if layer_id in self._tap_layers: + taps.append(mx.mean(hidden, axis=2)) + return hidden, mx.concatenate(taps, axis=-1) + + +class M3CompiledTailLayer: + """One native-width HC/attention boundary and one compiled M3 tail.""" + + __slots__ = ("_attention", "_attn_hc_pre", "_attn_norm", "_m3_tail") + + def __init__(self, layer: Any, m3_tail: Callable) -> None: + attention = getattr(layer, "attn", None) + hc = getattr(layer, "attn_hc", None) + norm = getattr(layer, "attn_norm", None) + pre = getattr(hc, "pre", None) + if not callable(attention) or not callable(pre) or not callable(norm): + raise M3TargetContractError( + "compiled M3 layer is missing its stock attention boundary" + ) + if not callable(m3_tail): + raise M3TargetContractError( + "compiled M3 layer requires a callable width-three tail" + ) + self._attention = attention + self._attn_hc_pre = pre + self._attn_norm = norm + self._m3_tail = m3_tail + + def __call__(self, hidden: mx.array, input_ids: mx.array, cache) -> mx.array: + attention_in, post, comb = self._attn_hc_pre(hidden) + attention_in = self._attn_norm(attention_in) + attention_out = self._attention(attention_in, mask=None, cache=cache) + return self._m3_tail( + attention_out, + hidden, + post, + comb, + input_ids, + ) + + +def build_m3_compiled_tail_layer( + layer: Any, + m3_tail: Callable, +) -> M3CompiledTailLayer: + return M3CompiledTailLayer(layer, m3_tail) + + +@dataclass(frozen=True, slots=True) +class FixedM3TargetRoute: + """Explicit width table: physical M3 or the prebound native base route.""" + + base: Callable + m3_body: _M3TargetBody + logits_from_hc_hidden: Callable + + def __call__(self, owner: Any, input_ids: mx.array, cache=None): + if int(input_ids.shape[1]) == _ROWS: + return self.m3_body(input_ids, cache) + return self.base(owner, input_ids, cache) + + def forward( + self, + input_ids: mx.array, + cache=None, + ) -> tuple[mx.array, mx.array, mx.array]: + hidden, taps = self.m3_body(input_ids, cache) + return self.logits_from_hc_hidden(hidden), hidden, taps + + +def build_0731_m3_target_route( + model: Any, + *, + full_layer_routes: Sequence[Callable], + base_route: Callable, +) -> FixedM3TargetRoute: + """Build the sole fixed-M3 route; nothing is attached here.""" + + contract = validate_0731_m3_target(model) + routes = tuple(full_layer_routes) + if len(routes) != contract.layers or not all(callable(route) for route in routes): + raise _contract_error("requires exactly 43 prebound full-layer routes") + if not callable(base_route): + raise _contract_error("prebound native base target route is absent") + return FixedM3TargetRoute( + base=base_route, + m3_body=_M3TargetBody(model.model, routes, contract.target_layer_ids), + logits_from_hc_hidden=model.logits_from_hc_hidden, + ) diff --git a/mtplx/deepseek_v4_0731_m3_wob.py b/mtplx/deepseek_v4_0731_m3_wob.py new file mode 100644 index 00000000..a3f8ddfc --- /dev/null +++ b/mtplx/deepseek_v4_0731_m3_wob.py @@ -0,0 +1,416 @@ +"""Construction-bound official-wheel fixed-M3 affine-Q6 0731 ``attn.wo_b``. + +The pinned full 0731 artifact stores all 43 body ``wo_b`` projections as one +unbatched affine Q6/group-128 U32 matrix with BF16 scales and biases. This +module derives the K=8192 loop directly from MLX ``qmv_fast_impl``: each of the +three rows retains its own Q6 load, qdot accumulation, and simd reduction tree. +It neither installs itself nor provides an enabled-path fallback. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from functools import lru_cache +from typing import Any, Callable + +import mlx.core as mx + + +_M = 3 +_K = 8192 +_N = 4096 +_BITS = 6 +_GROUP_SIZE = 128 + + +class M3WOBContractError(ValueError): + """A projection cannot use the fixed physical-M3 0731 wo_b primitive.""" + + +@dataclass(frozen=True, slots=True) +class M3WOBContract: + """Exact packed affine-Q6 BF16 storage for the 0731 body ``wo_b`` matrix.""" + + k: int = _K + n: int = _N + bits: int = _BITS + group_size: int = _GROUP_SIZE + dtype: mx.Dtype = mx.bfloat16 + + def __post_init__(self) -> None: + if (self.k, self.n, self.bits, self.group_size, self.dtype) != ( + _K, + _N, + _BITS, + _GROUP_SIZE, + mx.bfloat16, + ): + raise M3WOBContractError( + "fixed M3 wo_b supports only BF16 affine Q6/G128 K=8192 N=4096" + ) + + @property + def packed_cols(self) -> int: + return self.k * self.bits // 32 + + @property + def metadata_cols(self) -> int: + return self.k // self.group_size + + @property + def weight_shape(self) -> tuple[int, int]: + return (self.n, self.packed_cols) + + @property + def metadata_shape(self) -> tuple[int, int]: + return (self.n, self.metadata_cols) + + +def _shape(value: Any) -> tuple[int, ...] | None: + try: + return tuple(int(dimension) for dimension in value.shape) + except (AttributeError, TypeError, ValueError): + return None + + +def validate_wob_projection( + projection: Any, contract: M3WOBContract | None = None +) -> M3WOBContract: + """Fail at construction unless every immutable stock-storage fact matches.""" + + bound_contract = contract or M3WOBContract() + weight = getattr(projection, "weight", None) + scales = getattr(projection, "scales", None) + biases = getattr(projection, "biases", None) + if _shape(weight) != bound_contract.weight_shape: + if _shape(weight) and len(_shape(weight) or ()) != 2: + raise M3WOBContractError("fixed M3 wo_b requires an unbatched RHS") + raise M3WOBContractError("fixed M3 wo_b packed RHS shape does not match K=8192") + if ( + int(getattr(projection, "bits", 0) or 0) != _BITS + or int(getattr(projection, "group_size", 0) or 0) != _GROUP_SIZE + or str(getattr(projection, "mode", "")).lower() != "affine" + or getattr(projection, "bias", None) is not None + or _shape(scales) != bound_contract.metadata_shape + or _shape(biases) != bound_contract.metadata_shape + or getattr(weight, "dtype", None) != mx.uint32 + or getattr(scales, "dtype", None) != mx.bfloat16 + or getattr(biases, "dtype", None) != mx.bfloat16 + ): + raise M3WOBContractError( + "fixed M3 wo_b requires unbatched affine Q6/G128 U32 weights and " + "BF16 scales/biases" + ) + return bound_contract + + +def m3_wob_metal_source() -> str: + """Return the fixed K=8192 Q6 body with three independent stock M1 trees.""" + + return r""" +constexpr uint M = 3; +constexpr uint K = 8192; +constexpr uint N = 4096; +constexpr uint GS = 128; +constexpr uint PACKS_PER_THREAD = 2; +constexpr uint PACK_FACTOR = 4; +constexpr uint BYTES_PER_PACK = 3; +constexpr uint VALUES_PER_THREAD = PACK_FACTOR * PACKS_PER_THREAD; +constexpr uint BLOCK_SIZE = VALUES_PER_THREAD * 32; +constexpr uint RESULTS_PER_SIMDGROUP = 4; + +uint simd_gid = simdgroup_index_in_threadgroup; +uint simd_lid = thread_index_in_simdgroup; +uint out_row = threadgroup_position_in_grid.x * 8 + simd_gid * RESULTS_PER_SIMDGROUP; +uint group_count = K / GS; +uint row_bytes = K * 6 / 8; +const device uchar* ws = (const device uchar*)w + out_row * row_bytes + simd_lid * PACKS_PER_THREAD * BYTES_PER_PACK; +const device T* sc = scales + out_row * group_count + simd_lid / (GS / VALUES_PER_THREAD); +const device T* bs = biases + out_row * group_count + simd_lid / (GS / VALUES_PER_THREAD); +const device T* x0 = x + simd_lid * VALUES_PER_THREAD; +const device T* x1 = x0 + K; +const device T* x2 = x1 + K; +device T* y0 = y + out_row; +device T* y1 = y0 + N; +device T* y2 = y1 + N; +float result0[RESULTS_PER_SIMDGROUP] = {0.0f}; +float result1[RESULTS_PER_SIMDGROUP] = {0.0f}; +float result2[RESULTS_PER_SIMDGROUP] = {0.0f}; + +for (uint k = 0; k < K; k += BLOCK_SIZE) { + thread float x0_thread[VALUES_PER_THREAD]; + thread float x1_thread[VALUES_PER_THREAD]; + thread float x2_thread[VALUES_PER_THREAD]; + float sum0 = 0.0f, sum1 = 0.0f, sum2 = 0.0f; + for (uint i = 0; i < VALUES_PER_THREAD; i += 4) { + sum0 += x0[i] + x0[i + 1] + x0[i + 2] + x0[i + 3]; + sum1 += x1[i] + x1[i + 1] + x1[i + 2] + x1[i + 3]; + sum2 += x2[i] + x2[i + 1] + x2[i + 2] + x2[i + 3]; + x0_thread[i] = x0[i]; x0_thread[i + 1] = x0[i + 1] / 64.0f; + x0_thread[i + 2] = x0[i + 2] / 16.0f; x0_thread[i + 3] = x0[i + 3] / 4.0f; + x1_thread[i] = x1[i]; x1_thread[i + 1] = x1[i + 1] / 64.0f; + x1_thread[i + 2] = x1[i + 2] / 16.0f; x1_thread[i + 3] = x1[i + 3] / 4.0f; + x2_thread[i] = x2[i]; x2_thread[i + 1] = x2[i + 1] / 64.0f; + x2_thread[i + 2] = x2[i + 2] / 16.0f; x2_thread[i + 3] = x2[i + 3] / 4.0f; + } + for (uint row = 0; row < RESULTS_PER_SIMDGROUP; ++row) { + const device uchar* wl = ws + row * row_bytes; + const device T* sl = sc + row * group_count; + const device T* bl = bs + row * group_count; + thread uchar w_thread[PACKS_PER_THREAD * BYTES_PER_PACK]; + for (uint i = 0; i < PACKS_PER_THREAD * BYTES_PER_PACK; ++i) w_thread[i] = wl[i]; + float s = sl[0], b = bl[0]; + float dot0 = 0.0f, dot1 = 0.0f, dot2 = 0.0f; + const thread uchar* wp = w_thread; + const thread float* xp0 = x0_thread; const thread float* xp1 = x1_thread; const thread float* xp2 = x2_thread; + for (uint i = 0; i < VALUES_PER_THREAD / 4; ++i) { + xp0 += 4 * i; xp1 += 4 * i; xp2 += 4 * i; wp += 3 * i; + dot0 += (wp[0] & 0x3f) * xp0[0]; dot1 += (wp[0] & 0x3f) * xp1[0]; dot2 += (wp[0] & 0x3f) * xp2[0]; + dot0 += (wp[0] & 0xc0) * xp0[1]; dot1 += (wp[0] & 0xc0) * xp1[1]; dot2 += (wp[0] & 0xc0) * xp2[1]; + dot0 += (wp[1] & 0x0f) * (xp0[1] * 256.0f); dot1 += (wp[1] & 0x0f) * (xp1[1] * 256.0f); dot2 += (wp[1] & 0x0f) * (xp2[1] * 256.0f); + dot0 += (wp[1] & 0xf0) * xp0[2]; dot1 += (wp[1] & 0xf0) * xp1[2]; dot2 += (wp[1] & 0xf0) * xp2[2]; + dot0 += (wp[2] & 0x03) * (xp0[2] * 256.0f); dot1 += (wp[2] & 0x03) * (xp1[2] * 256.0f); dot2 += (wp[2] & 0x03) * (xp2[2] * 256.0f); + dot0 += (wp[2] & 0xfc) * xp0[3]; dot1 += (wp[2] & 0xfc) * xp1[3]; dot2 += (wp[2] & 0xfc) * xp2[3]; + } + result0[row] += s * dot0 + sum0 * b; + result1[row] += s * dot1 + sum1 * b; + result2[row] += s * dot2 + sum2 * b; + } + ws += BLOCK_SIZE * BYTES_PER_PACK / PACK_FACTOR; + sc += BLOCK_SIZE / GS; bs += BLOCK_SIZE / GS; + x0 += BLOCK_SIZE; x1 += BLOCK_SIZE; x2 += BLOCK_SIZE; +} +for (uint row = 0; row < RESULTS_PER_SIMDGROUP; ++row) { + float r0 = simd_sum(result0[row]); + float r1 = simd_sum(result1[row]); + float r2 = simd_sum(result2[row]); + if (simd_lid == 0) { y0[row] = T(r0); y1[row] = T(r1); y2[row] = T(r2); } +} +""" + + +@lru_cache(maxsize=1) +def _build_wob_kernel(): + return mx.fast.metal_kernel( + name="mtplx_dsv4_0731_official_m3_wob_q6", + input_names=["x", "w", "scales", "biases"], + output_names=["y"], + header="using namespace metal;", + source=m3_wob_metal_source(), + ensure_row_contiguous=False, + ) + + +class BoundM3WOB: + """Prebound direct fixed-M3 wo_b callable; no execution-path checks exist.""" + + __slots__ = ( + "_kernel", + "biases", + "grid", + "input_shape", + "output_shape", + "scales", + "threadgroup", + "weight", + ) + + def __init__(self, projection: Any, contract: M3WOBContract): + validate_wob_projection(projection, contract) + self.weight = projection.weight + self.scales = projection.scales + self.biases = projection.biases + self._kernel = _build_wob_kernel() + self.input_shape = (1, _M, _K) + self.output_shape = (1, _M, _N) + self.grid = ((_N // 8) * 64, 1, 1) + self.threadgroup = (64, 1, 1) + + def __call__(self, x: mx.array) -> mx.array: + (out,) = self._kernel( + inputs=[x.reshape(_M, _K), self.weight, self.scales, self.biases], + template=[("T", mx.bfloat16)], + grid=self.grid, + threadgroup=self.threadgroup, + output_shapes=[(_M, _N)], + output_dtypes=[mx.bfloat16], + ) + return out.reshape(self.output_shape) + + +def bind_m3_wob(projection: Any) -> Callable[[mx.array], mx.array]: + """Validate then bind the sole supported physical-M3 wo_b contract.""" + + return BoundM3WOB(projection, M3WOBContract()) + + +class _PreboundWOBRoute: + """The sole hot choice: captured stock phase versus physical M3.""" + + __slots__ = ("candidate", "stock") + + def __init__( + self, + stock: Callable[[mx.array], mx.array], + candidate: Callable[[mx.array], mx.array], + ) -> None: + self.stock = stock + self.candidate = candidate + + def __call__(self, x: mx.array) -> mx.array: + if int(x.shape[1]) == _M: + return self.candidate(x) + return self.stock(x) + + +def prebind_wob_route( + stock: Callable[[mx.array], mx.array], candidate: Callable[[mx.array], mx.array] +) -> Callable[[mx.array], mx.array]: + """Capture the M3/stock phase route; execution checks only logical shape.""" + + return _PreboundWOBRoute(stock, candidate) + + +@dataclass(frozen=True, slots=True) +class PreparedWOBM3Routes: + """Self-checked 43-layer ``wo_b`` bank awaiting atomic publication.""" + + attentions: tuple[Any, ...] + o_lora_impls: tuple[Any, ...] + stock_routes: tuple[Callable[[mx.array], mx.array], ...] + candidate_routes: tuple[Callable[[mx.array], mx.array], ...] + published_routes: tuple[Callable[[mx.array], mx.array], ...] + layer_count: int + q6_count: int + exact_selfchecked: int + o_lora_sink_count: int + + def publish(self) -> None: + try: + for attention, o_lora_impl, route in zip( + self.attentions, + self.o_lora_impls, + self.published_routes, + ): + attention.wo_b = route + o_lora_impl.wo_b = route + except Exception as publication_error: + try: + self.restore() + except Exception as restoration_error: + publication_error.add_note( + f"WOB rollback also raised: {restoration_error}" + ) + raise + + def restore(self) -> None: + errors = [] + for attention, o_lora_impl, stock in zip( + self.attentions, + self.o_lora_impls, + self.stock_routes, + ): + try: + attention.wo_b = stock + except Exception as exc: + errors.append(exc) + try: + o_lora_impl.wo_b = stock + except Exception as exc: + errors.append(exc) + if errors: + raise ExceptionGroup("WOB route restoration failed", errors) + + +def prepare_wob_m3( + layers: Any, + *, + exact_selfcheck: Callable[ + [Callable[[mx.array], mx.array], Callable[[mx.array], mx.array], int], bool + ], +) -> PreparedWOBM3Routes: + """Bind and self-check every ``wo_b`` route without publishing. + + The caller supplies an exact real-weight self-check. Every layer validates + before any custom candidate exists; every live gather-o-LoRA sink must + still alias that stock projection. Every candidate is then self-checked + before either reference changes. + """ + + layer_tuple = tuple(layers) + if len(layer_tuple) != 43: + raise M3WOBContractError("fixed M3 wo_b preparation requires exactly 43 layers") + if not callable(exact_selfcheck): + raise M3WOBContractError("exact WOB self-check is required") + + validated: list[tuple[Any, Any, Callable[[mx.array], mx.array]]] = [] + for index, layer in enumerate(layer_tuple): + try: + attention = layer.attn + stock = attention.wo_b + except AttributeError as exc: + raise M3WOBContractError(f"layer {index} has no attention wo_b") from exc + if not callable(stock): + raise M3WOBContractError(f"layer {index} attention wo_b is not callable") + try: + o_lora_impl = attention._o_lora_impl + o_lora_sink = o_lora_impl.wo_b + except AttributeError as exc: + raise M3WOBContractError( + f"layer {index} has no active o-LoRA wo_b sink" + ) from exc + if o_lora_sink is not stock: + raise M3WOBContractError(f"layer {index} active o-LoRA wo_b sink is stale") + validate_wob_projection(stock, M3WOBContract()) + validated.append((attention, o_lora_impl, stock)) + + staged: list[ + tuple[ + Any, + Any, + Callable[[mx.array], mx.array], + Callable[[mx.array], mx.array], + Callable[[mx.array], mx.array], + ] + ] = [] + rejected_selfcheck: int | None = None + for index, (attention, o_lora_impl, stock) in enumerate(validated): + candidate = bind_m3_wob(stock) + if not callable(candidate): + raise M3WOBContractError( + f"layer {index} fixed M3 candidate is not callable" + ) + try: + passed = exact_selfcheck(stock, candidate, index) + except Exception as exc: + raise M3WOBContractError( + f"layer {index} fixed M3 exact self-check raised" + ) from exc + if not passed and rejected_selfcheck is None: + rejected_selfcheck = index + staged.append( + ( + attention, + o_lora_impl, + stock, + candidate, + prebind_wob_route(stock, candidate), + ) + ) + + if rejected_selfcheck is not None: + raise M3WOBContractError( + f"layer {rejected_selfcheck} fixed M3 exact self-check failed" + ) + + return PreparedWOBM3Routes( + attentions=tuple(attention for attention, _, _, _, _ in staged), + o_lora_impls=tuple(o_lora_impl for _, o_lora_impl, _, _, _ in staged), + stock_routes=tuple(stock for _, _, stock, _, _ in staged), + candidate_routes=tuple(candidate for _, _, _, candidate, _ in staged), + published_routes=tuple(route for _, _, _, _, route in staged), + layer_count=len(staged), + q6_count=len(staged), + exact_selfchecked=len(staged), + o_lora_sink_count=len(staged), + ) diff --git a/mtplx/deepseek_v4_0731_m3_wqb_qnorm_rope.py b/mtplx/deepseek_v4_0731_m3_wqb_qnorm_rope.py new file mode 100644 index 00000000..2305cc58 --- /dev/null +++ b/mtplx/deepseek_v4_0731_m3_wqb_qnorm_rope.py @@ -0,0 +1,514 @@ +"""Fixed-M3 Q6 ``wq_b`` fused with row-owned Q-head norm and RoPE. + +One 256-threadgroup owns one of the 64 output heads and all three verifier +rows. It replays the official-wheel Q6/G128 affine QMV reduction for each +output component, stores the resulting BF16 head in threadgroup memory, then +replays the stock 32-lane per-head RMSNorm reduction and interleaved-RoPE +arithmetic. + +This is intentionally only a construction-bound micro candidate. It has a +small grid (64 threadgroups) and high register pressure, so it is neither a +throughput claim nor a parity claim until its guarded GPU bracket passes. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from functools import lru_cache +from typing import Any, Callable + +import numpy as np + +import mlx.core as mx + + +_M = 3 +_K = 1024 +_N = 32768 +_HEADS = 64 +_HEAD_DIM = 512 +_ROPE_DIM = 64 +_EPS = 1e-6 + +# Provenance of the recovered pre-geometry session snapshots. These identify +# the historical inputs, not this module after its staging publisher was added. +RECORDED_PRE_GEOMETRY_SOURCE_SHA256 = ( + "2eb4ce3d5bae9c9b71574d17fedfd37b6755d94299cb4ed6b01d2015c5f8f9a1" +) +RECORDED_PRE_GEOMETRY_TEST_SHA256 = ( + "0c551aa8d7f865454d3a9b6d22f46af5115e00a2fc48c5db82225516c2a77cbb" +) + + +class M3WQBNormRopeContractError(ValueError): + """The loaded projection cannot use the fixed 0731 micro candidate.""" + + +@dataclass(frozen=True, slots=True) +class M3WQBNormRopeContract: + """The immutable Q6 wq_b and post-projection head geometry.""" + + bits: int = 6 + group_size: int = 128 + k: int = _K + n: int = _N + heads: int = _HEADS + head_dim: int = _HEAD_DIM + rope_dim: int = _ROPE_DIM + eps: float = _EPS + dtype: mx.Dtype = mx.bfloat16 + + @property + def weight_shape(self) -> tuple[int, int]: + return (self.n, self.k * self.bits // 32) + + @property + def metadata_shape(self) -> tuple[int, int]: + return (self.n, self.k // self.group_size) + + def valid(self) -> bool: + return ( + (self.bits, self.group_size, self.k, self.n) == (6, 128, _K, _N) + and (self.heads, self.head_dim, self.rope_dim) + == ( + _HEADS, + _HEAD_DIM, + _ROPE_DIM, + ) + and self.eps == _EPS + and self.dtype == mx.bfloat16 + ) + + +def _shape(value: Any) -> tuple[int, ...] | None: + try: + return tuple(int(dimension) for dimension in value.shape) + except (AttributeError, TypeError, ValueError): + return None + + +def validate_0731_m3_wqb_qnorm_rope( + projection: Any, + contract: M3WQBNormRopeContract | None = None, +) -> M3WQBNormRopeContract: + """Validate fixed Q6/G128 storage once before publishing the callable.""" + + fixed = M3WQBNormRopeContract() if contract is None else contract + if not fixed.valid(): + raise M3WQBNormRopeContractError( + "candidate requires fixed 0731 Q6/G128 geometry" + ) + if ( + int(getattr(projection, "bits", 0) or 0) != fixed.bits + or int(getattr(projection, "group_size", 0) or 0) != fixed.group_size + or str(getattr(projection, "mode", "")).lower() != "affine" + or getattr(projection, "bias", None) is not None + ): + raise M3WQBNormRopeContractError("candidate requires affine Q6/G128 wq_b") + if ( + _shape(getattr(projection, "weight", None)) != fixed.weight_shape + or _shape(getattr(projection, "scales", None)) != fixed.metadata_shape + or _shape(getattr(projection, "biases", None)) != fixed.metadata_shape + or getattr(getattr(projection, "weight", None), "dtype", None) != mx.uint32 + or getattr(getattr(projection, "scales", None), "dtype", None) != mx.bfloat16 + or getattr(getattr(projection, "biases", None), "dtype", None) != mx.bfloat16 + ): + raise M3WQBNormRopeContractError( + "candidate requires exact packed wq_b Q6 storage" + ) + return fixed + + +def m3_wqb_qnorm_rope_metal_source(*, capture_projection: bool = False) -> str: + """Return the production source, or its test-only raw-projection variant.""" + + source = r""" +constexpr uint M = 3; +constexpr uint K = 1024; +constexpr uint N = 32768; +constexpr uint HEADS = 64; +constexpr uint HEAD_DIM = 512; +constexpr uint ROPE_DIM = 64; +constexpr float EPS = 1e-6f; +constexpr uint GS = 128; +constexpr uint PACKS_PER_THREAD = 2; +constexpr uint PACK_FACTOR = 4; +constexpr uint BYTES_PER_PACK = 3; +constexpr uint VALUES_PER_THREAD = PACK_FACTOR * PACKS_PER_THREAD; +constexpr uint BLOCK_SIZE = VALUES_PER_THREAD * 32; +constexpr uint RESULTS_PER_SIMDGROUP = 4; +constexpr uint SIMDGROUPS = 8; +constexpr uint ROWS_PER_TILE = SIMDGROUPS * RESULTS_PER_SIMDGROUP; +constexpr uint TILES_PER_HEAD = HEAD_DIM / ROWS_PER_TILE; +constexpr uint NORM_LANES = 32; +constexpr uint NORM_READS = 4; +constexpr uint NORM_BLOCKS = HEAD_DIM / (NORM_LANES * NORM_READS); + +uint tid = thread_index_in_threadgroup; +uint simd_lid = thread_index_in_simdgroup; +uint simd_gid = simdgroup_index_in_threadgroup; +uint head = threadgroup_position_in_grid.x; +uint row_bytes = K * 6 / 8; +uint group_count = K / GS; +threadgroup T q_shared[M][HEAD_DIM]; +threadgroup float norm_scale[M]; + +for (uint tile = 0; tile < TILES_PER_HEAD; ++tile) { + uint out_row = head * HEAD_DIM + tile * ROWS_PER_TILE + simd_gid * RESULTS_PER_SIMDGROUP; + const device uchar* ws = (const device uchar*)w + out_row * row_bytes + simd_lid * PACKS_PER_THREAD * BYTES_PER_PACK; + const device T* sc = scales + out_row * group_count + simd_lid / (GS / VALUES_PER_THREAD); + const device T* bs = biases + out_row * group_count + simd_lid / (GS / VALUES_PER_THREAD); + const device T* x0 = x + simd_lid * VALUES_PER_THREAD; + const device T* x1 = x0 + K; + const device T* x2 = x1 + K; + float result0[RESULTS_PER_SIMDGROUP] = {0.0f}; + float result1[RESULTS_PER_SIMDGROUP] = {0.0f}; + float result2[RESULTS_PER_SIMDGROUP] = {0.0f}; + + for (uint k = 0; k < K; k += BLOCK_SIZE) { + thread float x0_thread[VALUES_PER_THREAD]; + thread float x1_thread[VALUES_PER_THREAD]; + thread float x2_thread[VALUES_PER_THREAD]; + float sum0 = 0.0f, sum1 = 0.0f, sum2 = 0.0f; + for (uint i = 0; i < VALUES_PER_THREAD; i += 4) { + sum0 += x0[i] + x0[i + 1] + x0[i + 2] + x0[i + 3]; + sum1 += x1[i] + x1[i + 1] + x1[i + 2] + x1[i + 3]; + sum2 += x2[i] + x2[i + 1] + x2[i + 2] + x2[i + 3]; + x0_thread[i] = x0[i]; x0_thread[i + 1] = x0[i + 1] / 64.0f; + x0_thread[i + 2] = x0[i + 2] / 16.0f; x0_thread[i + 3] = x0[i + 3] / 4.0f; + x1_thread[i] = x1[i]; x1_thread[i + 1] = x1[i + 1] / 64.0f; + x1_thread[i + 2] = x1[i + 2] / 16.0f; x1_thread[i + 3] = x1[i + 3] / 4.0f; + x2_thread[i] = x2[i]; x2_thread[i + 1] = x2[i + 1] / 64.0f; + x2_thread[i + 2] = x2[i + 2] / 16.0f; x2_thread[i + 3] = x2[i + 3] / 4.0f; + } + for (uint output_row = 0; output_row < RESULTS_PER_SIMDGROUP; ++output_row) { + const device uchar* wl = ws + output_row * row_bytes; + const device T* sl = sc + output_row * group_count; + const device T* bl = bs + output_row * group_count; + thread uchar w_thread[PACKS_PER_THREAD * BYTES_PER_PACK]; + for (uint i = 0; i < PACKS_PER_THREAD * BYTES_PER_PACK; ++i) w_thread[i] = wl[i]; + float s = sl[0], b = bl[0]; + float dot0 = 0.0f, dot1 = 0.0f, dot2 = 0.0f; + const thread uchar* wp = w_thread; + const thread float* xp0 = x0_thread; const thread float* xp1 = x1_thread; const thread float* xp2 = x2_thread; + for (uint i = 0; i < VALUES_PER_THREAD / 4; ++i) { + xp0 += 4 * i; xp1 += 4 * i; xp2 += 4 * i; wp += 3 * i; + dot0 += (wp[0] & 0x3f) * xp0[0]; dot1 += (wp[0] & 0x3f) * xp1[0]; dot2 += (wp[0] & 0x3f) * xp2[0]; + dot0 += (wp[0] & 0xc0) * xp0[1]; dot1 += (wp[0] & 0xc0) * xp1[1]; dot2 += (wp[0] & 0xc0) * xp2[1]; + dot0 += (wp[1] & 0x0f) * (xp0[1] * 256.0f); dot1 += (wp[1] & 0x0f) * (xp1[1] * 256.0f); dot2 += (wp[1] & 0x0f) * (xp2[1] * 256.0f); + dot0 += (wp[1] & 0xf0) * xp0[2]; dot1 += (wp[1] & 0xf0) * xp1[2]; dot2 += (wp[1] & 0xf0) * xp2[2]; + dot0 += (wp[2] & 0x03) * (xp0[2] * 256.0f); dot1 += (wp[2] & 0x03) * (xp1[2] * 256.0f); dot2 += (wp[2] & 0x03) * (xp2[2] * 256.0f); + dot0 += (wp[2] & 0xfc) * xp0[3]; dot1 += (wp[2] & 0xfc) * xp1[3]; dot2 += (wp[2] & 0xfc) * xp2[3]; + } + result0[output_row] += s * dot0 + sum0 * b; + result1[output_row] += s * dot1 + sum1 * b; + result2[output_row] += s * dot2 + sum2 * b; + } + ws += BLOCK_SIZE * BYTES_PER_PACK / PACK_FACTOR; sc += BLOCK_SIZE / GS; bs += BLOCK_SIZE / GS; + x0 += BLOCK_SIZE; x1 += BLOCK_SIZE; x2 += BLOCK_SIZE; + } + for (uint output_row = 0; output_row < RESULTS_PER_SIMDGROUP; ++output_row) { + uint d = tile * ROWS_PER_TILE + simd_gid * RESULTS_PER_SIMDGROUP + output_row; + float r0 = simd_sum(result0[output_row]); float r1 = simd_sum(result1[output_row]); float r2 = simd_sum(result2[output_row]); + if (simd_lid == 0) { + q_shared[0][d] = T(r0); q_shared[1][d] = T(r1); q_shared[2][d] = T(r2); + /* CAPTURE_PROJECTION */ + } + } +} +threadgroup_barrier(mem_flags::mem_threadgroup); +for (uint row = 0; row < M; ++row) { + // Clone the 0.31.2 FP32 row_reduce_simple tree for a 512-element row: + // 32 lanes, four strided 128-value blocks, and four contiguous reads/lane. + if (simd_gid == 0) { + float sum = 0.0f; + for (uint group = 0; group < NORM_BLOCKS; ++group) { + uint d = group * NORM_LANES * NORM_READS + simd_lid * NORM_READS; + sum = float(q_shared[row][d]) * float(q_shared[row][d]) + sum; + sum = float(q_shared[row][d + 1]) * float(q_shared[row][d + 1]) + sum; + sum = float(q_shared[row][d + 2]) * float(q_shared[row][d + 2]) + sum; + sum = float(q_shared[row][d + 3]) * float(q_shared[row][d + 3]) + sum; + } + sum = simd_sum(sum); + if (simd_lid == 0) { + float mean = sum * (1.0f / float(HEAD_DIM)); + norm_scale[row] = metal::precise::rsqrt(mean + EPS); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float scale = norm_scale[row]; + const device float* c = cos + row * (ROPE_DIM / 2); + const device float* s = sin + row * (ROPE_DIM / 2); + device T* out = output + row * N + head * HEAD_DIM; + for (uint d = tid; d < HEAD_DIM; d += 256) { + if (d < HEAD_DIM - ROPE_DIM) { + out[d] = T(float(q_shared[row][d]) * scale); + continue; + } + uint r = d - (HEAD_DIM - ROPE_DIM); + if ((r & 1) != 0) continue; + uint pair = r / 2; + T normalized0 = T(float(q_shared[row][d]) * scale); + T normalized1 = T(float(q_shared[row][d + 1]) * scale); + float x0 = float(normalized0); float x1 = float(normalized1); + // The eager stock graph materializes each FP32 product in a separate + // binary kernel before its add/subtract. ``precise`` preserves those + // intermediate roundings inside this one-launch candidate. + precise float rope0_lhs = x0 * c[pair]; + precise float rope0_rhs = x1 * s[pair]; + precise float rope1_lhs = x0 * s[pair]; + precise float rope1_rhs = x1 * c[pair]; + out[d] = T(rope0_lhs - rope0_rhs); + out[d + 1] = T(rope1_lhs + rope1_rhs); + } + threadgroup_barrier(mem_flags::mem_threadgroup); +} +""" + capture = """ + device T* q0 = projection + head * HEAD_DIM + d; + device T* q1 = q0 + N; + device T* q2 = q1 + N; + q0[0] = T(r0); q1[0] = T(r1); q2[0] = T(r2);""" + return source.replace( + "/* CAPTURE_PROJECTION */", capture if capture_projection else "" + ) + + +@lru_cache(maxsize=1) +def _build_kernel(): + return mx.fast.metal_kernel( + name="mtplx_dsv4_0731_m3_q6_wqb_qnorm_rope", + input_names=["x", "w", "scales", "biases", "cos", "sin"], + output_names=["output"], + header="using namespace metal;", + source=m3_wqb_qnorm_rope_metal_source(), + ensure_row_contiguous=False, + ) + + +@lru_cache(maxsize=1) +def _build_debug_kernel(): + """Build the test-only twin that exposes the raw fused BF16 projection.""" + + return mx.fast.metal_kernel( + name="mtplx_dsv4_0731_m3_q6_wqb_qnorm_rope_debug", + input_names=["x", "w", "scales", "biases", "cos", "sin"], + output_names=["projection", "output"], + header="using namespace metal;", + source=m3_wqb_qnorm_rope_metal_source(capture_projection=True), + ensure_row_contiguous=False, + ) + + +class BoundM3WQBNormRope: + """Prepared one-launch callable; all geometry checks occurred before binding.""" + + __slots__ = ( + "_kernel", + "biases", + "contract", + "grid", + "input_shape", + "output_shape", + "scales", + "threadgroup", + "weight", + ) + + def __init__(self, projection: Any, contract: M3WQBNormRopeContract) -> None: + self.contract = contract + self.weight = projection.weight + self.scales = projection.scales + self.biases = projection.biases + self._kernel = _build_kernel() + self.input_shape = (1, _M, _K) + self.output_shape = (1, _M, _HEADS, _HEAD_DIM) + self.grid = (_HEADS * 256, 1, 1) + self.threadgroup = (256, 1, 1) + + def __call__(self, x: mx.array, cos: mx.array, sin: mx.array) -> mx.array: + (output,) = self._kernel( + inputs=[ + x.reshape(_M, _K), + self.weight, + self.scales, + self.biases, + cos.reshape(_M, _ROPE_DIM // 2), + sin.reshape(_M, _ROPE_DIM // 2), + ], + template=[("T", mx.bfloat16)], + grid=self.grid, + threadgroup=self.threadgroup, + output_shapes=[(_M, _N)], + output_dtypes=[mx.bfloat16], + ) + return output.reshape(self.output_shape) + + +def build_0731_m3_wqb_qnorm_rope( + projection: Any, +) -> Callable[[mx.array, mx.array, mx.array], mx.array]: + """Bind the fixed Q6/G128 M3 candidate without a hot-path fallback.""" + + contract = validate_0731_m3_wqb_qnorm_rope(projection) + return BoundM3WQBNormRope(projection, contract) + + +class _PreboundWQBQHeadRoute: + """Construction-selected fixed-M3 route or the captured stock route.""" + + __slots__ = ("candidate", "stock") + + def __init__( + self, + stock: Callable[[mx.array, mx.array, mx.array], mx.array], + candidate: Callable[[mx.array, mx.array, mx.array], mx.array], + ) -> None: + self.stock = stock + self.candidate = candidate + + def __call__(self, qr: mx.array, cos: mx.array, sin: mx.array) -> mx.array: + if int(qr.shape[1]) == _M: + return self.candidate(qr, cos, sin) + return self.stock(qr, cos, sin) + + +def prebind_wqb_qhead_route( + stock: Callable[[mx.array, mx.array, mx.array], mx.array], + candidate: Callable[[mx.array, mx.array, mx.array], mx.array], +) -> Callable[[mx.array, mx.array, mx.array], mx.array]: + return _PreboundWQBQHeadRoute(stock, candidate) + + +@dataclass(frozen=True, slots=True) +class PreparedWQBQHeadM3Routes: + """Self-checked 43-layer route bank awaiting atomic publication.""" + + attentions: tuple[Any, ...] + stock_routes: tuple[Callable[[mx.array, mx.array, mx.array], mx.array], ...] + candidate_routes: tuple[Callable[[mx.array, mx.array, mx.array], mx.array], ...] + published_routes: tuple[Callable[[mx.array, mx.array, mx.array], mx.array], ...] + q6_count: int + exact_selfchecked: int + + def publish(self) -> None: + try: + for attention, route in zip(self.attentions, self.published_routes): + attention._q_projection_qhead_route = route + except Exception as publication_error: + try: + self.restore() + except Exception as restoration_error: + publication_error.add_note( + f"WQB-qhead rollback also raised: {restoration_error}" + ) + raise + + def restore(self) -> None: + errors = [] + for attention, stock in zip(self.attentions, self.stock_routes): + try: + attention._q_projection_qhead_route = stock + except Exception as exc: + errors.append(exc) + if errors: + raise ExceptionGroup("WQB-qhead route restoration failed", errors) + + +def prepare_wqb_qhead_m3( + layers: Any, + *, + exact_selfcheck: Callable[ + [ + Callable[[mx.array, mx.array, mx.array], mx.array], + Callable[[mx.array, mx.array, mx.array], mx.array], + int, + ], + bool, + ], +) -> PreparedWQBQHeadM3Routes: + """Validate, build, and exact-self-check all routes without publishing.""" + + layer_tuple = tuple(layers) + if len(layer_tuple) != 43: + raise M3WQBNormRopeContractError( + "fixed M3 wq_b plus q-head preparation requires exactly 43 layers" + ) + if not callable(exact_selfcheck): + raise M3WQBNormRopeContractError("exact WQB-qhead self-check is required") + + validated = [] + for layer_index, layer in enumerate(layer_tuple): + try: + attention = layer.attn + projection = attention.wq_b + stock_route = attention._q_projection_qhead_route + except AttributeError as exc: + raise M3WQBNormRopeContractError( + f"layer {layer_index} lacks the q-projection/post route" + ) from exc + if not callable(stock_route): + raise M3WQBNormRopeContractError( + f"layer {layer_index} q-projection/post stock route is not callable" + ) + validate_0731_m3_wqb_qnorm_rope(projection) + validated.append((attention, projection, stock_route)) + + staged = [] + rejected_selfcheck = None + for layer_index, (attention, projection, stock_route) in enumerate(validated): + candidate = build_0731_m3_wqb_qnorm_rope(projection) + try: + passed = exact_selfcheck(stock_route, candidate, layer_index) + except Exception as exc: + raise M3WQBNormRopeContractError( + f"layer {layer_index} fused M3 exact self-check raised" + ) from exc + if not passed and rejected_selfcheck is None: + rejected_selfcheck = layer_index + staged.append( + ( + attention, + stock_route, + candidate, + prebind_wqb_qhead_route(stock_route, candidate), + ) + ) + + if rejected_selfcheck is not None: + raise M3WQBNormRopeContractError( + f"layer {rejected_selfcheck} fused M3 exact self-check failed" + ) + + return PreparedWQBQHeadM3Routes( + attentions=tuple(attention for attention, _, _, _ in staged), + stock_routes=tuple(stock for _, stock, _, _ in staged), + candidate_routes=tuple(candidate for _, _, candidate, _ in staged), + published_routes=tuple(route for _, _, _, route in staged), + q6_count=len(staged), + exact_selfchecked=len(staged), + ) + + +def q_head_norm_rope_cpu_oracle( + q: np.ndarray, cos: np.ndarray, sin: np.ndarray, *, eps: float, rope_dim: int +) -> np.ndarray: + """Independent CPU oracle for row-owned per-head norm and RoPE semantics.""" + + values = np.asarray(q, dtype=np.float32) + c = np.asarray(cos, dtype=np.float32) + s = np.asarray(sin, dtype=np.float32) + normalized = values * np.reciprocal( + np.sqrt(np.mean(np.square(values), axis=-1, keepdims=True) + float(eps)) + ) + out = normalized.copy() + tail = normalized[..., -int(rope_dim) :].reshape(*normalized.shape[:-1], -1, 2) + rotated = np.empty_like(tail) + rotated[..., 0] = ( + tail[..., 0] * c[None, :, None, :] - tail[..., 1] * s[None, :, None, :] + ) + rotated[..., 1] = ( + tail[..., 0] * s[None, :, None, :] + tail[..., 1] * c[None, :, None, :] + ) + out[..., -int(rope_dim) :] = rotated.reshape(*normalized.shape[:-1], int(rope_dim)) + return out diff --git a/mtplx/deepseek_v4_0731_moe.py b/mtplx/deepseek_v4_0731_moe.py new file mode 100644 index 00000000..a279f0e9 --- /dev/null +++ b/mtplx/deepseek_v4_0731_moe.py @@ -0,0 +1,217 @@ +"""Receipt-backed routed-MoE primitives for DeepSeek-V4-Flash-0731. + +This module owns only the measured target lane: affine Q2/group-128 routed +gate/up packing and the fixed M1 top-6 row-owned reduction. Construction +validates storage once; installed callables perform no environment reads, +eligibility checks, counters, or fallback routing. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from functools import lru_cache +from typing import Callable + +import mlx.core as mx +import mlx.nn as nn +from mlx_lm.models.switch_layers import QuantizedSwitchLinear + +from .moe_packed_projections import PackedSwitchGLU, _pack_pair + + +@dataclass(frozen=True, slots=True) +class RoutedQ2Contract: + bits: int + group_size: int + hidden_size: int + width: int + experts: int + + +class DeepseekV40731PackedQ2SwitchGLU(PackedSwitchGLU): + """Fixed M1/M3 unsorted gather path for the pinned top-6 lane.""" + + def __call__(self, x: mx.array, indices: mx.array) -> mx.array: + x = mx.expand_dims(x, (-2, -3)) + packed = self.gate_up_proj.gather(x, indices, False) + gate, up = mx.split(packed, [self._split_at], axis=-1) + routed = self.down_proj( + self.activation(up, gate), + indices, + sorted_indices=False, + ) + return routed.squeeze(-2) + + +def validate_routed_q2_pair( + gate: nn.Module, + up: nn.Module, + *, + hidden_size: int, + width: int, + experts: int, +) -> RoutedQ2Contract: + """Validate the exact affine-Q2 expert-bank layout before packing.""" + + hidden = int(hidden_size) + intermediate = int(width) + expert_count = int(experts) + expected_weight = (expert_count, intermediate, hidden // 16) + expected_meta = (expert_count, intermediate, hidden // 128) + valid = all( + ( + hidden > 0, + intermediate > 0, + expert_count > 0, + hidden % 128 == 0, + isinstance(gate, QuantizedSwitchLinear), + isinstance(up, QuantizedSwitchLinear), + int(getattr(gate, "bits", 0) or 0) == 2, + int(getattr(up, "bits", 0) or 0) == 2, + int(getattr(gate, "group_size", 0) or 0) == 128, + int(getattr(up, "group_size", 0) or 0) == 128, + str(getattr(gate, "mode", "")) == "affine", + str(getattr(up, "mode", "")) == "affine", + getattr(getattr(gate, "weight", None), "dtype", None) == mx.uint32, + getattr(getattr(up, "weight", None), "dtype", None) == mx.uint32, + tuple(getattr(getattr(gate, "weight", None), "shape", ())) + == expected_weight, + tuple(getattr(getattr(up, "weight", None), "shape", ())) == expected_weight, + tuple(getattr(getattr(gate, "scales", None), "shape", ())) == expected_meta, + tuple(getattr(getattr(up, "scales", None), "shape", ())) == expected_meta, + tuple(getattr(getattr(gate, "biases", None), "shape", ())) == expected_meta, + tuple(getattr(getattr(up, "biases", None), "shape", ())) == expected_meta, + getattr(getattr(gate, "scales", None), "dtype", None) == mx.bfloat16, + getattr(getattr(up, "scales", None), "dtype", None) == mx.bfloat16, + getattr(getattr(gate, "biases", None), "dtype", None) == mx.bfloat16, + getattr(getattr(up, "biases", None), "dtype", None) == mx.bfloat16, + "bias" not in gate, + "bias" not in up, + ) + ) + if not valid: + raise ValueError("DeepSeek-V4-0731 routed affine Q2/group-128 contract failed") + return RoutedQ2Contract( + bits=2, + group_size=128, + hidden_size=hidden, + width=intermediate, + experts=expert_count, + ) + + +def build_routed_q2_pair( + switch_mlp: nn.Module, + *, + hidden_size: int, + width: int, + experts: int, +) -> DeepseekV40731PackedQ2SwitchGLU: + """Pack gate/up output rows without changing any affine input group.""" + + validate_routed_q2_pair( + switch_mlp.gate_proj, + switch_mlp.up_proj, + hidden_size=hidden_size, + width=width, + experts=experts, + ) + packed = _pack_pair(switch_mlp.gate_proj, switch_mlp.up_proj, axis=1) + if isinstance(packed, str): + raise ValueError(f"DeepSeek-V4-0731 routed Q2 pair failed: {packed}") + gate_up, split_at = packed + if int(split_at) != int(width): + raise ValueError("DeepSeek-V4-0731 routed Q2 pair split is invalid") + return DeepseekV40731PackedQ2SwitchGLU( + gate_up, + switch_mlp.down_proj, + switch_mlp.activation, + split_at, + ) + + +@lru_cache(maxsize=None) +def _row_owned_combine_m1_kernel(): + source = r""" + uint column = thread_position_in_grid.x; + int hidden = int(HIDDEN_size); + if (int(column) >= hidden) { return; } + + bfloat accumulator = bfloat(0.0f); + _Pragma("unroll") + for (int expert = 0; expert < 6; ++expert) { + uint routed_index = uint(expert * hidden) + column; + bfloat score = bfloat(route_weights[expert]); + bfloat product = bfloat( + float(routed[routed_index]) * float(score)); + accumulator = bfloat(float(accumulator) + float(product)); + } + combined[column] = accumulator; + """ + return mx.fast.metal_kernel( + name="mtplx_dsv4_0731_row_owned_combine_top6_bf16", + input_names=["routed", "route_weights", "HIDDEN_size"], + output_names=["combined"], + source=source, + ensure_row_contiguous=True, + ) + + +def build_row_owned_combine_m1( + *, + hidden_size: int, + top_k: int, +) -> Callable[[mx.array, mx.array], mx.array]: + """Build the measured one-output-owner BF16 top-6 reduction.""" + + hidden = int(hidden_size) + if hidden != 4096 or int(top_k) != 6: + raise ValueError("DeepSeek-V4-0731 row-owned combine geometry is invalid") + if not mx.metal.is_available(): + raise ValueError("DeepSeek-V4-0731 row-owned combine requires Metal") + kernel = _row_owned_combine_m1_kernel() + + def combine(routed: mx.array, route_weights: mx.array) -> mx.array: + (output,) = kernel( + inputs=[ + routed.reshape(6, hidden), + route_weights.reshape(6), + hidden, + ], + grid=(hidden, 1, 1), + threadgroup=(128, 1, 1), + output_shapes=[(1, hidden)], + output_dtypes=[mx.bfloat16], + ) + return output + + return combine + + +def exact_selfcheck_row_owned_combine_m1( + combine: Callable[[mx.array, mx.array], mx.array], +) -> None: + """Execute the fixed BF16 reduction oracle once before publication.""" + + if not callable(combine): + raise ValueError("DeepSeek-V4-0731 row-owned combine is not callable") + routed = ( + ((mx.arange(6 * 4096, dtype=mx.float32) % 29 - 14) / 16.0) + .astype(mx.bfloat16) + .reshape(1, 6, 4096) + ) + route_weights = mx.array( + [[0.03125, 0.09375, 0.15625, 0.21875, 0.28125, 0.34375]], + dtype=mx.float32, + ) + expected = mx.zeros((1, 4096), dtype=mx.bfloat16) + weights_bf16 = route_weights.astype(mx.bfloat16) + for expert in range(6): + product = (routed[:, expert] * weights_bf16[:, expert : expert + 1]).astype( + mx.bfloat16 + ) + expected = (expected + product).astype(mx.bfloat16) + actual = combine(routed, route_weights) + mx.eval(actual, expected) + if not bool(mx.array_equal(actual, expected)): + raise ValueError("DeepSeek-V4-0731 row-owned combine exact self-check failed") diff --git a/mtplx/deepseek_v4_attention_island.py b/mtplx/deepseek_v4_attention_island.py index b0e54e74..f4631f8f 100644 --- a/mtplx/deepseek_v4_attention_island.py +++ b/mtplx/deepseek_v4_attention_island.py @@ -20,6 +20,8 @@ import mlx.core as mx from .attention_context import current_attention_phase, current_model_forward_kind +from .deepseek_v4_0731_moe import DeepseekV40731PackedQ2SwitchGLU +from .moe_packed_projections import PackedSwitchGLU from .models import deepseek_v4 as D @@ -52,9 +54,7 @@ class _Projection: input_dim: int -def _projection_contract( - module: Any, label: str, *, expected_bits: int -) -> _Projection: +def _projection_contract(module: Any, label: str, *, expected_bits: int) -> _Projection: """Validate one stock affine projection once and bind its array leaves.""" weight = getattr(module, "weight", None) @@ -115,9 +115,7 @@ def _qmm(x: mx.array, projection: _Projection) -> mx.array: ) -def _gather_qmm( - x: mx.array, indices: mx.array, projection: _Projection -) -> mx.array: +def _gather_qmm(x: mx.array, indices: mx.array, projection: _Projection) -> mx.array: return mx.gather_qmm( x, projection.weight, @@ -159,9 +157,7 @@ def _route( indices = mx.argpartition(-biased, kth=topk - 1, axis=-1)[..., :topk] route_weights = mx.take_along_axis(scores, indices, axis=-1) if score_func != "softmax": - route_weights = route_weights / mx.sum( - route_weights, axis=-1, keepdims=True - ) + route_weights = route_weights / mx.sum(route_weights, axis=-1, keepdims=True) return indices, route_weights * route_scale @@ -169,40 +165,43 @@ def _moe( x: mx.array, indices: mx.array, route_weights: mx.array, - routed_gate: _Projection, - routed_up: _Projection, + routed_gate: _Projection | None, + routed_up: _Projection | None, routed_down: _Projection, shared_gate: _Projection, shared_up: _Projection, shared_down: _Projection, *, + routed_gate_up: _Projection | None, + routed_combine: Callable | None, routed_limit: float, shared_limit: float, ) -> mx.array: """Stock unsorted Q2/Q4 arithmetic for the production top-6 tiny-M shape.""" gathered_x = mx.expand_dims(x, (-2, -3)) - up = _gather_qmm(gathered_x, indices, routed_up) - gate = _gather_qmm(gathered_x, indices, routed_gate) + if routed_gate_up is None: + up = _gather_qmm(gathered_x, indices, routed_up) # type: ignore[arg-type] + gate = _gather_qmm(gathered_x, indices, routed_gate) # type: ignore[arg-type] + else: + packed = _gather_qmm(gathered_x, indices, routed_gate_up) + gate, up = mx.split(packed, [routed_gate_up.output_dim // 2], axis=-1) if routed_limit > 0: up = mx.clip(up, -routed_limit, routed_limit) gate = mx.minimum(gate, routed_limit) routed = _gather_qmm(D.nn.silu(gate) * up, indices, routed_down) routed = routed.squeeze(-2) - routed = ( - routed * route_weights[..., None].astype(routed.dtype) - ).sum(axis=-2) + if routed_combine is None: + routed = (routed * route_weights[..., None].astype(routed.dtype)).sum(axis=-2) + else: + routed = routed_combine(routed, route_weights) shared_gate_out = _qmm(x, shared_gate) shared_up_out = _qmm(x, shared_up) if shared_limit > 0: - shared_up_out = mx.clip( - shared_up_out, -shared_limit, shared_limit - ) + shared_up_out = mx.clip(shared_up_out, -shared_limit, shared_limit) shared_gate_out = mx.minimum(shared_gate_out, shared_limit) - shared = _qmm( - D.nn.silu(shared_gate_out) * shared_up_out, shared_down - ) + shared = _qmm(D.nn.silu(shared_gate_out) * shared_up_out, shared_down) return routed + shared @@ -218,13 +217,15 @@ def _island_impl( norm_weight: mx.array, router_weight: mx.array, router_auxiliary: mx.array, - routed_gate: _Projection, - routed_up: _Projection, + routed_gate: _Projection | None, + routed_up: _Projection | None, routed_down: _Projection, shared_gate: _Projection, shared_up: _Projection, shared_down: _Projection, *, + routed_gate_up: _Projection | None, + routed_combine: Callable | None, hc: int, iters: int, hc_eps: float, @@ -283,6 +284,8 @@ def normalise(comb): shared_gate, shared_up, shared_down, + routed_gate_up=routed_gate_up, + routed_combine=routed_combine, routed_limit=routed_limit, shared_limit=shared_limit, ).reshape(shape) @@ -296,8 +299,10 @@ def _attention_island_tape( *, width: int, hash_router: bool, - routed_gate: _Projection, - routed_up: _Projection, + routed_gate: _Projection | None, + routed_up: _Projection | None, + routed_gate_up: _Projection | None, + routed_combine: Callable | None, routed_down: _Projection, shared_gate: _Projection, shared_up: _Projection, @@ -313,17 +318,29 @@ def _attention_island_tape( routed_limit: float, shared_limit: float, ) -> Callable: + paired = routed_gate_up is not None + if paired: + effective_gate = effective_up = effective_pair = routed_gate_up + else: + assert routed_gate is not None and routed_up is not None + effective_gate = routed_gate + effective_up = routed_up + # Preserve one compiled signature; this final leaf is unused when stock. + effective_pair = routed_gate projections = ( - routed_gate, - routed_up, + effective_gate, + effective_up, routed_down, shared_gate, shared_up, shared_down, + effective_pair, ) key = ( int(width), bool(hash_router), + paired, + routed_combine is not None, tuple((p.bits, p.group_size) for p in projections), int(hc), int(iters), @@ -374,6 +391,9 @@ def impl( sd_weight, sd_scales, sd_biases, + rgu_weight, + rgu_scales, + rgu_biases, ): arrays = ( (rg_weight, rg_scales, rg_biases), @@ -382,11 +402,13 @@ def impl( (sg_weight, sg_scales, sg_biases), (su_weight, su_scales, su_biases), (sd_weight, sd_scales, sd_biases), + (rgu_weight, rgu_scales, rgu_biases), ) bound = tuple( _Projection(*leaves, *spec) for leaves, spec in zip(arrays, specs, strict=True) ) + brg, bru, brd, bsg, bsu, bsd, brgu = bound return _island_impl( attn_out, attn_residual, @@ -399,7 +421,14 @@ def impl( norm_weight, router_weight, router_auxiliary, - *bound, + None if paired else brg, + None if paired else bru, + brd, + bsg, + bsu, + bsd, + routed_gate_up=brgu if paired else None, + routed_combine=routed_combine, hc=hc, iters=iters, hc_eps=hc_eps, @@ -438,54 +467,103 @@ def __call__(self, attn_out, attn_residual, attn_post, attn_comb, input_ids): def _bind_attention_island_layer( - layer: D.DeepseekV4DecoderLayer, *, width: int + layer: D.DeepseekV4DecoderLayer, + *, + width: int, + allowed_widths: tuple[int, ...] = _WIDTHS, + shared_bits: int = 4, + routed_pair: bool = False, + routed_combine: Callable | None = None, + routed_switch: Any | None = None, ) -> _BoundAttentionIslandLayer: """Validate and bind one layer; its hot call performs no discovery.""" - if int(width) not in _WIDTHS: + if int(width) not in allowed_widths: raise AttentionIslandError(f"unsupported verifier width {width}") + if int(shared_bits) not in {4, 8}: + raise AttentionIslandError(f"unsupported shared-expert bits {shared_bits}") if type(layer) is not D.DeepseekV4DecoderLayer: raise AttentionIslandError("requires an exact DeepseekV4DecoderLayer") ffn = layer.ffn if type(ffn) is not D.DeepseekV4MoE: raise AttentionIslandError("requires the stock DeepSeek-V4 MoE topology") - switch = ffn.switch_mlp + switch = ffn.switch_mlp if routed_switch is None else routed_switch shared = ffn.shared_experts + if routed_pair: + if not isinstance( + switch, + (DeepseekV40731PackedQ2SwitchGLU, PackedSwitchGLU), + ): + raise AttentionIslandError("requires the packed routed Q2 topology") + elif isinstance(switch, PackedSwitchGLU): + raise AttentionIslandError("requires the stock routed Q2 topology") if type(switch.activation) is not D.ClampedSwiGLU: raise AttentionIslandError("requires the exact clamped SwiGLU activation") - projections = tuple( - _projection_contract(module, label, expected_bits=bits) - for module, label, bits in ( - (switch.gate_proj, "routed gate", 2), - (switch.up_proj, "routed up", 2), - (switch.down_proj, "routed down", 2), - (shared.gate_proj, "shared gate", 4), - (shared.up_proj, "shared up", 4), - (shared.down_proj, "shared down", 4), + if routed_pair: + rgu = _projection_contract( + switch.gate_up_proj, + "packed routed gate/up", + expected_bits=2, + ) + rd = _projection_contract(switch.down_proj, "routed down", expected_bits=2) + if int(switch._split_at) * 2 != rgu.output_dim: + raise AttentionIslandError("packed routed gate/up split is inconsistent") + rg = ru = None + input_dim = rgu.input_dim + intermediate = int(switch._split_at) + effective_routed = (rgu, rgu, rd, rgu) + else: + rg, ru, rd = tuple( + _projection_contract(module, label, expected_bits=2) + for module, label in ( + (switch.gate_proj, "routed gate"), + (switch.up_proj, "routed up"), + (switch.down_proj, "routed down"), + ) + ) + rgu = None + input_dim = rg.input_dim + intermediate = rg.output_dim + effective_routed = (rg, ru, rd, rg) + sg, su, sd = tuple( + _projection_contract(module, label, expected_bits=int(shared_bits)) + for module, label in ( + (shared.gate_proj, "shared gate"), + (shared.up_proj, "shared up"), + (shared.down_proj, "shared down"), ) ) - rg, ru, rd, sg, su, sd = projections if ( - rg.input_dim != ru.input_dim - or rg.output_dim != ru.output_dim - or rd.input_dim != rg.output_dim - or rd.output_dim != rg.input_dim - or sg.input_dim != rg.input_dim + ( + not routed_pair + and (rg.input_dim != ru.input_dim or rg.output_dim != ru.output_dim) + ) + or rd.input_dim != intermediate + or rd.output_dim != input_dim + or sg.input_dim != input_dim or sg.output_dim != su.output_dim - or su.input_dim != rg.input_dim + or su.input_dim != input_dim or sd.input_dim != sg.output_dim - or sd.output_dim != rg.input_dim + or sd.output_dim != input_dim ): raise AttentionIslandError("routed/shared projection geometry is inconsistent") hc = layer.ffn_hc fn_t, base, scale_vec = hc._static() router = ffn.gate + if routed_combine is not None and ( + int(width) != 1 or int(router.topk) != 6 or rd.output_dim != 4096 + ): + raise AttentionIslandError( + "row-owned combine requires the exact AR M1/top-6/hidden-4096 route" + ) auxiliary = router.tid2eid if router.hash else router.e_score_correction_bias tape = _attention_island_tape( width=width, hash_router=bool(router.hash), routed_gate=rg, routed_up=ru, + routed_gate_up=rgu, + routed_combine=routed_combine, routed_down=rd, shared_gate=sg, shared_up=su, @@ -508,7 +586,11 @@ def _bind_attention_island_layer( layer.ffn_norm.weight, router.weight, auxiliary, - *(leaf for p in projections for leaf in (p.weight, p.scales, p.biases)), + *( + leaf + for p in effective_routed[:3] + (sg, su, sd, effective_routed[3]) + for leaf in (p.weight, p.scales, p.biases) + ), ) return _BoundAttentionIslandLayer(tape, leaves, width) @@ -583,7 +665,10 @@ def select_deepseek_v4_attention_island_arm(model: Any, enabled: bool) -> None: """Select a preinstalled bracket arm outside measured generation.""" selector = getattr(model, "_mtplx_dsv4_attention_island_selector", None) - if type(selector) is not _AttentionIslandArmSelector or selector._model is not model: + if ( + type(selector) is not _AttentionIslandArmSelector + or selector._model is not model + ): raise AttentionIslandError("attention-island arm selector is not installed") selector.select(enabled) diff --git a/mtplx/deepseek_v4_dspark_generation.py b/mtplx/deepseek_v4_dspark_generation.py new file mode 100644 index 00000000..1fd0aed2 --- /dev/null +++ b/mtplx/deepseek_v4_dspark_generation.py @@ -0,0 +1,176 @@ +"""DeepSeek-V4-0731 adapter for MTPLX's generic block speculation engine.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import mlx.core as mx + + +@dataclass(frozen=True) +class DeepseekV4DSparkBackend: + """Construction-selected DSpark operations; target policy stays generic.""" + + dspark: Any + embed_tokens: Any + lm_head: Any + target_forward: Callable[..., Any] + backend_id: str = "deepseek_v4_dspark_0731" + supported_depths: tuple[int, ...] = (2,) + prefill_chunk_size: int = 128 + # Absolute zero selects DSpark's attention-only prefill branch. A + # one-token prompt therefore needs one explicit target seed before the + # primary-inclusive fixed-K loop can start from main position one. + minimum_proposal_target_position: int = 2 + + @classmethod + def bind(cls, model: Any) -> "DeepseekV4DSparkBackend": + """Validate ownership once and bind direct hot-path callables.""" + dspark = getattr(model, "_dspark", None) + inner = getattr(model, "model", None) + embed_tokens = getattr(inner, "embed_tokens", None) + lm_head = getattr(model, "lm_head", None) + if ( + dspark is None + or not callable(embed_tokens) + or not callable(lm_head) + or not callable(model) + ): + raise ValueError("DeepSeek-V4 DSpark backend cannot bind model ownership") + if len(tuple(getattr(dspark, "stages", ()))) != 3: + raise ValueError("DeepSeek-V4 DSpark backend requires three owned stages") + return cls( + dspark=dspark, + embed_tokens=embed_tokens, + lm_head=lm_head, + target_forward=model, + ) + + def make_cache(self, rt: Any) -> Any: + del rt + return self.dspark.make_cache() + + def prefill(self, rt: Any, hidden: mx.array, cache: Any) -> None: + del rt + self.dspark.prefill(hidden, cache) + + def prefill_chunk( + self, + rt: Any, + hidden: mx.array, + cache: Any, + *, + start_pos: int, + ) -> None: + """Install one construction-fixed 0731 sliding-window prompt chunk.""" + del rt + start_pos = int(start_pos) + if start_pos == 0: + self.dspark.prefill(hidden, cache) + else: + self.dspark.commit_main(hidden, cache, start_pos=start_pos) + prefill_length = start_pos + int(hidden.shape[1]) + for entry in cache: + entry.prefill_length = prefill_length + + def propose( + self, + rt: Any, + hidden: mx.array, + token_id: int, + primary_token_id: int, + cache: Any, + *, + start_pos: int, + width: int, + ) -> mx.array: + del rt + # A logical MTP-N request drafts only N future tokens. The target has + # already sampled the primary token from its carried logits, so that + # authoritative id seeds DSpark's sequential Markov recurrence. The + # neural block still evaluates primary + N rows in parallel; the ids-only + # API returns [last_target, primary, future...], and only future rows + # leave this adapter. + # + # The generic engine names ``start_pos`` as the next target position. + # DSpark names it as the position of ``hidden``/``token_id`` (upstream + # passes ``checkpoint.len - 1`` after committing that row), so translate + # ownership once at this adapter boundary. Target verification and + # commit continue to use the unshifted next-target position. + main_position = int(start_pos) - 1 + draft_ids = self.dspark.forward( + hidden, + mx.array([int(token_id)], dtype=mx.int32), + self.embed_tokens, + self.lm_head, + cache, + start_pos=main_position, + greedy=True, + ids_only_width=int(width) + 1, + forced_first_token_ids=mx.array([int(primary_token_id)], dtype=mx.int32), + ) + return draft_ids[:, 2:] + + def commit( + self, + rt: Any, + hidden: mx.array, + cache: Any, + *, + start_pos: int, + ) -> None: + del rt + self.dspark.commit_main(hidden, cache, start_pos=int(start_pos)) + + def cache_roots(self, cache: Any) -> list[mx.array]: + return [ + ring + for entry in cache + if (ring := getattr(entry, "ring", None)) is not None + ] + + def bind_target_forward(self, rt: Any) -> Callable[..., Any]: + """Bind the construction-certified target route before generation.""" + del rt + return self.target_forward + + def snapshot(self, cache: Any) -> tuple[tuple[mx.array | None, int], ...]: + """Capture DSpark-owned stage rings before a speculative proposal.""" + return tuple( + ( + None if entry.ring is None else entry.ring[...], + int(entry.prefill_length), + ) + for entry in cache + ) + + def restore( + self, + cache: Any, + snapshot: tuple[tuple[mx.array | None, int], ...], + ) -> None: + """Restore only the proposal backend's ring ownership.""" + for entry, (ring, prefill_length) in zip(cache, snapshot): + entry.ring = ring + entry.prefill_length = int(prefill_length) + + def rollback_target(self, target_cache: Any, rejected_rows: int) -> None: + # Eligibility is proven when the DeepSeek-V4 backend is installed: all + # 43 target cache entries are exact trimmable DeepseekV4Cache instances. + # The enabled path executes that route directly without probing/fallback. + for entry in target_cache: + entry.trim(int(rejected_rows)) + + +def generate_deepseek_v4_dspark(rt: Any, prompt_ids: list[int], **kwargs: Any): + """Compatibility wrapper for callers that used the old dedicated runner.""" + from .native_block_speculation import generate_native_block_speculative + + return generate_native_block_speculative( + rt, + DeepseekV4DSparkBackend.bind(rt.model), + prompt_ids, + **kwargs, + ) diff --git a/mtplx/generation.py b/mtplx/generation.py index 7313d719..ac1e283a 100644 --- a/mtplx/generation.py +++ b/mtplx/generation.py @@ -156,7 +156,9 @@ def _resolve_runtime_base_hidden_variant( requested: str | None, ) -> str: if requested in {None, "auto", "contract"}: - return str(getattr(rt.contract, "base_hidden_variant", "post_norm") or "post_norm") + return str( + getattr(rt.contract, "base_hidden_variant", "post_norm") or "post_norm" + ) return str(requested) @@ -1981,8 +1983,7 @@ def _detect_repeated_token_suffix( repeats = 1 cursor = token_count - block_tokens while ( - cursor >= block_tokens - and tokens[cursor - block_tokens : cursor] == block + cursor >= block_tokens and tokens[cursor - block_tokens : cursor] == block ): repeats += 1 cursor -= block_tokens @@ -2116,9 +2117,7 @@ def emit_chunk(chunk_len: int, chunk_elapsed: float, started: float) -> None: else None ) cumulative_tok_s = ( - float(new_done) / elapsed - if elapsed > 0.0 and new_done > 0 - else None + float(new_done) / elapsed if elapsed > 0.0 and new_done > 0 else None ) chunk_callback( { @@ -2213,9 +2212,7 @@ def append_history( [int(token) for token in suffix[1:]], window_start=1, ) - target_forward_time += _maybe_repage_target_prefill_cache( - rt, restored.cache - ) + target_forward_time += _maybe_repage_target_prefill_cache(rt, restored.cache) _check_splice_consumed() return ( suffix_logits[:, -1, :], @@ -2224,9 +2221,8 @@ def append_history( mtp_history_time, ) - capture_boundaries = ( - gdn_boundary_sink is not None - and _cache_has_recurrent_entries(restored.cache) + capture_boundaries = gdn_boundary_sink is not None and _cache_has_recurrent_entries( + restored.cache ) if len(suffix) > 1: body = suffix[:-1] @@ -2293,9 +2289,7 @@ def append_history( cached_tokens + end, restored.cache, hidden_last=( - hidden_chunk[:, -1:, :] - if hidden_chunk is not None - else None + hidden_chunk[:, -1:, :] if hidden_chunk is not None else None ), ) if hidden_chunk is not None: @@ -2686,8 +2680,7 @@ def _near_debug(reason: str) -> None: cache_restore_time_s = time.perf_counter() - restore_started if prefix_restore is None: _near_debug( - "restore_failed:" - + str(getattr(session_bank, "last_miss_reason", None)) + "restore_failed:" + str(getattr(session_bank, "last_miss_reason", None)) ) continue _near_debug("served") @@ -2728,22 +2721,22 @@ def _near_debug(reason: str) -> None: ) if committed_history_required and mtp_history_cache is None: continue - if ( - boundary_restore - and committed_history_required - and boundary_hidden is None - ): + if boundary_restore and committed_history_required and boundary_hidden is None: # Without the boundary's hidden state the committed MTP history # cannot resume exactly at b; running a seed forward instead would # advance the recurrent state twice. Fail closed to the next # candidate (or cold). continue cache_source = str(getattr(entry, "cache_source", "ram") or "ram") - ssd_cache_hit = bool(getattr(entry, "ssd_cache_hit", False)) or cache_source == "ssd" + ssd_cache_hit = ( + bool(getattr(entry, "ssd_cache_hit", False)) or cache_source == "ssd" + ) ssd_restore_s = float(getattr(entry, "ssd_restore_s", 0.0) or 0.0) ssd_cached_tokens = restore_point if ssd_cache_hit else 0 total_cache_restore_time_s = ( - cache_restore_time_s + ssd_restore_s if ssd_cache_hit else cache_restore_time_s + cache_restore_time_s + ssd_restore_s + if ssd_cache_hit + else cache_restore_time_s ) _check_postcommit_abort(abort_check) @@ -2773,7 +2766,9 @@ def _near_debug(reason: str) -> None: repair_time = time.perf_counter() - started _check_postcommit_abort(abort_check) restore_kind_base = ( - "block_prefix" if int(entry.prefix_len) - matched > max_gap else "near_prefix" + "block_prefix" + if int(entry.prefix_len) - matched > max_gap + else "near_prefix" ) if restore_point < matched: restore_kind_base = f"{restore_kind_base}_boundary" @@ -2847,9 +2842,7 @@ def _near_debug(reason: str) -> None: restore_served=served_truth, ) suffix_boundary_sink: list[tuple[int, Any, Any]] | None = ( - list(inherited_boundaries) - if _gdn_boundary_capture_enabled() - else None + list(inherited_boundaries) if _gdn_boundary_capture_enabled() else None ) suffix_logits, suffix_hidden, suffix_time, mtp_history_time = ( _prefill_restored_prompt_suffix( @@ -2974,7 +2967,10 @@ def _thin_gdn_boundary_records( newest = ordered[-1] oldest = ordered[0] newest_pos = int(newest[0]) - kept: dict[int, tuple[int, Any, Any]] = {int(newest[0]): newest, int(oldest[0]): oldest} + kept: dict[int, tuple[int, Any, Any]] = { + int(newest[0]): newest, + int(oldest[0]): oldest, + } # Walk from the tail toward the head (distance from newest increasing). # Keep the first record past each doubling floor — one keeper per # distance scale, geometric spacing by construction regardless of how @@ -3029,9 +3025,7 @@ def _capture_gdn_boundary( hidden_leaf = None if hidden_last is not None: hidden_leaf = detach_array_leaf(hidden_last, mode="contiguous_eval") - sink.append( - (int(tokens_done), snapshot_untrimmable_cache(cache), hidden_leaf) - ) + sink.append((int(tokens_done), snapshot_untrimmable_cache(cache), hidden_leaf)) cap = _gdn_boundary_max_count() if len(sink) > cap: sink[:] = _thin_gdn_boundary_records(sink, cap) @@ -3143,7 +3137,9 @@ def _store_on_prefill_min_suffix() -> int: return 1024 -def _debug_prefix_divergence(rt: MTPLXRuntime, prompt_ids: list[int], session_bank: Any) -> None: +def _debug_prefix_divergence( + rt: MTPLXRuntime, prompt_ids: list[int], session_bank: Any +) -> None: """Env-gated diagnostic: report where the prompt diverges from each bank entry. For every bank entry that shares a non-trivial prefix with the incoming @@ -3184,9 +3180,9 @@ def _decode(ids: list[int]) -> str: f"[mtplx] prefix-diverge: entry_len={entry_len} matched={matched} " f"prompt_len={len(prompt)}\n" f" entry [{lo}:{matched + 40}]: " - f"{_decode(toks[lo:matched + 40])!r}\n" + f"{_decode(toks[lo : matched + 40])!r}\n" f" prompt[{lo}:{matched + 40}]: " - f"{_decode(prompt[lo:matched + 40])!r}", + f"{_decode(prompt[lo : matched + 40])!r}", file=sys.stderr, ) except Exception as exc: # diagnostic only - never break the request @@ -3341,7 +3337,9 @@ def _maybe_store_prefix_snapshot(state: PromptState) -> None: put_timing: dict[str, object] = {} entry = session_bank.put( runtime=rt, - token_ids=list(bank_key_ids if bank_key_ids is not None else prompt_ids), + token_ids=list( + bank_key_ids if bank_key_ids is not None else prompt_ids + ), cache=state.trunk_cache, logits=state.logits, hidden=state.hidden, @@ -3354,7 +3352,9 @@ def _maybe_store_prefix_snapshot(state: PromptState) -> None: policy_fingerprint=policy_fingerprint, mtp_history_snapshot=mtp_snapshot, snapshot_epoch=len(prompt_ids), - mtp_snapshot_epoch=len(prompt_ids) if mtp_snapshot is not None else None, + mtp_snapshot_epoch=len(prompt_ids) + if mtp_snapshot is not None + else None, gdn_boundaries=list(getattr(state, "gdn_boundaries", None) or []), timing_out=put_timing, ) @@ -3402,7 +3402,9 @@ def _emit_prefill_complete(state: PromptState) -> PromptState: "new_prefill_tokens": new_tokens, "elapsed_s": elapsed, "prompt_eval_time_s": compute_elapsed, - "prefill_tok_s": compute_tok_s if compute_tok_s is not None else wall_tok_s, + "prefill_tok_s": compute_tok_s + if compute_tok_s is not None + else wall_tok_s, "prefill_compute_tok_s": compute_tok_s, "prefill_wall_tok_s": wall_tok_s, "cache_hit": bool(state.cache_hit), @@ -3438,9 +3440,7 @@ def _emit_prefill_complete(state: PromptState) -> PromptState: if callable(longest_prefix): exact_entry = longest_prefix(bank_match_ids) if exact_entry is not None: - exact_prefix_len = int( - getattr(exact_entry, "prefix_len", 0) or 0 - ) + exact_prefix_len = int(getattr(exact_entry, "prefix_len", 0) or 0) except Exception: exact_prefix_len = 0 @@ -3498,9 +3498,7 @@ def _emit_prefill_complete(state: PromptState) -> PromptState: chunk_callback=prefill_callback, chunk_started_s=prefill_started_s, matched_ceiling=( - vision_restore_spans[0][0] - if vision_restore_spans - else None + vision_restore_spans[0][0] if vision_restore_spans else None ), cache_factory=restore_cache_factory, ) @@ -3557,30 +3555,34 @@ def _emit_prefill_complete(state: PromptState) -> PromptState: flush=True, ) if not suffix: - repage_time = _maybe_repage_target_prefill_cache( - rt, restored.cache + repage_time = _maybe_repage_target_prefill_cache(rt, restored.cache) + return _emit_prefill_complete( + PromptState( + trunk_cache=restored.cache, + logits=restored.logits, + hidden=restored.hidden, + committed_mtp_cache=restored.mtp_history_cache, + token_prefix=tuple(int(token) for token in prompt_ids), + prompt_eval_time_s=repage_time, + cache_restore_time_s=restore_elapsed_s, + mtp_history_policy=mtp_history_policy, + mtp_history_window_tokens=mtp_history_window_tokens, + cached_tokens=restored.entry.prefix_len, + suffix_tokens=0, + cache_hit=True, + cache_source=getattr(restored, "cache_source", "ram"), + ssd_cache_hit=bool(getattr(restored, "ssd_cache_hit", False)), + ssd_cached_tokens=int( + getattr(restored, "ssd_cached_tokens", 0) or 0 + ), + ssd_restore_s=float( + getattr(restored, "ssd_restore_s", 0.0) or 0.0 + ), + restore_mode=restored.restore_mode, + gdn_boundaries=inherited_boundaries, + restore_served=exact_served, + ) ) - return _emit_prefill_complete(PromptState( - trunk_cache=restored.cache, - logits=restored.logits, - hidden=restored.hidden, - committed_mtp_cache=restored.mtp_history_cache, - token_prefix=tuple(int(token) for token in prompt_ids), - prompt_eval_time_s=repage_time, - cache_restore_time_s=restore_elapsed_s, - mtp_history_policy=mtp_history_policy, - mtp_history_window_tokens=mtp_history_window_tokens, - cached_tokens=restored.entry.prefix_len, - suffix_tokens=0, - cache_hit=True, - cache_source=getattr(restored, "cache_source", "ram"), - ssd_cache_hit=bool(getattr(restored, "ssd_cache_hit", False)), - ssd_cached_tokens=int(getattr(restored, "ssd_cached_tokens", 0) or 0), - ssd_restore_s=float(getattr(restored, "ssd_restore_s", 0.0) or 0.0), - restore_mode=restored.restore_mode, - gdn_boundaries=inherited_boundaries, - restore_served=exact_served, - )) _check_postcommit_abort(abort_check) _emit_prefill_restore_progress( @@ -3591,9 +3593,7 @@ def _emit_prefill_complete(state: PromptState) -> PromptState: started_s=prefill_started_s, cache_source=getattr(restored, "cache_source", "ram"), ssd_cache_hit=bool(getattr(restored, "ssd_cache_hit", False)), - ssd_cached_tokens=int( - getattr(restored, "ssd_cached_tokens", 0) or 0 - ), + ssd_cached_tokens=int(getattr(restored, "ssd_cached_tokens", 0) or 0), ssd_restore_s=float(getattr(restored, "ssd_restore_s", 0.0) or 0.0), ssd_suffix_tokens=len(suffix), ) @@ -3632,32 +3632,36 @@ def _emit_prefill_complete(state: PromptState) -> PromptState: stable_prefix_len=stable_prefix_len, ) ) - return _emit_prefill_complete(PromptState( - trunk_cache=restored.cache, - logits=suffix_logits, - hidden=suffix_hidden, - committed_mtp_cache=restored.mtp_history_cache, - token_prefix=tuple(int(token) for token in prompt_ids), - prompt_eval_time_s=suffix_time + mtp_history_time, - prompt_mtp_history_time_s=mtp_history_time, - cache_restore_time_s=restore_elapsed_s, - mtp_history_policy=mtp_history_policy, - mtp_history_window_tokens=mtp_history_window_tokens, - cached_tokens=restored.entry.prefix_len, - suffix_tokens=len(suffix), - cache_hit=True, - cache_source=getattr(restored, "cache_source", "ram"), - ssd_cache_hit=bool(getattr(restored, "ssd_cache_hit", False)), - ssd_cached_tokens=int(getattr(restored, "ssd_cached_tokens", 0) or 0), - ssd_restore_s=float(getattr(restored, "ssd_restore_s", 0.0) or 0.0), - restore_mode=restored.restore_mode, - gdn_boundaries=( - suffix_boundary_sink - if suffix_boundary_sink is not None - else inherited_boundaries - ), - restore_served=exact_served, - )) + return _emit_prefill_complete( + PromptState( + trunk_cache=restored.cache, + logits=suffix_logits, + hidden=suffix_hidden, + committed_mtp_cache=restored.mtp_history_cache, + token_prefix=tuple(int(token) for token in prompt_ids), + prompt_eval_time_s=suffix_time + mtp_history_time, + prompt_mtp_history_time_s=mtp_history_time, + cache_restore_time_s=restore_elapsed_s, + mtp_history_policy=mtp_history_policy, + mtp_history_window_tokens=mtp_history_window_tokens, + cached_tokens=restored.entry.prefix_len, + suffix_tokens=len(suffix), + cache_hit=True, + cache_source=getattr(restored, "cache_source", "ram"), + ssd_cache_hit=bool(getattr(restored, "ssd_cache_hit", False)), + ssd_cached_tokens=int( + getattr(restored, "ssd_cached_tokens", 0) or 0 + ), + ssd_restore_s=float(getattr(restored, "ssd_restore_s", 0.0) or 0.0), + restore_mode=restored.restore_mode, + gdn_boundaries=( + suffix_boundary_sink + if suffix_boundary_sink is not None + else inherited_boundaries + ), + restore_served=exact_served, + ) + ) near_prompt_state = _restore_near_prefix_prompt_state( rt, @@ -3747,7 +3751,9 @@ def _emit_prefill_complete(state: PromptState) -> PromptState: keep = min(len(history_token_ids), mtp_history_window_tokens) dropped = len(history_token_ids) - keep mtp_history_position_base = ( - dropped + 1 if mtp_position_mode == "absolute" else max(0, dropped) + dropped + 1 + if mtp_position_mode == "absolute" + else max(0, dropped) ) history_token_ids = history_token_ids[-keep:] history_hidden = history_hidden[:, -keep:, :] @@ -3756,7 +3762,9 @@ def _emit_prefill_complete(state: PromptState) -> PromptState: if vision_splice is not None: pad_id = vision_splice.image_pad_token_id rows_before = sum( - 1 for token in prompt_ids[:history_window_start] if token == pad_id + 1 + for token in prompt_ids[:history_window_start] + if token == pad_id ) if any(token == pad_id for token in history_token_ids): from mtplx.vision.splice import spliced_embeddings_for_window @@ -3776,7 +3784,8 @@ def _emit_prefill_complete(state: PromptState) -> PromptState: mtp_hidden_variant=mtp_hidden_variant, position_offset=( mtp_history_position_base - if mtp_position_mode == "absolute" or mtp_history_policy == "last_window" + if mtp_position_mode == "absolute" + or mtp_history_policy == "last_window" else None ), input_embeddings=history_embeddings, @@ -3804,23 +3813,25 @@ def _emit_prefill_complete(state: PromptState) -> PromptState: stable_prefix_len=stable_prefix_len, ) prompt_eval_time = target_time - return _emit_prefill_complete(PromptState( - trunk_cache=cache, - logits=logits, - hidden=hidden, - committed_mtp_cache=mtp_history_cache, - token_prefix=tuple(int(token) for token in prompt_ids), - prompt_eval_time_s=prompt_eval_time, - prompt_mtp_history_time_s=prompt_history_time, - mtp_history_policy=mtp_history_policy, - mtp_history_window_tokens=mtp_history_window_tokens, - mtp_history_position_base=mtp_history_position_base, - suffix_tokens=len(prompt_ids), - cache_miss_reason=getattr(session_bank, "last_miss_reason", None) - if session_bank is not None - else None, - gdn_boundaries=list(gdn_boundary_sink or []), - )) + return _emit_prefill_complete( + PromptState( + trunk_cache=cache, + logits=logits, + hidden=hidden, + committed_mtp_cache=mtp_history_cache, + token_prefix=tuple(int(token) for token in prompt_ids), + prompt_eval_time_s=prompt_eval_time, + prompt_mtp_history_time_s=prompt_history_time, + mtp_history_policy=mtp_history_policy, + mtp_history_window_tokens=mtp_history_window_tokens, + mtp_history_position_base=mtp_history_position_base, + suffix_tokens=len(prompt_ids), + cache_miss_reason=getattr(session_bank, "last_miss_reason", None) + if session_bank is not None + else None, + gdn_boundaries=list(gdn_boundary_sink or []), + ) + ) def _decode(tokenizer, tokens: list[int]) -> str: @@ -3954,11 +3965,9 @@ def _sample_from_logits( def _greedy_draft_token_and_top2(logits: mx.array) -> tuple[int, float, float]: """Materialize one greedy token and its FP32 top-two values together.""" - row = ( - logits[:, -1, :][0] - if logits.ndim == 3 - else logits.reshape(-1) - ).astype(mx.float32) + row = (logits[:, -1, :][0] if logits.ndim == 3 else logits.reshape(-1)).astype( + mx.float32 + ) token_id = mx.argmax(row, axis=-1) top2_values = mx.topk(row, k=2) _eval(token_id, top2_values) @@ -4424,9 +4433,9 @@ def chain_fn(hidden_states, first_token_ids, level_keys): ) cdf = mx.cumsum(q_norm, axis=-1) u = mx.random.uniform(key=level_keys[level - 1]) - pick = mx.minimum( - (cdf <= u).sum(), int(top_idx.shape[0]) - 1 - ).astype(mx.int32) + pick = mx.minimum((cdf <= u).sum(), int(top_idx.shape[0]) - 1).astype( + mx.int32 + ) next_tok = top_idx[pick].reshape(1, 1) q_ids.append(top_idx) q_probs.append(q_norm) @@ -4528,8 +4537,8 @@ def _prefill( cache = _make_target_prefill_cache(rt) target_forward_time = 0.0 final_logits_only = _final_logits_prefill_enabled() - capture_boundaries = ( - gdn_boundary_sink is not None and _cache_has_recurrent_entries(cache) + capture_boundaries = gdn_boundary_sink is not None and _cache_has_recurrent_entries( + cache ) if len(prompt_ids) > 1: @@ -4633,8 +4642,8 @@ def _prefill_committed_mtp_history_streaming( target_forward_time = 0.0 prompt_history_time = 0.0 final_logits_only = _final_logits_prefill_enabled() - capture_boundaries = ( - gdn_boundary_sink is not None and _cache_has_recurrent_entries(cache) + capture_boundaries = gdn_boundary_sink is not None and _cache_has_recurrent_entries( + cache ) body = prompt_ids[:-1] history_start_token_index = 1 @@ -4644,7 +4653,9 @@ def _prefill_committed_mtp_history_streaming( window = max(1, int(history_window_tokens)) history_start_token_index = max(1, len(prompt_ids) - window) mtp_history_position_base = ( - history_start_token_index if use_absolute_positions else max(0, history_start_token_index - 1) + history_start_token_index + if use_absolute_positions + else max(0, history_start_token_index - 1) ) cursor = 0 @@ -4677,9 +4688,7 @@ def _prefill_committed_mtp_history_streaming( chunk_size=prefill_chunk_size, ) if capture_boundaries - else _iter_prefill_chunk_spans( - len(body), chunk_size=prefill_chunk_size - ) + else _iter_prefill_chunk_spans(len(body), chunk_size=prefill_chunk_size) ) for start, end in mtp_streaming_spans: _check_postcommit_abort(abort_check) @@ -4727,14 +4736,14 @@ def _prefill_committed_mtp_history_streaming( if chunk_callback is not None: try: now = time.perf_counter() - phase_start = chunk_started_s if chunk_started_s is not None else started + phase_start = ( + chunk_started_s if chunk_started_s is not None else started + ) chunk_elapsed = max(0.0, now - started) elapsed = max(0.0, now - phase_start) tokens_done = int(cursor + chunk_len) chunk_tok_s = ( - float(chunk_len) / chunk_elapsed - if chunk_elapsed > 0.0 - else None + float(chunk_len) / chunk_elapsed if chunk_elapsed > 0.0 else None ) cumulative_tok_s = ( float(tokens_done) / elapsed @@ -4777,10 +4786,7 @@ def _prefill_committed_mtp_history_streaming( if vision_splice is not None and pad_prefix_counts is not None: window_start = token_start_index + slice_start window_end = window_start + len(sliced_token_ids) - if ( - pad_prefix_counts[window_end] - > pad_prefix_counts[window_start] - ): + if pad_prefix_counts[window_end] > pad_prefix_counts[window_start]: from mtplx.vision.splice import ( spliced_embeddings_for_window, ) @@ -4810,9 +4816,7 @@ def _prefill_committed_mtp_history_streaming( ) _check_postcommit_abort(abort_check) cursor += chunk_len - boundary_hidden = ( - hidden_chunk[:, -1:, :] if hidden_chunk is not None else None - ) + boundary_hidden = hidden_chunk[:, -1:, :] if hidden_chunk is not None else None del hidden_chunk del logits_chunk target_forward_time += _prefill_chunk_cache_cleanup(rt) @@ -5126,11 +5130,7 @@ def generate_ar( if prefill_callback is not None: try: elapsed = max(0.0, time.perf_counter() - prefill_started_s) - tok_s = ( - (len(prompt_ids) / elapsed) - if elapsed > 0 and prompt_ids - else None - ) + tok_s = (len(prompt_ids) / elapsed) if elapsed > 0 and prompt_ids else None prefill_callback( { "phase": "completed", @@ -5299,9 +5299,8 @@ def emit_token(token: int) -> None: }, } ) - _steer_active = ( - (_loop_guard is not None and _loop_guard.armed) - or (_thinking_guard is not None and _thinking_guard.steering_active) + _steer_active = (_loop_guard is not None and _loop_guard.armed) or ( + _thinking_guard is not None and _thinking_guard.steering_active ) logits_row = logits[0] if constraint is not None: @@ -6198,7 +6197,9 @@ def generate_mtpk( if getattr(rt, "backend_id", None) == "gemma4_assistant": from .backends.gemma4_assistant import generate_gemma4_assistant - runtime_block_size = int(getattr(getattr(rt, "config", None), "draft_block_size", 0) or 0) + runtime_block_size = int( + getattr(getattr(rt, "config", None), "draft_block_size", 0) or 0 + ) requested_block_size = int(speculative_depth or 0) effective_block_size = ( runtime_block_size @@ -6227,6 +6228,129 @@ def generate_mtpk( repetition_stop=repetition_stop, requested_speculative_depth=requested_block_size, ) + block_speculative_backend = getattr(rt, "block_speculative_backend", None) + if block_speculative_backend is not None: + session_state_requested = bool( + session_bank is not None + or capture_final_state + or commit_prompt_state_to_bank + or commit_prompt_state_keep_live_ref + ) + unsupported = [ + label + for selected, label in ( + (base_hidden_variant is not None, "base hidden variant"), + (mtp_hidden_variant is not None, "MTP hidden variant"), + (mtp_cache_policy != "persistent", "mtp_cache_policy"), + (mtp_history_policy != "cycle", "mtp_history_policy"), + (draft_margin_threshold is not None, "draft margin threshold"), + (min_speculative_depth != 1, "minimum speculative depth"), + (verify_strategy != "batched", "verify strategy"), + (verify_core != "stock", "verify core"), + (draft_core != "stock", "draft core"), + (mtp_corrector is not None, "MTP corrector"), + (adaptive_policy is not None, "adaptive policy"), + ( + online_hidden_corrector_alpha != 0.0, + "online hidden corrector", + ), + (online_hidden_corrector_decay != 0.8, "online corrector decay"), + (online_hidden_corrector_warmup != 1, "online corrector warmup"), + ( + online_hidden_corrector_max_feed_depth is not None, + "online corrector feed depth", + ), + (online_hidden_corrector_key != "global", "online corrector key"), + (online_correction_cache, "online correction cache"), + ( + online_correction_cache_min_depth != 1, + "online correction cache depth", + ), + ( + online_correction_cache_key != "local_prefix", + "online correction cache key", + ), + (prompt_correction_cache, "prompt correction cache"), + ( + prompt_correction_cache_min_depth != 2, + "prompt correction cache depth", + ), + (adapter_ensemble_q, "adapter ensemble"), + (adapter_ensemble_epsilon != 0.5, "adapter ensemble epsilon"), + (adapter_ensemble_min_depth != 2, "adapter ensemble depth"), + (mtp_topk_reranker is not None, "MTP top-k reranker"), + (session_bank is not None, "session bank"), + (session_state_requested and session_id is not None, "session id"), + ( + session_state_requested and session_restore_mode != "clone", + "session restore mode", + ), + ( + session_state_requested and session_template_hash is not None, + "session template hash", + ), + ( + session_state_requested and session_draft_head_identity is not None, + "session draft-head identity", + ), + ( + session_state_requested and session_policy_fingerprint is not None, + "session policy fingerprint", + ), + (capture_final_state, "final state capture"), + (commit_prompt_state_to_bank, "prompt state bank commit"), + ( + commit_prompt_state_keep_live_ref, + "prompt state live reference", + ), + (trace_label is not None, "decode trace label"), + (trace_metadata is not None, "decode trace metadata"), + ( + bool( + getattr( + rt, + "block_speculative_decode_trace_requested", + False, + ) + ), + "decode trace output", + ), + (repetition_stop, "repetition stop"), + (loop_guard, "loop guard"), + (thinking_guard is not None, "thinking guard"), + (vision_splice is not None, "vision input"), + (constraint is not None, "constrained decoding"), + (adaptive_width_policy is not None, "adaptive width policy"), + ) + if selected + ] + if unsupported: + backend_id = str( + getattr(block_speculative_backend, "backend_id", "native block backend") + ) + raise ValueError( + f"{backend_id} does not yet support: {', '.join(unsupported)}" + ) + from .native_block_speculation import generate_native_block_speculative + + return generate_native_block_speculative( + rt, + block_speculative_backend, + prompt_ids, + abort_check=abort_check, + max_tokens=max_tokens, + sampler=sampler, + speculative_depth=speculative_depth, + seed=seed, + stop_token_ids=stop_token_ids, + draft_sampler=draft_sampler, + token_callback=token_callback, + prefill_callback=prefill_callback, + constraint=constraint, + vision_splice=vision_splice, + adaptive_policy=adaptive_policy, + adaptive_width_policy=adaptive_width_policy, + ) if not rt.mtp_enabled: raise RuntimeError("generate_mtpk requires an MTP-enabled runtime") base_hidden_variant = _resolve_runtime_base_hidden_variant(rt, base_hidden_variant) @@ -6316,6 +6440,7 @@ def generate_mtpk( from .context_copy import ( context_copy_target_prefix_enabled as _cc_tp_enabled_early, ) + _penalty_bearing_request = bool(sampler.presence_penalty) or bool( sampler.frequency_penalty ) @@ -6333,7 +6458,8 @@ def generate_mtpk( # the batch scheduler's dense host fallback). exact_a3b_target_prefix_factory = ( rt.a3b_compiled_target_prefix_factory - if target_prefix_verify and constraint is None + if target_prefix_verify + and constraint is None and not _ccopy_takes_over_lane and not _penalty_bearing_request else None @@ -6472,12 +6598,8 @@ def record_adaptive_width_event( adaptive_width_policy.stop_after_d1, adaptive_width_policy.stop_after_d2, ) - adaptive_width_d1_threshold = float( - adaptive_width_policy.d1_margin_threshold - ) - adaptive_width_d2_threshold = float( - adaptive_width_policy.d2_margin_threshold - ) + adaptive_width_d1_threshold = float(adaptive_width_policy.d1_margin_threshold) + adaptive_width_d2_threshold = float(adaptive_width_policy.d2_margin_threshold) adaptive_width_max_depth = int(adaptive_width_policy.max_speculative_depth) capture_forward_routes = adaptive_width_policy.target_routes @@ -6771,7 +6893,9 @@ def record_adaptive_width_event( # every verified MTP position by its growing in-block prefix (per-position / # vLLM-exact). Counts are rebuilt from `tokens` at each sample point — simple # and drift-proof; an incremental counter is a documented perf follow-up. - _penalties_active = bool(sampler.presence_penalty) or bool(sampler.frequency_penalty) + _penalties_active = bool(sampler.presence_penalty) or bool( + sampler.frequency_penalty + ) # Loop Guard: loop-armed DRY-style steering (see mtplx/loop_guard.py). # Disarmed = zero distribution impact (identity transform, fast paths kept). # Armed = target distributions get sparse anti-cycle penalties per position; @@ -6890,7 +7014,9 @@ def _steer_overlay(working: Sequence[int]) -> dict[int, float] | None: int( position_base_env if position_base_env is not None - else (len(prompt_state.token_prefix) if mtp_position_mode == "absolute" else 0) + else ( + len(prompt_state.token_prefix) if mtp_position_mode == "absolute" else 0 + ) ), ) # Validate env spelling before a long generation starts. @@ -6904,7 +7030,9 @@ def _steer_overlay(working: Sequence[int]) -> dict[int, float] | None: def mtp_position_offset_for_cache(mtp_cache) -> int | None: position_base = mtp_position_base - if mtp_cache is mtp_history_cache and _mtp_history_uses_committed_cache(mtp_history_policy): + if mtp_cache is mtp_history_cache and _mtp_history_uses_committed_cache( + mtp_history_policy + ): position_base = mtp_history_position_base return _mtp_position_offset( _mtp_cache_offset(mtp_cache), @@ -7476,10 +7604,17 @@ def emit_new_tokens() -> None: # ---- context-copy (prompt-lookup) drafting: always on (kill switch # MTPLX_CONTEXT_COPY=0); any temperature, no repetition penalties, on # capture-commit verify strategies ---- - from .context_copy import (NgramIndex, block_for_ext, context_copy_block_k, - context_copy_enabled, context_copy_min_ext, - context_copy_ng_max, context_copy_ng_min, - context_copy_target_prefix_enabled) + from .context_copy import ( + NgramIndex, + block_for_ext, + context_copy_block_k, + context_copy_enabled, + context_copy_min_ext, + context_copy_ng_max, + context_copy_ng_min, + context_copy_target_prefix_enabled, + ) + # Temperature is supported through the same probability-ratio acceptance # as the MTP path: the copy block is a point-mass proposal, so a copied # token is accepted with the target's own shaped probability and a @@ -7492,7 +7627,10 @@ def emit_new_tokens() -> None: # depth-1 draft; block rounds stay capture_commit-only -- their T+1-row # forwards are not AR-exact). With whole-MoE installed the compiled # route is kept and the flag is inert, recorded via disabled_reason. - _ccopy_capture_lane = verify_strategy in {"capture_commit", "graphbank_capture_commit"} + _ccopy_capture_lane = verify_strategy in { + "capture_commit", + "graphbank_capture_commit", + } _ccopy_tp_requested = ( context_copy_target_prefix_enabled() and verify_strategy == "target_prefix" ) @@ -7502,7 +7640,10 @@ def emit_new_tokens() -> None: ccopy_active = ( context_copy_enabled() and not _penalties_active - and (_ccopy_capture_lane or (_ccopy_tp_requested and not _ccopy_whole_moe_conflict)) + and ( + _ccopy_capture_lane + or (_ccopy_tp_requested and not _ccopy_whole_moe_conflict) + ) ) ccopy_rounds = ccopy_drafted = ccopy_accepted = 0 ccopy_probes = ccopy_blocks_accepted = ccopy_suspensions = 0 @@ -7513,9 +7654,9 @@ def emit_new_tokens() -> None: # excludes draft substitution. The compiled route is kept. ccopy_disabled_reason = "whole_moe_keeps_compiled_route" ccopy_ema, ccopy_seen, ccopy_suspend_until = 0.5, 0, 0 - ccopy_backoff = 64 # doubles on each suspension (self-repetitive novel text would - # otherwise re-trigger copy rounds after every backoff and pay - # the probe cost recurrently); a paying round resets it. + ccopy_backoff = 64 # doubles on each suspension (self-repetitive novel text would + # otherwise re-trigger copy rounds after every backoff and pay + # the probe cost recurrently); a paying round resets it. # Draft-source streak state (target_prefix takeover lane): the copy match # feeds the depth-1 DRAFT instead of a block round, so every forward stays # on the lane's proven 2-row verify geometry -- bit-exact by construction. @@ -7636,9 +7777,7 @@ def emit_new_tokens() -> None: sampler, rng, token_counts=Counter(tokens) if _penalties_active else None, - penalty_overlay=( - _steer_overlay(tokens) if _steer_active else None - ), + penalty_overlay=(_steer_overlay(tokens) if _steer_active else None), ) if first_primary_sample_time_s == 0.0: # First primary token sampled: any lazy tail forced by @@ -7822,7 +7961,12 @@ def emit_new_tokens() -> None: "correction": None, } # ---- context-copy round: verbatim block from context, no MTP compute this cycle ---- - if ccopy_active and _ccopy_capture_lane and cycle_depth >= 1 and len(tokens) >= ccopy_suspend_until: + if ( + ccopy_active + and _ccopy_capture_lane + and cycle_depth >= 1 + and len(tokens) >= ccopy_suspend_until + ): _cc_hist = prompt_ids + tokens ccopy_probes += 1 # Prompt-only contract: candidates whose continuation starts at the @@ -7833,7 +7977,7 @@ def emit_new_tokens() -> None: _cc_block: list[int] = [] if _cc_pos is not None and _cc_ext >= ccopy_min_ext: _cc_klen = block_for_ext(_cc_ext, ccopy_k) - _cc_block = [int(t) for t in prompt_ids[_cc_pos:_cc_pos + _cc_klen]] + _cc_block = [int(t) for t in prompt_ids[_cc_pos : _cc_pos + _cc_klen]] _cc_block = _cc_block[: max(1, max_tokens - len(tokens))] if constraint is not None: # Truncate the copy proposal at the first grammar-illegal @@ -7923,9 +8067,13 @@ def emit_new_tokens() -> None: _cc_ok = True if _cc_nacc < len(_cc_block): from .gdn_capture import commit_captured_prefix + started_commit = time.perf_counter() _cc_ok = commit_captured_prefix( - cache, _cc_captures, keep_tokens=_cc_m, verified_tokens=_cc_T, + cache, + _cc_captures, + keep_tokens=_cc_m, + verified_tokens=_cc_T, ) capture_commit_time += time.perf_counter() - started_commit if not _cc_ok: @@ -7959,15 +8107,25 @@ def emit_new_tokens() -> None: continue _cc_round_pos = len(tokens) _cc_acc = _cc_block[:_cc_nacc] - _cc_stop_idx = next((i for i, t in enumerate(_cc_acc) - if _is_stop(int(t), stop_token_ids)), None) + _cc_stop_idx = next( + ( + i + for i, t in enumerate(_cc_acc) + if _is_stop(int(t), stop_token_ids) + ), + None, + ) if _cc_stop_idx is not None: - _cc_acc = _cc_acc[:_cc_stop_idx + 1] + _cc_acc = _cc_acc[: _cc_stop_idx + 1] tokens.extend(_cc_acc) _cc_finished = _cc_stop_idx is not None - if constraint is not None and _cc_correction is not None and ( - constraint.validate_prefix([*_cc_acc, int(_cc_correction)]) - != len(_cc_acc) + 1 + if ( + constraint is not None + and _cc_correction is not None + and ( + constraint.validate_prefix([*_cc_acc, int(_cc_correction)]) + != len(_cc_acc) + 1 + ) ): # Grammar-illegal residual: drop it; the next cycle's # masked primary resamples the position, which preserves @@ -7992,7 +8150,7 @@ def emit_new_tokens() -> None: ccopy_ema = 0.7 * ccopy_ema + 0.3 * (_cc_nacc / len(_cc_block)) ccopy_seen += 1 if _cc_nacc / len(_cc_block) >= 0.5: - ccopy_backoff = 64 # copy is paying again: full retry rate + ccopy_backoff = 64 # copy is paying again: full retry rate if ccopy_seen >= 4 and ccopy_ema < 0.35: # acceptance collapsed (novel region with incidental repeats): # suspend copy rounds and let the MTP head work; retry with @@ -8019,7 +8177,10 @@ def emit_new_tokens() -> None: # Committed-history MTP caches pair every committed token with the # hidden state of the token before it, including (previous hidden, # primary), which the drafting path would normally have added. - if _mtp_history_uses_committed_cache(mtp_history_policy) and mtp_cache is not None: + if ( + _mtp_history_uses_committed_cache(mtp_history_policy) + and mtp_cache is not None + ): _cc_committed_toks = [primary] + _cc_acc _cc_hiddens = mx.concatenate( [hidden, _cc_hidden[:, : len(_cc_acc), :]], axis=1 @@ -8028,7 +8189,7 @@ def emit_new_tokens() -> None: mtp_cache, _cc_hiddens, _cc_committed_toks ) logits = _cc_logits[:, _cc_m - 1, :] - hidden = _cc_hidden[:, _cc_m - 1:_cc_m, :] + hidden = _cc_hidden[:, _cc_m - 1 : _cc_m, :] append_event(event) emit_new_tokens() if _cc_finished: @@ -8227,8 +8388,11 @@ def emit_new_tokens() -> None: ): stored = device_core["state_signature"] diffs = [ - (i, stored[i] if i < len(stored) else None, - live_signature[i] if i < len(live_signature) else None) + ( + i, + stored[i] if i < len(stored) else None, + live_signature[i] if i < len(live_signature) else None, + ) for i in range(max(len(stored), len(live_signature))) if (stored[i] if i < len(stored) else None) != (live_signature[i] if i < len(live_signature) else None) @@ -8277,9 +8441,7 @@ def emit_new_tokens() -> None: for depth_index, (draft_token, draft_q) in enumerate( zip(core_tokens, core_qs) ): - draft_probs.append( - draft_q if sampler.temperature > 0 else None - ) + draft_probs.append(draft_q if sampler.temperature > 0 else None) drafted += 1 drafted_by_depth[depth_index] += 1 event["drafts"].append( @@ -8496,13 +8658,11 @@ def emit_new_tokens() -> None: need_draft_distribution = ( sampler.temperature > 0 and not target_prefix_verify ) - draft_token, draft_q, adaptive_width_stop = ( - cycle_draft_reader( - draft_logits, - depth_index=depth_index, - need_distribution=need_draft_distribution, - decision_margins=adaptive_width_decision_margins, - ) + draft_token, draft_q, adaptive_width_stop = cycle_draft_reader( + draft_logits, + depth_index=depth_index, + need_distribution=need_draft_distribution, + decision_margins=adaptive_width_decision_margins, ) elapsed_draft = time.perf_counter() - started draft_time += elapsed_draft @@ -8664,9 +8824,7 @@ def emit_new_tokens() -> None: bonus_distribution_row_needed = ( not omit_speculative_bonus and len(tokens) + 1 < max_tokens ) - target_distribution_rows_needed = 1 + int( - bonus_distribution_row_needed - ) + target_distribution_rows_needed = 1 + int(bonus_distribution_row_needed) verified_token_count = 2 verify_input_array = mx.concatenate( (mx.array([[primary]]), device_draft_token.reshape(1, 1)), @@ -8841,17 +8999,13 @@ def emit_new_tokens() -> None: _guarded_rows = [] for _row_index in range(int(target_distribution_rows)): _row = target_distribution_logits[:, _row_index, :].reshape(-1) - _row_overlay = _steer_overlay( - [*tokens, *draft_tokens[:_row_index]] - ) + _row_overlay = _steer_overlay([*tokens, *draft_tokens[:_row_index]]) if _row_overlay: _row = apply_penalties_mlx( _row, None, penalty_overlay=_row_overlay ) _guarded_rows.append(_row) - target_distribution_logits = mx.stack(_guarded_rows, axis=0)[ - None, ... - ] + target_distribution_logits = mx.stack(_guarded_rows, axis=0)[None, ...] sampled_target_ids = sample_token_ids_from_mlx_logits( target_distribution_logits, sampler, @@ -9026,7 +9180,10 @@ def emit_new_tokens() -> None: target_distribution_logits, sampler, ) - if target_distribution_batch is not None or target_distributions is not None: + if ( + target_distribution_batch is not None + or target_distributions is not None + ): target_distribution_materialized_rows += int(target_distribution_rows) target_distribution_materialized_windows += 1 event["target_distribution_materialized"] = { @@ -9232,10 +9389,7 @@ def emit_new_tokens() -> None: event["drafts"][depth_index]["online_correction_cache"][ "stored_token" ] = cached_target - if ( - sampler.temperature > 0 - or a3b_target_prefix_route is not None - ) and ( + if (sampler.temperature > 0 or a3b_target_prefix_route is not None) and ( constraint is None or constraint.validate_prefix( [*draft_tokens[:depth_index], int(correction)] @@ -9421,9 +9575,8 @@ def emit_new_tokens() -> None: continue started_bonus = time.perf_counter() bonus_target_distribution_time = 0.0 - if ( - target_prefix_tokens is not None - and len(target_prefix_tokens) > len(draft_tokens) + if target_prefix_tokens is not None and len(target_prefix_tokens) > len( + draft_tokens ): bonus = int(target_prefix_tokens[len(draft_tokens)]) elif target_distribution_batch is not None and not lazy_bonus_verify: @@ -9449,9 +9602,7 @@ def emit_new_tokens() -> None: rng, token_counts=Counter(tokens) if _penalties_active else None, penalty_overlay=( - _steer_overlay(tokens) - if _steer_active - else None + _steer_overlay(tokens) if _steer_active else None ), ) if sampler.temperature > 0: @@ -9462,7 +9613,9 @@ def emit_new_tokens() -> None: target_distribution_materialized_windows += 1 lazy_target_distribution_window_counted = True target_distribution_materialized_rows += 1 - verify_target_distribution_time += bonus_target_distribution_time + verify_target_distribution_time += ( + bonus_target_distribution_time + ) verify_logits_eval_time += bonus_target_distribution_time verify_eval_time += bonus_target_distribution_time verify_time += bonus_target_distribution_time @@ -9472,7 +9625,9 @@ def emit_new_tokens() -> None: materialized.get("mode", "") ).startswith("lazy"): materialized["mode"] = "lazy_accept_bonus_path" - materialized["rows"] = int(materialized.get("rows") or 0) + 1 + materialized["rows"] = ( + int(materialized.get("rows") or 0) + 1 + ) materialized["time_s"] = float( materialized.get("time_s") or 0.0 ) + float(bonus_target_distribution_time) @@ -9491,7 +9646,9 @@ def emit_new_tokens() -> None: ) elapsed_bonus = max( 0.0, - time.perf_counter() - started_bonus - bonus_target_distribution_time, + time.perf_counter() + - started_bonus + - bonus_target_distribution_time, ) bonus_time += elapsed_bonus _add_timing(event, "bonus_sample", elapsed_bonus) @@ -9911,9 +10068,7 @@ def emit_new_tokens() -> None: stats = GenerationStats( mode="mtpk", constraint_active=constraint is not None, - constraint_completed=( - constraint.completed if constraint is not None else None - ), + constraint_completed=(constraint.completed if constraint is not None else None), constraint_masked_steps=( constraint.masked_steps if constraint is not None else 0 ), @@ -9978,7 +10133,9 @@ def emit_new_tokens() -> None: ssd_cache_hit=prompt_state.ssd_cache_hit, ssd_cached_tokens=prompt_state.ssd_cached_tokens, ssd_restore_s=prompt_state.ssd_restore_s, - ssd_suffix_tokens=prompt_state.suffix_tokens if prompt_state.ssd_cache_hit else 0, + ssd_suffix_tokens=prompt_state.suffix_tokens + if prompt_state.ssd_cache_hit + else 0, cache_miss_reason=prompt_state.cache_miss_reason, session_restore_mode=prompt_state.restore_mode, session_prompt_prefix_bank_commit=prompt_prefix_bank_commit, diff --git a/mtplx/models/deepseek_v4.py b/mtplx/models/deepseek_v4.py index abd59f1a..91115ec6 100644 --- a/mtplx/models/deepseek_v4.py +++ b/mtplx/models/deepseek_v4.py @@ -271,9 +271,11 @@ import math import os +from contextlib import contextmanager +from contextvars import ContextVar from dataclasses import dataclass, field, replace from functools import lru_cache -from typing import List, Optional, Tuple +from typing import Iterator, List, Optional, Tuple import mlx.core as mx import mlx.nn as nn @@ -286,11 +288,7 @@ # Default per-layer compress ratios for DeepSeek-V4-Flash (43 body layers; the # 44th entry is the dropped MTP layer). 0 = pure sliding-window; 4 = overlapping # compressor + indexer; 128 = non-overlapping compressor + strided index. -_DEFAULT_COMPRESS_RATIOS = ( - [0, 0] - + [4, 128] * 20 - + [4, 0] -) +_DEFAULT_COMPRESS_RATIOS = [0, 0] + [4, 128] * 20 + [4, 0] # How many token positions a :class:`DeepseekV4Cache` can un-decode (``trim``). # Speculative decode only ever rewinds the rejected tail of one verify batch, so @@ -336,8 +334,7 @@ def _o_lora_mode_from_env() -> str: return "cached" if raw not in _O_LORA_MODES: raise ValueError( - "MTPLX_DSV4_O_LORA must be one of " - f"{', '.join(_O_LORA_MODES)}; got {raw!r}" + f"MTPLX_DSV4_O_LORA must be one of {', '.join(_O_LORA_MODES)}; got {raw!r}" ) return raw @@ -368,8 +365,7 @@ def _attn_mode_from_env() -> str: return "fused" if raw not in _ATTN_MODES: raise ValueError( - "MTPLX_DSV4_ATTN must be one of " - f"{', '.join(_ATTN_MODES)}; got {raw!r}" + f"MTPLX_DSV4_ATTN must be one of {', '.join(_ATTN_MODES)}; got {raw!r}" ) return raw @@ -432,6 +428,25 @@ def _attn_mode_from_env() -> str: #: ``MTPLX_DSV4_SINKHORN_KERNEL`` at import; tests set the module attribute. _SINKHORN_KERNEL = _env_flag("MTPLX_DSV4_SINKHORN_KERNEL", False) +# The retained 0731 K2 stack is selected by an explicit runtime construction +# option. A context-local selector keeps that request out of ambient process +# policy and confines it to the model constructors invoked by one ``load``. +_DEEPSEEK_V4_0731_K2_CONSTRUCTION: ContextVar[bool] = ContextVar( + "deepseek_v4_0731_k2_construction", + default=False, +) + + +@contextmanager +def deepseek_v4_0731_k2_construction() -> Iterator[None]: + """Select the pinned Sinkhorn route for one model construction only.""" + + token = _DEEPSEEK_V4_0731_K2_CONSTRUCTION.set(True) + try: + yield + finally: + _DEEPSEEK_V4_0731_K2_CONSTRUCTION.reset(token) + #: Escape hatch restoring the pre-fix all-fp32 activation path (rope output, #: compressed KV rows and the attention probability block). The reference keeps @@ -654,14 +669,12 @@ def _validate_loaded_moe_tail_contract(model, config: dict) -> dict: layers = list(getattr(model, "layers", ())) if len(layers) != _MOE_TAIL_BODY_LAYERS: raise ValueError( - "MTPLX_DSV4_MOE_TAIL requires exactly 43 body layers; " - f"got {len(layers)}" + f"MTPLX_DSV4_MOE_TAIL requires exactly 43 body layers; got {len(layers)}" ) mtp_blocks = list(getattr(model, "mtp_blocks", ())) if len(mtp_blocks) != _MOE_TAIL_MTP_BLOCKS: raise ValueError( - "MTPLX_DSV4_MOE_TAIL requires exactly one MTP block; " - f"got {len(mtp_blocks)}" + f"MTPLX_DSV4_MOE_TAIL requires exactly one MTP block; got {len(mtp_blocks)}" ) args = getattr(model, "args", None) if args is None: @@ -729,7 +742,11 @@ def _validate_loaded_moe_tail_contract(model, config: dict) -> dict: "up_proj": ((256, 2048, 256), (256, 2048, 64), 64), "down_proj": ((256, 4096, 128), (256, 4096, 32), 64), } - for projection, (weight_shape, scale_shape, group_size) in routed_contract.items(): + for projection, ( + weight_shape, + scale_shape, + group_size, + ) in routed_contract.items(): stem = f"model.layers.{layer_id}.ffn.switch_mlp.{projection}" expected_spec = { "bits": 2, @@ -857,7 +874,9 @@ def _moe_tail_metal_kernel(): return _MOE_TAIL_KERNEL -def _moe_tail_apply(kernel, routed: mx.array, weights: mx.array, shared: mx.array) -> mx.array: +def _moe_tail_apply( + kernel, routed: mx.array, weights: mx.array, shared: mx.array +) -> mx.array: """Dispatch the precompiled fixed tail; ``rows`` is the only varying value.""" rows = int(routed.shape[0]) n_elements = rows * _MOE_TAIL_HIDDEN @@ -883,19 +902,23 @@ def _verify_moe_tail_exact(kernel) -> None: return for rows in (1, 4): n = rows * _MOE_TAIL_TOPK * _MOE_TAIL_HIDDEN - routed = ((mx.arange(n, dtype=mx.float32) % 29 - 14) / 7).reshape( - rows, _MOE_TAIL_TOPK, _MOE_TAIL_HIDDEN - ).astype(mx.bfloat16) - weights = ((mx.arange(rows * _MOE_TAIL_TOPK, dtype=mx.float32) % 13 - 6) / 5) + routed = ( + ((mx.arange(n, dtype=mx.float32) % 29 - 14) / 7) + .reshape(rows, _MOE_TAIL_TOPK, _MOE_TAIL_HIDDEN) + .astype(mx.bfloat16) + ) + weights = (mx.arange(rows * _MOE_TAIL_TOPK, dtype=mx.float32) % 13 - 6) / 5 weights = weights.reshape(rows, _MOE_TAIL_TOPK).astype(mx.bfloat16) - shared = ((mx.arange(rows * _MOE_TAIL_HIDDEN, dtype=mx.float32) % 31 - 15) / 11) + shared = (mx.arange(rows * _MOE_TAIL_HIDDEN, dtype=mx.float32) % 31 - 15) / 11 shared = shared.reshape(rows, _MOE_TAIL_HIDDEN).astype(mx.bfloat16) stock = _stock_moe_tail_combine(routed, weights, shared) fused = _moe_tail_apply(kernel, routed, weights, shared) mx.eval(stock, fused) if not mx.array_equal(stock, fused): max_abs = float( - mx.max(mx.abs(stock.astype(mx.float32) - fused.astype(mx.float32))).item() + mx.max( + mx.abs(stock.astype(mx.float32) - fused.astype(mx.float32)) + ).item() ) raise RuntimeError( "MTPLX_DSV4_MOE_TAIL failed exact Metal self-check at " @@ -1076,9 +1099,7 @@ def __init__(self, attention: "DeepseekV4Attention", quant: tuple) -> None: self.per_group_input = per_group_input self.weight = weight.reshape(groups, rank, -1) self.scales = scales.reshape(groups, rank, -1) - self.biases = ( - None if biases is None else biases.reshape(groups, rank, -1) - ) + self.biases = None if biases is None else biases.reshape(groups, rank, -1) self.group_size = int(group_size) self.bits = int(bits) self.mode = mode @@ -1099,9 +1120,7 @@ def __call__(self, o: mx.array) -> mx.array: mode=self.mode, ) return self.wo_b( - out.swapaxes(0, 1).reshape( - batch, sequence, self.groups * self.rank - ) + out.swapaxes(0, 1).reshape(batch, sequence, self.groups * self.rank) ) @@ -1341,16 +1360,12 @@ def __init__(self, attention: "DeepseekV4Attention", weight: mx.array) -> None: self.per_group_input = int( attention.n_heads * attention.head_dim // attention.n_groups ) - self.weight = weight.reshape( - self.groups, self.rank, self.per_group_input - ) + self.weight = weight.reshape(self.groups, self.rank, self.per_group_input) self.wo_b = attention.wo_b def __call__(self, o: mx.array) -> mx.array: batch, sequence, _ = o.shape - grouped = o.reshape( - batch, sequence, self.groups, self.per_group_input - ) + grouped = o.reshape(batch, sequence, self.groups, self.per_group_input) out = mx.einsum("bsgp,grp->bsgr", grouped, self.weight) return self.wo_b(out.reshape(batch, sequence, self.groups * self.rank)) @@ -1420,7 +1435,9 @@ class ModelArgs(BaseModelArgs): index_n_heads: int = 64 index_head_dim: int = 128 index_topk: int = 512 - compress_ratios: List[int] = field(default_factory=lambda: list(_DEFAULT_COMPRESS_RATIOS)) + compress_ratios: List[int] = field( + default_factory=lambda: list(_DEFAULT_COMPRESS_RATIOS) + ) compress_rope_theta: float = 160000.0 # hyper-connections hc_mult: int = 4 @@ -1500,9 +1517,12 @@ def _yarn_inv_freq( half = dim // 2 freqs = 1.0 / (base ** (mx.arange(0, dim, 2, dtype=mx.float32) / dim)) if original_seq_len and original_seq_len > 0: + def correction_dim(num_rot): - return dim * math.log(original_seq_len / (num_rot * 2 * math.pi)) / ( - 2 * math.log(base) + return ( + dim + * math.log(original_seq_len / (num_rot * 2 * math.pi)) + / (2 * math.log(base)) ) low = max(math.floor(correction_dim(beta_fast)), 0) @@ -1539,7 +1559,7 @@ def _hadamard_rotate(x: mx.array) -> mx.array: b = y[:, :, 1] y = mx.stack([a + b, a - b], axis=2).reshape(-1, n) stride *= 2 - return (y * (n ** -0.5)).reshape(x.shape) + return (y * (n**-0.5)).reshape(x.shape) def _topk_mask(key: mx.array, k_row: mx.array, k_max: int) -> mx.array: @@ -1566,11 +1586,13 @@ def _topk_mask(key: mx.array, k_row: mx.array, k_max: int) -> mx.array: else: ranked = mx.sort(mx.topk(key, k_max, axis=-1), axis=-1)[..., ::-1] kth = mx.clip(k_row - 1, 0, ranked.shape[-1] - 1) - thr = mx.take_along_axis(ranked, kth, axis=-1) # k_row-th largest + thr = mx.take_along_axis(ranked, kth, axis=-1) # k_row-th largest gt = key > thr eq = key == thr n_gt = mx.sum(gt.astype(mx.int32), axis=-1, keepdims=True) - tie_rank = mx.cumsum(eq.astype(mx.int32), axis=-1) - 1 # rank among equals, index order + tie_rank = ( + mx.cumsum(eq.astype(mx.int32), axis=-1) - 1 + ) # rank among equals, index order return gt | (eq & (tie_rank < (k_row - n_gt))) @@ -1616,11 +1638,11 @@ def _sinkhorn_ops(comb: mx.array, iters: int, eps: float) -> mx.array: the path both :func:`hc_split_sinkhorn` and :func:`_hc_pre_impl` take when the kernel is off; it is exactly the loop these two functions used to inline. """ - comb = mx.softmax(comb, axis=-1) + eps # row-softmax - comb = comb / (comb.sum(axis=-2, keepdims=True) + eps) # column normalise + comb = mx.softmax(comb, axis=-1) + eps # row-softmax + comb = comb / (comb.sum(axis=-2, keepdims=True) + eps) # column normalise for _ in range(iters - 1): - comb = comb / (comb.sum(axis=-1, keepdims=True) + eps) # row normalise - comb = comb / (comb.sum(axis=-2, keepdims=True) + eps) # column normalise + comb = comb / (comb.sum(axis=-1, keepdims=True) + eps) # row normalise + comb = comb / (comb.sum(axis=-2, keepdims=True) + eps) # column normalise return comb @@ -1744,12 +1766,18 @@ def _install_sinkhorn_normaliser(hc: int, iters: int, eps: float): geometry fails here, before measured generation, instead of silently taking a differently-shaped kernel or falling back. """ + def stock(comb: mx.array) -> mx.array: return _sinkhorn_ops(comb, iters, eps) - if not _SINKHORN_KERNEL: + explicit_0731_k2 = _DEEPSEEK_V4_0731_K2_CONSTRUCTION.get() + if not (explicit_0731_k2 or _SINKHORN_KERNEL): return False, stock if not mx.metal.is_available() or mx.default_device() != mx.gpu: + if explicit_0731_k2: + raise ValueError( + "DeepSeek-V4-0731 K2 construction requires an MLX Metal GPU" + ) return False, stock if (hc, iters, eps) != (4, 20, 1e-6): raise ValueError( @@ -1792,7 +1820,9 @@ def hc_split_sinkhorn( return pre, post, comb -def _hc_pre_impl(x, fn_t, base, scale_vec, hc: int, iters: int, eps: float, normalise=None): +def _hc_pre_impl( + x, fn_t, base, scale_vec, hc: int, iters: int, eps: float, normalise=None +): """:meth:`HyperConnection.pre` as one pure function of arrays. Identical arithmetic to ``_mixes`` + :func:`hc_split_sinkhorn` + the weighted @@ -1825,8 +1855,10 @@ def _hc_pre_impl(x, fn_t, base, scale_vec, hc: int, iters: int, eps: float, norm post = 2.0 * mx.sigmoid(t[..., hc : 2 * hc]) comb = t[..., 2 * hc :].reshape(*t.shape[:-1], hc, hc) # [..., j, k] if normalise is None: + def normalise(c): return _sinkhorn_ops(c, iters, eps) + comb = normalise(comb) y = mx.sum(pre[..., None] * xf, axis=-2) # [..., dim] @@ -1838,8 +1870,8 @@ def _hc_post_impl(x, residual, post, comb): dtype = x.dtype xf = x.astype(mx.float32) rf = residual.astype(mx.float32) - term = post[..., None] * xf[..., None, :] # [..., hc, dim] - mixed = mx.einsum("...jk,...jd->...kd", comb, rf) # sum_j comb[j,k] res[j] + term = post[..., None] * xf[..., None, :] # [..., hc, dim] + mixed = mx.einsum("...jk,...jd->...kd", comb, rf) # sum_j comb[j,k] res[j] return (term + mixed).astype(dtype) @@ -1880,16 +1912,16 @@ def _hc_compiled(kind: str, *consts): hc, iters, eps, sinkhorn_kernel = consts if sinkhorn_kernel: + def normalise(comb): return _sinkhorn_kernel_apply(comb, hc, iters, eps) else: + def normalise(comb): return _sinkhorn_ops(comb, iters, eps) def impl(x, fn_t, base, scale_vec): - return _hc_pre_impl( - x, fn_t, base, scale_vec, hc, iters, eps, normalise - ) + return _hc_pre_impl(x, fn_t, base, scale_vec, hc, iters, eps, normalise) elif kind == "post": impl = _hc_post_impl elif kind == "head": @@ -2005,6 +2037,7 @@ def post(self, x: mx.array, residual: mx.array, post: mx.array, comb: mx.array): impl = _hc_compiled("post") if _hc_use_compile(residual) else _hc_post_impl return impl(x, residual, post, comb) + class HeadHC(nn.Module): """Final head hyper-connection collapse (``ParallelHead.hc_head``, model.py L728). @@ -2091,8 +2124,12 @@ def __init__( # Compressor rope uses the compress theta + YaRN (reference passes the # compressor its own freqs_cis; window w gets position w*ratio). self._inv_freq = _yarn_inv_freq( - self.rope_head_dim, args.compress_rope_theta, args.original_seq_len, - args.rope_factor, args.beta_fast, args.beta_slow, + self.rope_head_dim, + args.compress_rope_theta, + args.original_seq_len, + args.rope_factor, + args.beta_fast, + args.beta_slow, ) def _overlap_transform( @@ -2111,14 +2148,16 @@ def _overlap_transform( """ b, nwin, r, _ = t.shape d = self.head_dim - cur = t[..., d:] # [b, nwin, ratio, d] (current, d: half) - prev_half = t[..., :d] # [b, nwin, ratio, d] (:d half) + cur = t[..., d:] # [b, nwin, ratio, d] (current, d: half) + prev_half = t[..., :d] # [b, nwin, ratio, d] (:d half) if prev is None: seed = mx.full((b, 1, r, d), value, dtype=t.dtype) else: - seed = prev[..., :d][:, None] # [b, 1, ratio, d] - prev_shift = mx.concatenate([seed, prev_half[:, :-1]], axis=1) # w -> window w-1 - return mx.concatenate([prev_shift, cur], axis=2) # [b, nwin, 2*ratio, d] + seed = prev[..., :d][:, None] # [b, 1, ratio, d] + prev_shift = mx.concatenate( + [seed, prev_half[:, :-1]], axis=1 + ) # w -> window w-1 + return mx.concatenate([prev_shift, cur], axis=2) # [b, nwin, 2*ratio, d] def _pool( self, kv: mx.array, score: mx.array, first_window: int, out_dtype @@ -2140,9 +2179,11 @@ def _pool( """ nwin = kv.shape[1] rd = self.rope_head_dim - pooled = mx.sum(kv * mx.softmax(score, axis=2), axis=2) # [b, nwin, d] + pooled = mx.sum(kv * mx.softmax(score, axis=2), axis=2) # [b, nwin, d] pooled = self.norm(pooled.astype(out_dtype)) - win_pos = (mx.arange(nwin, dtype=mx.float32) + float(first_window)) * self.compress_ratio + win_pos = ( + mx.arange(nwin, dtype=mx.float32) + float(first_window) + ) * self.compress_ratio ang = win_pos[:, None] * self._inv_freq[None, :] cos, sin = mx.cos(ang), mx.sin(ang) head = pooled[..., :-rd] @@ -2165,10 +2206,12 @@ def __call__(self, x: mx.array) -> mx.array: if nwin == 0: return mx.zeros((b, 0, d), dtype=out_dtype) xf = x.astype(mx.float32) - kv = self.wkv(xf)[:, :cutoff].reshape(b, nwin, ratio, -1) # [b,nwin,ratio,coff*d] + kv = self.wkv(xf)[:, :cutoff].reshape( + b, nwin, ratio, -1 + ) # [b,nwin,ratio,coff*d] score = self.wgate(xf)[:, :cutoff].reshape(b, nwin, ratio, -1) + self.ape if self.overlap: - kv = self._overlap_transform(kv, 0.0) # [b,nwin,2*ratio,d] + kv = self._overlap_transform(kv, 0.0) # [b,nwin,2*ratio,d] score = self._overlap_transform(score, float("-inf")) return self._pool(kv, score, 0, out_dtype) @@ -2194,8 +2237,8 @@ def step(self, x: mx.array, state: "CompressorState", offset: int) -> mx.array: d = self.head_dim out_dtype = _store_dtype(x.dtype) xf = x.astype(mx.float32) - kv_rows = self.wkv(xf) # [b, s, coff*d] - ape_idx = (mx.arange(s) + offset) % ratio # slot of each token + kv_rows = self.wkv(xf) # [b, s, coff*d] + ape_idx = (mx.arange(s) + offset) % ratio # slot of each token score_rows = self.wgate(xf) + self.ape[ape_idx] # Rollback journal: the projected rows are per-position pure functions, so # keeping the most recent few is all a rewind needs to rebuild the frontier @@ -2218,7 +2261,7 @@ def step(self, x: mx.array, state: "CompressorState", offset: int) -> mx.array: score_slots = self._overlap_transform( score_w, float("-inf"), state.prev_score ) - state.prev_kv = kv_w[:, -1] # [b, ratio, coff*d] + state.prev_kv = kv_w[:, -1] # [b, ratio, coff*d] state.prev_score = score_w[:, -1] else: kv_slots, score_slots = kv_w, score_w @@ -2277,15 +2320,21 @@ def __init__(self, args: ModelArgs, compress_ratio: int): self.index_topk = args.index_topk self.compress_ratio = compress_ratio self.q_lora_rank = args.q_lora_rank - self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.head_dim, bias=False) + self.wq_b = nn.Linear( + self.q_lora_rank, self.n_heads * self.head_dim, bias=False + ) self.weights_proj = nn.Linear(self.dim, self.n_heads, bias=False) - self.softmax_scale = self.head_dim ** -0.5 + self.softmax_scale = self.head_dim**-0.5 self.compressor = Compressor(args, compress_ratio, self.head_dim, rotate=True) # The reference hands the indexer the *attention layer's* freqs_cis # (model.py L494); on a ratio-4 layer that is compress_rope_theta + YaRN. self._inv_freq = _yarn_inv_freq( - self.rope_head_dim, args.compress_rope_theta, args.original_seq_len, - args.rope_factor, args.beta_fast, args.beta_slow, + self.rope_head_dim, + args.compress_rope_theta, + args.original_seq_len, + args.rope_factor, + args.beta_fast, + args.beta_slow, ) def scores( @@ -2316,8 +2365,8 @@ def scores( ) q = _hadamard_rotate(q.astype(mx.float32)) weights = self.weights_proj(x).astype(mx.float32) * ( - self.softmax_scale * self.n_heads ** -0.5 - ) # [b, s, n_heads] + self.softmax_scale * self.n_heads**-0.5 + ) # [b, s, n_heads] score = mx.einsum("bshd,btd->bsht", q, rows.astype(mx.float32)) return mx.sum(mx.maximum(score, 0.0) * weights[..., None], axis=2) # [b,s,t] @@ -2332,7 +2381,9 @@ def __call__( # Causality: window c holds tokens [c*ratio, (c+1)*ratio), so query p may use # it once p has completed it — the same rule the dense mask uses. - causal = (mx.arange(n_comp)[None, :] < ((positions[:, None] + 1) // ratio))[None] + causal = (mx.arange(n_comp)[None, :] < ((positions[:, None] + 1) // ratio))[ + None + ] key = mx.where(causal, score, mx.array(-float("inf"), mx.float32)) k_row = mx.minimum( mx.sum(causal.astype(mx.int32), axis=-1, keepdims=True), self.index_topk @@ -2386,11 +2437,11 @@ def __init__( if self.ratio <= 0 else (2 if self.overlap else 1) * self.ratio + self.rollback_capacity ) - self.cur_kv: Optional[mx.array] = None # [b, offset % ratio, coff*head_dim] - self.cur_score: Optional[mx.array] = None # same, post-``ape`` - self.prev_kv: Optional[mx.array] = None # [b, ratio, coff*head_dim] (overlap) + self.cur_kv: Optional[mx.array] = None # [b, offset % ratio, coff*head_dim] + self.cur_score: Optional[mx.array] = None # same, post-``ape`` + self.prev_kv: Optional[mx.array] = None # [b, ratio, coff*head_dim] (overlap) self.prev_score: Optional[mx.array] = None - self.tail_kv: Optional[mx.array] = None # [b, <=rollback_rows, coff*head_dim] + self.tail_kv: Optional[mx.array] = None # [b, <=rollback_rows, coff*head_dim] self.tail_score: Optional[mx.array] = None self.n_emitted = 0 @@ -2408,15 +2459,17 @@ def push_rollback_rows(self, kv: mx.array, score: mx.array) -> None: """Append this step's freshly projected rows to the bounded journal.""" if self.rollback_rows <= 0 or kv.shape[1] == 0: return - self.tail_kv = kv if self.tail_kv is None else mx.concatenate( - [self.tail_kv, kv], axis=1 + self.tail_kv = ( + kv if self.tail_kv is None else mx.concatenate([self.tail_kv, kv], axis=1) ) - self.tail_score = score if self.tail_score is None else mx.concatenate( - [self.tail_score, score], axis=1 + self.tail_score = ( + score + if self.tail_score is None + else mx.concatenate([self.tail_score, score], axis=1) ) if self.tail_kv.shape[1] > self.rollback_rows: - self.tail_kv = self.tail_kv[:, -self.rollback_rows:] - self.tail_score = self.tail_score[:, -self.rollback_rows:] + self.tail_kv = self.tail_kv[:, -self.rollback_rows :] + self.tail_score = self.tail_score[:, -self.rollback_rows :] def rollback(self, n: int, new_offset: int) -> None: """Rewind ``n`` token positions; ``new_offset`` is the resulting offset. @@ -2448,12 +2501,12 @@ def rollback(self, n: int, new_offset: int) -> None: self.tail_kv = self.tail_kv[:, :kept] self.tail_score = self.tail_score[:, :kept] self.n_emitted = int(new_offset) // self.ratio - self.cur_kv = None if r == 0 else self.tail_kv[:, kept - r:] - self.cur_score = None if r == 0 else self.tail_score[:, kept - r:] + self.cur_kv = None if r == 0 else self.tail_kv[:, kept - r :] + self.cur_score = None if r == 0 else self.tail_score[:, kept - r :] if self.overlap and self.n_emitted > 0: lo = kept - r - self.ratio - self.prev_kv = self.tail_kv[:, lo: lo + self.ratio] - self.prev_score = self.tail_score[:, lo: lo + self.ratio] + self.prev_kv = self.tail_kv[:, lo : lo + self.ratio] + self.prev_score = self.tail_score[:, lo : lo + self.ratio] else: self.prev_kv = None self.prev_score = None @@ -2520,8 +2573,8 @@ def __init__( self.head_dim = int(head_dim) self.rollback_capacity = max(0, int(rollback_capacity)) self.offset = 0 - self.window: Optional[mx.array] = None # [b, L, head_dim] - self.window_start = 0 # abs position of window[:, 0] + self.window: Optional[mx.array] = None # [b, L, head_dim] + self.window_start = 0 # abs position of window[:, 0] self.compressed: Optional[mx.array] = None # [b, n_comp, head_dim] overlap = self.compress_ratio == 4 self.comp = CompressorState( @@ -2543,7 +2596,9 @@ def n_compressed(self) -> int: @property def n_index_compressed(self) -> int: - return 0 if self.index_compressed is None else int(self.index_compressed.shape[1]) + return ( + 0 if self.index_compressed is None else int(self.index_compressed.shape[1]) + ) def update_window(self, kv: mx.array): """Append ``kv`` (positions ``offset..offset+s-1``) and return the rows this @@ -2780,12 +2835,14 @@ def __init__(self, args: ModelArgs, layer_id: int): self.window_size = args.window_size self.eps = args.rms_norm_eps self.compress_ratio = args.compress_ratios[layer_id] - self.softmax_scale = self.head_dim ** -0.5 + self.softmax_scale = self.head_dim**-0.5 self.attn_sink = mx.zeros((self.n_heads,)) self.wq_a = nn.Linear(self.dim, self.q_lora_rank, bias=False) self.q_norm = nn.RMSNorm(self.q_lora_rank, eps=self.eps) - self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.head_dim, bias=False) + self.wq_b = nn.Linear( + self.q_lora_rank, self.n_heads * self.head_dim, bias=False + ) self.wkv = nn.Linear(self.dim, self.head_dim, bias=False) self.kv_norm = nn.RMSNorm(self.head_dim, eps=self.eps) # o-LoRA: grouped down-projection (block matmul) then a dense up-projection. @@ -2818,8 +2875,12 @@ def __init__(self, args: ModelArgs, layer_id: int): # ratio==0 layers use base rope_theta with no YaRN. if self.compress_ratio: inv = _yarn_inv_freq( - self.rope_head_dim, args.compress_rope_theta, args.original_seq_len, - args.rope_factor, args.beta_fast, args.beta_slow, + self.rope_head_dim, + args.compress_rope_theta, + args.original_seq_len, + args.rope_factor, + args.beta_fast, + args.beta_slow, ) else: inv = _yarn_inv_freq(self.rope_head_dim, args.rope_theta, 0, 1.0, 32, 1) @@ -2917,9 +2978,7 @@ def install_o_lora_route(self, mode: str | None = None) -> dict: "direct": direct, "groups": int(self.n_groups), "rank": int(self.o_lora_rank), - "per_group_input": int( - self.n_heads * self.head_dim // self.n_groups - ), + "per_group_input": int(self.n_heads * self.head_dim // self.n_groups), } def _o_lora_dense(self, o: mx.array) -> mx.array: @@ -2978,7 +3037,7 @@ def _attn_mask( if kv_pos is not None: i = q_pos[:, None] j = kv_pos[None, :] - parts.append(((j <= i) & (j > i - self.window_size))[None]) # [1, s, n_win] + parts.append(((j <= i) & (j > i - self.window_size))[None]) # [1, s, n_win] elif n_win: parts.append(mx.ones((1, s, n_win), dtype=mx.bool_)) if n_comp: @@ -3107,16 +3166,28 @@ def __call__(self, x: mx.array, mask=None, cache=None) -> mx.array: qr = self.q_norm(self.wq_a(x)) q = self.wq_b(qr).reshape(b, s, self.n_heads, self.head_dim) # per-head RMS-like normalisation (no learned weight), reference L498 - q = q * mx.rsqrt(mx.mean(mx.square(q.astype(mx.float32)), axis=-1, keepdims=True) + self.eps) + q = q * mx.rsqrt( + mx.mean(mx.square(q.astype(mx.float32)), axis=-1, keepdims=True) + self.eps + ) q = q.astype(x.dtype) q = mx.concatenate( - [q[..., :-rd], _apply_interleaved_rope(q[..., -rd:], cos[None, :, None, :], sin[None, :, None, :])], + [ + q[..., :-rd], + _apply_interleaved_rope( + q[..., -rd:], cos[None, :, None, :], sin[None, :, None, :] + ), + ], axis=-1, ) kv = self.kv_norm(self.wkv(x)) # [b, s, head_dim] (single shared KV — MQA) kv = mx.concatenate( - [kv[..., :-rd], _apply_interleaved_rope(kv[..., -rd:], cos[None, :, :], sin[None, :, :])], + [ + kv[..., :-rd], + _apply_interleaved_rope( + kv[..., -rd:], cos[None, :, :], sin[None, :, :] + ), + ], axis=-1, ) @@ -3130,7 +3201,9 @@ def __call__(self, x: mx.array, mask=None, cache=None) -> mx.array: kvc = self.compressor(x) # [b, n_comp, head_dim] n_comp = kvc.shape[1] if n_comp: - full_kv = mx.concatenate([kv, kvc], axis=1) # [b, s+n_comp, head_dim] + full_kv = mx.concatenate( + [kv, kvc], axis=1 + ) # [b, s+n_comp, head_dim] if self._indexer_active(n_comp): # No cache to keep, so the indexer's compressor only runs when # its rows are actually about to be scored. @@ -3151,8 +3224,10 @@ def __call__(self, x: mx.array, mask=None, cache=None) -> mx.array: win_kv, win_start = cache.update_window(kv) n_comp = cache.n_compressed n_win = int(win_kv.shape[1]) - full_kv = win_kv if not n_comp else mx.concatenate( - [win_kv, cache.compressed], axis=1 + full_kv = ( + win_kv + if not n_comp + else mx.concatenate([win_kv, cache.compressed], axis=1) ) if self._indexer_active(n_comp): assert cache.n_index_compressed == n_comp, ( @@ -3163,23 +3238,28 @@ def __call__(self, x: mx.array, mask=None, cache=None) -> mx.array: # s == 1: update_window already dropped every row outside the query's # window, so that half needs no mask (the compressed half still does once # the indexer is filtering). - kv_pos = None if s == 1 else mx.arange( - win_start, win_start + win_kv.shape[1] + kv_pos = ( + None if s == 1 else mx.arange(win_start, win_start + win_kv.shape[1]) ) cache.advance(s) - q_t = q.transpose(0, 2, 1, 3) # [b, h, s, head_dim] + q_t = q.transpose(0, 2, 1, 3) # [b, h, s, head_dim] # ``full_kv`` is [b, s+n_comp, head_dim] and shared over heads (MQA). # q and the KV block always carry the same dtype (both follow x, or both # follow the fp32 escape hatch), so either one names the score dtype. add = self._attn_mask( positions, kv_pos, n_win, n_comp, ratio, q_t.dtype, comp_sel=comp_sel ) - o = self._attend(q_t, full_kv, add) # [b, h, s, head_dim] - o = o.transpose(0, 2, 1, 3) # [b, s, h, head_dim] + o = self._attend(q_t, full_kv, add) # [b, h, s, head_dim] + o = o.transpose(0, 2, 1, 3) # [b, s, h, head_dim] # de-rotate the tail dims (reference L534, inverse rope) o = mx.concatenate( - [o[..., :-rd], _apply_interleaved_rope(o[..., -rd:], cos[None, :, None, :], -sin[None, :, None, :])], + [ + o[..., :-rd], + _apply_interleaved_rope( + o[..., -rd:], cos[None, :, None, :], -sin[None, :, None, :] + ), + ], axis=-1, ) o = o.reshape(b, s, self.n_heads * self.head_dim) @@ -3254,8 +3334,8 @@ def __init__(self, limit: float = 0.0): def __call__(self, x: mx.array, gate: mx.array) -> mx.array: if self.limit > 0: - x = mx.clip(x, -self.limit, self.limit) # up: two-sided - gate = mx.minimum(gate, self.limit) # gate: upper tail only + x = mx.clip(x, -self.limit, self.limit) # up: two-sided + gate = mx.minimum(gate, self.limit) # gate: upper tail only return super().__call__(x, gate) @@ -3302,7 +3382,9 @@ def __call__(self, x: mx.array, input_ids: Optional[mx.array] = None): indices = self.tid2eid[input_ids.reshape(-1)] # [n, topk] else: biased = scores + self.e_score_correction_bias - indices = mx.argpartition(-biased, kth=self.topk - 1, axis=-1)[..., : self.topk] + indices = mx.argpartition(-biased, kth=self.topk - 1, axis=-1)[ + ..., : self.topk + ] weights = mx.take_along_axis(scores, indices, axis=-1) if self.score_func != "softmax": weights = weights / (mx.sum(weights, axis=-1, keepdims=True)) @@ -3471,8 +3553,8 @@ def __call__( sibling appended-layer backends make (GLM's modulo-into-layers, Hy3's single NextN layer reused at every depth). """ - e = self.enorm(embed_tokens(input_ids)) # [b, s, dim] - x = self.hnorm(h) # [b, s, hc, dim] + e = self.enorm(embed_tokens(input_ids)) # [b, s, dim] + x = self.hnorm(h) # [b, s, hc, dim] x = self.e_proj(e)[:, :, None, :] + self.h_proj(x) x = super().__call__(x, mask=None, cache=cache, input_ids=input_ids) logits = lm_head(self.norm(self.hc_head(x))) @@ -3492,7 +3574,9 @@ def get_dspark_topk_idxs( main = mx.arange(min(int(window_size), int(start_pos) + 1), dtype=mx.int32) draft = int(window_size) + mx.arange(int(block_size), dtype=mx.int32) row = mx.concatenate([main, draft]) - return mx.broadcast_to(row[None, None, :], (int(batch_size), int(block_size), row.shape[0])) + return mx.broadcast_to( + row[None, None, :], (int(batch_size), int(block_size), row.shape[0]) + ) class DeepseekV4DSparkCache: @@ -3515,8 +3599,12 @@ def prefill(self, main_kv: mx.array) -> None: else: last = main_kv[:, -win:] cutoff = seqlen % win - self.ring = last if cutoff == 0 else mx.concatenate( - [last[:, win - cutoff:], last[:, : win - cutoff]], axis=1 + self.ring = ( + last + if cutoff == 0 + else mx.concatenate( + [last[:, win - cutoff :], last[:, : win - cutoff]], axis=1 + ) ) self.prefill_length = int(seqlen) @@ -3541,9 +3629,7 @@ def commit_main(self, start_pos: int, main_kv: mx.array) -> None: ) remaining = rows - first if remaining: - ring = mx.concatenate( - [main_kv[:, first:], ring[:, remaining:]], axis=1 - ) + ring = mx.concatenate([main_kv[:, first:], ring[:, remaining:]], axis=1) self.ring = ring def replace_main(self, start_pos: int, main_kv: mx.array) -> None: @@ -3566,7 +3652,10 @@ def _kv(self, x: mx.array, positions: mx.array) -> mx.array: cos, sin = self._rope_tables(positions) kv = self.kv_norm(self.wkv(x)) return mx.concatenate( - [kv[..., :-rd], _apply_interleaved_rope(kv[..., -rd:], cos[None], sin[None])], + [ + kv[..., :-rd], + _apply_interleaved_rope(kv[..., -rd:], cos[None], sin[None]), + ], axis=-1, ) @@ -3589,7 +3678,9 @@ def __call__( raise ValueError("DSpark decode requires one complete five-token block") cache.replace_main(start_pos, main_kv) block = int(x.shape[1]) - positions = mx.arange(int(start_pos) + main_len, int(start_pos) + main_len + block) + positions = mx.arange( + int(start_pos) + main_len, int(start_pos) + main_len + block + ) cos, sin = self._rope_tables(positions) rd = self.rope_head_dim @@ -3600,7 +3691,12 @@ def __call__( ) q = q.astype(x.dtype) q = mx.concatenate( - [q[..., :-rd], _apply_interleaved_rope(q[..., -rd:], cos[None, :, None], sin[None, :, None])], + [ + q[..., :-rd], + _apply_interleaved_rope( + q[..., -rd:], cos[None, :, None], sin[None, :, None] + ), + ], axis=-1, ) draft_kv = self._kv(x, positions) @@ -3612,7 +3708,12 @@ def __call__( o = self._attend(q.transpose(0, 2, 1, 3), visible_kv, None) o = o.transpose(0, 2, 1, 3) o = mx.concatenate( - [o[..., :-rd], _apply_interleaved_rope(o[..., -rd:], cos[None, :, None], -sin[None, :, None])], + [ + o[..., :-rd], + _apply_interleaved_rope( + o[..., -rd:], cos[None, :, None], -sin[None, :, None] + ), + ], axis=-1, ) return self._o_lora(o.reshape(b, block, self.n_heads * self.head_dim)) @@ -3673,14 +3774,17 @@ def __init__(self, args: ModelArgs, stage_id: int): self.confidence_head = None if stage_id == 0: self.main_proj = nn.Linear( - args.hidden_size * len(args.dspark_target_layer_ids), args.hidden_size, + args.hidden_size * len(args.dspark_target_layer_ids), + args.hidden_size, bias=False, ) self.main_norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) if stage_id == _DSPARK_STAGE_COUNT - 1: self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) self.hc_head = HeadHC(args.hidden_size, args.hc_mult, args.hc_eps) - self.markov_head = DSparkMarkovHead(args.vocab_size, args.dspark_markov_rank) + self.markov_head = DSparkMarkovHead( + args.vocab_size, args.dspark_markov_rank + ) self.confidence_head = DSparkConfidenceHead( args.hidden_size, args.dspark_markov_rank ) @@ -3701,7 +3805,12 @@ def prefill(self, h: mx.array, cache, main_x: mx.array) -> mx.array: return h def __call__( - self, h: mx.array, *, start_pos: int, cache=None, input_ids=None, + self, + h: mx.array, + *, + start_pos: int, + cache=None, + input_ids=None, main_x=None, ) -> mx.array: if int(start_pos) == 0: @@ -3753,7 +3862,10 @@ def _validate_dspark_manifest(args: ModelArgs) -> None: raise ValueError("DSpark-0731 requires num_nextn_predict_layers=1") if int(args.dspark_noise_token_id or 0) != _DSPARK_NOISE_TOKEN_ID: raise ValueError("DSpark-0731 requires dspark_noise_token_id=128799") - if tuple(int(x) for x in (args.dspark_target_layer_ids or ())) != _DSPARK_TARGET_LAYER_IDS: + if ( + tuple(int(x) for x in (args.dspark_target_layer_ids or ())) + != _DSPARK_TARGET_LAYER_IDS + ): raise ValueError("DSpark-0731 requires target taps (40, 41, 42)") if args.num_hidden_layers <= _DSPARK_TARGET_LAYER_IDS[-1]: raise ValueError("DSpark-0731 target taps are absent from this trunk") @@ -3762,7 +3874,9 @@ def _validate_dspark_manifest(args: ModelArgs) -> None: if int(args.dspark_markov_rank or 0) != _DSPARK_MARKOV_RANK: raise ValueError("DSpark-0731 requires dspark_markov_rank=256") ratios = list(args.compress_ratios) - for layer_id in range(args.num_hidden_layers, args.num_hidden_layers + _DSPARK_STAGE_COUNT): + for layer_id in range( + args.num_hidden_layers, args.num_hidden_layers + _DSPARK_STAGE_COUNT + ): if layer_id < len(ratios) and int(ratios[layer_id]) != 0: raise ValueError("DSpark-0731 stages require uncompressed attention") @@ -3790,13 +3904,18 @@ def __init__(self, args: ModelArgs): self.block_size = _DSPARK_BLOCK_SIZE self.noise_token_id = _DSPARK_NOISE_TOKEN_ID self.target_layer_ids = _DSPARK_TARGET_LAYER_IDS - self.stages = [DeepseekV4DSparkStage(args, i) for i in range(_DSPARK_STAGE_COUNT)] + self.stages = [ + DeepseekV4DSparkStage(args, i) for i in range(_DSPARK_STAGE_COUNT) + ] def draft_input_ids(self, target_ids: mx.array) -> mx.array: if target_ids.ndim != 1: raise ValueError("DSpark target ids must be a [batch] tensor") - noise = mx.full((target_ids.shape[0], self.block_size), self.noise_token_id, - dtype=target_ids.dtype) + noise = mx.full( + (target_ids.shape[0], self.block_size), + self.noise_token_id, + dtype=target_ids.dtype, + ) return mx.concatenate([target_ids[:, None], noise[:, 1:]], axis=1) def make_cache(self) -> list: @@ -3837,8 +3956,13 @@ def commit_main(self, main_hidden: mx.array, caches, *, start_pos: int) -> None: cache.commit_main(start_pos, stage.attn._kv(main_x, positions)) def finish( - self, logits: mx.array, hidden: mx.array, target_ids: mx.array, - *, greedy: bool = False, key=None, + self, + logits: mx.array, + hidden: mx.array, + target_ids: mx.array, + *, + greedy: bool = False, + key=None, ) -> Tuple[mx.array, mx.array, mx.array]: """Apply the sequential Markov recurrence and return fp32 confidence.""" final = self.stages[-1] @@ -3850,7 +3974,11 @@ def finish( biased_rows = [] markov_embeds = [] previous = target_ids - keys = [None] * self.block_size if key is None else list(mx.random.split(key, self.block_size)) + keys = ( + [None] * self.block_size + if key is None + else list(mx.random.split(key, self.block_size)) + ) for i in range(self.block_size): bias, markov_embed = final.markov_head(previous) row = logits[:, i] + bias @@ -3929,14 +4057,21 @@ def forward( if caches is None: caches = self.make_cache() if len(caches) != _DSPARK_STAGE_COUNT: - raise ValueError("DSpark requires one cache owned by each of its three stages") + raise ValueError( + "DSpark requires one cache owned by each of its three stages" + ) main_x = self.stages[0].fuse_main(main_hidden) ids = self.draft_input_ids(target_ids) h = embed_tokens(ids) - h = mx.broadcast_to(h[:, :, None, :], (*h.shape[:2], self.args.hc_mult, h.shape[-1])) + h = mx.broadcast_to( + h[:, :, None, :], (*h.shape[:2], self.args.hc_mult, h.shape[-1]) + ) for stage, cache in zip(self.stages, caches): h = stage( - h, start_pos=start_pos, cache=cache, input_ids=ids, + h, + start_pos=start_pos, + cache=cache, + input_ids=ids, main_x=main_x, ) if int(start_pos) == 0: @@ -3975,7 +4110,9 @@ class _DSparkTargetRoute: def __call__(self, owner, inputs: mx.array, cache): h = owner.model.embed_tokens(inputs) - h = mx.broadcast_to(h[:, :, None, :], (*h.shape[:2], owner.args.hc_mult, h.shape[-1])) + h = mx.broadcast_to( + h[:, :, None, :], (*h.shape[:2], owner.args.hc_mult, h.shape[-1]) + ) if cache is None: cache = [None] * len(owner.model.layers) taps = [] @@ -4058,7 +4195,9 @@ def __init__(self, args: ModelArgs): DeepseekV4MTP(args, args.num_hidden_layers + i) for i in range(max(int(args.num_nextn_predict_layers or 0), 0)) ] - self._target_hidden_route = _DSparkTargetRoute() if self._dspark else _LegacyTargetRoute() + self._target_hidden_route = ( + _DSparkTargetRoute() if self._dspark else _LegacyTargetRoute() + ) # Construction-time performance installers may replace this with a # typed phase/width router. The stock callable is explicit and direct; # decoder layers never probe candidate eligibility or fall back. @@ -4107,7 +4246,7 @@ def __call__( if emit_logits: source = h if logits_keep is not None: - source = h[:, -max(1, int(logits_keep)):] + source = h[:, -max(1, int(logits_keep)) :] logits = self.logits_from_hc_hidden(source) if not return_hidden: return logits @@ -4348,11 +4487,34 @@ def sanitize(self, weights: dict) -> dict: if target.endswith(source_suffix): target = target[: -len(source_suffix)] + target_suffix break + # Flash-0731 preserves o-LoRA's explicit group axis in storage: + # [o_groups, o_lora_rank, packed_input] + # QuantizedLinear owns the identical row-major matrix as + # [o_groups * o_lora_rank, packed_input]. + # Collapse the two logical row axes once at load; the byte order and + # therefore the group/rank ownership are unchanged. Older exports + # already use the 2-D form and pass through untouched. + if ( + (target.startswith("model.layers.") or target.startswith("mtp.")) + and any( + target.endswith(f".attn.wo_a.{field}") + for field in ("weight", "scales", "biases") + ) + and value.ndim == 3 + ): + expected = (self.args.o_groups, self.args.o_lora_rank) + if tuple(value.shape[:2]) != expected: + raise ValueError( + f"invalid grouped 0731 o-LoRA storage for {target}: " + f"expected leading axes {expected}, got {tuple(value.shape)}" + ) + value = value.reshape(expected[0] * expected[1], value.shape[-1]) translated[target] = value weights = translated if self._dspark is not None: missing = [ - stage_id for stage_id in range(_DSPARK_STAGE_COUNT) + stage_id + for stage_id in range(_DSPARK_STAGE_COUNT) if not any(str(k).startswith(f"mtp.{stage_id}.") for k in weights) ] if missing: @@ -4492,9 +4654,7 @@ def __init__(self, blocks): } -def _require_o_lora_array( - value, *, label: str, shape: tuple[int, int], dtype -) -> None: +def _require_o_lora_array(value, *, label: str, shape: tuple[int, int], dtype) -> None: if tuple(getattr(value, "shape", ())) != shape: raise ValueError( f"{label} shape {tuple(getattr(value, 'shape', ()))} does not match {shape}" @@ -4517,9 +4677,7 @@ def _require_canonical_quantized_linear( ): observed = getattr(linear, attribute, None) if observed != expected: - raise ValueError( - f"{label} {attribute}={observed!r}, expected {expected!r}" - ) + raise ValueError(f"{label} {attribute}={observed!r}, expected {expected!r}") weight = getattr(linear, "weight", None) scales = getattr(linear, "scales", None) @@ -4582,9 +4740,7 @@ def _require_canonical_quantized_linear( ) -def _require_canonical_dense_linear( - linear, *, label: str, shape: tuple[int, int] -): +def _require_canonical_dense_linear(linear, *, label: str, shape: tuple[int, int]): if not isinstance(linear, nn.Linear) or isinstance(linear, nn.QuantizedLinear): raise ValueError(f"{label} must be a dense nn.Linear, not QuantizedLinear") weight = getattr(linear, "weight", None) @@ -4597,9 +4753,7 @@ def _require_canonical_dense_linear( if getattr(linear, "bias", None) is not None: raise ValueError(f"{label} additive bias must be absent") accidental = [ - attribute - for attribute in _O_LORA_QUANT_FIELDS - if hasattr(linear, attribute) + attribute for attribute in _O_LORA_QUANT_FIELDS if hasattr(linear, attribute) ] if accidental: raise ValueError( @@ -4765,15 +4919,13 @@ def install_deepseek_v4_o_lora_routes( """ trunk = [layer.attn for layer in model.layers] mtp = [block.attn for block in model.mtp_blocks] - selected = ( - str(mode) - if mode is not None - else _o_lora_mode_from_env() - ) + selected = str(mode) if mode is not None else _o_lora_mode_from_env() if selected not in _O_LORA_MODES: raise ValueError(f"unsupported o-LoRA route {selected!r}") if not canonical_mixed_route: - reports = [attention.install_o_lora_route(selected) for attention in trunk + mtp] + reports = [ + attention.install_o_lora_route(selected) for attention in trunk + mtp + ] return { "mode": selected, "module_count": len(reports), @@ -4800,15 +4952,11 @@ def install_deepseek_v4_o_lora_routes( else None ) body_route_type = ( - _DirectGatherOLoraWideM4 - if selected == "gather_qmm" - else _DirectCachedOLora + _DirectGatherOLoraWideM4 if selected == "gather_qmm" else _DirectCachedOLora ) if selected == "gather_qmm": body_impls = [ - body_route_type( - attention, quant, activation_dtype=activation_dtype - ) + body_route_type(attention, quant, activation_dtype=activation_dtype) for attention, quant in zip(trunk, body_quant) ] else: @@ -4848,9 +4996,7 @@ def install_deepseek_v4_o_lora_routes( callable_census = { "body_route_objects": len(body_impls), "body_route_kind": ( - "gather_qmm_m4_wide_direct" - if selected == "gather_qmm" - else "cached_direct" + "gather_qmm_m4_wide_direct" if selected == "gather_qmm" else "cached_direct" ), "body_callable_class": body_route_type.__name__, "mtp_route_objects": len(mtp_impls), diff --git a/mtplx/native_block_speculation.py b/mtplx/native_block_speculation.py new file mode 100644 index 00000000..833f5e3b --- /dev/null +++ b/mtplx/native_block_speculation.py @@ -0,0 +1,536 @@ +"""Generic serial-exact greedy speculation for fixed-block proposal backends. + +The backend owns only model-specific proposal/cache operations. This module +owns the target protocol: first-token gating, serial-M1 draft verification, +callbacks, and common generation statistics. +""" + +from __future__ import annotations + +import time +from collections.abc import Callable +from typing import Any, Protocol + +import mlx.core as mx +import numpy as np + +from .attention_context import attention_phase +from .sampling import SamplerConfig + + +class NativeBlockSpeculativeBackend(Protocol): + """Construction-installed operations for one native block proposer.""" + + backend_id: str + supported_depths: tuple[int, ...] + minimum_proposal_target_position: int + prefill_chunk_size: int + + def make_cache(self, rt: Any) -> Any: ... + + def prefill(self, rt: Any, hidden: mx.array, cache: Any) -> None: ... + + def prefill_chunk( + self, + rt: Any, + hidden: mx.array, + cache: Any, + *, + start_pos: int, + ) -> None: ... + + def propose( + self, + rt: Any, + hidden: mx.array, + token_id: int, + primary_token_id: int, + cache: Any, + *, + start_pos: int, + width: int, + ) -> mx.array: ... + + def commit( + self, + rt: Any, + hidden: mx.array, + cache: Any, + *, + start_pos: int, + ) -> None: ... + + def cache_roots(self, cache: Any) -> list[mx.array]: ... + + def bind_target_forward(self, rt: Any) -> Callable[..., Any]: ... + + def snapshot(self, cache: Any) -> Any: ... + + def restore(self, cache: Any, snapshot: Any) -> None: ... + + def rollback_target(self, target_cache: Any, rejected_rows: int) -> None: ... + + +def _validate_request( + backend: NativeBlockSpeculativeBackend, + sampler: SamplerConfig, + speculative_depth: int, + *, + draft_sampler: SamplerConfig | None, + constraint: Any | None, + vision_splice: Any | None, + adaptive_policy: Any | None, + adaptive_width_policy: Any | None, +) -> None: + """Validate the fixed greedy lane once, before prompt execution.""" + if ( + float(sampler.temperature) > 0.0 + or float(sampler.presence_penalty) != 0.0 + or float(sampler.frequency_penalty) != 0.0 + ): + raise ValueError( + f"{backend.backend_id} currently requires greedy target sampling" + ) + if draft_sampler is not None and ( + float(draft_sampler.temperature) > 0.0 + or float(draft_sampler.presence_penalty) != 0.0 + or float(draft_sampler.frequency_penalty) != 0.0 + ): + raise ValueError( + f"{backend.backend_id} currently requires greedy draft sampling" + ) + if int(speculative_depth) not in backend.supported_depths: + supported = ", ".join(str(value) for value in backend.supported_depths) + raise ValueError( + f"{backend.backend_id} measured proposal width must be one of: {supported}" + ) + if constraint is not None: + raise ValueError( + f"{backend.backend_id} does not yet support constrained decoding" + ) + if vision_splice is not None: + raise ValueError(f"{backend.backend_id} does not support vision input") + if adaptive_policy is not None or adaptive_width_policy is not None: + raise ValueError( + f"{backend.backend_id} width is fixed for each benchmark request" + ) + + +def generate_native_block_speculative( + rt: Any, + backend: NativeBlockSpeculativeBackend, + prompt_ids: list[int], + *, + abort_check: Callable[[], bool] | None, + max_tokens: int, + sampler: SamplerConfig, + speculative_depth: int, + seed: int, + stop_token_ids: set[int] | None, + draft_sampler: SamplerConfig | None, + token_callback: Callable[[list[int]], None] | None, + prefill_callback: Callable[[dict[str, Any]], None] | None, + constraint: Any | None, + vision_splice: Any | None, + adaptive_policy: Any | None, + adaptive_width_policy: Any | None, +): + """Run a construction-installed native block proposer against the target. + + Depth two means two genuinely future DSpark drafts. The already sampled + target primary plus accepted drafts advance the target through serial M1 + calls. The target remains the sole token and state authority; a rejected + draft never enters its cache. + """ + from .generation import ( + GenerationOutput, + GenerationStats, + _attach_runtime_diagnostics, + _decode, + _default_stop_tokens, + _eval, + _finish_reason_from_tokens, + _final_logits_prefill_enabled, + _generation_rate_fields, + _is_stop, + _iter_prefill_chunk_spans, + _make_target_prefill_cache, + _mean_accept_probability_by_depth, + _runtime_counter_snapshot, + _strip_terminal_stop, + ) + + _validate_request( + backend, + sampler, + speculative_depth, + draft_sampler=draft_sampler, + constraint=constraint, + vision_splice=vision_splice, + adaptive_policy=adaptive_policy, + adaptive_width_policy=adaptive_width_policy, + ) + # Bind the construction-certified target callable once. The loop invokes + # this exact callable for every serial M1 target row. + target_forward = backend.bind_target_forward(rt) + if not prompt_ids: + raise ValueError("prompt_ids must not be empty") + if max_tokens < 0: + raise ValueError("max_tokens must be non-negative") + + counter_start = _runtime_counter_snapshot(rt) + stop_token_ids = ( + _default_stop_tokens(rt.tokenizer) if stop_token_ids is None else stop_token_ids + ) + del seed + started_all = time.perf_counter() + prefill_started = started_all + tokens: list[int] = [] + accepted_drafts = 0 + drafted_tokens = 0 + rejected_drafts = 0 + # With hot event dictionaries disabled, this is the aggregate count of + # main-loop proposal/verification cycles. It is the denominator consumed + # by acceptance reporting, not a physical target-forward count. It includes + # terminal/tail cycles but excludes the one-time position seed. + verify_calls = 0 + prompt_target_time = 0.0 + prompt_proposal_time = 0.0 + prompt_eval_time = 0.0 + accepted_by_depth = [0] * int(speculative_depth) + drafted_by_depth = [0] * int(speculative_depth) + target_position = len(prompt_ids) + last_target_token = int(prompt_ids[-1]) + + def finish() -> GenerationOutput: + elapsed = max(0.0, time.perf_counter() - started_all) + stats = GenerationStats( + mode="mtpk", + generated_tokens=len(tokens), + elapsed_s=elapsed, + **_generation_rate_fields( + generated_tokens=len(tokens), + elapsed_s=elapsed, + prompt_eval_time_s=prompt_eval_time, + ), + accepted_drafts=accepted_drafts, + rejected_drafts=rejected_drafts, + drafted_tokens=drafted_tokens, + verify_time_s=0.0, + verify_forward_time_s=0.0, + verify_eval_time_s=0.0, + verify_joint_eval_time_s=0.0, + draft_time_s=0.0, + target_forward_time_s=0.0, + prompt_eval_time_s=prompt_eval_time, + prompt_tps=( + len(prompt_ids) / prompt_eval_time if prompt_eval_time > 0 else 0.0 + ), + prompt_target_prefill_time_s=prompt_target_time, + prompt_mtp_history_time_s=prompt_proposal_time, + prompt_target_prefill_tok_s=( + len(prompt_ids) / prompt_target_time if prompt_target_time > 0 else 0.0 + ), + accept_time_s=0.0, + rollback_time_s=0.0, + repair_time_s=0.0, + commit_time_s=0.0, + peak_memory_bytes=mx.get_peak_memory(), + speculative_depth=int(speculative_depth), + requested_speculative_depth=int(speculative_depth), + accepted_by_depth=accepted_by_depth, + drafted_by_depth=drafted_by_depth, + mean_accept_probability_by_depth=_mean_accept_probability_by_depth( + [float(value) for value in accepted_by_depth], drafted_by_depth + ), + verify_calls=verify_calls, + events=[], + ) + _attach_runtime_diagnostics(stats, rt, counter_start) + return GenerationOutput( + tokens=tokens, + text=_decode(rt.tokenizer, _strip_terminal_stop(tokens, stop_token_ids)), + stats=stats, + finish_reason=_finish_reason_from_tokens( + tokens, stop_token_ids=stop_token_ids, max_tokens=max_tokens + ), + ) + + if prefill_callback is not None: + try: + prefill_callback( + { + "phase": "started", + "tokens_done": 0, + "tokens_total": len(prompt_ids), + "cached_tokens": 0, + "new_prefill_tokens": len(prompt_ids), + "elapsed_s": 0.0, + "started_s": prefill_started, + } + ) + except Exception: + pass + + target_cache = _make_target_prefill_cache(rt) + proposal_cache = backend.make_cache(rt) + logits = None + current_hidden = None + prompt_length = len(prompt_ids) + body_length = prompt_length - 1 + final_logits_only = _final_logits_prefill_enabled() + # Match AR's cache-only body geometry exactly, then run its mandatory + # one-token logits/hidden tail. DSpark retains its separately measured + # 128-row ring geometry below. + target_spans = [ + *_iter_prefill_chunk_spans(body_length), + (body_length, prompt_length), + ] + proposal_chunk_size = int(backend.prefill_chunk_size) + proposal_start = 0 + proposal_remainder = None + + def prefill_proposal_chunk(hidden: mx.array) -> bool: + nonlocal prompt_eval_time, prompt_proposal_time, proposal_start + if abort_check is not None and abort_check(): + return False + proposal_chunk_started = time.perf_counter() + backend.prefill_chunk( + rt, + hidden, + proposal_cache, + start_pos=proposal_start, + ) + proposal_roots = backend.cache_roots(proposal_cache) + if proposal_roots: + _eval(*proposal_roots) + prompt_proposal_time += max(0.0, time.perf_counter() - proposal_chunk_started) + prompt_eval_time = prompt_target_time + prompt_proposal_time + proposal_start += int(hidden.shape[1]) + return True + + for chunk_start, chunk_end in target_spans: + if abort_check is not None and abort_check(): + return finish() + final_chunk = chunk_end == prompt_length + target_chunk_started = time.perf_counter() + with attention_phase("prefill"): + chunk_logits, chunk_hidden = target_forward( + mx.array([prompt_ids[chunk_start:chunk_end]]), + cache=target_cache, + return_hidden=True, + emit_logits=final_chunk or not final_logits_only, + logits_keep=1 if final_chunk and final_logits_only else None, + ) + if chunk_logits is not None: + _eval(chunk_logits, chunk_hidden) + else: + _eval(chunk_hidden) + prompt_target_time += max(0.0, time.perf_counter() - target_chunk_started) + prompt_eval_time = prompt_target_time + prompt_proposal_time + if final_chunk: + logits = chunk_logits[:, -1, :] + current_hidden = chunk_hidden[:, -1:] + + # Stream every complete proposal block now. Only the cross-target-span + # remainder is retained, so long prompts never accumulate a full hidden + # sequence in the scheduler. + hidden_offset = 0 + hidden_width = int(chunk_hidden.shape[1]) + if proposal_remainder is not None: + needed = proposal_chunk_size - int(proposal_remainder.shape[1]) + taken = min(needed, hidden_width) + proposal_remainder = mx.concatenate( + [proposal_remainder, chunk_hidden[:, :taken]], axis=1 + ) + hidden_offset = taken + if int(proposal_remainder.shape[1]) == proposal_chunk_size: + if not prefill_proposal_chunk(proposal_remainder): + return finish() + proposal_remainder = None + while hidden_offset + proposal_chunk_size <= hidden_width: + proposal_chunk = chunk_hidden[ + :, hidden_offset : hidden_offset + proposal_chunk_size + ] + if not prefill_proposal_chunk(proposal_chunk): + return finish() + hidden_offset += proposal_chunk_size + if hidden_offset < hidden_width: + proposal_remainder = chunk_hidden[:, hidden_offset:] + if final_chunk and proposal_remainder is not None: + if not prefill_proposal_chunk(proposal_remainder): + return finish() + proposal_remainder = None + if prefill_callback is not None: + try: + prefill_elapsed = max(0.0, time.perf_counter() - prefill_started) + compute_tok_s = ( + len(prompt_ids) / prompt_eval_time if prompt_eval_time > 0 else None + ) + wall_tok_s = ( + len(prompt_ids) / prefill_elapsed if prefill_elapsed > 0 else None + ) + prefill_callback( + { + "phase": "completed", + "tokens_total": len(prompt_ids), + "new_prefill_tokens": len(prompt_ids), + "cached_tokens": 0, + "elapsed_s": prefill_elapsed, + "prompt_eval_time_s": prompt_eval_time, + "prefill_tok_s": ( + compute_tok_s if compute_tok_s is not None else wall_tok_s + ), + "prefill_compute_tok_s": compute_tok_s, + "prefill_wall_tok_s": wall_tok_s, + "cache_hit": False, + } + ) + except Exception: + pass + + def emit(block: list[int]) -> None: + if token_callback is None: + return + visible = [token for token in block if not _is_stop(token, stop_token_ids)] + if visible: + token_callback(visible) + + # Some native proposers overload absolute position zero as their prefill + # signal. Prime such a backend once with ordinary target M1 so every + # later proposal sees the real position of its carried hidden. This is a + # construction-selected initial phase, outside the measured fixed-K loop. + minimum_proposal_position = int( + getattr(backend, "minimum_proposal_target_position", 1) + ) + seed_terminal = False + while ( + target_position < minimum_proposal_position + and len(tokens) < max_tokens + and not seed_terminal + ): + if abort_check is not None and abort_check(): + break + current_top = int(mx.argmax(logits[0], axis=-1).item()) + with attention_phase("ar_decode"): + seed_logits, seed_hidden = target_forward( + mx.array([[current_top]]), cache=target_cache, return_hidden=True + ) + _eval(seed_logits, seed_hidden) + current_hidden = seed_hidden[:, -1:] + backend.commit(rt, current_hidden, proposal_cache, start_pos=target_position) + roots = backend.cache_roots(proposal_cache) + if roots: + _eval(*roots) + logits = seed_logits[:, -1, :] + last_target_token = current_top + tokens.append(current_top) + emit([current_top]) + target_position += 1 + seed_terminal = _is_stop(current_top, stop_token_ids) + + while len(tokens) < max_tokens and not seed_terminal: + if abort_check is not None and abort_check(): + break + # The primary is sampled from carried target logits and is therefore + # authoritative, never a DSpark acceptance gate. DSpark still proposes + # K future rows together, but target ownership advances only through + # serial M1 calls. A draft is fed only after the preceding M1 logits + # accept it, so target cache state is exact by construction. + current_top = int(mx.argmax(logits[0], axis=-1).item()) + remaining = max_tokens - len(tokens) + width = min(int(speculative_depth) + 1, remaining) + if _is_stop(current_top, stop_token_ids): + width = 1 + proposal_snapshot = None + future = None + if width > 1: + proposal_snapshot = backend.snapshot(proposal_cache) + future = backend.propose( + rt, + current_hidden, + last_target_token, + current_top, + proposal_cache, + start_pos=target_position, + width=width - 1, + ) + + # Build the authoritative primary M1 before forcing proposal IDs to the + # host. The two graphs are independent given carried hidden/current_top, + # so they share one evaluation boundary without changing target math. + with attention_phase("ar_decode"): + row_logits, row_hidden = target_forward( + mx.array([[current_top]], dtype=mx.int32), + cache=target_cache, + return_hidden=True, + ) + if future is None: + _eval(row_logits, row_hidden) + else: + _eval(future, row_logits, row_hidden) + + accepted_hidden = [row_hidden[:, -1:]] + next_logits = row_logits[:, -1, :] + future_tokens: list[int] = [] + if future is not None: + future_tokens = [int(token) for token in np.asarray(future)[0]] + for index, token in enumerate(future_tokens): + drafted_by_depth[index] += 1 + if _is_stop(token, stop_token_ids): + future_tokens = future_tokens[: index + 1] + width = index + 2 + break + drafted_tokens += len(future_tokens) + proposal_tokens = [current_top, *future_tokens] + for token in future_tokens: + previous_top = int(mx.argmax(next_logits[0], axis=-1).item()) + if token != previous_top: + break + with attention_phase("ar_decode"): + row_logits, row_hidden = target_forward( + mx.array([[token]], dtype=mx.int32), + cache=target_cache, + return_hidden=True, + ) + _eval(row_logits, row_hidden) + next_logits = row_logits[:, -1, :] + accepted_hidden.append(row_hidden[:, -1:]) + if _is_stop(token, stop_token_ids): + break + verify_calls += 1 + accepted = len(accepted_hidden) + verify_hidden = mx.concatenate(accepted_hidden, axis=1) + + rejected_suffix = len(proposal_tokens) - accepted + + accepted_drafts += max(0, accepted - 1) + rejected_drafts += rejected_suffix + for index in range(1, accepted): + accepted_by_depth[index - 1] += 1 + + # Proposal cache state is backend-owned. Discard its speculative state + # and install only the target-verified causal prefix in one commit. + if proposal_snapshot is not None: + backend.restore(proposal_cache, proposal_snapshot) + backend.commit( + rt, + verify_hidden[:, :accepted], + proposal_cache, + start_pos=target_position, + ) + roots = backend.cache_roots(proposal_cache) + if roots: + _eval(*roots) + logits = next_logits + current_hidden = verify_hidden[:, accepted - 1 : accepted] + + committed = proposal_tokens[:accepted] + tokens.extend(committed) + emit(committed) + last_target_token = committed[-1] + target_position += accepted + if _is_stop(committed[-1], stop_token_ids): + break + + return finish() diff --git a/mtplx/runtime.py b/mtplx/runtime.py index fa17c8da..43f36b14 100644 --- a/mtplx/runtime.py +++ b/mtplx/runtime.py @@ -10,7 +10,7 @@ import re import subprocess import sys -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import dataclass, field from pathlib import Path from typing import TYPE_CHECKING, Any @@ -32,6 +32,72 @@ if TYPE_CHECKING: from .a3b_compiled_target_prefix import A3BCompiledTargetPrefixFactory + from .native_block_speculation import NativeBlockSpeculativeBackend + + +DEEPSEEK_V4_DSPARK_BACKEND_ID = "deepseek_v4_dspark_0731" + + +def build_mtpk_request_kwargs( + rt: Any, + *, + common: Mapping[str, Any], + legacy_defaults: Mapping[str, Any], + explicit_legacy: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + """Build one request while preserving explicit DSpark incompatibilities.""" + request = dict(common) + backend = getattr(rt, "block_speculative_backend", None) + if getattr(backend, "backend_id", None) == DEEPSEEK_V4_DSPARK_BACKEND_ID: + request.update(explicit_legacy or {}) + else: + request.update(legacy_defaults) + return request + + +def _snapshot_deepseek_v4_o_lora_routes(model: Any) -> tuple[tuple[Any, str, Any], ...]: + """Capture every trunk and DSpark route before explicit K2 publication.""" + + attentions = tuple(layer.attn for layer in model.layers) + tuple( + block.attn for block in model.mtp_blocks + ) + try: + return tuple( + (attention, str(attention.o_lora_mode), attention._o_lora_impl) + for attention in attentions + ) + except AttributeError as exc: + raise ValueError("DeepSeek-V4-0731 K2 o-LoRA ownership is incomplete") from exc + + +def _restore_deepseek_v4_o_lora_routes( + state: tuple[tuple[Any, str, Any], ...] | None, +) -> list[Exception]: + failures: list[Exception] = [] + for attention, mode, implementation in state or (): + try: + attention.o_lora_mode = mode + except Exception as exc: + failures.append(exc) + try: + attention._o_lora_impl = implementation + except Exception as exc: + failures.append(exc) + return failures + + +def _rollback_deepseek_v4_0731_k2( + prepared: tuple[Any, ...] | None, + o_lora_state: tuple[tuple[Any, str, Any], ...] | None, +) -> list[Exception]: + failures: list[Exception] = [] + for route_bank in reversed(prepared or ()): + try: + route_bank.restore() + except Exception as exc: + failures.append(exc) + failures.extend(_restore_deepseek_v4_o_lora_routes(o_lora_state)) + return failures def _detect_total_system_memory_bytes() -> int | None: @@ -94,6 +160,10 @@ class MTPLXRuntime: deepseek_v4_o_lora_report: dict[str, Any] | None = None deepseek_v4_attn_proj_wide_m3_report: dict[str, Any] | None = None deepseek_v4_attention_island_report: dict[str, Any] | None = None + deepseek_v4_dspark_enabled: bool = False + deepseek_v4_0731_k2_receipt: dict[str, Any] | None = None + block_speculative_backend: NativeBlockSpeculativeBackend | None = None + block_speculative_decode_trace_requested: bool = False a3b_compiled_target_prefix_factory: A3BCompiledTargetPrefixFactory | None = None a3b_whole_moe_installed: bool = False qwen_row_owned_router_report: dict[str, Any] = field(default_factory=dict) @@ -102,15 +172,21 @@ class MTPLXRuntime: init=False, repr=False, ) - _a3b_whole_moe_request_geometry_keys: dict[ - tuple[int, str, str], str - ] = field(default_factory=dict, init=False, repr=False) + _a3b_whole_moe_request_geometry_keys: dict[tuple[int, str, str], str] = field( + default_factory=dict, init=False, repr=False + ) diagnostic_counters: dict[str, int] = field(default_factory=dict) - _forward_ar_supports_emit_logits: bool | None = field(default=None, init=False, repr=False) - _forward_ar_supports_logits_keep: bool | None = field(default=None, init=False, repr=False) + _forward_ar_supports_emit_logits: bool | None = field( + default=None, init=False, repr=False + ) + _forward_ar_supports_logits_keep: bool | None = field( + default=None, init=False, repr=False + ) def _count(self, key: str, amount: int = 1) -> None: - self.diagnostic_counters[key] = int(self.diagnostic_counters.get(key, 0)) + int(amount) + self.diagnostic_counters[key] = int(self.diagnostic_counters.get(key, 0)) + int( + amount + ) @staticmethod def _sequence_len(input_ids: Any) -> int: @@ -162,7 +238,9 @@ def forward_ar( logits_keep: int | None = None, input_embeddings=None, ): - self._count("forward_ar_hidden_calls" if return_hidden else "forward_ar_plain_calls") + self._count( + "forward_ar_hidden_calls" if return_hidden else "forward_ar_plain_calls" + ) if not self.mtp_enabled and return_hidden: raise RuntimeError("return_hidden requires an MTP-patched runtime") if input_embeddings is not None and not self.mtp_enabled: @@ -206,9 +284,7 @@ def forward_ar( # unprimed cache: seeding the compiled graph from its None KV # leaves throws, and its shape differs from a single-token decode # step, forcing a retrace. Prefill stays eager. - compiled = ( - self._compiled_ar_forward(cache) if sequence_len == 1 else None - ) + compiled = self._compiled_ar_forward(cache) if sequence_len == 1 else None if compiled is not None: # Engagement proof: arm A (flag off) must report 0 here, # arm B (on) > 0 — the A/B credits nothing without it. @@ -343,7 +419,9 @@ def draft_mtp( else str(mtp_hidden_variant) ) resolved_concat_order = ( - self.contract.concat_order if concat_order in {None, "auto", "contract"} else concat_order + self.contract.concat_order + if concat_order in {None, "auto", "contract"} + else concat_order ) with mtp_adapter_depth(self.model, mtp_depth): kwargs = { @@ -380,7 +458,9 @@ def update_mtp_cache( else str(mtp_hidden_variant) ) resolved_concat_order = ( - self.contract.concat_order if concat_order in {None, "auto", "contract"} else concat_order + self.contract.concat_order + if concat_order in {None, "auto", "contract"} + else concat_order ) update = getattr(self.model, "mtp_update_cache", None) if update is not None: @@ -570,6 +650,7 @@ def load( gemma4_target_distribution_mode: str | None = None, proj_quant: str | None = None, proj_requant: str | None = None, + deepseek_v4_0731_k2: bool = False, ) -> MTPLXRuntime: """Load an MLX model and optionally inject native MTP support. @@ -579,7 +660,20 @@ def load( to the trunk only, before MTP injection, so a draft head's precision is never reduced. """ + if deepseek_v4_0731_k2 and not mtp: + raise ValueError("DeepSeek-V4-0731 K2 construction requires mtp=True") path = Path(model_path) + k2_config = None + if deepseek_v4_0731_k2: + from .deepseek_v4_0731_full_install import ( + validate_full_0731_dspark_artifact, + ) + + k2_config = load_config(path) + # This validates config/index bytes, source revision, topology, and the + # exact DSpark manifest before any paired or ordinary loader can + # construct a model. + validate_full_0731_dspark_artifact(path, k2_config) from .gemma4_pair import resolve_gemma4_pair_paths gemma4_pair = resolve_gemma4_pair_paths(path) @@ -592,9 +686,7 @@ def load( ) metadata = gemma4_pair["metadata"] - benchmark = ( - metadata.get("benchmark") if isinstance(metadata, dict) else {} - ) + benchmark = metadata.get("benchmark") if isinstance(metadata, dict) else {} draft_block_size = DEFAULT_DRAFT_BLOCK_SIZE if isinstance(benchmark, dict): try: @@ -618,7 +710,7 @@ def load( runtime.bundle_path = path return runtime path = Path(gemma4_pair["target_model"]) - config = load_config(path) + config = k2_config if k2_config is not None else load_config(path) from .a3b_whole_moe import validate_a3b_whole_moe_load_options validate_a3b_whole_moe_load_options( @@ -661,6 +753,11 @@ def load( tokenizer = _load_tokenizer_resilient(path, config) model, _loaded_config = load_model(path) + elif deepseek_v4_0731_k2: + from .models.deepseek_v4 import deepseek_v4_0731_k2_construction + + with deepseek_v4_0731_k2_construction(): + model, tokenizer = _load_base_model(path, config) else: model, tokenizer = _load_base_model(path, config) import os as _os @@ -674,16 +771,21 @@ def load( touched = quantize_projections(model, proj_quant) logger.info( "[proj-quant] quantized %d trunk *_proj modules to %s", - len(touched), proj_quant, + len(touched), + proj_quant, ) if proj_requant: touched = requantize_projections(model, proj_requant) logger.info( "[proj-quant] requantized %d trunk *_proj modules to %s", - len(touched), proj_requant, + len(touched), + proj_requant, ) deepseek_v4_attn_proj_wide_m3_report = None - if str((config or {}).get("model_type") or "").lower() == "deepseek_v4": + if ( + str((config or {}).get("model_type") or "").lower() == "deepseek_v4" + and not deepseek_v4_0731_k2 + ): from .models.deepseek_v4 import configure_deepseek_v4_moe_tail configure_deepseek_v4_moe_tail(model, config) @@ -710,20 +812,38 @@ def load( .with_config_defaults(config) ) mtp_enabled = False + block_speculative_backend = None if mtp: - from .deepseek_mtp_patch import inject_deepseek_mtp_support, is_deepseek_mtp_config + from .deepseek_mtp_patch import ( + inject_deepseek_mtp_support, + is_deepseek_mtp_config, + ) from .glm_mtp_patch import inject_glm_mtp_support, is_glm_mtp_config from .mimo_mtp_patch import inject_mimo_mtp_support, is_mimo_mtp_config - from .nemotron_h_mtp_patch import inject_nemotron_h_mtp_support, is_nemotron_h_mtp_config + from .nemotron_h_mtp_patch import ( + inject_nemotron_h_mtp_support, + is_nemotron_h_mtp_config, + ) from .step3p5_mtp_patch import inject_step3p5_mtp_support from .hy_v3_mtp_patch import inject_hy_v3_mtp_support, is_hy_v3_mtp_config from .models.deepseek_v4 import ( + _config_has_dspark_signature, inject_deepseek_v4_mtp_support, is_deepseek_v4_mtp_config, ) from .qwen3_5_mtp_patch import inject_qwen3_5_mtp_support - if is_deepseek_v4_mtp_config(config): + if _config_has_dspark_signature(config): + if not deepseek_v4_0731_k2: + from .deepseek_v4_dspark_generation import DeepseekV4DSparkBackend + + # The 0731 artifact owns a fixed three-stage proposer rather than + # the preview model's reusable one-block MTP head. Bind every + # callable and ownership invariant once, before the runtime is + # published; this route never enters legacy MTP validation. + block_speculative_backend = DeepseekV4DSparkBackend.bind(model) + mtp_enabled = True + elif is_deepseek_v4_mtp_config(config): # Native draft head: the block binds through the ordinary load path # and the model already carries the runtime surface, so this only # publishes it. Placed ahead of is_deepseek_mtp_config defensively -- @@ -747,22 +867,23 @@ def load( mtp_enabled = inject_deepseek_mtp_support(model, path, config, contract) else: mtp_enabled = inject_mtp_support(model, path, config, contract) - if mtp_enabled: - if not validate_mtp_support(model): + if block_speculative_backend is None and not deepseek_v4_0731_k2: + if mtp_enabled: + if not validate_mtp_support(model): + raise RuntimeError(f"MTP injection failed for {path}") + elif mtp_weights_present_on_disk(path, config): + # MTP weights ship with the model but injection could not use + # them: a genuine failure the operator should see. raise RuntimeError(f"MTP injection failed for {path}") - elif mtp_weights_present_on_disk(path, config): - # MTP weights ship with the model but injection could not use - # them: a genuine failure the operator should see. - raise RuntimeError(f"MTP injection failed for {path}") - else: - # The config declares MTP layers but no MTP weights are present on - # disk (e.g. a quant conversion that dropped the draft head). - # Degrade to autoregressive rather than failing the load. - logger.warning( - "[MTP] %s declares MTP layer(s) but ships no MTP weights; " - "serving autoregressive (no speculative draft head).", - path, - ) + else: + # The config declares MTP layers but no MTP weights are present on + # disk (e.g. a quant conversion that dropped the draft head). + # Degrade to autoregressive rather than failing the load. + logger.warning( + "[MTP] %s declares MTP layer(s) but ships no MTP weights; " + "serving autoregressive (no speculative draft head).", + path, + ) compiled_target_factory = None whole_moe_plan = None selfcheck_report = None @@ -840,7 +961,9 @@ def load( "whole-MoE target M2 requires the accepted row-owned router/combine route" ) if router_plan is not None: - router_report = install_qwen_row_owned_routers(router_plan, selfcheck_report) + router_report = install_qwen_row_owned_routers( + router_plan, selfcheck_report + ) logger.info("[qwen-row-owned-router] %s", router_report) if whole_moe_plan is not None: selfcheck_report = run_a3b_whole_moe_selfcheck( @@ -848,9 +971,7 @@ def load( selfcheck_report, ) if postconv_plan is not None: - postconv_factory = install_a3b_gdn_postconv( - postconv_plan, selfcheck_report - ) + postconv_factory = install_a3b_gdn_postconv(postconv_plan, selfcheck_report) from .gdn_capture import gdn_postconv_stats logger.info("[a3b-gdn-postconv] %s", gdn_postconv_stats()) @@ -872,13 +993,20 @@ def load( raise RuntimeError("merge_mtp_adapter requires mtp_adapter") deepseek_v4_o_lora_report = None deepseek_v4_attention_island_report = None + k2_o_lora_state = None if str(config.get("model_type") or "").lower() == "deepseek_v4": from .models.deepseek_v4 import ( _o_lora_mode_from_env, install_deepseek_v4_o_lora_routes, ) - selected_o_lora_mode = _o_lora_mode_from_env() + selected_o_lora_mode = ( + "gather_qmm" if deepseek_v4_0731_k2 else _o_lora_mode_from_env() + ) + if deepseek_v4_0731_k2: + # Snapshot now, but leave O-Lora and the target/FFN routes stock while + # their candidates run construction-time preparation and self-checks. + k2_o_lora_state = _snapshot_deepseek_v4_o_lora_routes(model) # The canonical mixed route hard-validates the exact DeepSeek-V4-Flash # topology (43 body layers, rank-1024 Q4/g64 wo_a/wo_b, one dense-BF16 # MTP block) and refuses anything else. That strictness is correct for @@ -888,27 +1016,31 @@ def load( # route — which is bit-identical on the canonical artifact anyway # (test_cached_dequant_is_bit_identical). canonical_mixed_route = bool( - mtp_enabled and selected_o_lora_mode == "gather_qmm" + mtp_enabled + and block_speculative_backend is None + and not deepseek_v4_0731_k2 + and selected_o_lora_mode == "gather_qmm" ) if not mtp_enabled: # An artifact that declared but did not ship MTP weights already # degraded to AR above. It has no dense MTP module to validate or # route, so bind the trunk's explicit stock/cached construction. selected_o_lora_mode = "cached" - deepseek_v4_o_lora_report = install_deepseek_v4_o_lora_routes( - model, - mode=selected_o_lora_mode, - canonical_mixed_route=canonical_mixed_route, - ) - logger.info("[deepseek-v4-o-lora] %s", deepseek_v4_o_lora_report) + if not deepseek_v4_0731_k2: + deepseek_v4_o_lora_report = install_deepseek_v4_o_lora_routes( + model, + mode=selected_o_lora_mode, + canonical_mixed_route=canonical_mixed_route, + ) + logger.info("[deepseek-v4-o-lora] %s", deepseek_v4_o_lora_report) from .deepseek_v4_attention_island import ( deepseek_v4_attention_island_enabled, install_deepseek_v4_attention_island, ) - if deepseek_v4_attention_island_enabled(): - deepseek_v4_attention_island_report = ( - install_deepseek_v4_attention_island(model, config) + if not deepseek_v4_0731_k2 and deepseek_v4_attention_island_enabled(): + deepseek_v4_attention_island_report = install_deepseek_v4_attention_island( + model, config ) logger.info( "[deepseek-v4-attention-island] %s", @@ -925,27 +1057,91 @@ def load( fused_report = _laguna_install_fused(model) if fused_report: logger.info("[laguna-fused] %s", fused_report) + deepseek_v4_0731_k2_receipt = None + k2_prepared = None + if deepseek_v4_0731_k2: + from .deepseek_v4_0731_dspark_ffn import ( + prepare_dspark_q3_packed_gate_up_m5, + ) + from .deepseek_v4_0731_full_install import ( + prepare_full_0731_dspark_compiled_tail_q2_pair, + ) + from .deepseek_v4_0731_m3_wob import prepare_wob_m3 + from .deepseek_v4_0731_m3_wqb_qnorm_rope import prepare_wqb_qhead_m3 + from .deepseek_v4_dspark_generation import DeepseekV4DSparkBackend + + try: + target_prepared = prepare_full_0731_dspark_compiled_tail_q2_pair( + model, + config, + path, + prepare_wqb_qhead=prepare_wqb_qhead_m3, + prepare_wob=prepare_wob_m3, + ) + ffn_prepared = prepare_dspark_q3_packed_gate_up_m5(model) + k2_prepared = (target_prepared, ffn_prepared) + deepseek_v4_o_lora_report = install_deepseek_v4_o_lora_routes( + model, + mode="gather_qmm", + canonical_mixed_route=False, + ) + logger.info("[deepseek-v4-o-lora] %s", deepseek_v4_o_lora_report) + target_prepared.publish() + ffn_prepared.publish() + block_speculative_backend = DeepseekV4DSparkBackend.bind(model) + except Exception as failure: + rollback_failures = _rollback_deepseek_v4_0731_k2( + k2_prepared, + k2_o_lora_state, + ) + if rollback_failures: + raise ExceptionGroup( + "DeepSeek-V4-0731 K2 publication and rollback failed", + [failure, *rollback_failures], + ) from failure + raise + deepseek_v4_0731_k2_receipt = { + "target": target_prepared.receipt, + "dspark_ffn": ffn_prepared.receipt, + } runtime_class = ( - LagunaARRuntime - if _is_laguna_s_2_1_mlx_4bit_config(config) - else MTPLXRuntime - ) - runtime = runtime_class( - model, - tokenizer, - path, - mtp_enabled, - contract, - mtp_adapter_path=adapter_path, - mtp_adapter_metadata=adapter_metadata, - mtp_adapter_merge_report=adapter_merge_report, - deepseek_v4_o_lora_report=deepseek_v4_o_lora_report, - deepseek_v4_attn_proj_wide_m3_report=deepseek_v4_attn_proj_wide_m3_report, - deepseek_v4_attention_island_report=deepseek_v4_attention_island_report, - a3b_compiled_target_prefix_factory=compiled_target_factory, - a3b_whole_moe_installed=False, - qwen_row_owned_router_report=router_report, + LagunaARRuntime if _is_laguna_s_2_1_mlx_4bit_config(config) else MTPLXRuntime ) + try: + runtime = runtime_class( + model, + tokenizer, + path, + mtp_enabled, + contract, + mtp_adapter_path=adapter_path, + mtp_adapter_metadata=adapter_metadata, + mtp_adapter_merge_report=adapter_merge_report, + deepseek_v4_o_lora_report=deepseek_v4_o_lora_report, + deepseek_v4_attn_proj_wide_m3_report=deepseek_v4_attn_proj_wide_m3_report, + deepseek_v4_attention_island_report=deepseek_v4_attention_island_report, + deepseek_v4_dspark_enabled=block_speculative_backend is not None, + deepseek_v4_0731_k2_receipt=deepseek_v4_0731_k2_receipt, + block_speculative_backend=block_speculative_backend, + block_speculative_decode_trace_requested=bool( + block_speculative_backend is not None + and os.environ.get("MTPLX_DECODE_TRACE_JSONL") + ), + a3b_compiled_target_prefix_factory=compiled_target_factory, + a3b_whole_moe_installed=False, + qwen_row_owned_router_report=router_report, + ) + except Exception as failure: + rollback_failures = _rollback_deepseek_v4_0731_k2( + k2_prepared, + k2_o_lora_state, + ) + if rollback_failures: + raise ExceptionGroup( + "DeepSeek-V4-0731 K2 runtime publication and rollback failed", + [failure, *rollback_failures], + ) from failure + raise if whole_moe_plan is not None: if compiled_target_factory is None: from .a3b_whole_moe import A3BWholeMoeConfigError @@ -1100,10 +1296,18 @@ def _load_tokenizer_resilient(model_path: Path, config: dict[str, Any]) -> Any: from transformers import PreTrainedTokenizerFast tcfg_path = model_path / "tokenizer_config.json" - tcfg = json.loads(tcfg_path.read_text(encoding="utf-8")) if tcfg_path.exists() else {} + tcfg = ( + json.loads(tcfg_path.read_text(encoding="utf-8")) if tcfg_path.exists() else {} + ) passthrough = { key: tcfg[key] - for key in ("bos_token", "eos_token", "pad_token", "unk_token", "additional_special_tokens") + for key in ( + "bos_token", + "eos_token", + "pad_token", + "unk_token", + "additional_special_tokens", + ) if key in tcfg } hf_tokenizer = PreTrainedTokenizerFast( @@ -1158,10 +1362,7 @@ def _mtp_alias_load_path(path: Path, config: dict[str, Any] | None) -> Path: import importlib.util def _mlx_lm_has(model_type_name: str) -> bool: - return ( - importlib.util.find_spec(f"mlx_lm.models.{model_type_name}") - is not None - ) + return importlib.util.find_spec(f"mlx_lm.models.{model_type_name}") is not None if _mlx_lm_has(model_type) or not _mlx_lm_has(base_type): return path diff --git a/mtplx/server/openai.py b/mtplx/server/openai.py index 9aeb8dc0..810a84b9 100644 --- a/mtplx/server/openai.py +++ b/mtplx/server/openai.py @@ -222,7 +222,7 @@ def _safe_stdout_print(*values: Any, **kwargs: Any) -> bool: is_background_request, system_prompt_hash, ) - from mtplx.runtime import load + from mtplx.runtime import build_mtpk_request_kwargs, load from mtplx.session_bank import CacheMissReason, common_prefix_len from mtplx.cache_state import restore_cache, snapshot_cache @@ -248,6 +248,7 @@ def _missing_runtime(*_args: Any, **_kwargs: Any) -> Any: restore_cache = _missing_runtime snapshot_cache = _missing_runtime load = _missing_runtime + build_mtpk_request_kwargs = _missing_runtime class PostcommitAbort(RuntimeError): pass @@ -1762,9 +1763,26 @@ def _validate_mtp_batch_settings(args: argparse.Namespace) -> None: _require_mlx_lm_arrays_cache_fix() +def _validate_deepseek_v4_0731_k2_entrypoint(args: argparse.Namespace) -> None: + if not bool(getattr(args, "deepseek_v4_0731_k2", False)): + return + cli_flags = getattr(args, "_cli_flags", set()) or set() + if "depth" not in cli_flags or int(getattr(args, "depth", 3)) != 2: + raise ValueError( + "DeepSeek-V4-0731 K2 requires explicit --depth 2 before model load" + ) + if ( + getattr(args, "load_mtp", True) is False + or str(getattr(args, "generation_mode", "mtp")) != "mtp" + or bool(getattr(args, "stock_ar", False)) + ): + raise ValueError("DeepSeek-V4-0731 K2 requires MTP generation") + + class ServerState: def __init__(self, args: argparse.Namespace) -> None: _validate_mtp_batch_settings(args) + _validate_deepseek_v4_0731_k2_entrypoint(args) self.args = args try: args.paged_kv_quantization = normalize_paged_kv_quantization( @@ -1886,6 +1904,9 @@ def __init__(self, args: argparse.Namespace) -> None: ) _startup_line(" Model load in progress (this may take a minute).") load_heartbeat = _startup_heartbeat("Model still loading") + construction_options = {} + if bool(getattr(args, "deepseek_v4_0731_k2", False)): + construction_options["deepseek_v4_0731_k2"] = True try: self.runtime = self.model_scheduler.submit_foreground( load, @@ -1903,6 +1924,7 @@ def __init__(self, args: argparse.Namespace) -> None: args, startup_backend, ), + **construction_options, batch_key="startup.load", ).result() except BaseException as exc: @@ -2807,7 +2829,9 @@ def _admit_pending(self, generator: Any, config_dict: dict[str, Any]) -> None: self._active[int(uid)] = job self._condition.notify_all() - def _commit_prompt_boundary(self, job: _BatchedARJob, generator: Any, uid: int) -> None: + def _commit_prompt_boundary( + self, job: _BatchedARJob, generator: Any, uid: int + ) -> None: """Store a batched row's PROMPT-ONLY state at its first generation step. At the first response, the step that produced token 1 has just @@ -13744,7 +13768,9 @@ def _mtplx_dashboard_snapshot(state: "ServerState") -> dict[str, Any]: "model_id": state.model_id, # Always present, so a client can tell "no retrieval configured" apart # from "this build has no retrieval support". - "retrieval": retrieval.status() if retrieval is not None else {"enabled": False, "models": []}, + "retrieval": retrieval.status() + if retrieval is not None + else {"enabled": False, "models": []}, "profile": state.profile.to_dict() if hasattr(state.profile, "to_dict") else {"name": getattr(state.profile, "name", "unknown")}, @@ -17407,7 +17433,10 @@ def session_restore() -> Any | None: outcome: dict[str, Any] = {"hit": False} started = time.perf_counter() try: - if len(tokens) <= min_restore_tokens or not _boundary_true_restore_enabled(): + if ( + len(tokens) <= min_restore_tokens + or not _boundary_true_restore_enabled() + ): return None candidates = candidates_fn( tokens, @@ -17420,9 +17449,7 @@ def session_restore() -> Any | None: policy_fingerprint=policy_fingerprint, min_restore_tokens=min_restore_tokens, ) - for entry, matched in sorted( - candidates, key=lambda item: -int(item[1]) - ): + for entry, matched in sorted(candidates, key=lambda item: -int(item[1])): restored = restore_fn( rt, entry, @@ -17478,9 +17505,7 @@ def session_restore() -> Any | None: outcome["error"] = f"{type(exc).__name__}: {exc}" return None finally: - outcome.setdefault( - "restore_s", round(time.perf_counter() - started, 6) - ) + outcome.setdefault("restore_s", round(time.perf_counter() - started, 6)) request_observability["mtp_batch_session_restore"] = outcome def session_commit( @@ -17606,8 +17631,7 @@ def _run_mtp_batch_generation_dispatched( ), "mtp_disabled_reason": None, "mtp_batch_session_cache_bypass": ( - kwargs.get("session_bank") is not None - and session_restore_hook is None + kwargs.get("session_bank") is not None and session_restore_hook is None ), } ) @@ -18146,6 +18170,37 @@ def record_tokens(new_tokens: list[int]) -> None: "generation_mode": effective_mode, **(request_observability or {}), } + cli_flags = getattr(state.args, "_cli_flags", set()) or set() + explicit_legacy = {} + for flag, key, value, baseline in ( + ( + "verify-strategy", + "verify_strategy", + state.args.verify_strategy, + "capture_commit", + ), + ( + "verify-core", + "verify_core", + state.args.verify_core, + "linear-gdn-from-conv-tape", + ), + ): + if flag in cli_flags or value != baseline: + explicit_legacy[key] = value + selected_legacy_kwargs = build_mtpk_request_kwargs( + state.runtime, + common={}, + legacy_defaults={ + "mtp_hidden_variant": "post_norm", + "mtp_history_policy": "committed", + "verify_strategy": state.args.verify_strategy, + "verify_core": state.args.verify_core, + "trace_label": trace_label, + "trace_metadata": trace_metadata, + }, + explicit_legacy=explicit_legacy, + ) for attempt in range(max_attempts): generation_seed, seed_is_explicit = _resolve_seed(state, seed) lock_started = time.perf_counter() @@ -18235,40 +18290,34 @@ def record_tokens(new_tokens: list[int]) -> None: # Retries and tool-loop redispatches replay the # full prompt, so the image rows must rewind. vision_splice.reset() - out = generate_mtpk( - state.runtime, - prompt_ids, - constraint=constraint, - vision_splice=vision_splice, - abort_check=( + request_kwargs = { + "constraint": constraint, + "vision_splice": vision_splice, + "abort_check": ( (lambda: bool(cancel_event.is_set())) if cancel_event is not None else None ), - max_tokens=response_max, - sampler=sampler, - draft_sampler=effective_draft_sampler, - speculative_depth=effective_depth, - seed=generation_seed, - mtp_hidden_variant="post_norm", - mtp_cache_policy="persistent", - mtp_history_policy="committed", - verify_strategy=state.args.verify_strategy, - verify_core=state.args.verify_core, - draft_core=str( + "max_tokens": response_max, + "sampler": sampler, + "draft_sampler": effective_draft_sampler, + "speculative_depth": effective_depth, + "seed": generation_seed, + "mtp_cache_policy": "persistent", + "draft_core": str( getattr(state.args, "draft_core", None) or "stock" ), - token_callback=record_tokens, - session_bank=session_bank, - session_id=session_id, - session_restore_mode=_session_bank_restore_mode( + "token_callback": record_tokens, + "session_bank": session_bank, + "session_id": session_id, + "session_restore_mode": _session_bank_restore_mode( session_restore_mode ), - session_template_hash=session_template_hash, - session_draft_head_identity=session_draft_head_identity, - session_policy_fingerprint=session_policy_fingerprint, - capture_final_state=session_bank is not None, - commit_prompt_state_to_bank=( + "session_template_hash": session_template_hash, + "session_draft_head_identity": session_draft_head_identity, + "session_policy_fingerprint": session_policy_fingerprint, + "capture_final_state": session_bank is not None, + "commit_prompt_state_to_bank": ( commit_prompt_prefix_to_bank and session_bank is not None and session_id is not None @@ -18276,45 +18325,45 @@ def record_tokens(new_tokens: list[int]) -> None: # Prompt-prefix commits happen before decode mutates # the same KV/MTP cache objects. They must snapshot or # skip, not live-lease the mutable prompt cache. - commit_prompt_state_keep_live_ref=False, - trace_label=trace_label, - trace_metadata=trace_metadata, - prefill_callback=prefill_callback, - adaptive_policy=adaptive_policy, - repetition_stop=uncapped_repetition_stop, - loop_guard=_loop_guard_enabled(), - thinking_guard=thinking_guard_config, - online_correction_cache=bool( + "commit_prompt_state_keep_live_ref": False, + "prefill_callback": prefill_callback, + "adaptive_policy": adaptive_policy, + "repetition_stop": uncapped_repetition_stop, + "loop_guard": _loop_guard_enabled(), + "thinking_guard": thinking_guard_config, + "online_correction_cache": bool( state.args.online_correction_cache ), - online_correction_cache_min_depth=int( + "online_correction_cache_min_depth": int( state.args.online_correction_cache_min_depth ), - online_correction_cache_key=str( + "online_correction_cache_key": str( state.args.online_correction_cache_key ), - prompt_correction_cache=bool( + "prompt_correction_cache": bool( state.args.prompt_correction_cache ), - prompt_correction_cache_min_depth=int( + "prompt_correction_cache_min_depth": int( state.args.prompt_correction_cache_min_depth ), - online_hidden_corrector_alpha=float( + "online_hidden_corrector_alpha": float( state.args.online_hidden_corrector_alpha ), - online_hidden_corrector_decay=float( + "online_hidden_corrector_decay": float( state.args.online_hidden_corrector_decay ), - online_hidden_corrector_warmup=int( + "online_hidden_corrector_warmup": int( state.args.online_hidden_corrector_warmup ), - online_hidden_corrector_max_feed_depth=( + "online_hidden_corrector_max_feed_depth": ( state.args.online_hidden_corrector_max_feed_depth ), - online_hidden_corrector_key=str( + "online_hidden_corrector_key": str( state.args.online_hidden_corrector_key ), - ) + **selected_legacy_kwargs, + } + out = generate_mtpk(state.runtime, prompt_ids, **request_kwargs) except PostcommitAbort: # abort_check tripped inside the prefill: the client disconnected # mid-prompt-processing. Reuse the exact cancellation path client @@ -23495,7 +23544,9 @@ async def rerank(request: RerankRequest) -> dict[str, Any]: detail="no reranking model is configured; start MTPLX with --reranker-model", ) if not request.query or not str(request.query).strip(): - raise HTTPException(status_code=400, detail="query must be a non-empty string") + raise HTTPException( + status_code=400, detail="query must be a non-empty string" + ) documents = _as_text_list(request.documents, field="documents") try: scores, spec, prompt_tokens = await asyncio.to_thread( @@ -26244,9 +26295,7 @@ def finish_orphan_stream_guards() -> list[str]: stream_orphan_tool_markup_suppressed = True remainder = guard.take_orphan_remainder() if remainder: - deferred_orphan_stream_remainders.append( - (field, remainder) - ) + deferred_orphan_stream_remainders.append((field, remainder)) if not flushed: continue chunks.extend( @@ -28747,6 +28796,14 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: help="Load and inject the native MTP sidecar. Disable only for stock AR diagnostics.", ) parser.add_argument("--depth", type=int, default=3) + parser.add_argument( + "--deepseek-v4-0731-k2", + action="store_true", + help=( + "Select the exact construction-bound DeepSeek-V4-Flash-0731 " + "DSpark K2 stack. Requires explicit --depth 2 and MTP." + ), + ) parser.add_argument( "--max-response-tokens", "--max-tokens", diff --git a/scripts/deepseek_v4_0731_k2_bench.py b/scripts/deepseek_v4_0731_k2_bench.py new file mode 100644 index 00000000..62cc9bb5 --- /dev/null +++ b/scripts/deepseek_v4_0731_k2_bench.py @@ -0,0 +1,1114 @@ +"""Official-wheel bracket for the isolated 0731 scheduler evaluation boundary. + +Each clean source tree is one arm. The harness derives ``lazy_joint_eval`` or +``materialize_first`` from the reviewed scheduler source, loads the unchanged +generic ``mtp=True`` runtime once, proves unmeasured prompt-cache parity, and +then runs the same AR/K2 primers and five measured repetitions. There is no +runtime or environment arm selector. +""" + +from __future__ import annotations + +import argparse +import hashlib +from importlib import metadata +import json +from pathlib import Path +from statistics import median +import subprocess +import sys +from types import ModuleType +from typing import Any, Callable +from urllib.parse import urlsplit + +import numpy as np + + +PROMPT_TEXT = "Explain why speculative decoding can preserve greedy output." +PROMPT_TOKEN_COUNT = 9 +REPETITIONS = 5 +EXPECTED_MLX_VERSION = "0.32.0" +EXPECTED_MLX_CORE_SHA256 = ( + "f96aede5d6eee539d4826a52690914e79794e2ad2c691935d02dca6b0c421c56" +) +EXPECTED_MLX_LIB_SHA256 = ( + "1876795e05b3434925e745fbf6e9f0c8c0446b666224c9d881609ab353e94e51" +) +EXPECTED_MLX_METALLIB_SHA256 = ( + "1518c08860738b08dc4563ddcf380a08dec4e6ad146c0d54888790e80656e9e3" +) +EXPECTED_MODEL_CONFIG_SHA256 = ( + "44735712733fcf8f299bdf1faa1d87fac88f1917efe1d3876d6d4c582f79a68f" +) +EXPECTED_MODEL_INDEX_SHA256 = ( + "f1332b2b209769c2db335954c2651652a8048e7d7dbf60296c2f2c0198715861" +) +EXPECTED_MODEL_METADATA_REVISION = "10001e0065f8394e03e968e652cbbe7cd2ca122c" +_SCHEDULER_SOURCE = "mtplx/native_block_speculation.py" +_BRACKET_SOURCE_PATHS = ( + "mtplx/models/deepseek_v4.py", + "mtplx/deepseek_v4_dspark_generation.py", + _SCHEDULER_SOURCE, + "mtplx/runtime.py", + "mtplx/generation.py", + "mtplx/sampling.py", + "scripts/deepseek_v4_guard_window.py", + "scripts/deepseek_v4_0731_k2_bench.py", +) +_REQUIRED_IMPORTED_MODULES = ( + "mtplx.models.deepseek_v4", + "mtplx.deepseek_v4_dspark_generation", + "mtplx.native_block_speculation", + "mtplx.runtime", + "mtplx.generation", + "mtplx.sampling", +) +_ARM_EVENTS = { + "lazy_joint_eval": [ + "proposal_graph", + "target_row_graph", + "joint_eval", + "draft_materialize", + ], + "materialize_first": [ + "proposal_graph", + "proposal_eval", + "draft_materialize", + "target_row_graph", + ], +} +_ARM_LABELS = frozenset(_ARM_EVENTS) +EXPECTED_NORMALIZED_SCHEDULER_SHA256 = ( + "10f7a52f59044ca7e7600156626b28826773886657e68201644f8b50385ba2e1" +) +EXPECTED_SCHEDULER_BOUNDARY_PATCH_SHA256 = ( + "f09d68378f940eb948a58cf4f9b24e90bfb9d40119483348b3e6f5d8b849205e" +) +_LAZY_BOUNDARY_BLOCK = """ # Build the authoritative primary M1 before forcing proposal IDs to the + # host. The two graphs are independent given carried hidden/current_top, + # so they share one evaluation boundary without changing target math. + with attention_phase("ar_decode"): + row_logits, row_hidden = target_forward( + mx.array([[current_top]], dtype=mx.int32), + cache=target_cache, + return_hidden=True, + ) + if future is None: + _eval(row_logits, row_hidden) + else: + _eval(future, row_logits, row_hidden) + + accepted_hidden = [row_hidden[:, -1:]] + next_logits = row_logits[:, -1, :] + future_tokens: list[int] = [] + if future is not None: + future_tokens = [int(token) for token in np.asarray(future)[0]] + for index, token in enumerate(future_tokens): + drafted_by_depth[index] += 1 + if _is_stop(token, stop_token_ids): + future_tokens = future_tokens[: index + 1] + width = index + 2 + break + drafted_tokens += len(future_tokens) +""" +_MATERIALIZE_BOUNDARY_BLOCK = """ # Settle and materialize proposal IDs before constructing target row zero. + future_tokens: list[int] = [] + if future is not None: + _eval(future) + future_tokens = [int(token) for token in np.asarray(future)[0]] + for index, token in enumerate(future_tokens): + drafted_by_depth[index] += 1 + if _is_stop(token, stop_token_ids): + future_tokens = future_tokens[: index + 1] + width = index + 2 + break + drafted_tokens += len(future_tokens) + + with attention_phase("ar_decode"): + row_logits, row_hidden = target_forward( + mx.array([[current_top]], dtype=mx.int32), + cache=target_cache, + return_hidden=True, + ) + _eval(row_logits, row_hidden) + + accepted_hidden = [row_hidden[:, -1:]] + next_logits = row_logits[:, -1, :] +""" + + +def _sha256(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def _canonical_bytes(value: Any) -> bytes: + return json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ).encode() + + +def attest_official_mlx(mx: Any, distribution: Any) -> dict[str, Any]: + """Attest an installed wheel and reject editable/source import overlays.""" + + package_version = str(distribution.version) + if package_version != EXPECTED_MLX_VERSION: + raise ValueError( + f"expected MLX {EXPECTED_MLX_VERSION}, got {package_version or 'unknown'}" + ) + installer = str(distribution.read_text("INSTALLER") or "").strip() + if not installer: + raise ValueError("MLX distribution has no INSTALLER attestation") + distribution_root = Path(distribution.locate_file("")).resolve() + core_path = Path(getattr(mx, "__file__", "")).resolve() + if not core_path.is_file(): + raise ValueError(f"MLX core module is unreadable: {core_path}") + try: + core_path.relative_to(distribution_root) + except ValueError as exc: + raise ValueError( + f"MLX core import is outside installed distribution: {core_path}" + ) from exc + core_sha = _sha256(core_path) + if core_sha != EXPECTED_MLX_CORE_SHA256: + raise ValueError(f"MLX core SHA mismatch: {core_sha}") + + libmlx_path = core_path.parent / "lib" / "libmlx.dylib" + metallib_path = core_path.parent / "lib" / "mlx.metallib" + if not libmlx_path.is_file() or not metallib_path.is_file(): + raise ValueError("MLX dylib/metallib wheel artifacts are missing") + libmlx_sha = _sha256(libmlx_path) + metallib_sha = _sha256(metallib_path) + if libmlx_sha != EXPECTED_MLX_LIB_SHA256: + raise ValueError(f"MLX libmlx SHA mismatch: {libmlx_sha}") + if metallib_sha != EXPECTED_MLX_METALLIB_SHA256: + raise ValueError(f"MLX metallib SHA mismatch: {metallib_sha}") + + direct_url_text = distribution.read_text("direct_url.json") + direct_url = None + if direct_url_text: + direct_url = json.loads(direct_url_text) + url = str(direct_url.get("url") or "") + is_wheel_archive = urlsplit(url).path.lower().endswith(".whl") + if ( + direct_url.get("dir_info") is not None + or direct_url.get("vcs_info") is not None + or not is_wheel_archive + ): + raise ValueError( + "MLX source/direct overlay is not an official installed wheel" + ) + + return { + "version": package_version, + "core_path": str(core_path), + "core_sha256": core_sha, + "libmlx": {"path": str(libmlx_path), "sha256": libmlx_sha}, + "metallib": {"path": str(metallib_path), "sha256": metallib_sha}, + "distribution_root": str(distribution_root), + "installer": installer, + "direct_url": direct_url, + } + + +def attest_model(model_path: Path) -> dict[str, Any]: + root = model_path.expanduser().resolve() + config_path = root / "config.json" + index_path = root / "model.safetensors.index.json" + config_sha = _sha256(config_path) + index_sha = _sha256(index_path) + if config_sha != EXPECTED_MODEL_CONFIG_SHA256: + raise ValueError(f"model config SHA mismatch: {config_sha}") + if index_sha != EXPECTED_MODEL_INDEX_SHA256: + raise ValueError(f"model index SHA mismatch: {index_sha}") + metadata_root = root / ".cache" / "huggingface" / "download" + metadata_paths = { + "config": metadata_root / "config.json.metadata", + "index": metadata_root / "model.safetensors.index.json.metadata", + } + model_metadata = {} + for name, path in metadata_paths.items(): + try: + revision = path.read_text(encoding="utf-8").splitlines()[0] + except (OSError, IndexError) as exc: + raise ValueError(f"model {name} metadata is unreadable: {exc}") from exc + if revision != EXPECTED_MODEL_METADATA_REVISION: + raise ValueError(f"model {name} metadata revision mismatch: {revision!r}") + model_metadata[name] = { + "path": str(path), + "sha256": _sha256(path), + "revision": revision, + } + return { + "path": str(root), + "config_path": str(config_path), + "config_sha256": config_sha, + "index_path": str(index_path), + "index_sha256": index_sha, + "metadata": model_metadata, + } + + +def attest_git(repo: Path) -> dict[str, Any]: + """Require a clean committed tree before any MLX import can occur.""" + + root = repo.resolve() + + def git(*arguments: str) -> str: + return subprocess.run( + ["git", "-C", str(root), *arguments], + check=True, + capture_output=True, + text=True, + ).stdout.strip() + + status_text = git("status", "--porcelain=v1", "--untracked-files=all") + status = status_text.splitlines() if status_text else [] + if status: + raise RuntimeError( + "scheduler bracket requires a clean committed worktree: " + + "; ".join(status) + ) + head_tree_rows = git("ls-tree", "-r", "--full-tree", "HEAD").splitlines() + head_tree_files = {} + for row in head_tree_rows: + metadata_text, path = row.split("\t", 1) + mode, object_type, object_id = metadata_text.split(" ", 2) + head_tree_files[path] = { + "mode": mode, + "type": object_type, + "object": object_id, + } + head_python_sha256 = { + path: _sha256(root / path) + for path in head_tree_files + if path.startswith("mtplx/") and path.endswith(".py") + } + return { + "repository": str(root), + "commit": git("rev-parse", "HEAD"), + "head_tree": git("rev-parse", "HEAD^{tree}"), + "head_tree_files": head_tree_files, + "head_tree_files_sha256": hashlib.sha256( + _canonical_bytes(head_tree_files) + ).hexdigest(), + "head_python_sha256": head_python_sha256, + "head_python_set_sha256": hashlib.sha256( + _canonical_bytes(head_python_sha256) + ).hexdigest(), + "dirty": False, + "status": [], + } + + +def attest_sources(repo: Path) -> dict[str, Any]: + root = repo.resolve() + files = {relative: _sha256(root / relative) for relative in _BRACKET_SOURCE_PATHS} + importable_mtplx_files = { + str(path.relative_to(root)): _sha256(path) + for path in sorted((root / "mtplx").rglob("*.py")) + } + return { + "files": files, + "source_set_sha256": hashlib.sha256(_canonical_bytes(files)).hexdigest(), + "importable_mtplx_files": importable_mtplx_files, + "importable_mtplx_set_sha256": hashlib.sha256( + _canonical_bytes(importable_mtplx_files) + ).hexdigest(), + } + + +def _scheduler_patch_digest() -> str: + payload = ( + _LAZY_BOUNDARY_BLOCK.encode() + b"\0" + _MATERIALIZE_BOUNDARY_BLOCK.encode() + ) + return hashlib.sha256(payload).hexdigest() + + +def _normalize_scheduler_source(source: str) -> tuple[str, str]: + """Normalize only the exact reviewed lazy/materialize boundary motion.""" + + lazy_count = source.count(_LAZY_BOUNDARY_BLOCK) + materialize_count = source.count(_MATERIALIZE_BOUNDARY_BLOCK) + if (lazy_count, materialize_count) == (1, 0): + label = "lazy_joint_eval" + normalized = source + elif (lazy_count, materialize_count) == (0, 1): + label = "materialize_first" + normalized = source.replace( + _MATERIALIZE_BOUNDARY_BLOCK, + _LAZY_BOUNDARY_BLOCK, + 1, + ) + else: + raise ValueError("scheduler source is not an exact sanctioned bracket arm") + normalized_sha = hashlib.sha256(normalized.encode()).hexdigest() + if normalized_sha != EXPECTED_NORMALIZED_SCHEDULER_SHA256: + raise ValueError( + "scheduler source contains changes outside reviewed boundary motion" + ) + patch_sha = _scheduler_patch_digest() + if patch_sha != EXPECTED_SCHEDULER_BOUNDARY_PATCH_SHA256: + raise RuntimeError("scheduler reviewed boundary patch constant is invalid") + return label, normalized_sha + + +def _classify_scheduler_source(source: str) -> tuple[str, list[str]]: + label, _normalized_sha = _normalize_scheduler_source(source) + return label, list(_ARM_EVENTS[label]) + + +def attest_scheduler_arm(repo: Path) -> dict[str, Any]: + path = repo.resolve() / _SCHEDULER_SOURCE + source_sha = _sha256(path) + source = path.read_bytes().decode("utf-8") + label, normalized_sha = _normalize_scheduler_source(source) + events = list(_ARM_EVENTS[label]) + return { + "label": label, + "source_path": _SCHEDULER_SOURCE, + "source_sha256": source_sha, + "arm_id": f"{label}:{source_sha}", + "normalized_source_sha256": normalized_sha, + "reviewed_boundary_patch_sha256": _scheduler_patch_digest(), + "sanctioned_event_sequence": events, + } + + +def attest_imported_mtplx_modules( + repo: Path, + *, + git_identity: dict[str, Any], + source_identity: dict[str, Any], + modules: dict[str, ModuleType] | None = None, +) -> dict[str, Any]: + """Bind every imported MTPLX Python module to the reviewed worktree.""" + + root = repo.resolve() + observed_modules = sys.modules if modules is None else modules + files = {} + for name, module in sorted(observed_modules.items()): + if name != "mtplx" and not name.startswith("mtplx."): + continue + path_text = getattr(module, "__file__", None) + if not path_text: + continue + path = Path(path_text).resolve() + try: + relative = path.relative_to(root) + except ValueError as exc: + raise RuntimeError( + f"imported reviewed module {name} is outside worktree: {path}" + ) from exc + relative_text = str(relative) + actual_sha = _sha256(path) + head_sha = git_identity.get("head_python_sha256", {}).get(relative_text) + if actual_sha != head_sha: + raise RuntimeError( + f"imported reviewed module {name} does not match preflight HEAD" + ) + source_sha = source_identity.get("importable_mtplx_files", {}).get( + relative_text + ) + if actual_sha != source_sha: + raise RuntimeError( + f"imported reviewed module {name} does not match source attestation" + ) + files[name] = { + "path": relative_text, + "sha256": actual_sha, + "head_sha256": head_sha, + "reviewed_source_sha256": source_sha, + } + missing = sorted(set(_REQUIRED_IMPORTED_MODULES) - set(files)) + if missing: + raise RuntimeError(f"reviewed MTPLX modules were not imported: {missing}") + return { + "files": files, + "module_set_sha256": hashlib.sha256(_canonical_bytes(files)).hexdigest(), + "preflight_head_bound": True, + "preflight_sources_bound": True, + } + + +def _state_manifest(value: Any) -> Any: + if value is None or isinstance(value, (bool, int, float, str)): + return value + if isinstance(value, bytes): + return { + "kind": "bytes", + "nbytes": len(value), + "sha256": hashlib.sha256(value).hexdigest(), + } + if hasattr(value, "shape") and hasattr(value, "dtype"): + array = np.asarray(value) + payload = array.tobytes(order="C") + return { + "kind": "array", + "shape": [int(dimension) for dimension in array.shape], + "dtype": str(array.dtype), + "nbytes": len(payload), + "sha256": hashlib.sha256(payload).hexdigest(), + } + if isinstance(value, (list, tuple)): + return { + "kind": type(value).__name__, + "items": [_state_manifest(item) for item in value], + } + if isinstance(value, dict): + return { + "kind": "dict", + "items": { + str(key): _state_manifest(item) + for key, item in sorted(value.items(), key=lambda row: str(row[0])) + }, + } + try: + attributes = vars(value) + except TypeError as exc: + raise TypeError(f"unsupported state value: {type(value)!r}") from exc + return { + "kind": "object", + "class": f"{type(value).__module__}.{type(value).__qualname__}", + "attributes": { + key: _state_manifest(item) for key, item in sorted(attributes.items()) + }, + } + + +def _metadata_manifest(value: Any) -> Any: + if isinstance(value, list): + return [_metadata_manifest(item) for item in value] + if isinstance(value, dict): + return { + key: _metadata_manifest(item) + for key, item in value.items() + if not (value.get("kind") in {"array", "bytes"} and key == "sha256") + } + return value + + +def _array_totals(value: Any) -> tuple[int, int]: + if isinstance(value, list): + rows = [_array_totals(item) for item in value] + return sum(row[0] for row in rows), sum(row[1] for row in rows) + if isinstance(value, dict): + if value.get("kind") == "array": + return 1, int(value["nbytes"]) + rows = [_array_totals(item) for item in value.values()] + return sum(row[0] for row in rows), sum(row[1] for row in rows) + return 0, 0 + + +def _state_receipt(value: Any) -> dict[str, Any]: + manifest = _state_manifest(value) + metadata_manifest = _metadata_manifest(manifest) + array_count, array_bytes = _array_totals(manifest) + return { + "state_sha256": hashlib.sha256(_canonical_bytes(manifest)).hexdigest(), + "metadata_sha256": hashlib.sha256( + _canonical_bytes(metadata_manifest) + ).hexdigest(), + "array_count": array_count, + "array_bytes": array_bytes, + } + + +class _BackendCapture: + def __init__(self, backend: Any, proposal_caches: list[Any]): + self._backend = backend + self._proposal_caches = proposal_caches + + def make_cache(self, rt: Any) -> Any: + cache = self._backend.make_cache(rt) + self._proposal_caches.append(cache) + return cache + + def __getattr__(self, name: str) -> Any: + return getattr(self._backend, name) + + +def prove_prefill_state( + runtime: Any, + prompt_ids: list[int], + *, + generate_ar: Callable[..., Any], + generate_mtpk: Callable[..., Any], + sampler: Any, +) -> dict[str, Any]: + """Compare state after three target rows and one complete K2 cycle.""" + + backend = runtime.block_speculative_backend + target_caches: list[Any] = [] + proposal_caches: list[Any] = [] + original_make_cache = runtime.make_cache + absent = object() + original_instance_make_cache = vars(runtime).get("make_cache", absent) + + def capture_target_cache() -> Any: + cache = original_make_cache() + target_caches.append(cache) + return cache + + runtime.make_cache = capture_target_cache + runtime.block_speculative_backend = _BackendCapture(backend, proposal_caches) + try: + ar_output = generate_ar( + runtime, + list(prompt_ids), + **_generation_kwargs(4, sampler), + ) + if len(target_caches) != 1: + raise RuntimeError("AR state proof did not create exactly one target cache") + ar_target_cache = target_caches[0] + k2_output = generate_mtpk( + runtime, + list(prompt_ids), + speculative_depth=2, + **_generation_kwargs(3, sampler), + ) + if len(target_caches) != 2 or len(proposal_caches) != 1: + raise RuntimeError("K2 state proof did not expose exact cache ownership") + k2_target_cache = target_caches[1] + proposal_snapshot = backend.snapshot(proposal_caches[0]) + finally: + if original_instance_make_cache is absent: + del runtime.make_cache + else: + runtime.make_cache = original_instance_make_cache + runtime.block_speculative_backend = backend + + if len(ar_output.tokens) != 4 or len(k2_output.tokens) != 3: + raise RuntimeError("state proof did not emit the required AR/K2 control rows") + if list(ar_output.tokens[:3]) != list(k2_output.tokens): + raise RuntimeError("state proof AR/K2 target token prefix is not exact") + k2_stats = k2_output.stats + drafted_by_depth = [int(value) for value in k2_stats.drafted_by_depth] + complete_k2_cycle = ( + int(k2_stats.verify_calls) >= 1 + and len(drafted_by_depth) >= 2 + and drafted_by_depth[0] >= 1 + and drafted_by_depth[1] >= 1 + ) + if not complete_k2_cycle: + raise RuntimeError("state proof did not execute one complete K2 cycle") + ar_target = _state_receipt(ar_target_cache) + k2_target = _state_receipt(k2_target_cache) + target_equal = ar_target == k2_target + if not target_equal: + raise RuntimeError("AR/K2 target cache state is not bit-exact after K2 cycle") + return { + "measured": False, + "ar_max_tokens": 4, + "k2_max_tokens": 3, + "target_rows_consumed": 3, + "semantic_boundary": "three_serial_target_rows_after_prompt", + "complete_k2_cycle": True, + "k2_verify_calls": int(k2_stats.verify_calls), + "k2_drafted_by_depth": drafted_by_depth, + "target_token_prefix": [int(token) for token in k2_output.tokens], + "wrappers_restored_before_primers": True, + "ar_target": ar_target, + "k2_target": k2_target, + "target_state_equal": True, + "proposal_snapshot": _state_receipt(proposal_snapshot), + } + + +def _generation_kwargs(max_tokens: int, sampler: Any) -> dict[str, Any]: + return { + "max_tokens": int(max_tokens), + "sampler": sampler, + "seed": 0, + "stop_token_ids": set(), + } + + +def _measurement(output: Any, mx: Any) -> dict[str, Any]: + stats = output.stats + decode_tok_s = float(stats.decode_tok_s) + end_to_end_tok_s = float(stats.end_to_end_tok_s) + if decode_tok_s <= 0.0 or end_to_end_tok_s <= 0.0: + raise RuntimeError("measured throughput must be positive") + return { + "tokens": [int(token) for token in output.tokens], + "generated_tokens": len(output.tokens), + "decode_tok_s": decode_tok_s, + "end_to_end_tok_s": end_to_end_tok_s, + "prompt_eval_time_s": float(stats.prompt_eval_time_s), + "prompt_target_prefill_time_s": float( + getattr(stats, "prompt_target_prefill_time_s", 0.0) + ), + "prompt_mtp_history_time_s": float( + getattr(stats, "prompt_mtp_history_time_s", 0.0) + ), + "prompt_target_prefill_tok_s": float( + getattr(stats, "prompt_target_prefill_tok_s", 0.0) + ), + "accepted_drafts": int(getattr(stats, "accepted_drafts", 0)), + "rejected_drafts": int(getattr(stats, "rejected_drafts", 0)), + "drafted_tokens": int(getattr(stats, "drafted_tokens", 0)), + "accepted_by_depth": [ + int(value) for value in getattr(stats, "accepted_by_depth", []) + ], + "drafted_by_depth": [ + int(value) for value in getattr(stats, "drafted_by_depth", []) + ], + "verify_calls": int(getattr(stats, "verify_calls", 0)), + "peak_memory_bytes": int(mx.get_peak_memory()), + "active_memory_bytes": int(mx.get_active_memory()), + } + + +def _acceptance_signature(measurement: dict[str, Any]) -> dict[str, Any]: + return { + key: measurement[key] + for key in ( + "accepted_drafts", + "rejected_drafts", + "drafted_tokens", + "accepted_by_depth", + "drafted_by_depth", + "verify_calls", + ) + } + + +def run_benchmark( + args: argparse.Namespace, + *, + mx: Any, + runtime_load: Callable[..., Any], + generate_ar: Callable[..., Any], + generate_mtpk: Callable[..., Any], + sampler_factory: Callable[..., Any], + imported_modules_attestation: Callable[[], dict[str, Any]], + post_run_git_attestation: Callable[[], dict[str, Any]], + mlx_identity: dict[str, Any], + model_identity: dict[str, Any], + git_identity: dict[str, Any], + source_identity: dict[str, Any], + scheduler_arm: dict[str, Any], + guard_attestation: dict[str, Any], +) -> dict[str, Any]: + """Run one source-derived arm with one load and fixed repetitions.""" + + if int(args.max_tokens) <= 0: + raise ValueError("--max-tokens must be positive") + if git_identity.get("dirty") is not False: + raise ValueError("benchmark provenance must be a clean HEAD tree") + _label, scheduler_sha = _arm_from_receipt({"scheduler_arm": scheduler_arm}) + if source_identity.get("files", {}).get(_SCHEDULER_SOURCE) != scheduler_sha: + raise ValueError("scheduler arm hash does not match source provenance") + + runtime = runtime_load(args.model, mtp=True) + backend = getattr(runtime, "block_speculative_backend", None) + if getattr(backend, "backend_id", None) != "deepseek_v4_dspark_0731": + raise ValueError("loaded runtime has no native DeepSeek-V4 DSpark backend") + if getattr(runtime, "deepseek_v4_0731_k2_receipt", None) is not None: + raise ValueError("scheduler bracket must keep explicit 0731 kernels stock") + imported_modules_pre_run = imported_modules_attestation() + if not ( + imported_modules_pre_run.get("preflight_head_bound") is True + and imported_modules_pre_run.get("preflight_sources_bound") is True + ): + raise RuntimeError("imported MTPLX modules lack preflight source binding") + prompt_ids = [int(token) for token in runtime.tokenizer.encode(PROMPT_TEXT)] + if len(prompt_ids) != PROMPT_TOKEN_COUNT: + raise RuntimeError( + "fixed prompt tokenizer drift: expected " + f"{PROMPT_TOKEN_COUNT} tokens, got {len(prompt_ids)}" + ) + + sampler = sampler_factory(temperature=0.0, top_p=1.0, top_k=0) + kwargs = _generation_kwargs(args.max_tokens, sampler) + state_proof = prove_prefill_state( + runtime, + prompt_ids, + generate_ar=generate_ar, + generate_mtpk=generate_mtpk, + sampler=sampler, + ) + + generate_ar(runtime, list(prompt_ids), **kwargs) + generate_mtpk( + runtime, + list(prompt_ids), + speculative_depth=2, + **kwargs, + ) + + samples = [] + for repetition in range(1, REPETITIONS + 1): + mx.reset_peak_memory() + ar_output = generate_ar(runtime, list(prompt_ids), **kwargs) + ar_measurement = _measurement(ar_output, mx) + + mx.reset_peak_memory() + k2_output = generate_mtpk( + runtime, + list(prompt_ids), + speculative_depth=2, + **kwargs, + ) + k2_measurement = _measurement(k2_output, mx) + samples.append( + { + "repetition": repetition, + "ar": ar_measurement, + "k2": k2_measurement, + "exact_vs_ar": ar_measurement["tokens"] == k2_measurement["tokens"], + "acceptance_signature": _acceptance_signature(k2_measurement), + } + ) + + ar_reference = samples[0]["ar"]["tokens"] + k2_reference = samples[0]["k2"]["tokens"] + acceptance_reference = samples[0]["acceptance_signature"] + ar_deterministic = all(row["ar"]["tokens"] == ar_reference for row in samples) + k2_deterministic = all(row["k2"]["tokens"] == k2_reference for row in samples) + acceptance_deterministic = all( + row["acceptance_signature"] == acceptance_reference for row in samples + ) + exact_all_samples = all(row["exact_vs_ar"] for row in samples) + gates = { + "state_proof_target_equal": state_proof["target_state_equal"], + "tokens_exact_all_samples": exact_all_samples, + "ar_deterministic": ar_deterministic, + "k2_deterministic": k2_deterministic, + "acceptance_signature_identical_all_samples": acceptance_deterministic, + } + passed = all(gates.values()) + ar_decode_median = float(median(row["ar"]["decode_tok_s"] for row in samples)) + k2_decode_median = float(median(row["k2"]["decode_tok_s"] for row in samples)) + ar_end_to_end_median = float( + median(row["ar"]["end_to_end_tok_s"] for row in samples) + ) + k2_end_to_end_median = float( + median(row["k2"]["end_to_end_tok_s"] for row in samples) + ) + imported_modules = imported_modules_attestation() + if not ( + imported_modules.get("preflight_head_bound") is True + and imported_modules.get("preflight_sources_bound") is True + ): + raise RuntimeError("post-run MTPLX imports lack preflight source binding") + for name, identity in imported_modules_pre_run["files"].items(): + if imported_modules.get("files", {}).get(name) != identity: + raise RuntimeError("imported MTPLX module identity changed during bracket") + post_run_git = post_run_git_attestation() + if post_run_git != git_identity: + raise RuntimeError("repository provenance changed during scheduler bracket") + + return { + "schema_version": 2, + "kind": "deepseek_v4_0731_scheduler_boundary_benchmark", + "scheduler_arm": scheduler_arm, + "single_model_load": True, + "baseline": "generic_mtp_true_stock", + "load_kwargs": {"mtp": True}, + "prompt": { + "text": PROMPT_TEXT, + "token_ids": prompt_ids, + "tokens": len(prompt_ids), + }, + "max_tokens": int(args.max_tokens), + "repetitions": REPETITIONS, + "speculative_depth": 2, + "sampling": { + "temperature": 0.0, + "top_p": 1.0, + "top_k": 0, + "seed": 0, + "stop_token_ids": [], + }, + "provenance": { + "mlx": mlx_identity, + "model": model_identity, + "git": git_identity, + "git_post_run": post_run_git, + "sources": source_identity, + "imported_mtplx_modules_pre_run": imported_modules_pre_run, + "imported_mtplx_modules": imported_modules, + }, + "guard_attestation": guard_attestation, + "state_proof": state_proof, + "primers": { + "ar": {"executed": True, "measured": False}, + "k2": { + "executed": True, + "measured": False, + "speculative_depth": 2, + }, + }, + "measurements": { + "samples": samples, + "summary": { + "ar": { + "median_decode_tok_s": ar_decode_median, + "median_end_to_end_tok_s": ar_end_to_end_median, + }, + "k2": { + "median_decode_tok_s": k2_decode_median, + "median_end_to_end_tok_s": k2_end_to_end_median, + }, + "k2_over_ar_decode_ratio": k2_decode_median / ar_decode_median, + "k2_over_ar_end_to_end_ratio": ( + k2_end_to_end_median / ar_end_to_end_median + ), + }, + }, + "deterministic": { + "ar": ar_deterministic, + "k2": k2_deterministic, + "acceptance_signature": acceptance_deterministic, + }, + "exact_vs_ar": exact_all_samples and ar_deterministic and k2_deterministic, + "gates": gates, + "passed": passed, + } + + +def _arm_from_receipt(receipt: dict[str, Any]) -> tuple[str, str]: + arm = receipt.get("scheduler_arm") or {} + label = arm.get("label") + source_sha = arm.get("source_sha256") + if ( + label not in _ARM_LABELS + or arm.get("arm_id") != f"{label}:{source_sha}" + or arm.get("normalized_source_sha256") != EXPECTED_NORMALIZED_SCHEDULER_SHA256 + or arm.get("reviewed_boundary_patch_sha256") + != EXPECTED_SCHEDULER_BOUNDARY_PATCH_SHA256 + or arm.get("sanctioned_event_sequence") != _ARM_EVENTS.get(label) + ): + raise ValueError("receipt scheduler arm attribution is invalid") + if not isinstance(source_sha, str) or len(source_sha) != 64: + raise ValueError("receipt scheduler source hash is invalid") + return label, source_sha + + +def compare_receipts(first: dict[str, Any], second: dict[str, Any]) -> dict[str, Any]: + """Gate a clean lazy/materialize pair without importing MLX.""" + + receipts = [first, second] + arms = {_arm_from_receipt(receipt)[0]: receipt for receipt in receipts} + if set(arms) != _ARM_LABELS: + raise ValueError("comparison requires one lazy and one materialize receipt") + lazy = arms["lazy_joint_eval"] + materialize = arms["materialize_first"] + lazy_sha = lazy["scheduler_arm"]["source_sha256"] + materialize_sha = materialize["scheduler_arm"]["source_sha256"] + + source_keys = set(lazy["provenance"]["sources"]["files"]) | set( + materialize["provenance"]["sources"]["files"] + ) + source_differences = sorted( + key + for key in source_keys + if lazy["provenance"]["sources"]["files"].get(key) + != materialize["provenance"]["sources"]["files"].get(key) + ) + head_file_keys = set(lazy["provenance"]["git"]["head_tree_files"]) | set( + materialize["provenance"]["git"]["head_tree_files"] + ) + head_tree_differences = sorted( + key + for key in head_file_keys + if lazy["provenance"]["git"]["head_tree_files"].get(key) + != materialize["provenance"]["git"]["head_tree_files"].get(key) + ) + lazy_imports = lazy["provenance"]["imported_mtplx_modules"]["files"] + materialize_imports = materialize["provenance"]["imported_mtplx_modules"]["files"] + imported_names = set(lazy_imports) | set(materialize_imports) + imported_module_differences = sorted( + name + for name in imported_names + if lazy_imports.get(name) != materialize_imports.get(name) + ) + state_keys = ("ar_target", "k2_target", "proposal_snapshot") + state_equal = { + key: lazy["state_proof"][key] == materialize["state_proof"][key] + for key in state_keys + } + common_fields = ("load_kwargs", "prompt", "max_tokens", "repetitions", "sampling") + common_configuration = all(lazy[key] == materialize[key] for key in common_fields) + lazy_samples = lazy["measurements"]["samples"] + materialize_samples = materialize["measurements"]["samples"] + cross_arm_tokens = all( + left[lane]["tokens"] == right[lane]["tokens"] + for left, right in zip(lazy_samples, materialize_samples, strict=True) + for lane in ("ar", "k2") + ) + cross_arm_acceptance = all( + left["acceptance_signature"] == right["acceptance_signature"] + for left, right in zip(lazy_samples, materialize_samples, strict=True) + ) + mlx_fields = ("version", "core_sha256", "libmlx", "metallib") + same_mlx = all( + lazy["provenance"]["mlx"].get(key) == materialize["provenance"]["mlx"].get(key) + for key in mlx_fields + ) + model_fields = ("config_sha256", "index_sha256", "metadata") + same_model = all( + lazy["provenance"]["model"].get(key) + == materialize["provenance"]["model"].get(key) + for key in model_fields + ) + gates = { + "both_arms_passed": bool(lazy.get("passed") and materialize.get("passed")), + "source_hashes_distinct": lazy_sha != materialize_sha, + "head_trees_distinct": ( + lazy["provenance"]["git"]["head_tree"] + != materialize["provenance"]["git"]["head_tree"] + ), + "only_scheduler_source_differs": source_differences == [_SCHEDULER_SOURCE], + "only_scheduler_head_blob_differs": head_tree_differences + == [_SCHEDULER_SOURCE], + "only_scheduler_import_differs": imported_module_differences + == ["mtplx.native_block_speculation"], + "normalized_scheduler_source_identical": ( + lazy["scheduler_arm"]["normalized_source_sha256"] + == materialize["scheduler_arm"]["normalized_source_sha256"] + == EXPECTED_NORMALIZED_SCHEDULER_SHA256 + ), + "reviewed_boundary_patch_identical": ( + lazy["scheduler_arm"]["reviewed_boundary_patch_sha256"] + == materialize["scheduler_arm"]["reviewed_boundary_patch_sha256"] + == EXPECTED_SCHEDULER_BOUNDARY_PATCH_SHA256 + ), + "preflight_postrun_git_identical_within_arms": all( + receipt["provenance"]["git"] == receipt["provenance"]["git_post_run"] + for receipt in (lazy, materialize) + ), + "common_configuration": common_configuration, + "official_mlx_identical": same_mlx, + "model_identity_identical": same_model, + "ar_target_state_identical": state_equal["ar_target"], + "k2_target_state_identical": state_equal["k2_target"], + "proposal_snapshot_identical": state_equal["proposal_snapshot"], + "tokens_identical_cross_arm": cross_arm_tokens, + "acceptance_signature_identical_cross_arm": cross_arm_acceptance, + } + return { + "schema_version": 1, + "kind": "deepseek_v4_0731_scheduler_boundary_comparison", + "arms": { + "lazy_joint_eval": lazy["scheduler_arm"], + "materialize_first": materialize["scheduler_arm"], + }, + "source_differences": source_differences, + "head_tree_differences": head_tree_differences, + "imported_module_differences": imported_module_differences, + "state_digests_equal": state_equal, + "gates": gates, + "passed": all(gates.values()), + } + + +def write_receipt(receipt: dict[str, Any], output_path: Path) -> int: + path = output_path.expanduser().resolve() + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps(receipt, indent=2, sort_keys=True, allow_nan=False) + "\n", + encoding="utf-8", + ) + return 0 if receipt.get("passed", receipt.get("exact_vs_ar", False)) else 1 + + +def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + action = parser.add_mutually_exclusive_group(required=True) + action.add_argument("--model", type=Path) + action.add_argument( + "--compare", + nargs=2, + type=Path, + metavar=("LAZY_RECEIPT", "MATERIALIZE_RECEIPT"), + ) + parser.add_argument("--max-tokens", type=int, default=64) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args(argv) + if args.max_tokens <= 0: + parser.error("--max-tokens must be positive") + return args + + +def main(argv: list[str] | None = None) -> int: + args = _parse_args(argv) + if args.compare is not None: + first, second = ( + json.loads(path.expanduser().read_text(encoding="utf-8")) + for path in args.compare + ) + return write_receipt(compare_receipts(first, second), args.out) + + repo = Path(__file__).resolve().parents[1] + # These source/provenance checks intentionally precede the guard bridge and + # MLX imports. A dirty arm never initializes Metal or loads model weights. + git_identity = attest_git(repo) + source_identity = attest_sources(repo) + scheduler_arm = attest_scheduler_arm(repo) + + from deepseek_v4_guard_window import ( + WINDOW_PATH_ENV, + WINDOW_SHA256_ENV, + issue_guard_window, + load_verified_guard_window, + ) + + guard_path, guard_digest = issue_guard_window() + try: + guard_attestation = load_verified_guard_window( + environment={ + WINDOW_PATH_ENV: str(guard_path), + WINDOW_SHA256_ENV: guard_digest, + } + ) + import mlx.core as mx + + mlx_identity = attest_official_mlx(mx, metadata.distribution("mlx")) + model_identity = attest_model(args.model) + + from mtplx import deepseek_v4_dspark_generation as _adapter_module # noqa: F401 + from mtplx import generation as generation_module + from mtplx import native_block_speculation as _scheduler_module # noqa: F401 + from mtplx import runtime as runtime_module + from mtplx import sampling as sampling_module + from mtplx.models import deepseek_v4 as _model_module # noqa: F401 + + receipt = run_benchmark( + args, + mx=mx, + runtime_load=runtime_module.load, + generate_ar=generation_module.generate_ar, + generate_mtpk=generation_module.generate_mtpk, + sampler_factory=sampling_module.SamplerConfig, + imported_modules_attestation=lambda: attest_imported_mtplx_modules( + repo, + git_identity=git_identity, + source_identity=source_identity, + ), + post_run_git_attestation=lambda: attest_git(repo), + mlx_identity=mlx_identity, + model_identity=model_identity, + git_identity=git_identity, + source_identity=source_identity, + scheduler_arm=scheduler_arm, + guard_attestation=guard_attestation, + ) + return write_receipt(receipt, args.out) + finally: + try: + guard_path.unlink() + except FileNotFoundError: + pass + try: + guard_path.parent.rmdir() + except OSError: + pass + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/deepseek_v4_guard_window.py b/scripts/deepseek_v4_guard_window.py index a262f109..cfb75f2a 100755 --- a/scripts/deepseek_v4_guard_window.py +++ b/scripts/deepseek_v4_guard_window.py @@ -100,7 +100,9 @@ def _checked_attestation( ) if ( attestation.get("schema_version") != 1 - or any(isinstance(value, bool) or not isinstance(value, int) for value in integers) + or any( + isinstance(value, bool) or not isinstance(value, int) for value in integers + ) or not _valid_digest(attestation.get("nonce_sha256")) ): raise RuntimeError("repository guard attestation receipt is malformed") @@ -110,7 +112,10 @@ def _checked_attestation( raise RuntimeError("repository guard attestation expiry is malformed") lock_path = attestation.get("lock_path") resolved_lock = expected_lock.resolve(strict=True) - if not isinstance(lock_path, str) or Path(lock_path).resolve(strict=True) != resolved_lock: + if ( + not isinstance(lock_path, str) + or Path(lock_path).resolve(strict=True) != resolved_lock + ): raise RuntimeError( f"guard attested {lock_path!r}, expected lock {str(expected_lock)!r}" ) @@ -225,7 +230,9 @@ def load_verified_guard_window( try: document = json.loads(encoded) except (UnicodeDecodeError, json.JSONDecodeError) as error: - raise RuntimeError(f"verified guard window receipt is malformed: {error}") from error + raise RuntimeError( + f"verified guard window receipt is malformed: {error}" + ) from error if not isinstance(document, dict) or _canonical_json(document) != encoded: raise RuntimeError("verified guard window receipt is not canonical") attestation = document.get("attestation") diff --git a/tests/test_deepseek_v4_0731_dspark_ffn.py b/tests/test_deepseek_v4_0731_dspark_ffn.py new file mode 100644 index 00000000..b7ba7fee --- /dev/null +++ b/tests/test_deepseek_v4_0731_dspark_ffn.py @@ -0,0 +1,496 @@ +"""CPU construction gates for the retained DSpark native packed-Q3 lane.""" + +from __future__ import annotations + +import inspect +from types import SimpleNamespace + +import pytest + +pytest.importorskip("mlx.core") +import mlx.core as mx # noqa: E402 +import mlx.nn as nn # noqa: E402 +from mlx_lm.models.switch_layers import QuantizedSwitchLinear # noqa: E402 + +from mtplx.deepseek_v4_0731_dspark_ffn import ( # noqa: E402 + DSPARK_Q3_GATE_UP_GEOMETRY, + DeepseekV40731DSparkM5PackedSwitchGLU, + PreparedDSparkQ3PackedGateUpM5, + build_dspark_q3_packed_gate_up, + install_dspark_q3_packed_gate_up_m5, + prepare_dspark_q3_packed_gate_up_m5, + validate_dspark_q3_gate_up, +) +from mtplx.models.deepseek_v4 import ( # noqa: E402 + ClampedSwiGLU, + DeepseekV4DSpark, + DeepseekV4DSparkStage, + DeepseekV4MoE, + Model, + MoEGate, +) + + +@pytest.fixture(autouse=True) +def _cpu_default_device(): + previous = mx.default_device() + mx.set_default_device(mx.cpu) + try: + yield + finally: + mx.set_default_device(previous) + + +def _q3_pair(*, hidden_size: int = 128, width: int = 16, experts: int = 8): + gate = QuantizedSwitchLinear( + hidden_size, width, experts, bias=False, group_size=128, bits=3 + ) + up = QuantizedSwitchLinear( + hidden_size, width, experts, bias=False, group_size=128, bits=3 + ) + for projection in (gate, up): + projection.scales = projection.scales.astype(mx.bfloat16) + projection.biases = projection.biases.astype(mx.bfloat16) + mx.eval(gate.parameters(), up.parameters()) + return gate, up + + +def _q3_switch(*, hidden_size: int = 128, width: int = 128, experts: int = 8): + gate, up = _q3_pair( + hidden_size=hidden_size, + width=width, + experts=experts, + ) + down = QuantizedSwitchLinear( + width, hidden_size, experts, bias=False, group_size=128, bits=3 + ) + down.scales = down.scales.astype(mx.bfloat16) + down.biases = down.biases.astype(mx.bfloat16) + mx.eval(down.parameters()) + return SimpleNamespace( + gate_proj=gate, + up_proj=up, + down_proj=down, + activation=ClampedSwiGLU(10.0), + ) + + +def _projection_bytes(projection) -> int: + return sum(int(projection[name].nbytes) for name in ("weight", "scales", "biases")) + + +def _stage(stage_id: int, switch=None): + stage = DeepseekV4DSparkStage.__new__(DeepseekV4DSparkStage) + ffn = DeepseekV4MoE.__new__(DeepseekV4MoE) + gate = MoEGate.__new__(MoEGate) + object.__setattr__(gate, "topk", 6) + object.__setattr__(ffn, "gate", gate) + object.__setattr__(ffn, "switch_mlp", object() if switch is None else switch) + object.__setattr__(stage, "stage_id", stage_id) + object.__setattr__(stage, "block_size", 5) + object.__setattr__(stage, "ffn", ffn) + return stage + + +def _model_owner(stages=None): + stages = [_stage(index) for index in range(3)] if stages is None else list(stages) + dspark = object.__new__(DeepseekV4DSpark) + object.__setattr__(dspark, "stages", stages) + object.__setattr__(dspark, "block_size", 5) + owner = Model.__new__(Model) + object.__setattr__(owner, "_dspark", dspark) + object.__setattr__(owner, "mtp", SimpleNamespace(layers=stages)) + return owner + + +def test_retained_geometry_is_the_physical_five_row_dspark_layout(): + assert DSPARK_Q3_GATE_UP_GEOMETRY == { + "rows": 5, + "hidden_size": 4096, + "width": 2048, + "experts": 256, + "top_k": 6, + "bits": 3, + "group_size": 128, + "weight_shape": (256, 2048, 384), + "metadata_shape": (256, 2048, 32), + } + + +def test_loaded_storage_contract_accepts_only_affine_q3_group128_u32_bf16(): + gate, up = _q3_pair() + + contract = validate_dspark_q3_gate_up( + gate, + up, + hidden_size=128, + width=16, + experts=8, + top_k=6, + rows=5, + ) + + assert contract.bits == 3 + assert contract.group_size == 128 + assert contract.weight_shape == (8, 16, 12) + assert contract.metadata_shape == (8, 16, 1) + + +@pytest.mark.parametrize( + "mutator, match", + [ + (lambda gate, up: setattr(up, "bits", 4), "Q3"), + (lambda gate, up: setattr(up, "group_size", 64), "group-128"), + (lambda gate, up: setattr(up, "mode", "mxfp4"), "affine"), + (lambda gate, up: setattr(up, "biases", None), "biases"), + ], +) +def test_loaded_storage_contract_rejects_nonphysical_q3(mutator, match): + gate, up = _q3_pair() + mutator(gate, up) + + with pytest.raises(ValueError, match=match): + validate_dspark_q3_gate_up( + gate, + up, + hidden_size=128, + width=16, + experts=8, + top_k=6, + rows=5, + ) + + +def test_fixed_m5_pack_is_exact_one_dispatch_and_adds_no_resident_weight_bytes(): + switch = _q3_switch() + original_bytes = _projection_bytes(switch.gate_proj) + _projection_bytes( + switch.up_proj + ) + + packed = build_dspark_q3_packed_gate_up( + switch, + hidden_size=128, + width=128, + experts=8, + top_k=6, + rows=5, + ) + + assert type(packed) is DeepseekV40731DSparkM5PackedSwitchGLU + assert packed.down_proj is switch.down_proj + assert packed.activation is switch.activation + assert _projection_bytes(packed.gate_up_proj) == original_bytes + + x = (mx.arange(5 * 128).reshape(5, 128) % 17 - 8).astype(mx.bfloat16) + indices = mx.array([[0, 1, 2, 3, 4, 5]] * 5, dtype=mx.uint32) + expanded = mx.expand_dims(x, (-2, -3)) + stock = switch.down_proj( + switch.activation( + switch.up_proj(expanded, indices), + switch.gate_proj(expanded, indices), + ), + indices, + ).squeeze(-2) + candidate = packed(x, indices) + mx.eval(stock, candidate) + assert mx.array_equal(candidate, stock) + + class GatherSpy(nn.Module): + def __init__(self, projection): + super().__init__() + self.projection = projection + self.calls = 0 + self.sorted_indices = [] + + def gather(self, x, indices, sorted_indices): + self.calls += 1 + self.sorted_indices.append(sorted_indices) + return self.projection.gather(x, indices, sorted_indices) + + spy = GatherSpy(packed.gate_up_proj) + packed.gate_up_proj = spy + candidate = packed(x, indices) + mx.eval(candidate) + + assert spy.calls == 1 + assert spy.sorted_indices == [False] + assert tuple(candidate.shape) == (5, 6, 128) + source = inspect.getsource(DeepseekV40731DSparkM5PackedSwitchGLU.__call__) + assert "indices.size" not in source + assert "moe_force_unsorted" not in source + + +@pytest.mark.parametrize("stage_count", [2, 4]) +def test_installer_requires_exactly_three_dspark_stages(stage_count): + owner = _model_owner([_stage(index) for index in range(stage_count)]) + + with pytest.raises(ValueError, match="exactly three"): + prepare_dspark_q3_packed_gate_up_m5(owner) + + +def test_preparer_rejects_bare_arbitrary_stage_lists(): + with pytest.raises(ValueError, match="model owner"): + prepare_dspark_q3_packed_gate_up_m5([_stage(index) for index in range(3)]) + + +@pytest.mark.parametrize( + ("mutate", "match"), + [ + ( + lambda owner: owner._dspark.stages.__setitem__( + 1, SimpleNamespace(stage_id=1) + ), + "stage identity", + ), + ( + lambda owner: setattr(owner._dspark.stages[1], "stage_id", 2), + "stage order", + ), + ( + lambda owner: setattr(owner._dspark.stages[1], "ffn", SimpleNamespace()), + "FFN identity", + ), + ( + lambda owner: setattr( + owner._dspark.stages[1].ffn, "gate", SimpleNamespace(topk=6) + ), + "router identity", + ), + ( + lambda owner: setattr( + owner.mtp, "layers", list(reversed(owner._dspark.stages)) + ), + "model.mtp", + ), + ( + lambda owner: setattr(owner._dspark.stages[1].ffn.gate, "topk", 5), + "top-k=6", + ), + (lambda owner: setattr(owner._dspark, "block_size", 4), "M=5"), + ( + lambda owner: setattr(owner._dspark.stages[2], "block_size", 4), + "M=5", + ), + ], +) +def test_preparer_rejects_adversarial_dspark_ownership(mutate, match): + owner = _model_owner() + mutate(owner) + + with pytest.raises(ValueError, match=match): + prepare_dspark_q3_packed_gate_up_m5(owner) + + +def test_installer_validates_and_builds_every_stage_before_atomic_publication( + monkeypatch, +): + original = [object(), object(), object()] + stages = [_stage(index, switch) for index, switch in enumerate(original)] + owner = _model_owner(stages) + replacements = [object(), object(), object()] + calls = [] + + def validate(switch, **geometry): + calls.append(("validate", switch, geometry)) + + def build(switch, **geometry): + calls.append(("build", switch, geometry)) + return replacements[len([call for call in calls if call[0] == "build"]) - 1] + + monkeypatch.setattr( + "mtplx.deepseek_v4_0731_dspark_ffn._validate_dspark_q3_switch", + validate, + ) + monkeypatch.setattr( + "mtplx.deepseek_v4_0731_dspark_ffn.build_dspark_q3_packed_gate_up", + build, + ) + + prepared = prepare_dspark_q3_packed_gate_up_m5(owner) + + expected_geometry = { + "hidden_size": 4096, + "width": 2048, + "experts": 256, + "top_k": 6, + "rows": 5, + } + assert isinstance(prepared, PreparedDSparkQ3PackedGateUpM5) + assert [stage.ffn.switch_mlp for stage in stages] == original + assert [call[0] for call in calls] == ["validate"] * 3 + ["build"] * 3 + assert all(call[2] == expected_geometry for call in calls) + assert prepared.receipt == { + "candidate": "dspark-native-packed-q3-gate-up-m5", + "stages": 3, + "geometry": DSPARK_Q3_GATE_UP_GEOMETRY, + "gate_up_dispatches_per_stage": 1, + "stock_gate_up_dispatches_per_stage": 2, + "explicit_dequantize": False, + "resident_weight_bytes_added": 0, + } + prepared.publish() + assert [stage.ffn.switch_mlp for stage in stages] == replacements + prepared.restore() + assert [stage.ffn.switch_mlp for stage in stages] == original + + +def test_installer_leaves_every_original_stage_unchanged_when_build_fails(monkeypatch): + original = [object(), object(), object()] + stages = [_stage(index, switch) for index, switch in enumerate(original)] + owner = _model_owner(stages) + build_calls = 0 + + monkeypatch.setattr( + "mtplx.deepseek_v4_0731_dspark_ffn._validate_dspark_q3_switch", + lambda *_args, **_kwargs: None, + ) + + def build(*_args, **_kwargs): + nonlocal build_calls + build_calls += 1 + if build_calls == 2: + raise ValueError("stage two failed") + return object() + + monkeypatch.setattr( + "mtplx.deepseek_v4_0731_dspark_ffn.build_dspark_q3_packed_gate_up", + build, + ) + + with pytest.raises(ValueError, match="stage two failed"): + prepare_dspark_q3_packed_gate_up_m5(owner) + + assert [stage.ffn.switch_mlp for stage in stages] == original + + +def test_prepared_publication_restores_every_stage_if_one_assignment_fails( + monkeypatch, +): + original = [object(), object(), object()] + replacements = [object(), object(), object()] + stages = [_stage(index, switch) for index, switch in enumerate(original)] + owner = _model_owner(stages) + monkeypatch.setattr( + "mtplx.deepseek_v4_0731_dspark_ffn._validate_dspark_q3_switch", + lambda *_args, **_kwargs: None, + ) + built = iter(replacements) + monkeypatch.setattr( + "mtplx.deepseek_v4_0731_dspark_ffn.build_dspark_q3_packed_gate_up", + lambda *_args, **_kwargs: next(built), + ) + prepared = prepare_dspark_q3_packed_gate_up_m5(owner) + original_setattr = DeepseekV4MoE.__setattr__ + + def guarded_setattr(self, name, value): + if self is stages[1].ffn and name == "switch_mlp" and value is replacements[1]: + raise RuntimeError("publication failed") + return original_setattr(self, name, value) + + monkeypatch.setattr(DeepseekV4MoE, "__setattr__", guarded_setattr) + + with pytest.raises(RuntimeError, match="publication failed"): + prepared.publish() + + assert [stage.ffn.switch_mlp for stage in stages] == original + + +def test_prepared_restore_attempts_every_stage_and_groups_setter_failures(monkeypatch): + originals = [object(), object(), object()] + replacements = [object(), object(), object()] + stages = [ + _stage(index, replacement) for index, replacement in enumerate(replacements) + ] + prepared = PreparedDSparkQ3PackedGateUpM5( + stages=tuple(stages), + originals=tuple(originals), + replacements=tuple(replacements), + receipt={}, + ) + original_setattr = DeepseekV4MoE.__setattr__ + attempts = [] + + def guarded_setattr(self, name, value): + if name == "switch_mlp": + stage_index = next( + index for index, stage in enumerate(stages) if stage.ffn is self + ) + if value is originals[stage_index]: + attempts.append(stage_index) + if stage_index in {0, 1}: + raise RuntimeError(f"stage {stage_index} restoration failed") + return original_setattr(self, name, value) + + monkeypatch.setattr(DeepseekV4MoE, "__setattr__", guarded_setattr) + + with pytest.raises(ExceptionGroup, match="DSpark FFN restoration failed") as exc: + prepared.restore() + + assert attempts == [0, 1, 2] + assert [str(error) for error in exc.value.exceptions] == [ + "stage 0 restoration failed", + "stage 1 restoration failed", + ] + assert stages[0].ffn.switch_mlp is replacements[0] + assert stages[1].ffn.switch_mlp is replacements[1] + assert stages[2].ffn.switch_mlp is originals[2] + + +def test_publication_keeps_original_error_and_notes_grouped_rollback_failure( + monkeypatch, +): + originals = [object(), object(), object()] + replacements = [object(), object(), object()] + stages = [_stage(index, original) for index, original in enumerate(originals)] + prepared = PreparedDSparkQ3PackedGateUpM5( + stages=tuple(stages), + originals=tuple(originals), + replacements=tuple(replacements), + receipt={}, + ) + original_setattr = DeepseekV4MoE.__setattr__ + publication_error = RuntimeError("stage 1 publication failed") + restore_attempts = [] + + def guarded_setattr(self, name, value): + if name == "switch_mlp": + stage_index = next( + index for index, stage in enumerate(stages) if stage.ffn is self + ) + if stage_index == 1 and value is replacements[1]: + raise publication_error + if value is originals[stage_index]: + restore_attempts.append(stage_index) + if stage_index == 0: + raise RuntimeError("stage 0 rollback failed") + return original_setattr(self, name, value) + + monkeypatch.setattr(DeepseekV4MoE, "__setattr__", guarded_setattr) + + with pytest.raises(RuntimeError, match="stage 1 publication failed") as exc: + prepared.publish() + + assert exc.value is publication_error + assert restore_attempts == [0, 1, 2] + assert exc.value.__notes__ == [ + "DSpark FFN publication rollback also failed: " + "DSpark FFN restoration failed (1 sub-exception)" + ] + + +def test_convenience_installer_prepares_then_publishes(monkeypatch): + owner = _model_owner() + replacements = [object(), object(), object()] + monkeypatch.setattr( + "mtplx.deepseek_v4_0731_dspark_ffn._validate_dspark_q3_switch", + lambda *_args, **_kwargs: None, + ) + monkeypatch.setattr( + "mtplx.deepseek_v4_0731_dspark_ffn.build_dspark_q3_packed_gate_up", + lambda switch, **_geometry: replacements.pop(0), + ) + + receipt = install_dspark_q3_packed_gate_up_m5(owner) + + assert receipt["candidate"] == "dspark-native-packed-q3-gate-up-m5" + assert all(stage.ffn.switch_mlp is not None for stage in owner._dspark.stages) diff --git a/tests/test_deepseek_v4_0731_full_install.py b/tests/test_deepseek_v4_0731_full_install.py new file mode 100644 index 00000000..8c922d54 --- /dev/null +++ b/tests/test_deepseek_v4_0731_full_install.py @@ -0,0 +1,663 @@ +"""Construction gates for the single receipt-backed 0731 target stack.""" + +from __future__ import annotations + +import hashlib +import inspect +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +pytest.importorskip("mlx.core") +import mlx.core as mx # noqa: E402 + +from mtplx import deepseek_v4_0731_full_install as full + + +@pytest.fixture(autouse=True) +def _cpu_default_device(): + previous = mx.default_device() + mx.set_default_device(mx.cpu) + try: + yield + finally: + mx.set_default_device(previous) + + +@pytest.fixture +def full_artifact(tmp_path: Path, monkeypatch): + config = { + "model_type": "deepseek_v4", + "hidden_size": 4096, + "num_hidden_layers": 43, + "num_attention_heads": 64, + "num_key_value_heads": 1, + "head_dim": 512, + "n_routed_experts": 256, + "num_experts_per_tok": 6, + "moe_intermediate_size": 2048, + "n_shared_experts": 1, + "swiglu_limit": 10.0, + "num_nextn_predict_layers": 1, + "dspark_block_size": 5, + "dspark_noise_token_id": 128799, + "dspark_target_layer_ids": [40, 41, 42], + "dspark_markov_rank": 256, + } + config_bytes = json.dumps(config, sort_keys=True).encode() + index_bytes = b'{"metadata":{},"weight_map":{"model.layers.0":"a.safetensors"}}' + (tmp_path / "config.json").write_bytes(config_bytes) + (tmp_path / "model.safetensors.index.json").write_bytes(index_bytes) + metadata_root = tmp_path / ".cache/huggingface/download" + metadata_root.mkdir(parents=True) + for name in ("config.json.metadata", "model.safetensors.index.json.metadata"): + (metadata_root / name).write_text( + f"{full.EXPECTED_SOURCE_REVISION}\nblob\ntimestamp\n", + encoding="utf-8", + ) + monkeypatch.setattr( + full, + "EXPECTED_FULL_CONFIG_SHA256", + hashlib.sha256(config_bytes).hexdigest(), + ) + monkeypatch.setattr( + full, + "EXPECTED_FULL_INDEX_SHA256", + hashlib.sha256(index_bytes).hexdigest(), + ) + return tmp_path, config + + +def _model(): + switches = [ + SimpleNamespace(gate_proj=object(), up_proj=object()) for _ in range(43) + ] + + class WOBProjection(SimpleNamespace): + def __call__(self, value): + return value + + def attention(): + def qhead_stock(qr, _cos, _sin): + return qr + + wq_b = SimpleNamespace( + bits=6, + group_size=128, + mode="affine", + bias=None, + weight=SimpleNamespace(shape=(32768, 192), dtype=mx.uint32), + scales=SimpleNamespace(shape=(32768, 8), dtype=mx.bfloat16), + biases=SimpleNamespace(shape=(32768, 8), dtype=mx.bfloat16), + ) + wo_b = WOBProjection( + bits=6, + group_size=128, + mode="affine", + bias=None, + weight=SimpleNamespace(shape=(4096, 1536), dtype=mx.uint32), + scales=SimpleNamespace(shape=(4096, 64), dtype=mx.bfloat16), + biases=SimpleNamespace(shape=(4096, 64), dtype=mx.bfloat16), + ) + return SimpleNamespace( + wq_b=wq_b, + _q_projection_qhead_route=qhead_stock, + wo_b=wo_b, + _o_lora_impl=SimpleNamespace(wo_b=wo_b), + ) + + layers = [ + SimpleNamespace(ffn=SimpleNamespace(switch_mlp=switch), attn=attention()) + for switch in switches + ] + stages = (object(), object(), object()) + stock_calls = [] + + def stock(owner, input_ids, cache=None): + stock_calls.append((owner, input_ids, cache)) + return "stock-hidden", "stock-taps" + + model = SimpleNamespace( + model=SimpleNamespace(layers=layers), + mtp=stages, + _dspark=SimpleNamespace(target_layer_ids=(40, 41, 42), stages=stages), + _target_hidden_route=stock, + ) + return model, layers, switches, stock_calls + + +def test_artifact_contract_pins_config_index_and_both_hf_revisions(full_artifact): + path, config = full_artifact + + contract = full.validate_full_0731_dspark_artifact(path, config) + + assert contract.layers == 43 + assert contract.target_layer_ids == (40, 41, 42) + assert contract.stage_count == 3 + assert contract.source_revision == full.EXPECTED_SOURCE_REVISION + assert contract.config_sha256 == full.EXPECTED_FULL_CONFIG_SHA256 + assert contract.index_sha256 == full.EXPECTED_FULL_INDEX_SHA256 + + (path / "model.safetensors.index.json").write_bytes(b"wrong") + with pytest.raises(ValueError, match="index SHA-256"): + full.validate_full_0731_dspark_artifact(path, config) + + +def test_artifact_contract_rejects_null_or_wrong_topology_and_revision(full_artifact): + path, config = full_artifact + with pytest.raises(ValueError, match="dspark_block_size"): + full.validate_full_0731_dspark_artifact( + path, + {**config, "dspark_block_size": None}, + ) + + metadata = ( + path / ".cache/huggingface/download/model.safetensors.index.json.metadata" + ) + metadata.write_text("wrong-revision\nblob\ntimestamp\n", encoding="utf-8") + with pytest.raises(ValueError, match="metadata revision"): + full.validate_full_0731_dspark_artifact(path, config) + + +def _install_with_spies( + full_artifact, + monkeypatch, + *, + prepare_only=False, + fail_wob_prepare=False, + fail_publish=None, + fail_restore=None, + prepare_wqb_override=None, + prepare_wob_override=None, +): + path, config = full_artifact + model, layers, switches, stock_calls = _model() + stock_route = model._target_hidden_route + qhead_stocks = tuple(layer.attn._q_projection_qhead_route for layer in layers) + wob_stocks = tuple(layer.attn.wo_b for layer in layers) + replacements = [object() for _ in layers] + validations = [] + bindings = [] + m3_layers = [] + row_owned = object() + prepared = [] + publication = [] + combine_selfchecks = [] + projection_selfchecks = [] + + monkeypatch.setattr( + full, + "validate_routed_q2_pair", + lambda gate, up, **kwargs: validations.append((gate, up, kwargs)), + ) + + def build_pair(switch, **kwargs): + index = switches.index(switch) + return replacements[index] + + monkeypatch.setattr(full, "build_routed_q2_pair", build_pair) + monkeypatch.setattr( + full, + "build_row_owned_combine_m1", + lambda **kwargs: row_owned, + ) + monkeypatch.setattr( + full, + "exact_selfcheck_row_owned_combine_m1", + lambda combine: combine_selfchecks.append(combine), + ) + + def bind_tail(layer, **kwargs): + index = layers.index(layer) + assert [owned.ffn.switch_mlp for owned in layers] == switches + assert kwargs.pop("routed_switch") is replacements[index] + route = ("tail", index, kwargs["width"]) + bindings.append((layer, kwargs, route)) + return route + + monkeypatch.setattr(full.AI, "_bind_attention_island_layer", bind_tail) + monkeypatch.setattr( + full, + "build_m3_compiled_tail_layer", + lambda layer, tail: ( + m3_layers.append((layer, tail)) or ("m3-layer", layers.index(layer)) + ), + ) + + class M3Route: + def __init__(self, base): + self.base = base + + def __call__(self, owner, input_ids, cache=None): + if tuple(input_ids.shape) == (1, 3): + return "m3-hidden", "m3-taps" + return self.base(owner, input_ids, cache) + + monkeypatch.setattr( + full, + "build_0731_m3_target_route", + lambda owner, *, full_layer_routes, base_route: M3Route(base_route), + ) + + class Prepared: + def __init__(self, label): + self.label = label + self.published_routes = tuple(object() for _ in range(43)) + self.q6_count = 43 + self.exact_selfchecked = 43 + self.o_lora_sink_count = 43 + + def publish(self): + publication.append(f"{self.label}.publish") + if fail_publish == self.label: + raise RuntimeError(f"{self.label} publication failed") + + def restore(self): + publication.append(f"{self.label}.restore") + if fail_restore == self.label: + raise RuntimeError(f"{self.label} restoration failed") + + def prepare_wqb(layer_bank, *, exact_selfcheck): + assert tuple(layer_bank) == tuple(layers) + assert model._target_hidden_route is stock_route + assert [layer.ffn.switch_mlp for layer in layers] == switches + projection_selfchecks.append(("wqb", exact_selfcheck)) + prepared.append("wqb") + return Prepared("wqb") + + def prepare_wob(layer_bank, *, exact_selfcheck): + assert tuple(layer_bank) == tuple(layers) + assert model._target_hidden_route is stock_route + assert [layer.ffn.switch_mlp for layer in layers] == switches + projection_selfchecks.append(("wob", exact_selfcheck)) + prepared.append("wob") + if fail_wob_prepare: + raise RuntimeError("wob self-check failed") + return Prepared("wob") + + result = None + error = None + try: + operation = ( + full.prepare_full_0731_dspark_compiled_tail_q2_pair + if prepare_only + else full.install_full_0731_dspark_compiled_tail_q2_pair + ) + result = operation( + model, + config, + path, + prepare_wqb_qhead=prepare_wqb_override or prepare_wqb, + prepare_wob=prepare_wob_override or prepare_wob, + ) + except Exception as exc: # asserted by the failure test + error = exc + return SimpleNamespace( + result=result, + error=error, + model=model, + layers=layers, + switches=switches, + replacements=replacements, + stock_calls=stock_calls, + stock_route=stock_route, + qhead_stocks=qhead_stocks, + wob_stocks=wob_stocks, + validations=validations, + bindings=bindings, + m3_layers=m3_layers, + prepared=prepared, + publication=publication, + combine_selfchecks=combine_selfchecks, + projection_selfchecks=projection_selfchecks, + row_owned=row_owned, + ) + + +def test_installer_publishes_only_the_receipt_backed_configuration( + full_artifact, + monkeypatch, +): + state = _install_with_spies(full_artifact, monkeypatch) + + assert state.error is None + assert len(state.validations) == 43 + assert len(state.bindings) == 86 + assert len(state.m3_layers) == 43 + assert state.prepared == ["wqb", "wob"] + assert [label for label, _check in state.projection_selfchecks] == ["wqb", "wob"] + assert all(callable(check) for _label, check in state.projection_selfchecks) + assert state.combine_selfchecks == [state.row_owned] + assert state.publication == ["wqb.publish", "wob.publish"] + assert [layer.ffn.switch_mlp for layer in state.layers] == state.replacements + assert all( + kwargs + == { + "width": 1, + "allowed_widths": (1,), + "shared_bits": 8, + "routed_pair": True, + "routed_combine": state.row_owned, + } + for _layer, kwargs, _route in state.bindings[::2] + ) + assert all( + kwargs + == { + "width": 3, + "allowed_widths": (3,), + "shared_bits": 8, + "routed_pair": True, + } + for _layer, kwargs, _route in state.bindings[1::2] + ) + assert state.result == { + "candidate": "mtplx-full-dspark-compiled-tail-packed-q2-pair-m1-m3", + "artifact_label": full.RECORDED_ARTIFACT_LABEL, + "validated_config_sha256": full.EXPECTED_FULL_CONFIG_SHA256, + "validated_index_sha256": full.EXPECTED_FULL_INDEX_SHA256, + "validated_metadata_revision": full.EXPECTED_SOURCE_REVISION, + "layers_installed": 43, + "decode_m": 1, + "fixed_k": 2, + "physical_target_rows": 3, + "m3_tail": "fixed-width3-compiled-tail", + "m3_wqb": { + "candidate": "official-wheel-custom-fixed-m3-wqb-qhead-fused", + "layers_installed": 43, + "q6_g128_layers": 43, + "exact_selfchecked_layers": 43, + "shape": [1, 3, 1024], + "output_shape": [1, 3, 64, 512], + }, + "m3_wob": { + "candidate": "official-wheel-custom-fixed-m3-affine-qmv", + "layers_installed": 43, + "q6_g128_layers": 43, + "exact_selfchecked_layers": 43, + "active_o_lora_sinks_installed": 43, + "shape": [1, 3, 8192], + "output_size": 4096, + }, + "row_owned_combine": True, + "non_m1_m3_route": "native-dspark", + "routed_bits": 2, + "routed_group_size": 128, + "routed_gate_up_paired": True, + "shared_bits": 8, + "target_taps": (40, 41, 42), + "dspark_stages": 3, + "stage_ownership": "native", + } + + +def test_failed_required_preparation_keeps_every_route_stock( + full_artifact, + monkeypatch, +): + state = _install_with_spies( + full_artifact, + monkeypatch, + fail_wob_prepare=True, + ) + + assert isinstance(state.error, RuntimeError) + assert "wob self-check failed" in str(state.error) + assert state.prepared == ["wqb", "wob"] + assert [layer.ffn.switch_mlp for layer in state.layers] == state.switches + assert state.publication == [] + ids = SimpleNamespace(shape=(1, 1)) + assert state.model._target_hidden_route(state.model, ids, "cache") == ( + "stock-hidden", + "stock-taps", + ) + + +def test_prepared_target_stack_publishes_and_restores_as_one_transaction( + full_artifact, + monkeypatch, +): + state = _install_with_spies(full_artifact, monkeypatch, prepare_only=True) + + assert state.error is None + assert [layer.ffn.switch_mlp for layer in state.layers] == state.switches + assert state.model._target_hidden_route is state.stock_route + assert state.publication == [] + assert state.result.receipt["layers_installed"] == 43 + + state.result.publish() + assert [layer.ffn.switch_mlp for layer in state.layers] == state.replacements + assert state.model._target_hidden_route is not state.stock_route + assert state.publication == ["wqb.publish", "wob.publish"] + + state.result.restore() + assert [layer.ffn.switch_mlp for layer in state.layers] == state.switches + assert state.model._target_hidden_route is state.stock_route + assert state.publication[-2:] == ["wob.restore", "wqb.restore"] + + +def test_publication_failure_restores_projection_q2_and_target_routes( + full_artifact, + monkeypatch, +): + state = _install_with_spies( + full_artifact, + monkeypatch, + fail_publish="wob", + ) + + assert isinstance(state.error, RuntimeError) + assert "wob publication failed" in str(state.error) + assert [layer.ffn.switch_mlp for layer in state.layers] == state.switches + assert state.model._target_hidden_route is state.stock_route + assert state.publication == [ + "wqb.publish", + "wob.publish", + "wob.restore", + "wqb.restore", + ] + + +def test_publication_rollback_attempts_every_restore_after_one_restore_fails( + full_artifact, + monkeypatch, +): + state = _install_with_spies( + full_artifact, + monkeypatch, + fail_publish="wob", + fail_restore="wob", + ) + + assert isinstance(state.error, RuntimeError) + assert "wob publication failed" in str(state.error) + assert [layer.ffn.switch_mlp for layer in state.layers] == state.switches + assert state.model._target_hidden_route is state.stock_route + assert state.publication == [ + "wqb.publish", + "wob.publish", + "wob.restore", + "wqb.restore", + ] + + +def test_wob_receipt_rejects_incomplete_exact_selfcheck(): + receipt = SimpleNamespace( + published_routes=tuple(object() for _ in range(43)), + q6_count=43, + exact_selfchecked=42, + o_lora_sink_count=43, + ) + + with pytest.raises(ValueError, match="WOB preparation is not 43/43 exact"): + full._require_wob_receipt(receipt) + + +def test_real_weight_projection_checks_require_three_exact_stock_m1_rows(): + qhead_rows = [] + + def qhead_stock(qr, cos, sin): + qhead_rows.append((tuple(qr.shape), tuple(cos.shape), tuple(sin.shape))) + return qr + + qhead_check = full._m3_wqb_qhead_exact_selfcheck() + assert qhead_check(qhead_stock, lambda qr, _cos, _sin: qr, 0) is True + assert qhead_rows == [((1, 1, 1024), (1, 32), (1, 32))] * 3 + assert qhead_check(qhead_stock, lambda qr, _cos, _sin: qr + 1, 0) is False + + wob_check = full._m3_wob_exact_selfcheck() + assert wob_check(lambda value: value, lambda value: value, 0) is True + assert wob_check(lambda value: value, lambda value: value + 1, 0) is False + + +@pytest.mark.parametrize("failing_projection", ["wqb", "wob"]) +def test_raw_projection_preparer_failure_keeps_every_live_route_stock( + full_artifact, + monkeypatch, + failing_projection, +): + from mtplx import deepseek_v4_0731_m3_wob as wob + from mtplx import deepseek_v4_0731_m3_wqb_qnorm_rope as wqb + + built = 0 + + def build_qhead(_projection): + nonlocal built + layer_index = built + built += 1 + if failing_projection == "wqb" and layer_index == 9: + return lambda qr, _cos, _sin: qr + 1 + return lambda qr, _cos, _sin: qr + + built_wob = 0 + + def build_wob(_projection): + nonlocal built_wob + layer_index = built_wob + built_wob += 1 + if failing_projection == "wob" and layer_index == 9: + return lambda value: value + 1 + return lambda value: value + + monkeypatch.setattr(wqb, "build_0731_m3_wqb_qnorm_rope", build_qhead) + monkeypatch.setattr(wob, "bind_m3_wob", build_wob) + + state = _install_with_spies( + full_artifact, + monkeypatch, + prepare_wqb_override=wqb.prepare_wqb_qhead_m3, + prepare_wob_override=wob.prepare_wob_m3, + ) + + expected_error = ( + wqb.M3WQBNormRopeContractError + if failing_projection == "wqb" + else wob.M3WOBContractError + ) + assert isinstance(state.error, expected_error) + assert "layer 9" in str(state.error) + assert [layer.ffn.switch_mlp for layer in state.layers] == state.switches + assert state.model._target_hidden_route is state.stock_route + assert state.publication == [] + assert all( + layer.attn._q_projection_qhead_route is stock + for layer, stock in zip(state.layers, state.qhead_stocks) + ) + assert all( + layer.attn.wo_b is stock and layer.attn._o_lora_impl.wo_b is stock + for layer, stock in zip(state.layers, state.wob_stocks) + ) + + +def test_installer_has_no_modes_or_optional_projection_fallbacks(): + signature = inspect.signature(full.install_full_0731_dspark_compiled_tail_q2_pair) + assert tuple(signature.parameters) == ( + "model", + "config", + "model_path", + "prepare_wqb_qhead", + "prepare_wob", + ) + assert signature.parameters["prepare_wqb_qhead"].default is inspect.Parameter.empty + assert signature.parameters["prepare_wob"].default is inspect.Parameter.empty + source = inspect.getsource(full) + assert "m3_tail_mode" not in source + assert "official-custom" not in source + assert "row-exact-control" not in source + assert "hybrid" not in source.lower() + m1_route_source = inspect.getsource(full._FullDSparkTargetRoute.__call__) + assert "tuple(" not in m1_route_source + assert "shape[1]" in m1_route_source + + +def test_bound_m1_body_preserves_tap_order_and_cache_ownership(monkeypatch): + class FakeMX: + @staticmethod + def broadcast_to(value, shape): + assert shape == (1, 1, 4, 4096) + return value + + @staticmethod + def mean(value, axis): + assert axis == 2 + return f"mean-{value}" + + @staticmethod + def concatenate(values, *, axis): + assert axis == -1 + return tuple(values) + + class Layer: + def __init__(self, layer_id): + self.attn_hc = SimpleNamespace(pre=lambda hidden: (hidden, "post", "comb")) + self.attn_norm = lambda hidden: hidden + self.cache_entries = [] + + def attention(hidden, **kwargs): + self.cache_entries.append(kwargs["cache"]) + return hidden + + self.attn = attention + + class Tail: + def __init__(self, layer_id): + self.layer_id = layer_id + + def __call__(self, *_args): + return f"h{self.layer_id}" + + class Embedded: + shape = (1, 1, 4096) + + def __getitem__(self, _key): + return self + + monkeypatch.setattr(full, "mx", FakeMX) + layer_tails = tuple((Layer(i), Tail(i)) for i in range(43)) + body = SimpleNamespace( + embed_tokens=lambda _ids: Embedded(), + hc_mult=4, + ) + candidate = full._BoundDSparkM1Body(body, layer_tails, (40, 41, 42)) + + class CacheEntries: + def __iter__(self): + return (f"cache-{i}" for i in range(43)) + + def __len__(self): + raise AssertionError("bound route must not revalidate cache length") + + source = inspect.getsource(full._BoundDSparkM1Body.__call__) + assert "len(entries)" not in source + assert "strict=True" not in source + assert "missed a DSpark tap" not in source + + hidden, taps = candidate(SimpleNamespace(shape=(1, 1)), CacheEntries()) + + assert hidden == "h42" + assert taps == ("mean-h40", "mean-h41", "mean-h42") + assert [layer.cache_entries for layer, _tail in layer_tails] == [ + [f"cache-{i}"] for i in range(43) + ] diff --git a/tests/test_deepseek_v4_0731_k2_bench.py b/tests/test_deepseek_v4_0731_k2_bench.py new file mode 100644 index 00000000..e42ac899 --- /dev/null +++ b/tests/test_deepseek_v4_0731_k2_bench.py @@ -0,0 +1,762 @@ +"""CPU-only contracts for the source-isolated 0731 scheduler bracket.""" + +from __future__ import annotations + +import argparse +from copy import deepcopy +import importlib.util +import json +from pathlib import Path +import subprocess +from types import ModuleType, SimpleNamespace + +import pytest + + +ROOT = Path(__file__).parents[1] +SCRIPT = ROOT / "scripts" / "deepseek_v4_0731_k2_bench.py" + + +def _load_harness(): + spec = importlib.util.spec_from_file_location("deepseek_v4_0731_k2_bench", SCRIPT) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class _FakeMX: + def __init__(self): + self.reset_calls = 0 + self.memory_reads = 0 + + def reset_peak_memory(self): + self.reset_calls += 1 + + def get_peak_memory(self): + self.memory_reads += 1 + return 1_000 + self.memory_reads + + def get_active_memory(self): + return 900 + self.memory_reads + + +def _stats(*, speculative: bool, signature_variant: int = 0): + return SimpleNamespace( + decode_tok_s=41.25 if speculative else 31.5, + end_to_end_tok_s=30.5 if speculative else 25.0, + prompt_eval_time_s=0.8, + prompt_target_prefill_time_s=0.5, + prompt_mtp_history_time_s=0.3 if speculative else 0.0, + prompt_target_prefill_tok_s=18.0, + accepted_drafts=(4 + signature_variant) if speculative else 0, + rejected_drafts=2 if speculative else 0, + drafted_tokens=(6 + signature_variant) if speculative else 0, + accepted_by_depth=([3 + signature_variant, 1] if speculative else []), + drafted_by_depth=([4 + signature_variant, 2] if speculative else []), + verify_calls=3 if speculative else 0, + ) + + +class _FakeBackend: + backend_id = "deepseek_v4_dspark_0731" + + def make_cache(self, _runtime): + return SimpleNamespace(ring=b"proposal-ring", prefill_length=9) + + def snapshot(self, cache): + return ((cache.ring, cache.prefill_length),) + + +class _FakeRuntime: + def __init__(self): + self.model = object() + self.tokenizer = SimpleNamespace( + encode=lambda text: list(range(11, 20)) if text else [] + ) + self.block_speculative_backend = _FakeBackend() + self.deepseek_v4_0731_k2_receipt = None + self.target_cache_calls = 0 + + def make_cache(self): + self.target_cache_calls += 1 + return SimpleNamespace( + offset=9, + state=(b"target-state",), + metadata_version="fake-v1", + ) + + +def _scheduler_arm( + label: str = "lazy_joint_eval", + source_sha: str = "a" * 64, + *, + normalized_sha: str = "10f7a52f59044ca7e7600156626b28826773886657e68201644f8b50385ba2e1", + patch_sha: str = "f09d68378f940eb948a58cf4f9b24e90bfb9d40119483348b3e6f5d8b849205e", +): + events = ( + ["proposal_graph", "target_row_graph", "joint_eval", "draft_materialize"] + if label == "lazy_joint_eval" + else [ + "proposal_graph", + "proposal_eval", + "draft_materialize", + "target_row_graph", + ] + ) + return { + "label": label, + "source_path": "mtplx/native_block_speculation.py", + "source_sha256": source_sha, + "arm_id": f"{label}:{source_sha}", + "normalized_source_sha256": normalized_sha, + "reviewed_boundary_patch_sha256": patch_sha, + "sanctioned_event_sequence": events, + } + + +def _identities(harness, scheduler_sha: str = "a" * 64): + sources = {path: "e" * 64 for path in harness._BRACKET_SOURCE_PATHS} + sources[harness._SCHEDULER_SOURCE] = scheduler_sha + importable = { + name.replace(".", "/") + ".py": ( + scheduler_sha if name == "mtplx.native_block_speculation" else "e" * 64 + ) + for name in harness._REQUIRED_IMPORTED_MODULES + } + return { + "mlx_identity": { + "version": "0.32.0", + "core_sha256": "1" * 64, + "libmlx": {"sha256": "2" * 64}, + "metallib": {"sha256": "3" * 64}, + }, + "model_identity": { + "config_sha256": "4" * 64, + "index_sha256": "5" * 64, + "metadata": {"revision": "6" * 40}, + }, + "git_identity": { + "commit": "7" * 40, + "head_tree": "8" * 40, + "head_tree_files": { + path: {"mode": "100644", "object": digest} + for path, digest in sources.items() + }, + "head_python_sha256": importable, + "dirty": False, + "status": [], + }, + "source_identity": { + "source_set_sha256": "9" * 64, + "files": sources, + "importable_mtplx_files": importable, + }, + "guard_attestation": { + "window_id": "b" * 64, + "attestation": {"lock_device": 1, "lock_inode": 2}, + "lock_identity": {"device": 1, "inode": 2}, + }, + } + + +def _run_fake_benchmark( + harness, + *, + ar_tokens=None, + k2_tokens=None, + k2_signature_variant=None, + post_run_git_mutator=None, +): + calls = [] + loads = [] + mx = _FakeMX() + runtime = _FakeRuntime() + original_backend = runtime.block_speculative_backend + ar_tokens = ar_tokens or (lambda _call: [101, 102, 103]) + k2_tokens = k2_tokens or (lambda _call: [101, 102, 103]) + k2_signature_variant = k2_signature_variant or (lambda _call: 0) + ar_calls = 0 + k2_calls = 0 + + def load(path, **kwargs): + loads.append((path, kwargs)) + return runtime + + def generate_ar(active_runtime, prompt_ids, **kwargs): + nonlocal ar_calls + ar_calls += 1 + cache = active_runtime.make_cache() + calls.append( + ( + "ar", + prompt_ids, + tuple(prompt_ids), + kwargs, + active_runtime.block_speculative_backend is original_backend, + "make_cache" in vars(active_runtime), + ) + ) + tokens = ( + [101, 102, 103, 104] if kwargs["max_tokens"] == 4 else ar_tokens(ar_calls) + ) + assert cache.offset == 9 + return SimpleNamespace(tokens=tokens, stats=_stats(speculative=False)) + + def generate_mtpk(active_runtime, prompt_ids, **kwargs): + nonlocal k2_calls + k2_calls += 1 + target_cache = active_runtime.make_cache() + active_runtime.block_speculative_backend.make_cache(active_runtime) + calls.append( + ( + "k2", + prompt_ids, + tuple(prompt_ids), + kwargs, + active_runtime.block_speculative_backend is original_backend, + "make_cache" in vars(active_runtime), + ) + ) + tokens = k2_tokens(k2_calls) + assert target_cache.offset == 9 + return SimpleNamespace( + tokens=tokens, + stats=_stats( + speculative=True, + signature_variant=k2_signature_variant(k2_calls), + ), + ) + + identities = _identities(harness) + post_run_git = deepcopy(identities["git_identity"]) + if post_run_git_mutator is not None: + post_run_git_mutator(post_run_git) + receipt = harness.run_benchmark( + argparse.Namespace( + model=Path("/model").resolve(), + max_tokens=64, + out=Path("/receipt.json"), + ), + mx=mx, + runtime_load=load, + generate_ar=generate_ar, + generate_mtpk=generate_mtpk, + sampler_factory=lambda **kwargs: SimpleNamespace(**kwargs), + imported_modules_attestation=lambda: { + "module_set_sha256": "c" * 64, + "files": { + name: { + "path": name.replace(".", "/") + ".py", + "sha256": ( + "a" * 64 + if name == "mtplx.native_block_speculation" + else "e" * 64 + ), + } + for name in harness._REQUIRED_IMPORTED_MODULES + }, + "preflight_head_bound": True, + "preflight_sources_bound": True, + }, + post_run_git_attestation=lambda: post_run_git, + scheduler_arm=_scheduler_arm(), + **identities, + ) + return receipt, runtime, original_backend, calls, loads, mx + + +def test_run_is_stock_one_load_with_unmeasured_state_proof_and_fixed_five(): + harness = _load_harness() + receipt, runtime, original_backend, calls, loads, mx = _run_fake_benchmark(harness) + + assert loads == [(Path("/model").resolve(), {"mtp": True})] + assert [call[0] for call in calls] == [ + "ar", + "k2", + "ar", + "k2", + *(lane for _ in range(5) for lane in ("ar", "k2")), + ] + assert [call[3]["max_tokens"] for call in calls[:2]] == [4, 3] + assert all(call[3]["max_tokens"] == 64 for call in calls[2:]) + assert all(call[3]["stop_token_ids"] == set() for call in calls) + assert len({id(call[1]) for call in calls}) == len(calls) + assert all(call[2] == tuple(range(11, 20)) for call in calls) + assert calls[0][4:] == (False, True) + assert calls[1][4:] == (False, True) + assert all(call[4:] == (True, False) for call in calls[2:]) + assert runtime.block_speculative_backend is original_backend + assert "make_cache" not in vars(runtime) + assert mx.reset_calls == 10 + + assert receipt["baseline"] == "generic_mtp_true_stock" + assert receipt["load_kwargs"] == {"mtp": True} + assert receipt["repetitions"] == harness.REPETITIONS == 5 + assert receipt["scheduler_arm"] == _scheduler_arm() + assert receipt["provenance"]["git"] == receipt["provenance"]["git_post_run"] + assert receipt["state_proof"]["measured"] is False + assert receipt["state_proof"]["complete_k2_cycle"] is True + assert receipt["state_proof"]["target_rows_consumed"] == 3 + assert receipt["state_proof"]["k2_drafted_by_depth"] == [4, 2] + assert receipt["state_proof"]["target_state_equal"] is True + assert receipt["state_proof"]["wrappers_restored_before_primers"] is True + assert receipt["state_proof"]["ar_target"] == receipt["state_proof"]["k2_target"] + assert len(receipt["state_proof"]["proposal_snapshot"]["state_sha256"]) == 64 + assert len(receipt["measurements"]["samples"]) == 5 + assert all( + row["acceptance_signature"] + == receipt["measurements"]["samples"][0]["acceptance_signature"] + for row in receipt["measurements"]["samples"] + ) + assert receipt["gates"] == { + "state_proof_target_equal": True, + "tokens_exact_all_samples": True, + "ar_deterministic": True, + "k2_deterministic": True, + "acceptance_signature_identical_all_samples": True, + } + assert receipt["passed"] is True + + +def test_state_proof_rejects_target_cache_drift_and_restores_wrappers(): + harness = _load_harness() + runtime = _FakeRuntime() + original_backend = runtime.block_speculative_backend + + def ar(active_runtime, _prompt, **_kwargs): + active_runtime.make_cache() + return SimpleNamespace(tokens=[101, 102, 103, 104]) + + def k2(active_runtime, _prompt, **_kwargs): + target_cache = active_runtime.make_cache() + target_cache.offset = 10 + active_runtime.block_speculative_backend.make_cache(active_runtime) + return SimpleNamespace(tokens=[101, 102, 103], stats=_stats(speculative=True)) + + with pytest.raises(RuntimeError, match="target cache state is not bit-exact"): + harness.prove_prefill_state( + runtime, + list(range(9)), + generate_ar=ar, + generate_mtpk=k2, + sampler=SimpleNamespace(), + ) + assert runtime.block_speculative_backend is original_backend + assert "make_cache" not in vars(runtime) + + +def test_state_proof_requires_both_k2_draft_depths(): + harness = _load_harness() + runtime = _FakeRuntime() + + def ar(active_runtime, _prompt, **_kwargs): + active_runtime.make_cache() + return SimpleNamespace(tokens=[101, 102, 103, 104]) + + def k2(active_runtime, _prompt, **_kwargs): + active_runtime.make_cache() + active_runtime.block_speculative_backend.make_cache(active_runtime) + stats = _stats(speculative=True) + stats.drafted_by_depth = [1, 0] + return SimpleNamespace(tokens=[101, 102, 103], stats=stats) + + with pytest.raises(RuntimeError, match="one complete K2 cycle"): + harness.prove_prefill_state( + runtime, + list(range(9)), + generate_ar=ar, + generate_mtpk=k2, + sampler=SimpleNamespace(), + ) + + +def test_all_samples_gate_tokens_determinism_and_acceptance_signature(): + harness = _load_harness() + receipt, *_ = _run_fake_benchmark( + harness, + ar_tokens=lambda call: [999] if call == 7 else [101, 102, 103], + k2_signature_variant=lambda call: 1 if call == 7 else 0, + ) + assert receipt["gates"]["ar_deterministic"] is False + assert receipt["gates"]["tokens_exact_all_samples"] is False + assert receipt["gates"]["acceptance_signature_identical_all_samples"] is False + assert receipt["passed"] is False + + +def test_post_run_git_must_match_exact_preflight_tree(): + harness = _load_harness() + + def change_tree(identity): + identity["head_tree"] = "0" * 40 + + with pytest.raises(RuntimeError, match="provenance changed during"): + _run_fake_benchmark(harness, post_run_git_mutator=change_tree) + + +def test_fixed_prompt_fails_before_state_proof(): + harness = _load_harness() + runtime = _FakeRuntime() + runtime.tokenizer.encode = lambda _text: list(range(8)) + with pytest.raises(RuntimeError, match="fixed prompt tokenizer drift"): + harness.run_benchmark( + argparse.Namespace(model=Path("/model"), max_tokens=64), + mx=_FakeMX(), + runtime_load=lambda *_args, **_kwargs: runtime, + generate_ar=lambda *_args, **_kwargs: None, + generate_mtpk=lambda *_args, **_kwargs: None, + sampler_factory=lambda **kwargs: SimpleNamespace(**kwargs), + imported_modules_attestation=lambda: { + "preflight_head_bound": True, + "preflight_sources_bound": True, + }, + post_run_git_attestation=lambda: {}, + scheduler_arm=_scheduler_arm(), + **_identities(harness), + ) + + +def test_write_receipt_returns_nonzero_for_failed_gate(tmp_path): + harness = _load_harness() + receipt = {"passed": False, "gates": {"tokens": False}} + output = tmp_path / "receipt.json" + assert harness.write_receipt(receipt, output) == 1 + assert json.loads(output.read_text()) == receipt + + +def test_scheduler_arm_is_derived_from_only_two_sanctioned_source_orders(): + harness = _load_harness() + actual = harness.attest_scheduler_arm(ROOT) + assert actual["label"] in harness._ARM_LABELS + assert actual["arm_id"] == f"{actual['label']}:{actual['source_sha256']}" + assert actual["source_sha256"] == harness._sha256( + ROOT / "mtplx/native_block_speculation.py" + ) + assert ( + actual["normalized_source_sha256"] + == harness.EXPECTED_NORMALIZED_SCHEDULER_SHA256 + ) + assert ( + actual["reviewed_boundary_patch_sha256"] + == harness.EXPECTED_SCHEDULER_BOUNDARY_PATCH_SHA256 + ) + + source = (ROOT / harness._SCHEDULER_SOURCE).read_text() + lazy_source = source.replace( + harness._MATERIALIZE_BOUNDARY_BLOCK, + harness._LAZY_BOUNDARY_BLOCK, + ) + materialize_source = lazy_source.replace( + harness._LAZY_BOUNDARY_BLOCK, + harness._MATERIALIZE_BOUNDARY_BLOCK, + ) + assert harness._classify_scheduler_source(lazy_source) == ( + "lazy_joint_eval", + harness._ARM_EVENTS["lazy_joint_eval"], + ) + assert harness._classify_scheduler_source(materialize_source) == ( + "materialize_first", + harness._ARM_EVENTS["materialize_first"], + ) + with pytest.raises(ValueError, match="outside reviewed boundary motion"): + harness._classify_scheduler_source(lazy_source + "# unrelated edit\n") + with pytest.raises(ValueError, match="exact sanctioned bracket arm"): + harness._classify_scheduler_source( + materialize_source.replace("Settle and materialize", "Changed materialize") + ) + + +def test_scheduler_arm_rejects_crlf_byte_drift(tmp_path): + harness = _load_harness() + source = (ROOT / harness._SCHEDULER_SOURCE).read_bytes() + scheduler_path = tmp_path / harness._SCHEDULER_SOURCE + scheduler_path.parent.mkdir(parents=True) + scheduler_path.write_bytes(source.replace(b"\n", b"\r\n")) + + with pytest.raises(ValueError, match="exact sanctioned bracket arm"): + harness.attest_scheduler_arm(tmp_path) + + +def _comparison_pair(harness): + lazy, *_ = _run_fake_benchmark(harness) + materialize = deepcopy(lazy) + materialize_sha = "d" * 64 + materialize["scheduler_arm"] = _scheduler_arm("materialize_first", materialize_sha) + materialize["provenance"]["git"] = { + **materialize["provenance"]["git"], + "commit": "e" * 40, + "head_tree": "f" * 40, + } + materialize["provenance"]["git"]["head_tree_files"] = deepcopy( + materialize["provenance"]["git"]["head_tree_files"] + ) + materialize["provenance"]["git"]["head_tree_files"][harness._SCHEDULER_SOURCE][ + "object" + ] = materialize_sha + materialize["provenance"]["git"]["head_python_sha256"][ + harness._SCHEDULER_SOURCE + ] = materialize_sha + materialize["provenance"]["sources"]["files"][harness._SCHEDULER_SOURCE] = ( + materialize_sha + ) + materialize["provenance"]["sources"]["source_set_sha256"] = "0" * 64 + materialize["provenance"]["sources"]["importable_mtplx_files"][ + harness._SCHEDULER_SOURCE + ] = materialize_sha + materialize["provenance"]["imported_mtplx_modules"]["files"][ + "mtplx.native_block_speculation" + ]["sha256"] = materialize_sha + materialize["provenance"]["git_post_run"] = deepcopy( + materialize["provenance"]["git"] + ) + return lazy, materialize + + +def test_comparator_gates_source_isolation_state_tokens_and_acceptance(): + harness = _load_harness() + lazy, materialize = _comparison_pair(harness) + comparison = harness.compare_receipts(materialize, lazy) + assert comparison["source_differences"] == [harness._SCHEDULER_SOURCE] + assert comparison["head_tree_differences"] == [harness._SCHEDULER_SOURCE] + assert comparison["imported_module_differences"] == [ + "mtplx.native_block_speculation" + ] + assert all(comparison["state_digests_equal"].values()) + assert all(comparison["gates"].values()) + assert comparison["passed"] is True + + materialize["state_proof"]["proposal_snapshot"]["state_sha256"] = "1" * 64 + materialize["measurements"]["samples"][0]["acceptance_signature"][ + "accepted_drafts" + ] += 1 + failed = harness.compare_receipts(lazy, materialize) + assert failed["gates"]["proposal_snapshot_identical"] is False + assert failed["gates"]["acceptance_signature_identical_cross_arm"] is False + assert failed["passed"] is False + + materialize = _comparison_pair(harness)[1] + materialize["provenance"]["git"]["head_tree_files"]["unrelated.txt"] = { + "mode": "100644", + "object": "2" * 40, + } + unrelated = harness.compare_receipts(lazy, materialize) + assert unrelated["gates"]["only_scheduler_head_blob_differs"] is False + + materialize = _comparison_pair(harness)[1] + materialize["provenance"]["imported_mtplx_modules"]["files"]["mtplx.runtime"][ + "sha256" + ] = "3" * 64 + imported = harness.compare_receipts(lazy, materialize) + assert imported["gates"]["only_scheduler_import_differs"] is False + + +def test_comparator_rejects_arm_label_not_bound_to_source_hash(): + harness = _load_harness() + lazy, materialize = _comparison_pair(harness) + materialize["scheduler_arm"]["arm_id"] = "materialize_first:wrong" + with pytest.raises(ValueError, match="arm attribution is invalid"): + harness.compare_receipts(lazy, materialize) + + +def test_git_attestation_requires_clean_committed_head_tree(tmp_path): + harness = _load_harness() + subprocess.run(["git", "init", "-q", str(tmp_path)], check=True) + subprocess.run( + ["git", "-C", str(tmp_path), "config", "user.email", "test@example.com"], + check=True, + ) + subprocess.run( + ["git", "-C", str(tmp_path), "config", "user.name", "Test"], check=True + ) + (tmp_path / "tracked.txt").write_text("clean\n") + subprocess.run(["git", "-C", str(tmp_path), "add", "tracked.txt"], check=True) + subprocess.run( + ["git", "-C", str(tmp_path), "commit", "-q", "-m", "clean"], check=True + ) + identity = harness.attest_git(tmp_path) + assert identity["dirty"] is False + assert len(identity["commit"]) == 40 + assert len(identity["head_tree"]) == 40 + + (tmp_path / "dirty.txt").write_text("dirty\n") + with pytest.raises(RuntimeError, match="clean committed worktree"): + harness.attest_git(tmp_path) + + +def test_imported_module_attestation_rejects_any_reviewed_overlay(tmp_path): + harness = _load_harness() + modules = {} + preflight_hashes = {} + for name in harness._REQUIRED_IMPORTED_MODULES: + path = tmp_path / (name.replace(".", "/") + ".py") + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(name) + module = ModuleType(name) + module.__file__ = str(path) + modules[name] = module + preflight_hashes[str(path.relative_to(tmp_path))] = harness._sha256(path) + git_identity = {"head_python_sha256": dict(preflight_hashes)} + source_identity = {"importable_mtplx_files": dict(preflight_hashes)} + identity = harness.attest_imported_mtplx_modules( + tmp_path, + git_identity=git_identity, + source_identity=source_identity, + modules=modules, + ) + assert set(identity["files"]) == set(harness._REQUIRED_IMPORTED_MODULES) + assert identity["preflight_head_bound"] is True + assert identity["preflight_sources_bound"] is True + + first_path = next(iter(preflight_hashes)) + git_identity["head_python_sha256"][first_path] = "0" * 64 + with pytest.raises(RuntimeError, match="does not match preflight HEAD"): + harness.attest_imported_mtplx_modules( + tmp_path, + git_identity=git_identity, + source_identity=source_identity, + modules=modules, + ) + git_identity["head_python_sha256"] = dict(preflight_hashes) + source_identity["importable_mtplx_files"][first_path] = "1" * 64 + with pytest.raises(RuntimeError, match="does not match source attestation"): + harness.attest_imported_mtplx_modules( + tmp_path, + git_identity=git_identity, + source_identity=source_identity, + modules=modules, + ) + source_identity["importable_mtplx_files"] = dict(preflight_hashes) + + outside = tmp_path.parent / "overlay.py" + outside.write_text("overlay") + modules[harness._REQUIRED_IMPORTED_MODULES[0]].__file__ = str(outside) + with pytest.raises(RuntimeError, match="outside worktree"): + harness.attest_imported_mtplx_modules( + tmp_path, + git_identity=git_identity, + source_identity=source_identity, + modules=modules, + ) + + +def test_official_mlx_attestation_rejects_editable_or_import_overlay(tmp_path): + harness = _load_harness() + site = tmp_path / "site-packages" + core = site / "mlx" / "core.cpython.so" + core.parent.mkdir(parents=True) + core.write_bytes(b"official wheel core") + libmlx = site / "mlx" / "lib" / "libmlx.dylib" + metallib = site / "mlx" / "lib" / "mlx.metallib" + libmlx.parent.mkdir() + libmlx.write_bytes(b"official wheel dylib") + metallib.write_bytes(b"official wheel metallib") + harness.EXPECTED_MLX_CORE_SHA256 = harness._sha256(core) + harness.EXPECTED_MLX_LIB_SHA256 = harness._sha256(libmlx) + harness.EXPECTED_MLX_METALLIB_SHA256 = harness._sha256(metallib) + + class Distribution: + version = "0.32.0" + + def __init__(self, direct_url=None): + self.direct_url = direct_url + + def locate_file(self, value): + return site / value + + def read_text(self, name): + if name == "INSTALLER": + return "uv\n" + if name == "direct_url.json": + return self.direct_url + return None + + mx = SimpleNamespace(__file__=str(core)) + identity = harness.attest_official_mlx(mx, Distribution()) + assert identity["version"] == "0.32.0" + assert identity["installer"] == "uv" + assert identity["core_path"] == str(core.resolve()) + assert identity["libmlx"]["sha256"] == harness._sha256(libmlx) + assert identity["metallib"]["sha256"] == harness._sha256(metallib) + + editable = json.dumps( + {"url": "file:///tmp/mlx-source", "dir_info": {"editable": True}} + ) + with pytest.raises(ValueError, match="source/direct overlay"): + harness.attest_official_mlx(mx, Distribution(editable)) + + outside = tmp_path / "mlx-overlay" / "core.cpython.so" + outside.parent.mkdir() + outside.write_bytes(b"overlay") + with pytest.raises(ValueError, match="outside installed distribution"): + harness.attest_official_mlx( + SimpleNamespace(__file__=str(outside)), Distribution() + ) + + +def test_model_attestation_requires_pinned_0731_hashes(tmp_path): + harness = _load_harness() + (tmp_path / "config.json").write_bytes(b"config") + (tmp_path / "model.safetensors.index.json").write_bytes(b"index") + harness.EXPECTED_MODEL_CONFIG_SHA256 = harness._sha256(tmp_path / "config.json") + harness.EXPECTED_MODEL_INDEX_SHA256 = harness._sha256( + tmp_path / "model.safetensors.index.json" + ) + metadata_root = tmp_path / ".cache" / "huggingface" / "download" + metadata_root.mkdir(parents=True) + for name in ("config.json.metadata", "model.safetensors.index.json.metadata"): + (metadata_root / name).write_text( + harness.EXPECTED_MODEL_METADATA_REVISION + "\nmetadata\n" + ) + + identity = harness.attest_model(tmp_path) + assert identity["config_sha256"] == harness.EXPECTED_MODEL_CONFIG_SHA256 + assert identity["index_sha256"] == harness.EXPECTED_MODEL_INDEX_SHA256 + assert {row["revision"] for row in identity["metadata"].values()} == { + harness.EXPECTED_MODEL_METADATA_REVISION + } + + (tmp_path / "config.json").write_bytes(b"drift") + with pytest.raises(ValueError, match="config SHA mismatch"): + harness.attest_model(tmp_path) + + +def test_cli_has_fixed_repetitions_and_source_only_arm(tmp_path): + harness = _load_harness() + args = harness._parse_args( + ["--model", str(tmp_path), "--out", str(tmp_path / "receipt.json")] + ) + assert args.model == tmp_path + assert not hasattr(args, "repetitions") + assert not hasattr(args, "mode") + assert harness.REPETITIONS == 5 + + compare = harness._parse_args( + [ + "--compare", + str(tmp_path / "lazy.json"), + str(tmp_path / "materialize.json"), + "--out", + str(tmp_path / "comparison.json"), + ] + ) + assert compare.compare == [tmp_path / "lazy.json", tmp_path / "materialize.json"] + + +def test_source_is_scheduler_only_clean_before_mlx_and_includes_guard_bridge(): + source = SCRIPT.read_text() + assert "os.environ" not in source + assert "MLX_DISPATCH_CENSUS" not in source + assert "prepare_dspark_q3_packed_gate_up_m5" not in source + assert "dspark-ffn-q3-m5" not in source + assert 'add_argument("--mode"' not in source + assert '"--repetitions"' not in source + assert "deepseek_v4_0731_k2=" not in source + assert "REPETITIONS = 5" in source + assert "attest_scheduler_arm(" in source + assert "scripts/deepseek_v4_guard_window.py" in source + assert source.index("git_identity = attest_git(repo)") < source.index( + "import mlx.core as mx" + ) + assert source.index("issue_guard_window()") < source.index("import mlx.core as mx") + assert source.index("load_verified_guard_window(") < source.index( + "import mlx.core as mx" + ) diff --git a/tests/test_deepseek_v4_0731_m3_target.py b/tests/test_deepseek_v4_0731_m3_target.py new file mode 100644 index 00000000..8a033f11 --- /dev/null +++ b/tests/test_deepseek_v4_0731_m3_target.py @@ -0,0 +1,229 @@ +"""CPU contracts for the single receipt-backed physical-M3 target route.""" + +from __future__ import annotations + +import inspect +from types import SimpleNamespace + +import pytest + +pytest.importorskip("mlx.core") +import mlx.core as mx # noqa: E402 + +from mtplx import deepseek_v4_0731_m3_target as target # noqa: E402 + + +@pytest.fixture(autouse=True) +def _cpu_default_device(): + previous = mx.default_device() + mx.set_default_device(mx.cpu) + try: + yield + finally: + mx.set_default_device(previous) + + +class _Body: + hc_mult = 4 + + def __init__(self): + self.layers = tuple(object() for _ in range(43)) + + @staticmethod + def embed_tokens(ids): + return mx.broadcast_to(ids[..., None].astype(mx.float32), (*ids.shape, 2)) + + +class _Model: + def __init__(self): + self.args = SimpleNamespace( + hidden_size=4096, + num_hidden_layers=43, + num_attention_heads=64, + num_key_value_heads=1, + head_dim=512, + ) + self.model = _Body() + self._dspark = SimpleNamespace( + target_layer_ids=(40, 41, 42), + stages=(object(), object(), object()), + ) + self.base_calls = [] + + def base(owner, input_ids, cache=None): + self.base_calls.append((owner, input_ids, cache)) + return "base-hidden", "base-taps" + + self._target_hidden_route = base + + @staticmethod + def logits_from_hc_hidden(hidden): + return mx.mean(hidden, axis=2) + + +def test_fixed_m3_route_runs_43_prebound_layers_once_and_captures_taps(): + model = _Model() + calls = [] + + def layer_route(hidden, input_ids, cache, *, layer_id): + calls.append((layer_id, tuple(input_ids.shape), cache)) + return hidden + 1 + + routes = tuple( + lambda hidden, input_ids, cache, layer_id=layer_id: layer_route( + hidden, + input_ids, + cache, + layer_id=layer_id, + ) + for layer_id in range(43) + ) + route = target.build_0731_m3_target_route( + model, + full_layer_routes=routes, + base_route=model._target_hidden_route, + ) + caches = tuple(object() for _ in range(43)) + + logits, hidden, taps = route.forward(mx.array([[7, 8, 9]]), caches) + mx.eval(logits, hidden, taps) + + assert hidden.shape == (1, 3, 4, 2) + assert logits.shape == (1, 3, 2) + assert taps.shape == (1, 3, 6) + assert [layer_id for layer_id, _, _ in calls] == list(range(43)) + assert [cache for _, _, cache in calls] == list(caches) + assert all(shape == (1, 3) for _, shape, _ in calls) + assert mx.array_equal(taps[..., :2], hidden[..., 0, :] - 2) + assert mx.array_equal(taps[..., 2:4], hidden[..., 0, :] - 1) + assert mx.array_equal(taps[..., 4:], hidden[..., 0, :]) + + +def test_fixed_m3_route_delegates_only_non_m3_shapes_to_prebound_base(): + model = _Model() + routes = tuple(lambda hidden, _ids, _cache: hidden for _ in range(43)) + route = target.build_0731_m3_target_route( + model, + full_layer_routes=routes, + base_route=model._target_hidden_route, + ) + ids = mx.array([[7]]) + + assert route(model, ids, "cache") == ("base-hidden", "base-taps") + assert model.base_calls == [(model, ids, "cache")] + + # Batch size is construction-owned; the hot route reads logical M only. + m3_ids = mx.array([[1, 2, 3], [4, 5, 6]]) + hidden, taps = route(model, m3_ids, (None,) * 43) + mx.eval(hidden, taps) + assert hidden.shape == (2, 3, 4, 2) + assert model.base_calls == [(model, ids, "cache")] + + +def test_bound_m3_body_directly_executes_without_hot_invariant_checks(): + source = inspect.getsource(target._M3TargetBody.__call__) + assert "physical-M3 target route requires" not in source + assert "len(entries)" not in source + assert "strict=True" not in source + assert "missed a DSpark tap" not in source + + model = _Model() + calls = [] + routes = tuple( + lambda hidden, _ids, cache, layer_id=layer_id: ( + calls.append((layer_id, cache)) or hidden + ) + for layer_id in range(43) + ) + route = target.build_0731_m3_target_route( + model, + full_layer_routes=routes, + base_route=model._target_hidden_route, + ) + + class CacheEntries: + def __iter__(self): + return iter(range(43)) + + def __len__(self): + raise AssertionError("bound route must not revalidate cache length") + + logits, hidden, taps = route.forward(mx.array([[7, 8]]), CacheEntries()) + mx.eval(logits, hidden, taps) + + assert hidden.shape == (1, 2, 4, 2) + assert taps.shape == (1, 2, 6) + assert calls == list(enumerate(range(43))) + + +def test_fixed_m3_builder_rejects_missing_routes_and_non_0731_owner(): + model = _Model() + with pytest.raises(target.M3TargetContractError, match="43 prebound"): + target.build_0731_m3_target_route( + model, + full_layer_routes=(), + base_route=model._target_hidden_route, + ) + + model._dspark.target_layer_ids = (39, 40, 41) + with pytest.raises(target.M3TargetContractError, match="target taps"): + target.build_0731_m3_target_route( + model, + full_layer_routes=tuple(lambda *_args: None for _ in range(43)), + base_route=model._target_hidden_route, + ) + + +def test_compiled_tail_layer_keeps_one_full_width_three_boundary(): + calls = [] + + class HC: + @staticmethod + def pre(hidden): + calls.append(("pre", tuple(hidden.shape))) + value = hidden[:, :, 0, :] + return value, value + 20, value[:, :, None, :] + 30 + + class Layer: + attn_hc = HC() + attn_norm = staticmethod(lambda value: value + 10) + + @staticmethod + def attn(value, *, mask, cache): + calls.append(("attention", tuple(value.shape), cache)) + return value + 100 + + def compiled_tail(attn_out, residual, post, comb, input_ids): + calls.append( + ( + "tail", + tuple(attn_out.shape), + tuple(residual.shape), + tuple(post.shape), + tuple(comb.shape), + tuple(input_ids.shape), + ) + ) + return attn_out[:, :, None, :] + mx.zeros_like(residual) + + route = target.build_m3_compiled_tail_layer(Layer(), compiled_tail) + hidden = mx.zeros((1, 3, 4, 2), dtype=mx.float32) + ids = mx.array([[7, 8, 9]]) + cache = object() + + got = route(hidden, ids, cache) + mx.eval(got) + + assert got.shape == (1, 3, 4, 2) + assert calls == [ + ("pre", (1, 3, 4, 2)), + ("attention", (1, 3, 2), cache), + ("tail", (1, 3, 2), (1, 3, 4, 2), (1, 3, 2), (1, 3, 1, 2), (1, 3)), + ] + + +def test_module_exposes_no_unmeasured_control_or_hybrid_modes(): + assert not hasattr(target, "RowExactControlLayer") + assert not hasattr(target, "M3HybridControlLayer") + assert not hasattr(target, "build_row_exact_control_layer") + assert not hasattr(target, "build_m3_hybrid_control_layer") diff --git a/tests/test_deepseek_v4_0731_m3_wob.py b/tests/test_deepseek_v4_0731_m3_wob.py new file mode 100644 index 00000000..6ab0448b --- /dev/null +++ b/tests/test_deepseek_v4_0731_m3_wob.py @@ -0,0 +1,260 @@ +"""CPU contracts for the receipt-backed fixed-M3 Q6/G128 WOB primitive.""" + +from __future__ import annotations + +import inspect +import os +from types import SimpleNamespace + +import pytest + +pytest.importorskip("mlx.core") +import mlx.core as mx # noqa: E402 + + +@pytest.fixture(autouse=True) +def _cpu_default_device(): + previous = mx.default_device() + mx.set_default_device(mx.cpu) + try: + yield + finally: + mx.set_default_device(previous) + + +class _CallableProjection(SimpleNamespace): + def __call__(self, value): + return ("stock", value) + + +def _projection(*, bits=6, group_size=128, mode="affine"): + return _CallableProjection( + bits=bits, + group_size=group_size, + mode=mode, + bias=None, + weight=SimpleNamespace(shape=(4096, 1536), dtype=mx.uint32), + scales=SimpleNamespace(shape=(4096, 64), dtype=mx.bfloat16), + biases=SimpleNamespace(shape=(4096, 64), dtype=mx.bfloat16), + ) + + +def _layers(): + layers = [] + for _ in range(43): + stock = _projection() + attention = SimpleNamespace( + wo_b=stock, + _o_lora_impl=SimpleNamespace(wo_b=stock), + ) + layers.append(SimpleNamespace(attn=attention)) + return layers + + +def test_contract_and_binder_are_only_fixed_m3_q6_g128(monkeypatch): + from mtplx import deepseek_v4_0731_m3_wob as candidate + + contract = candidate.M3WOBContract() + assert (contract.k, contract.n, contract.bits, contract.group_size) == ( + 8192, + 4096, + 6, + 128, + ) + candidate.validate_wob_projection(_projection(), contract) + with pytest.raises(candidate.M3WOBContractError, match="Q6/G128"): + candidate.validate_wob_projection(_projection(bits=8), contract) + + calls = [] + candidate._build_wob_kernel.cache_clear() + monkeypatch.setattr( + candidate.mx.fast, + "metal_kernel", + lambda **kwargs: calls.append(kwargs) or object(), + ) + bound = candidate.bind_m3_wob(_projection()) + assert bound.input_shape == (1, 3, 8192) + assert bound.output_shape == (1, 3, 4096) + assert bound.grid == ((4096 // 8) * 64, 1, 1) + assert bound.threadgroup == (64, 1, 1) + assert calls[0]["input_names"] == ["x", "w", "scales", "biases"] + assert calls[0]["output_names"] == ["y"] + + +def test_source_preserves_three_official_q6_m1_reduction_trees(): + from mtplx import deepseek_v4_0731_m3_wob as candidate + + source = candidate.m3_wob_metal_source() + for fragment in ( + "constexpr uint M = 3;", + "constexpr uint K = 8192;", + "constexpr uint N = 4096;", + "constexpr uint GS = 128;", + "float result0[RESULTS_PER_SIMDGROUP]", + "float result1[RESULTS_PER_SIMDGROUP]", + "float result2[RESULTS_PER_SIMDGROUP]", + "simd_sum(result0[row])", + "simd_sum(result1[row])", + "simd_sum(result2[row])", + "thread uchar w_thread[PACKS_PER_THREAD * BYTES_PER_PACK]", + "sum0 += x0[i] + x0[i + 1] + x0[i + 2] + x0[i + 3];", + ): + assert fragment in source + assert "return projection(x)" not in source + + +def test_preparation_is_atomic_reversible_and_rebinds_active_o_lora(monkeypatch): + from mtplx import deepseek_v4_0731_m3_wob as candidate + + layers = _layers() + stocks = tuple(layer.attn.wo_b for layer in layers) + built = [] + checked = [] + + def build(_stock): + index = len(built) + + def fixed(value): + return ("candidate", index, value) + + built.append(fixed) + return fixed + + def selfcheck(stock, fixed, index): + checked.append(index) + assert stock is stocks[index] + assert fixed is built[index] + assert tuple(layer.attn.wo_b for layer in layers) == stocks + assert tuple(layer.attn._o_lora_impl.wo_b for layer in layers) == stocks + return True + + monkeypatch.setattr(candidate, "bind_m3_wob", build) + prepared = candidate.prepare_wob_m3(layers, exact_selfcheck=selfcheck) + + assert checked == list(range(43)) + assert prepared.layer_count == 43 + assert prepared.q6_count == 43 + assert prepared.exact_selfchecked == 43 + assert prepared.o_lora_sink_count == 43 + assert tuple(layer.attn.wo_b for layer in layers) == stocks + assert tuple(layer.attn._o_lora_impl.wo_b for layer in layers) == stocks + + prepared.publish() + assert tuple(layer.attn.wo_b for layer in layers) == prepared.published_routes + assert tuple(layer.attn._o_lora_impl.wo_b for layer in layers) == ( + prepared.published_routes + ) + value = SimpleNamespace(shape=(4, 3, 8192)) + assert layers[5].attn.wo_b(value) == ("candidate", 5, value) + prepared.restore() + assert tuple(layer.attn.wo_b for layer in layers) == stocks + assert tuple(layer.attn._o_lora_impl.wo_b for layer in layers) == stocks + + +def test_stale_o_lora_sink_rejected_before_binding(monkeypatch): + from mtplx import deepseek_v4_0731_m3_wob as candidate + + layers = _layers() + stocks = tuple(layer.attn.wo_b for layer in layers) + layers[19].attn._o_lora_impl.wo_b = _projection() + binds = [] + monkeypatch.setattr(candidate, "bind_m3_wob", lambda stock: binds.append(stock)) + + with pytest.raises(candidate.M3WOBContractError, match="o-LoRA wo_b sink"): + candidate.prepare_wob_m3(layers, exact_selfcheck=lambda *_: True) + + assert binds == [] + assert tuple(layer.attn.wo_b for layer in layers) == stocks + + +def test_failed_selfcheck_checks_all_layers_and_publishes_nothing(monkeypatch): + from mtplx import deepseek_v4_0731_m3_wob as candidate + + layers = _layers() + stocks = tuple(layer.attn.wo_b for layer in layers) + checked = [] + monkeypatch.setattr( + candidate, + "bind_m3_wob", + lambda _stock: lambda value: ("candidate", value), + ) + + with pytest.raises(candidate.M3WOBContractError, match="layer 17.*failed"): + candidate.prepare_wob_m3( + layers, + exact_selfcheck=lambda _stock, _fixed, index: ( + checked.append(index) or index != 17 + ), + ) + + assert checked == list(range(43)) + assert tuple(layer.attn.wo_b for layer in layers) == stocks + assert tuple(layer.attn._o_lora_impl.wo_b for layer in layers) == stocks + + +def test_restore_attempts_every_attention_and_o_lora_owner_before_raising(): + from mtplx import deepseek_v4_0731_m3_wob as candidate + + class RejectingOwner: + def __init__(self, stock): + object.__setattr__(self, "stock", stock) + object.__setattr__(self, "wo_b", object()) + + def __setattr__(self, name, value): + if name == "wo_b" and value is self.stock: + raise RuntimeError("first WOB restore failed") + object.__setattr__(self, name, value) + + first_stock = object() + second_stock = object() + first_attention = RejectingOwner(first_stock) + first_o_lora = SimpleNamespace(wo_b=object()) + second_attention = SimpleNamespace(wo_b=object()) + second_o_lora = SimpleNamespace(wo_b=object()) + prepared = candidate.PreparedWOBM3Routes( + attentions=(first_attention, second_attention), + o_lora_impls=(first_o_lora, second_o_lora), + stock_routes=(first_stock, second_stock), + candidate_routes=(object(), object()), + published_routes=(object(), object()), + layer_count=2, + q6_count=2, + exact_selfchecked=2, + o_lora_sink_count=2, + ) + + with pytest.raises(ExceptionGroup, match="WOB route restoration"): + prepared.restore() + + assert first_o_lora.wo_b is first_stock + assert second_attention.wo_b is second_stock + assert second_o_lora.wo_b is second_stock + + +def test_fixed_route_has_only_logical_m_choice(): + from mtplx import deepseek_v4_0731_m3_wob as candidate + + calls = [] + + def stock(value): + calls.append(("stock", value)) + return "stock" + + def fixed(value): + calls.append(("fixed", value)) + return "fixed" + + route = candidate.prebind_wob_route(stock, fixed) + assert route(SimpleNamespace(shape=(2, 3, 8192))) == "fixed" + assert route(SimpleNamespace(shape=(1, 1, 8192))) == "stock" + source = inspect.getsource(type(route).__call__) + assert "shape[1]" in source + assert "try:" not in source + + +@pytest.mark.skipif( + mx.default_device() != mx.gpu or os.environ.get("MTPLX_GPU_TESTS") != "1", + reason="requires an explicitly guarded Metal GPU lane", +) +def test_gpu_wob_is_exact_to_three_stock_m1_calls(): + pytest.skip("GPU gate is intentionally deferred to the locked lane") diff --git a/tests/test_deepseek_v4_0731_m3_wqb_qnorm_rope.py b/tests/test_deepseek_v4_0731_m3_wqb_qnorm_rope.py new file mode 100644 index 00000000..863ab6d5 --- /dev/null +++ b/tests/test_deepseek_v4_0731_m3_wqb_qnorm_rope.py @@ -0,0 +1,297 @@ +"""CPU contracts for the receipt-backed pre-geometry fixed-M3 WQB fusion.""" + +from __future__ import annotations + +import inspect +import os +from types import SimpleNamespace + +import numpy as np +import pytest + +pytest.importorskip("mlx.core") +import mlx.core as mx # noqa: E402 + + +@pytest.fixture(autouse=True) +def _cpu_default_device(): + previous = mx.default_device() + mx.set_default_device(mx.cpu) + try: + yield + finally: + mx.set_default_device(previous) + + +def _projection(): + return SimpleNamespace( + bits=6, + group_size=128, + mode="affine", + bias=None, + weight=SimpleNamespace(shape=(32768, 192), dtype=mx.uint32), + scales=SimpleNamespace(shape=(32768, 8), dtype=mx.bfloat16), + biases=SimpleNamespace(shape=(32768, 8), dtype=mx.bfloat16), + ) + + +def _layers(): + layers = [] + for index in range(43): + + def stock(qr, cos, sin, *, layer_index=index): + return ("stock", layer_index, qr, cos, sin) + + layers.append( + SimpleNamespace( + attn=SimpleNamespace( + wq_b=_projection(), + _q_projection_qhead_route=stock, + ) + ) + ) + return layers + + +def test_contract_binds_only_pre_geometry_q6_g128_wqb(monkeypatch): + from mtplx import deepseek_v4_0731_m3_wqb_qnorm_rope as candidate + + calls = [] + candidate._build_kernel.cache_clear() + monkeypatch.setattr( + candidate.mx.fast, + "metal_kernel", + lambda **kwargs: calls.append(kwargs) or object(), + ) + + bound = candidate.build_0731_m3_wqb_qnorm_rope(_projection()) + + assert bound.input_shape == (1, 3, 1024) + assert bound.output_shape == (1, 3, 64, 512) + assert bound.grid == (64 * 256, 1, 1) + assert bound.threadgroup == (256, 1, 1) + assert calls[0]["input_names"] == [ + "x", + "w", + "scales", + "biases", + "cos", + "sin", + ] + assert calls[0]["output_names"] == ["output"] + assert calls[0]["ensure_row_contiguous"] is False + assert not hasattr(candidate, "m3_wqb_qhead_geometry_variants") + + +def test_contract_rejects_nonreceipt_storage(): + from mtplx import deepseek_v4_0731_m3_wqb_qnorm_rope as candidate + + projection = _projection() + projection.bits = 8 + with pytest.raises(candidate.M3WQBNormRopeContractError, match="Q6/G128"): + candidate.validate_0731_m3_wqb_qnorm_rope(projection) + + projection = _projection() + projection.weight.shape = (32768, 191) + with pytest.raises(candidate.M3WQBNormRopeContractError, match="packed wq_b"): + candidate.validate_0731_m3_wqb_qnorm_rope(projection) + + +def test_source_preserves_pre_geometry_qmv_norm_and_rope_arithmetic(): + from mtplx import deepseek_v4_0731_m3_wqb_qnorm_rope as candidate + + assert candidate.RECORDED_PRE_GEOMETRY_SOURCE_SHA256 == ( + "2eb4ce3d5bae9c9b71574d17fedfd37b6755d94299cb4ed6b01d2015c5f8f9a1" + ) + assert candidate.RECORDED_PRE_GEOMETRY_TEST_SHA256 == ( + "0c551aa8d7f865454d3a9b6d22f46af5115e00a2fc48c5db82225516c2a77cbb" + ) + source = candidate.m3_wqb_qnorm_rope_metal_source() + for fragment in ( + "constexpr uint M = 3;", + "constexpr uint K = 1024;", + "constexpr uint N = 32768;", + "constexpr uint HEADS = 64;", + "constexpr uint HEAD_DIM = 512;", + "constexpr uint ROPE_DIM = 64;", + "float result0[RESULTS_PER_SIMDGROUP] = {0.0f};", + "float result1[RESULTS_PER_SIMDGROUP] = {0.0f};", + "float result2[RESULTS_PER_SIMDGROUP] = {0.0f};", + "constexpr uint NORM_LANES = 32;", + "constexpr uint NORM_READS = 4;", + "metal::precise::rsqrt(mean + EPS);", + "precise float rope0_lhs = x0 * c[pair];", + "precise float rope1_rhs = x1 * c[pair];", + ): + assert fragment in source + assert ( + "geometry" + not in inspect.signature(candidate.m3_wqb_qnorm_rope_metal_source).parameters + ) + assert "parallel_norm" not in source + assert "shared_x" not in source + + +def test_bound_call_is_one_fixed_m3_launch_without_hot_validation(monkeypatch): + from mtplx import deepseek_v4_0731_m3_wqb_qnorm_rope as candidate + + launches = [] + + def fake_kernel(*, inputs, grid, threadgroup, output_shapes, output_dtypes, **_): + launches.append((inputs, grid, threadgroup, output_shapes, output_dtypes)) + return (mx.zeros(output_shapes[0], dtype=output_dtypes[0]),) + + candidate._build_kernel.cache_clear() + monkeypatch.setattr(candidate.mx.fast, "metal_kernel", lambda **_: fake_kernel) + bound = candidate.build_0731_m3_wqb_qnorm_rope(_projection()) + actual = bound( + mx.zeros((1, 3, 1024), dtype=mx.bfloat16), + mx.zeros((3, 32)), + mx.zeros((3, 32)), + ) + + assert len(launches) == 1 + inputs, grid, threadgroup, output_shapes, output_dtypes = launches[0] + assert tuple(inputs[0].shape) == (3, 1024) + assert tuple(inputs[4].shape) == tuple(inputs[5].shape) == (3, 32) + assert grid == (64 * 256, 1, 1) + assert threadgroup == (256, 1, 1) + assert output_shapes == [(3, 32768)] + assert output_dtypes == [mx.bfloat16] + assert actual.shape == (1, 3, 64, 512) + hot_source = inspect.getsource(type(bound).__call__) + assert "if " not in hot_source + assert "try:" not in hot_source + + +def test_preparation_selfchecks_43_before_reversible_publication(monkeypatch): + from mtplx import deepseek_v4_0731_m3_wqb_qnorm_rope as candidate + + layers = _layers() + stocks = tuple(layer.attn._q_projection_qhead_route for layer in layers) + built = [] + checked = [] + + def build(_projection): + index = len(built) + + def fused(qr, cos, sin): + return ("candidate", index, qr, cos, sin) + + built.append(fused) + return fused + + def selfcheck(stock, fused, index): + checked.append(index) + assert stock is stocks[index] + assert fused is built[index] + assert tuple(layer.attn._q_projection_qhead_route for layer in layers) == stocks + return True + + monkeypatch.setattr(candidate, "build_0731_m3_wqb_qnorm_rope", build) + prepared = candidate.prepare_wqb_qhead_m3( + layers, + exact_selfcheck=selfcheck, + ) + + assert checked == list(range(43)) + assert prepared.q6_count == 43 + assert prepared.exact_selfchecked == 43 + assert len(prepared.published_routes) == 43 + assert tuple(layer.attn._q_projection_qhead_route for layer in layers) == stocks + + prepared.publish() + assert tuple(layer.attn._q_projection_qhead_route for layer in layers) == ( + prepared.published_routes + ) + qr = SimpleNamespace(shape=(2, 3, 1024)) + assert layers[17].attn._q_projection_qhead_route(qr, "c", "s")[:2] == ( + "candidate", + 17, + ) + prepared.restore() + assert tuple(layer.attn._q_projection_qhead_route for layer in layers) == stocks + + +def test_preparation_failure_publishes_nothing(monkeypatch): + from mtplx import deepseek_v4_0731_m3_wqb_qnorm_rope as candidate + + layers = _layers() + stocks = tuple(layer.attn._q_projection_qhead_route for layer in layers) + monkeypatch.setattr( + candidate, + "build_0731_m3_wqb_qnorm_rope", + lambda _projection: lambda qr, cos, sin: (qr, cos, sin), + ) + + with pytest.raises( + candidate.M3WQBNormRopeContractError, + match="layer 9.*self-check failed", + ): + candidate.prepare_wqb_qhead_m3( + layers, + exact_selfcheck=lambda _stock, _fused, index: index != 9, + ) + + assert tuple(layer.attn._q_projection_qhead_route for layer in layers) == stocks + + +def test_restore_attempts_every_qhead_owner_before_raising(): + from mtplx import deepseek_v4_0731_m3_wqb_qnorm_rope as candidate + + good_stock = object() + + class RejectingOwner: + def __init__(self, stock): + object.__setattr__(self, "stock", stock) + object.__setattr__(self, "_q_projection_qhead_route", object()) + + def __setattr__(self, name, value): + if name == "_q_projection_qhead_route" and value is self.stock: + raise RuntimeError("first qhead restore failed") + object.__setattr__(self, name, value) + + first_stock = object() + first = RejectingOwner(first_stock) + second = SimpleNamespace(_q_projection_qhead_route=object()) + prepared = candidate.PreparedWQBQHeadM3Routes( + attentions=(first, second), + stock_routes=(first_stock, good_stock), + candidate_routes=(object(), object()), + published_routes=(object(), object()), + q6_count=2, + exact_selfchecked=2, + ) + + with pytest.raises(ExceptionGroup, match="qhead route restoration"): + prepared.restore() + + assert second._q_projection_qhead_route is good_stock + + +def test_cpu_oracle_keeps_norm_and_rope_owned_by_each_row(): + from mtplx.deepseek_v4_0731_m3_wqb_qnorm_rope import ( + q_head_norm_rope_cpu_oracle, + ) + + q = np.zeros((1, 3, 1, 4), dtype=np.float32) + q[..., 0] = 3.0 + q[..., 1] = 4.0 + q[..., 2] = 1.0 + q[..., 3] = 2.0 + cos = np.array([[1.0], [0.0], [-1.0]], dtype=np.float32) + sin = np.array([[0.0], [1.0], [0.0]], dtype=np.float32) + + got = q_head_norm_rope_cpu_oracle(q, cos, sin, eps=0.0, rope_dim=2) + scale = 1.0 / np.sqrt(7.5) + assert np.allclose(got[0, 0, 0], np.array([3, 4, 1, 2]) * scale) + assert np.allclose(got[0, 1, 0], np.array([3, 4, -2, 1]) * scale) + assert np.allclose(got[0, 2, 0], np.array([3, 4, -1, -2]) * scale) + + +@pytest.mark.skipif( + mx.default_device() != mx.gpu or os.environ.get("MTPLX_GPU_TESTS") != "1", + reason="requires an explicitly guarded Metal GPU lane", +) +def test_gpu_pre_geometry_candidate_is_exact_to_three_stock_m1_calls(): + pytest.skip("GPU gate is intentionally deferred to the locked lane") diff --git a/tests/test_deepseek_v4_0731_moe.py b/tests/test_deepseek_v4_0731_moe.py new file mode 100644 index 00000000..0bb8dd16 --- /dev/null +++ b/tests/test_deepseek_v4_0731_moe.py @@ -0,0 +1,165 @@ +"""CPU construction contracts for the receipt-backed 0731 routed-Q2 lane.""" + +from __future__ import annotations + +import inspect +from types import SimpleNamespace + +import pytest + +pytest.importorskip("mlx.core") +import mlx.core as mx # noqa: E402 +from mlx_lm.models.switch_layers import QuantizedSwitchLinear # noqa: E402 + +from mtplx.deepseek_v4_0731_moe import ( # noqa: E402 + DeepseekV40731PackedQ2SwitchGLU, + build_routed_q2_pair, + build_row_owned_combine_m1, + exact_selfcheck_row_owned_combine_m1, + validate_routed_q2_pair, +) + + +@pytest.fixture(autouse=True) +def _cpu_default_device(): + previous = mx.default_device() + mx.set_default_device(mx.cpu) + try: + yield + finally: + mx.set_default_device(previous) + + +def _q2(in_features: int = 256, out_features: int = 64, experts: int = 8): + projection = QuantizedSwitchLinear( + in_features, + out_features, + experts, + bias=False, + group_size=128, + bits=2, + ) + projection.scales = projection.scales.astype(mx.bfloat16) + projection.biases = projection.biases.astype(mx.bfloat16) + mx.eval(projection.parameters()) + return projection + + +def test_routed_q2_pair_requires_exact_affine_q2_group128_storage(): + gate = _q2() + up = _q2() + + contract = validate_routed_q2_pair( + gate, + up, + hidden_size=256, + width=64, + experts=8, + ) + + assert contract.bits == 2 + assert contract.group_size == 128 + assert contract.hidden_size == 256 + assert contract.width == 64 + assert contract.experts == 8 + + up.group_size = 64 + with pytest.raises(ValueError, match="affine Q2/group-128"): + validate_routed_q2_pair( + gate, + up, + hidden_size=256, + width=64, + experts=8, + ) + + +def test_routed_q2_pair_packs_only_output_rows_and_owns_fixed_unsorted_path(): + gate = _q2() + up = _q2() + down = object() + activation = object() + switch = SimpleNamespace( + gate_proj=gate, + up_proj=up, + down_proj=down, + activation=activation, + ) + + packed = build_routed_q2_pair( + switch, + hidden_size=256, + width=64, + experts=8, + ) + + assert type(packed) is DeepseekV40731PackedQ2SwitchGLU + assert packed._split_at == 64 + assert packed.down_proj is down + assert packed.activation is activation + assert packed.gate_up_proj.weight.shape == (8, 128, 16) + assert packed.gate_up_proj.scales.shape == (8, 128, 2) + assert packed.gate_up_proj.biases.shape == (8, 128, 2) + hot_source = inspect.getsource(type(packed).__call__) + assert "moe_force_unsorted_enabled" not in hot_source + assert "environ" not in hot_source + assert "try:" not in hot_source + + +def test_row_owned_m1_combine_binds_fixed_top6_geometry_without_gpu( + monkeypatch, +): + calls = [] + + def kernel(**kwargs): + calls.append(kwargs) + return (mx.zeros((1, 4096), dtype=mx.bfloat16),) + + import mtplx.deepseek_v4_0731_moe as moe + + monkeypatch.setattr(moe.mx.metal, "is_available", lambda: True) + monkeypatch.setattr(moe, "_row_owned_combine_m1_kernel", lambda: kernel) + combine = build_row_owned_combine_m1(hidden_size=4096, top_k=6) + + got = combine( + mx.zeros((1, 6, 4096), dtype=mx.bfloat16), + mx.zeros((1, 6), dtype=mx.float32), + ) + + assert got.shape == (1, 4096) + assert calls[0]["grid"] == (4096, 1, 1) + assert calls[0]["threadgroup"] == (128, 1, 1) + assert calls[0]["output_dtypes"] == [mx.bfloat16] + + +def test_row_owned_m1_combine_rejects_every_nonreceipt_geometry(monkeypatch): + import mtplx.deepseek_v4_0731_moe as moe + + monkeypatch.setattr(moe.mx.metal, "is_available", lambda: True) + with pytest.raises(ValueError, match="row-owned combine geometry"): + build_row_owned_combine_m1(hidden_size=4096, top_k=8) + with pytest.raises(ValueError, match="row-owned combine geometry"): + build_row_owned_combine_m1(hidden_size=4100, top_k=6) + + +def test_row_owned_m1_exact_selfcheck_executes_and_rejects_mismatch(): + calls = [] + + def exact(routed, route_weights): + calls.append((routed, route_weights)) + accumulator = mx.zeros((1, 4096), dtype=mx.bfloat16) + weights = route_weights.astype(mx.bfloat16) + for expert in range(6): + product = (routed[:, expert] * weights[:, expert : expert + 1]).astype( + mx.bfloat16 + ) + accumulator = (accumulator + product).astype(mx.bfloat16) + return accumulator + + exact_selfcheck_row_owned_combine_m1(exact) + assert len(calls) == 1 + + with pytest.raises(ValueError, match="exact self-check failed"): + exact_selfcheck_row_owned_combine_m1( + lambda routed, _weights: mx.zeros_like(routed[:, 0]) + ) diff --git a/tests/test_deepseek_v4_attention_island.py b/tests/test_deepseek_v4_attention_island.py index 3c714e85..b5c3c1de 100644 --- a/tests/test_deepseek_v4_attention_island.py +++ b/tests/test_deepseek_v4_attention_island.py @@ -2,6 +2,7 @@ from __future__ import annotations +import inspect from types import SimpleNamespace from pathlib import Path @@ -12,6 +13,7 @@ from mlx.utils import tree_flatten, tree_unflatten # noqa: E402 from mtplx import deepseek_v4_attention_island as AI # noqa: E402 +from mtplx.moe_packed_projections import PackedSwitchGLU # noqa: E402 from mtplx.models import deepseek_v4 as D # noqa: E402 @@ -93,9 +95,9 @@ def _quantized_layer( def _post_attention_inputs(args, width: int): mx.random.seed(90 + width) x = mx.random.normal((1, width, args.hidden_size)).astype(mx.bfloat16) - residual = mx.random.normal( - (1, width, args.hc_mult, args.hidden_size) - ).astype(mx.bfloat16) + residual = mx.random.normal((1, width, args.hc_mult, args.hidden_size)).astype( + mx.bfloat16 + ) post = mx.random.normal((1, width, args.hc_mult)).astype(mx.float32) comb = mx.softmax( mx.random.normal((1, width, args.hc_mult, args.hc_mult)), axis=-1 @@ -105,6 +107,96 @@ def _post_attention_inputs(args, width: int): return x, residual, post, comb, ids +def _fake_affine_projection(*, bits, input_dim, output_dim): + packed = input_dim // (32 // bits) + groups = input_dim // 128 + + def leaf(shape, dtype): + return SimpleNamespace(shape=shape, ndim=len(shape), dtype=dtype) + + return SimpleNamespace( + bits=bits, + group_size=128, + mode="affine", + bias=None, + weight=leaf((output_dim, packed), mx.uint32), + scales=leaf((output_dim, groups), mx.bfloat16), + biases=leaf((output_dim, groups), mx.bfloat16), + ) + + +def _synthetic_target_stack_layer(monkeypatch): + class FakeActivation: + limit = 10.0 + + class FakePackedSwitch(PackedSwitchGLU): + pass + + class FakeMoE: + pass + + class FakeLayer: + pass + + monkeypatch.setattr(AI.D, "ClampedSwiGLU", FakeActivation) + monkeypatch.setattr(AI.D, "DeepseekV4MoE", FakeMoE) + monkeypatch.setattr(AI.D, "DeepseekV4DecoderLayer", FakeLayer) + + gate_up = _fake_affine_projection( + bits=2, + input_dim=4096, + output_dim=4096, + ) + down = _fake_affine_projection( + bits=2, + input_dim=2048, + output_dim=4096, + ) + switch = FakePackedSwitch(gate_up, down, FakeActivation(), 2048) + + shared = SimpleNamespace( + gate_proj=_fake_affine_projection( + bits=8, + input_dim=4096, + output_dim=2048, + ), + up_proj=_fake_affine_projection( + bits=8, + input_dim=4096, + output_dim=2048, + ), + down_proj=_fake_affine_projection( + bits=8, + input_dim=2048, + output_dim=4096, + ), + limit=10.0, + ) + router = SimpleNamespace( + hash=False, + e_score_correction_bias=object(), + weight=object(), + topk=6, + score_func="sigmoid", + route_scale=1.5, + ) + ffn = FakeMoE() + ffn.switch_mlp = switch + ffn.shared_experts = shared + ffn.gate = router + layer = FakeLayer() + layer.ffn = ffn + layer.ffn_hc = SimpleNamespace( + _static=lambda: (object(), object(), object()), + hc=4, + _iters=20, + eps=1e-6, + _sinkhorn_kernel=True, + ) + layer.ffn_norm = SimpleNamespace(weight=object(), eps=1e-6) + return layer + + def _stock_post_attention(layer, x, residual, post, comb, ids): h = layer.attn_hc.post(x, residual, post, comb) ffn_residual = h @@ -156,8 +248,10 @@ def _shape_model(seed: int = 71): for layer_index, layer in enumerate(model.layers): for name in ("gate_proj", "up_proj", "down_proj"): projection = getattr(layer.ffn.switch_mlp, name) - group_size = 64 if layer_index == 2 and name == "gate_proj" else ( - 32 if name == "gate_proj" else 64 + group_size = ( + 64 + if layer_index == 2 and name == "gate_proj" + else (32 if name == "gate_proj" else 64) ) setattr( layer.ffn.switch_mlp, @@ -223,6 +317,75 @@ def test_attention_island_matches_eager_post_attention_chain( assert mx.array_equal(want, got) +def test_target_stack_binder_reaches_all_43_real_calls_with_paired_q2_shared_q8( + monkeypatch, +): + layer = _synthetic_target_stack_layer(monkeypatch) + tape_calls = [] + combine = object() + + def fake_tape(**kwargs): + tape_calls.append(kwargs) + return lambda *_args: None + + monkeypatch.setattr(AI, "_attention_island_tape", fake_tape) + + bound = [ + AI._bind_attention_island_layer( + layer, + width=1, + allowed_widths=(1,), + shared_bits=8, + routed_pair=True, + routed_combine=combine, + ) + for _ in range(43) + ] + + assert len(bound) == 43 + assert len(tape_calls) == 43 + assert all(call["routed_gate"] is None for call in tape_calls) + assert all(call["routed_up"] is None for call in tape_calls) + assert all(call["routed_gate_up"].bits == 2 for call in tape_calls) + assert all(call["routed_gate_up"].group_size == 128 for call in tape_calls) + assert all(call["shared_gate"].bits == 8 for call in tape_calls) + assert all(call["shared_up"].bits == 8 for call in tape_calls) + assert all(call["shared_down"].bits == 8 for call in tape_calls) + assert all(call["routed_combine"] is combine for call in tape_calls) + + +def test_target_stack_binder_accepts_staged_switch_without_mutating_live_layer( + monkeypatch, +): + layer = _synthetic_target_stack_layer(monkeypatch) + staged_switch = layer.ffn.switch_mlp + live_switch = object() + layer.ffn.switch_mlp = live_switch + monkeypatch.setattr(AI, "_attention_island_tape", lambda **_kwargs: lambda *_: None) + + bound = AI._bind_attention_island_layer( + layer, + width=3, + allowed_widths=(3,), + shared_bits=8, + routed_pair=True, + routed_switch=staged_switch, + ) + + assert callable(bound) + assert layer.ffn.switch_mlp is live_switch + + +def test_target_stack_binder_has_no_shared_pair_or_custom_q2_kernel_modes(): + parameters = inspect.signature(AI._bind_attention_island_layer).parameters + assert "allowed_widths" in parameters + assert "shared_bits" in parameters + assert "routed_pair" in parameters + assert "routed_combine" in parameters + assert "shared_pair" not in parameters + assert "routed_gate_up_kernel" not in parameters + + def test_attention_island_reuses_tapes_by_width_router_and_q2_layout(): _, score_a = _quantized_layer(31, routed_gate_group=32) _, score_b = _quantized_layer(32, routed_gate_group=32) @@ -267,7 +430,9 @@ def test_bound_width_body_matches_full_stock_body_cache_logits_and_argmax(width) assert mx.array_equal(want_hidden, got_hidden) assert mx.array_equal(want_logits, got_logits) - assert mx.array_equal(mx.argmax(want_logits, axis=-1), mx.argmax(got_logits, axis=-1)) + assert mx.array_equal( + mx.argmax(want_logits, axis=-1), mx.argmax(got_logits, axis=-1) + ) _assert_cache_equal(control_cache, candidate_cache) # The ratio-4 cache exposes only its logical rows, never physical padding. assert control_cache[1].compressed.shape[1] == control_cache[1].n_compressed @@ -357,9 +522,7 @@ def stock(ids, cache): assert route(SimpleNamespace(shape=(1, 3)), "cache") == "stock" monkeypatch.setattr(AI, "current_attention_phase", lambda: "decode_verify") for kind in ("repair", "other"): - monkeypatch.setattr( - AI, "current_model_forward_kind", lambda kind=kind: kind - ) + monkeypatch.setattr(AI, "current_model_forward_kind", lambda kind=kind: kind) assert route(SimpleNamespace(shape=(1, 3)), "cache") == "stock" monkeypatch.setattr(AI, "current_model_forward_kind", lambda: "target_verify") assert route(SimpleNamespace(shape=(1, 1)), "cache") == "stock" diff --git a/tests/test_deepseek_v4_dspark.py b/tests/test_deepseek_v4_dspark.py index 0933ae3a..7536f8c3 100644 --- a/tests/test_deepseek_v4_dspark.py +++ b/tests/test_deepseek_v4_dspark.py @@ -512,6 +512,18 @@ def test_official_hc_names_map_to_exact_installed_parameter_keys(tiny_model): assert expected <= installed +def test_dspark_grouped_o_lora_storage_flattens_at_load_boundary(tiny_model): + raw = { + "mtp.0.attn.wo_a.weight": mx.zeros((1, 8, 3), dtype=mx.uint32), + "mtp.1.attn.wo_a.weight": mx.zeros((1, 8, 3), dtype=mx.uint32), + "mtp.2.attn.wo_a.weight": mx.zeros((1, 8, 3), dtype=mx.uint32), + } + mapped = tiny_model.sanitize(raw) + assert all( + mapped[f"mtp.{stage}.attn.wo_a.weight"].shape == (8, 3) for stage in range(3) + ) + + def test_dspark_visibility_is_exact_and_includes_all_five_draft_rows(): got = np.array(D.get_dspark_topk_idxs(8, 2, 5, 3)) expected = [0, 1, 2, 3, 8, 9, 10, 11, 12] diff --git a/tests/test_deepseek_v4_dspark_generation.py b/tests/test_deepseek_v4_dspark_generation.py new file mode 100644 index 00000000..bcfda459 --- /dev/null +++ b/tests/test_deepseek_v4_dspark_generation.py @@ -0,0 +1,1131 @@ +import inspect +from pathlib import Path + +import numpy as np +import pytest + +pytest.importorskip("mlx.core") +import mlx.core as mx # noqa: E402 + +from mtplx import generation as generation_module # noqa: E402 +from mtplx import native_block_speculation as native_speculation # noqa: E402 +from mtplx.deepseek_v4_dspark_generation import ( # noqa: E402 + DeepseekV4DSparkBackend, +) +from mtplx.mtp_patch import MTPContract # noqa: E402 +from mtplx.models.deepseek_v4 import DeepseekV4DSparkCache # noqa: E402 +from mtplx.runtime import MTPLXRuntime # noqa: E402 +from mtplx.sampling import SamplerConfig # noqa: E402 + + +def generate_mtpk(rt, prompt_ids, **kwargs): + """Exercise the native scheduler without depending on outer dispatch.""" + return native_speculation.generate_native_block_speculative( + rt, + rt.block_speculative_backend, + prompt_ids, + abort_check=kwargs.pop("abort_check", None), + max_tokens=kwargs.pop("max_tokens"), + sampler=kwargs.pop("sampler"), + speculative_depth=kwargs.pop("speculative_depth"), + seed=kwargs.pop("seed", 0), + stop_token_ids=kwargs.pop("stop_token_ids", None), + draft_sampler=kwargs.pop("draft_sampler", None), + token_callback=kwargs.pop("token_callback", None), + prefill_callback=kwargs.pop("prefill_callback", None), + constraint=kwargs.pop("constraint", None), + vision_splice=kwargs.pop("vision_splice", None), + adaptive_policy=kwargs.pop("adaptive_policy", None), + adaptive_width_policy=kwargs.pop("adaptive_width_policy", None), + ) + + +class _Tokenizer: + eos_token_id = None + pad_token_id = None + + def decode(self, tokens): + return " ".join(str(int(token)) for token in tokens) + + +class _TargetCache: + def __init__(self): + self.offset = 0 + self.trimmed = [] + + def trim(self, n): + self.offset -= int(n) + self.trimmed.append(int(n)) + return int(n) + + +class _ProposalCache: + def __init__(self): + self.ring = None + self.prefill_length = 0 + + +class _FakeDSpark: + def __init__(self): + self.owner = None + self.stages = (object(), object(), object()) + + def make_cache(self): + cache = [_ProposalCache(), _ProposalCache(), _ProposalCache()] + self.owner.proposal_cache = cache + return cache + + def prefill(self, hidden, cache): + for entry in cache: + entry.ring = hidden[...] + entry.prefill_length = int(hidden.shape[1]) + self.owner.prefill_hidden = np.asarray(hidden).copy() + self.owner.proposal_prefill_spans.append((0, np.asarray(hidden).copy())) + + def forward( + self, + hidden, + token_ids, + embed_tokens, + lm_head, + cache, + *, + start_pos, + greedy, + ids_only_width, + forced_first_token_ids=None, + ): + del embed_tokens, lm_head, greedy + for entry in cache: + entry.ring = mx.full((1, 1, 1), 999.0) + ids, _logits, _confidence = self.owner.draft_deepseek_v4_dspark( + hidden, token_ids, cache, start_pos=start_pos + ) + if forced_first_token_ids is not None: + forced = int(np.asarray(forced_first_token_ids)[0]) + ids = mx.concatenate([ids[:, :1], mx.array([[forced]]), ids[:, 2:]], axis=1) + self.owner.forced_primary_ids.append(forced) + self.owner.proposal_inputs.append( + (int(np.asarray(token_ids)[0]), int(start_pos), int(ids_only_width)) + ) + self.owner.proposal_widths.append(int(ids_only_width)) + return ids[:, : 1 + int(ids_only_width)] + + def commit_main(self, hidden, cache, *, start_pos): + self.owner.proposal_prefill_spans.append( + (int(start_pos), np.asarray(hidden).copy()) + ) + self.owner.commit_ring_history.append( + [np.asarray(entry.ring).copy() for entry in cache] + ) + for entry in cache: + entry.ring = hidden[...] + entry.prefill_length = int(start_pos) + int(hidden.shape[1]) + self.owner.commits.append((int(start_pos), np.asarray(hidden).copy())) + + +class _DSparkModel: + def __init__(self): + self._dspark = _FakeDSpark() + self.model = type("_Body", (), {"embed_tokens": staticmethod(lambda x: x)})() + self.lm_head = lambda x: x + self.owner = None + + def __call__(self, *args, **kwargs): + return self.owner.target_forward(*args, **kwargs) + + +class _DSparkRuntime(MTPLXRuntime): + def __init__(self): + model = _DSparkModel() + super().__init__( + model=model, + tokenizer=_Tokenizer(), + model_path=Path("."), + mtp_enabled=True, + contract=MTPContract(), + ) + self.deepseek_v4_dspark_enabled = True + model._dspark.owner = self + model.owner = self + self.block_speculative_backend = DeepseekV4DSparkBackend.bind(model) + self.target_cache = _TargetCache() + self.prefill_hidden = None + self.proposal_cache = None + self.draft_widths = [] + self.proposal_widths = [] + self.commits = [] + self.commit_ring_history = [] + self.proposal_inputs = [] + self.proposal_prefill_spans = [] + self.forced_primary_ids = [] + self.target_forward_widths = [] + self.target_forward_inputs = [] + self.target_forward_options = [] + self.runtime_forward_calls = 0 + + def make_cache(self): + return [self.target_cache] + + def target_forward( + self, + input_ids, + cache=None, + return_hidden=False, + hidden_variant=None, + emit_logits=True, + logits_keep=None, + input_embeddings=None, + ): + del hidden_variant, input_embeddings + ids = np.asarray(input_ids, dtype=np.int32) + self.target_forward_widths.append(int(ids.shape[1])) + self.target_forward_inputs.append(ids.copy()) + self.target_forward_options.append( + { + "return_hidden": bool(return_hidden), + "emit_logits": bool(emit_logits), + "logits_keep": logits_keep, + } + ) + cache[0].offset += ids.shape[1] + vocab = 1024 + logits = np.full((1, ids.shape[1], vocab), -1000.0, dtype=np.float32) + for row, token in enumerate(ids[0]): + logits[0, row, int(token) + 1] = 1000.0 + if logits_keep is not None: + logits = logits[:, -int(logits_keep) :] + hidden = np.repeat(ids[..., None].astype(np.float32), 3, axis=-1) + result = mx.array(logits) if emit_logits else None + return (result, mx.array(hidden)) if return_hidden else result + + def forward_ar(self, *args, **kwargs): + self.runtime_forward_calls += 1 + self._count("forward_ar_hidden_calls") + return self.target_forward(*args, **kwargs) + + def make_deepseek_v4_dspark_cache(self): + return [object(), object(), object()] + + def prefill_deepseek_v4_dspark(self, hidden, cache): + del cache + self.prefill_hidden = np.asarray(hidden).copy() + + def draft_deepseek_v4_dspark(self, hidden, token_ids, cache, *, start_pos): + del hidden, cache, start_pos + token = int(np.asarray(token_ids)[0]) + # Primary plus two exact future drafts, followed by a deliberate miss. + # This exercises K2 proposal acceptance against serial target rows. + ids = mx.array([[token, token + 1, token + 2, token + 3, 999, 999]]) + return ids, mx.zeros((1, 5, 1024)), mx.ones((1, 5)) + + def draft_deepseek_v4_dspark_ids( + self, hidden, token_ids, cache, *, start_pos, width + ): + self.proposal_widths.append(int(width)) + ids, _logits, _confidence = self.draft_deepseek_v4_dspark( + hidden, token_ids, cache, start_pos=start_pos + ) + return ids[:, : 1 + int(width)] + + def commit_deepseek_v4_dspark(self, hidden, cache, *, start_pos): + del cache + self.commits.append((int(start_pos), np.asarray(hidden).copy())) + + +def test_dspark_fixed_k2_preserves_the_greedy_target_stream(): + rt = _DSparkRuntime() + callback = [] + out = generate_mtpk( + rt, + [9, 10], + max_tokens=8, + sampler=SamplerConfig(temperature=0.0, top_p=1.0, top_k=0), + speculative_depth=2, + stop_token_ids=set(), + token_callback=lambda block: callback.extend(block), + ) + assert out.tokens == list(range(11, 19)) + assert callback == out.tokens + assert out.stats.speculative_depth == 2 + assert out.stats.runtime_mtp_enabled is True + assert rt.prefill_hidden.tolist() == [[[9.0, 9.0, 9.0], [10.0, 10.0, 10.0]]] + assert all(hidden.shape[1] <= 3 for _, hidden in rt.commits) + assert all(float(hidden.max()) < 1000.0 for _, hidden in rt.commits) + # The tail token is accepted from the final verified block, so every emitted + # token is already represented in the target cache. + assert rt.target_cache.offset == 2 + len(out.tokens) + assert rt.proposal_widths + assert all(width <= 3 for width in rt.proposal_widths) + assert out.stats.events == [] + + +class _WidthDivergentTargetRuntime(_DSparkRuntime): + """Model the receipt-proven case where batched target rows are not serial M1.""" + + def __init__(self): + super().__init__() + self.batched_decode_calls = 0 + + def target_forward(self, input_ids, cache=None, **kwargs): + offset_before = cache[0].offset + logits, hidden = super().target_forward(input_ids, cache=cache, **kwargs) + width = int(np.asarray(input_ids).shape[1]) + if offset_before > 0 and width > 1: + self.batched_decode_calls += 1 + poisoned_logits = np.full(logits.shape, -1000.0, dtype=np.float32) + poisoned_logits[..., 777] = 1000.0 + logits = mx.array(poisoned_logits) + hidden = hidden + 1000.0 + return logits, hidden + + +def test_dspark_target_verification_is_serial_m1_when_physical_m3_diverges(): + rt = _WidthDivergentTargetRuntime() + + out = generate_mtpk( + rt, + [9, 10], + max_tokens=4, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + ) + + assert out.tokens == [11, 12, 13, 14] + assert rt.batched_decode_calls == 0 + assert rt.target_forward_widths == [1, 1, 1, 1, 1, 1] + assert all(float(hidden.max()) < 1000.0 for _, hidden in rt.commits) + assert rt.target_cache.offset == 2 + len(out.tokens) + + +def test_dspark_uses_one_sanctioned_scheduler_evaluation_boundary(monkeypatch): + events = [] + + class _OrderingRuntime(_DSparkRuntime): + future = None + + def target_forward(self, input_ids, cache=None, **kwargs): + if cache[0].offset >= 2 and int(np.asarray(input_ids).shape[1]) == 1: + events.append("target_row_zero_graph") + return super().target_forward(input_ids, cache=cache, **kwargs) + + rt = _OrderingRuntime() + original_propose = DeepseekV4DSparkBackend.propose + + def tracked_propose(self, *args, **kwargs): + events.append("proposal_graph") + rt.future = original_propose(self, *args, **kwargs) + return rt.future + + original_asarray = native_speculation.np.asarray + + def tracked_asarray(value, *args, **kwargs): + if value is rt.future: + events.append("proposal_materialized") + return original_asarray(value, *args, **kwargs) + + monkeypatch.setattr(DeepseekV4DSparkBackend, "propose", tracked_propose) + monkeypatch.setattr(native_speculation.np, "asarray", tracked_asarray) + + out = generate_mtpk( + rt, + [9, 10], + max_tokens=3, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + ) + + assert out.tokens == [11, 12, 13] + assert events[:3] in ( + ["proposal_graph", "target_row_zero_graph", "proposal_materialized"], + ["proposal_graph", "proposal_materialized", "target_row_zero_graph"], + ) + + +def test_dspark_uses_generic_backend_even_without_legacy_family_flag(): + rt = _DSparkRuntime() + rt.deepseek_v4_dspark_enabled = False + out = generate_mtpk( + rt, + [9, 10], + max_tokens=4, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + ) + assert out.tokens == [11, 12, 13, 14] + + +def test_dspark_depth_two_keeps_k2_proposal_with_serial_m1_target_rows(): + rt = _DSparkRuntime() + + out = generate_mtpk( + rt, + [9, 10], + max_tokens=4, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + ) + + assert out.tokens == [11, 12, 13, 14] + # Depth two still proposes primary + two future drafts together, while the + # target advances one exact serial row per emitted token. + assert rt.target_forward_widths == [1, 1, 1, 1, 1, 1] + assert rt.proposal_widths == [3] + assert out.stats.accepted_drafts == 2 + assert out.stats.drafted_tokens == 2 + assert out.stats.accepted_by_depth == [1, 1] + assert out.stats.repair_time_s == 0.0 + + +def test_dspark_proposal_restore_precedes_the_full_accepted_prefix_commit(): + """The proposal's poisoned rings never reach the accepted-prefix commit.""" + rt = _DSparkRuntime() + + out = generate_mtpk( + rt, + [9, 10], + max_tokens=3, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + ) + + assert out.stats.accepted_drafts == 2 + assert rt.target_cache.trimmed == [] + assert len(rt.commit_ring_history) == 1 + for ring in rt.commit_ring_history[0]: + np.testing.assert_array_equal(ring, rt.prefill_hidden) + assert float(np.max(ring)) != 999.0 + + +def test_dspark_single_token_prompt_uses_one_explicit_m1_seed_before_fixed_k2(): + rt = _DSparkRuntime() + + out = generate_mtpk( + rt, + [10], + max_tokens=4, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + ) + + assert out.tokens == [11, 12, 13, 14] + assert rt.target_forward_widths == [1, 1, 1, 1, 1] + assert rt.proposal_inputs == [(11, 1, 3)] + assert out.stats.verify_calls == 1 + assert out.stats.accepted_drafts == 2 + + +def test_dspark_backend_uses_construction_bound_model_operations(): + rt = _DSparkRuntime() + + def forbidden(*_args, **_kwargs): + raise AssertionError("enabled backend called a validating runtime wrapper") + + rt.make_deepseek_v4_dspark_cache = forbidden + rt.prefill_deepseek_v4_dspark = forbidden + rt.draft_deepseek_v4_dspark_ids = forbidden + rt.commit_deepseek_v4_dspark = forbidden + + out = generate_mtpk( + rt, + [9, 10], + max_tokens=4, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + ) + + assert out.tokens == [11, 12, 13, 14] + + +def test_dspark_backend_binds_the_target_callable_before_hot_decode(): + rt = _DSparkRuntime() + before = dict(rt.diagnostic_counters) + target_forward = rt.block_speculative_backend.bind_target_forward(rt) + + def forbidden(*_args, **_kwargs): + raise AssertionError("hot decode re-looked up the runtime target route") + + rt.forward_ar = forbidden + logits, hidden = target_forward( + mx.array([[10]]), cache=rt.make_cache(), return_hidden=True + ) + + assert logits.shape[1] == 1 + assert hidden.shape[1] == 1 + assert rt.target_forward_widths == [1] + assert rt.runtime_forward_calls == 0 + assert rt.diagnostic_counters == before + + +class _SecondProposalMissRuntime(_DSparkRuntime): + def draft_deepseek_v4_dspark(self, hidden, token_ids, cache, *, start_pos): + del hidden, cache, start_pos + token = int(np.asarray(token_ids)[0]) + ids = mx.array([[token, token + 1, 999, 999, 999, 999]]) + return ids, mx.zeros((1, 5, 1024)), mx.ones((1, 5)) + + +def test_dspark_rejected_drafts_never_enter_the_target_cache(): + rt = _SecondProposalMissRuntime() + + out = generate_mtpk( + rt, + [9, 10], + max_tokens=4, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + ) + + assert out.tokens == [11, 12, 13, 14] + assert rt.target_forward_widths == [1, 1, 1, 1, 1, 1] + assert rt.target_cache.trimmed == [] + assert out.stats.rejected_drafts >= 2 + + +def test_dspark_proposal_restore_precedes_the_primary_only_commit(): + """A rejected suffix restores proposal rings before committing target state.""" + rt = _SecondProposalMissRuntime() + + out = generate_mtpk( + rt, + [9, 10], + max_tokens=3, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + ) + + assert out.stats.accepted_drafts == 0 + assert rt.target_cache.trimmed == [] + assert len(rt.commit_ring_history) >= 1 + for ring in rt.commit_ring_history[0]: + np.testing.assert_array_equal(ring, rt.prefill_hidden) + assert float(np.max(ring)) != 999.0 + + +class _AcceptOneThenExactRuntime(_DSparkRuntime): + def draft_deepseek_v4_dspark(self, hidden, token_ids, cache, *, start_pos): + del hidden, cache, start_pos + token = int(np.asarray(token_ids)[0]) + if token == 10: + ids = mx.array([[token, token + 1, 999, 999, 999, 999]]) + else: + ids = mx.array([[token, token + 1, token + 2, token + 3, 999, 999]]) + return ids, mx.zeros((1, 5, 1024)), mx.ones((1, 5)) + + +def test_dspark_accept_one_resumes_at_the_target_correction_boundary(): + rt = _AcceptOneThenExactRuntime() + + out = generate_mtpk( + rt, + [9, 10], + max_tokens=4, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + ) + + assert out.tokens == [11, 12, 13, 14] + assert rt.target_forward_widths == [1, 1, 1, 1, 1, 1] + # The generic engine owns the next target position; DSpark RoPE/ring setup + # owns the carried hidden's position, exactly one row earlier. + assert rt.proposal_inputs == [(10, 1, 3), (11, 2, 3)] + assert out.stats.accepted_drafts == 2 + assert out.stats.rejected_drafts == 2 + + +class _AcceptTwoThenExactRuntime(_DSparkRuntime): + def draft_deepseek_v4_dspark(self, hidden, token_ids, cache, *, start_pos): + del hidden, cache, start_pos + token = int(np.asarray(token_ids)[0]) + if token == 10: + ids = mx.array([[token, token + 1, token + 2, 999, 999, 999]]) + else: + ids = mx.array([[token, token + 1, token + 2, token + 3, 999, 999]]) + return ids, mx.zeros((1, 5, 1024)), mx.ones((1, 5)) + + +def test_dspark_accept_two_resumes_at_the_target_correction_boundary(): + rt = _AcceptTwoThenExactRuntime() + + out = generate_mtpk( + rt, + [9, 10], + max_tokens=5, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + ) + + assert out.tokens == [11, 12, 13, 14, 15] + assert rt.target_forward_widths == [1, 1, 1, 1, 1, 1, 1] + assert rt.proposal_inputs == [(10, 1, 3), (12, 3, 3)] + assert out.stats.accepted_drafts == 3 + assert out.stats.rejected_drafts == 1 + + +class _FirstMissRuntime(_DSparkRuntime): + def draft_deepseek_v4_dspark(self, hidden, token_ids, cache, *, start_pos): + del hidden, token_ids, cache, start_pos + ids = mx.array([[0, 999, 999, 999, 999, 999]]) + return ids, mx.zeros((1, 5, 1024)), mx.ones((1, 5)) + + +def test_dspark_wrong_internal_primary_is_replaced_before_serial_target_rows(): + rt = _FirstMissRuntime() + out = generate_mtpk( + rt, + [9, 10], + max_tokens=3, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + ) + assert out.tokens == [11, 12, 13] + assert rt.forced_primary_ids == [11, 12] + assert rt.target_forward_widths == [1, 1, 1, 1, 1] + assert out.stats.verify_calls == 3 + + +def test_dspark_rejects_non_greedy_sampling_before_prefill(): + rt = _DSparkRuntime() + with pytest.raises(ValueError, match="greedy"): + generate_mtpk( + rt, + [9, 10], + max_tokens=4, + sampler=SamplerConfig(temperature=0.7), + speculative_depth=2, + stop_token_ids=set(), + ) + assert rt.target_cache.offset == 0 + + +@pytest.mark.parametrize("depth", [0, 1, 3, 4, 6]) +def test_dspark_rejects_unbenchmarked_widths_before_prefill(depth): + rt = _DSparkRuntime() + with pytest.raises(ValueError, match="width"): + generate_mtpk( + rt, + [9, 10], + max_tokens=4, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=depth, + stop_token_ids=set(), + ) + assert rt.target_cache.offset == 0 + + +@pytest.mark.parametrize( + "draft_sampler", + [ + SamplerConfig(temperature=0.7), + SamplerConfig(temperature=0.0, presence_penalty=1.0), + SamplerConfig(temperature=0.0, frequency_penalty=1.0), + ], +) +def test_dspark_rejects_unsupported_draft_sampling_before_prefill(draft_sampler): + rt = _DSparkRuntime() + with pytest.raises(ValueError, match="greedy"): + generate_mtpk( + rt, + [9, 10], + max_tokens=4, + sampler=SamplerConfig(temperature=0.0), + draft_sampler=draft_sampler, + speculative_depth=2, + stop_token_ids=set(), + ) + assert rt.target_forward_widths == [] + assert rt.prefill_hidden is None + + +def test_dspark_stop_primary_commits_one_target_row_without_drafting(): + rt = _DSparkRuntime() + callback = [] + out = generate_mtpk( + rt, + [9, 10], + max_tokens=4, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids={11}, + token_callback=lambda block: callback.extend(block), + ) + assert out.tokens == [11] + assert out.finish_reason == "stop" + assert callback == [] + assert rt.target_forward_widths == [1, 1, 1] + assert rt.proposal_widths == [] + assert rt.target_cache.offset == 3 + + +def test_dspark_accepted_stop_commits_only_the_terminal_prefix(): + rt = _DSparkRuntime() + callback = [] + out = generate_mtpk( + rt, + [9, 10], + max_tokens=4, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids={12}, + token_callback=lambda block: callback.extend(block), + ) + assert out.tokens == [11, 12] + assert out.finish_reason == "stop" + assert callback == [11] + assert rt.target_forward_widths == [1, 1, 1, 1] + assert rt.target_cache.trimmed == [] + assert out.stats.accepted_drafts == 1 + assert out.stats.drafted_tokens == 1 + + +class _RejectedStopRuntime(_DSparkRuntime): + def draft_deepseek_v4_dspark(self, hidden, token_ids, cache, *, start_pos): + del hidden, cache, start_pos + token = int(np.asarray(token_ids)[0]) + if token == 10: + ids = mx.array([[token, token + 1, 99, 999, 999, 999]]) + else: + ids = mx.array([[token, token + 1, token + 2, token + 3, 999, 999]]) + return ids, mx.zeros((1, 5, 1024)), mx.ones((1, 5)) + + +def test_dspark_rejected_stop_is_trimmed_and_never_emitted(): + rt = _RejectedStopRuntime() + out = generate_mtpk( + rt, + [9, 10], + max_tokens=4, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids={99}, + ) + assert out.tokens == [11, 12, 13, 14] + assert 99 not in out.tokens + assert rt.target_forward_widths == [1, 1, 1, 1, 1, 1] + assert rt.target_cache.trimmed == [] + assert out.stats.rejected_drafts == 1 + assert out.stats.accepted_drafts == 2 + + +@pytest.mark.parametrize( + ("max_tokens", "expected_tokens", "expected_widths", "proposal_widths"), + [ + (1, [11], [1, 1, 1], []), + (2, [11, 12], [1, 1, 1, 1], [2]), + ], +) +def test_dspark_generation_tail_uses_only_the_remaining_target_rows( + max_tokens, expected_tokens, expected_widths, proposal_widths +): + rt = _DSparkRuntime() + out = generate_mtpk( + rt, + [9, 10], + max_tokens=max_tokens, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + ) + assert out.tokens == expected_tokens + assert rt.target_forward_widths == expected_widths + assert rt.proposal_widths == proposal_widths + assert rt.target_cache.offset == 2 + max_tokens + + +def test_dspark_abort_before_target_prefill_does_no_model_work(): + rt = _DSparkRuntime() + out = generate_mtpk( + rt, + [9, 10], + abort_check=lambda: True, + max_tokens=4, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + ) + assert out.tokens == [] + assert rt.target_forward_widths == [] + assert rt.prefill_hidden is None + + +def test_dspark_abort_before_proposal_prefill_does_no_proposal_work(): + rt = _DSparkRuntime() + decisions = iter([False, False, True]) + out = generate_mtpk( + rt, + [9, 10], + abort_check=lambda: next(decisions), + max_tokens=4, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + ) + assert out.tokens == [] + assert rt.target_forward_widths == [1, 1] + assert rt.prefill_hidden is None + + +def test_dspark_prefill_callback_uses_standard_compute_and_wall_schema(): + rt = _DSparkRuntime() + callbacks = [] + generate_mtpk( + rt, + [9, 10], + max_tokens=1, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + prefill_callback=callbacks.append, + ) + assert [event["phase"] for event in callbacks] == ["started", "completed"] + assert set(callbacks[0]) == { + "phase", + "tokens_done", + "tokens_total", + "cached_tokens", + "new_prefill_tokens", + "elapsed_s", + "started_s", + } + assert set(callbacks[1]) == { + "phase", + "tokens_total", + "cached_tokens", + "new_prefill_tokens", + "elapsed_s", + "prompt_eval_time_s", + "prefill_tok_s", + "prefill_compute_tok_s", + "prefill_wall_tok_s", + "cache_hit", + } + assert callbacks[1]["prefill_compute_tok_s"] is not None + assert callbacks[1]["prefill_wall_tok_s"] is not None + + +@pytest.mark.parametrize("raise_on_call", [1, 2]) +def test_dspark_prefill_callback_failure_does_not_abort_generation(raise_on_call): + rt = _DSparkRuntime() + calls = 0 + + def failing_callback(_event): + nonlocal calls + calls += 1 + if calls == raise_on_call: + raise RuntimeError("dashboard callback failed") + + out = generate_mtpk( + rt, + [9, 10], + max_tokens=1, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + prefill_callback=failing_callback, + ) + assert out.tokens == [11] + + +def test_dspark_rejection_after_required_seed_restores_without_target_trim(): + rt = _SecondProposalMissRuntime() + out = generate_mtpk( + rt, + [10], + max_tokens=4, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + ) + assert out.tokens == [11, 12, 13, 14] + assert rt.target_forward_widths[:3] == [1, 1, 1] + assert rt.proposal_inputs[0] == (11, 1, 3) + assert rt.target_cache.trimmed == [] + assert out.stats.rejected_drafts >= 2 + + +def test_dspark_does_not_collect_per_cycle_timing_or_event_dicts(monkeypatch): + class _Clock: + def __init__(self): + self.calls = 0 + + def __call__(self): + self.calls += 1 + return float(self.calls) + + clock = _Clock() + monkeypatch.setattr(native_speculation.time, "perf_counter", clock) + + short = generate_mtpk( + _DSparkRuntime(), + [9, 10], + max_tokens=1, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + ) + short_calls = clock.calls + clock.calls = 0 + long = generate_mtpk( + _DSparkRuntime(), + [9, 10], + max_tokens=10, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + ) + assert clock.calls == short_calls + assert short.stats.events == [] + assert long.stats.events == [] + for stats in (short.stats, long.stats): + assert stats.verify_time_s == 0.0 + assert stats.verify_forward_time_s == 0.0 + assert stats.verify_eval_time_s == 0.0 + assert stats.verify_joint_eval_time_s == 0.0 + assert stats.draft_time_s == 0.0 + assert stats.target_forward_time_s == 0.0 + assert stats.accept_time_s == 0.0 + assert stats.rollback_time_s == 0.0 + assert stats.commit_time_s == 0.0 + + +def test_dspark_prefill_component_timing_accumulates_exact_chunk_boundaries( + monkeypatch, +): + class _Clock: + def __init__(self): + self.calls = 0 + + def __call__(self): + value = float(self.calls) + self.calls += 1 + return value + + clock = _Clock() + monkeypatch.setattr(native_speculation.time, "perf_counter", clock) + monkeypatch.setenv("MTPLX_SUSTAINED_PREFILL", "1") + monkeypatch.setenv("MTPLX_PREFILL_CHUNK_SIZE", "128") + out = generate_mtpk( + _DSparkRuntime(), + list(range(300)), + max_tokens=0, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + ) + assert out.stats.prompt_target_prefill_time_s == 4.0 + assert out.stats.prompt_mtp_history_time_s == 3.0 + assert out.stats.prompt_eval_time_s == 7.0 + assert out.stats.prompt_eval_time_s == ( + out.stats.prompt_target_prefill_time_s + out.stats.prompt_mtp_history_time_s + ) + assert out.stats.prompt_target_prefill_tok_s == 75.0 + + +def test_dspark_rollback_parameter_names_rejected_rows(): + parameters = inspect.signature(DeepseekV4DSparkBackend.rollback_target).parameters + assert "rejected_rows" in parameters + assert "verified_rows" not in parameters + + +@pytest.mark.parametrize("prompt_length", [9, 129, 256, 300, 385]) +def test_dspark_prefill_matches_ar_spans_and_streams_exact_dspark_chunks( + monkeypatch, prompt_length +): + monkeypatch.setenv("MTPLX_SUSTAINED_PREFILL", "1") + monkeypatch.setenv("MTPLX_PREFILL_CHUNK_SIZE", "128") + rt = _DSparkRuntime() + prompt = list(range(prompt_length)) + out = generate_mtpk( + rt, + prompt, + max_tokens=0, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + ) + assert out.tokens == [] + ar_spans = generation_module._iter_prefill_chunk_spans(prompt_length - 1) + assert rt.target_forward_widths == [ + *(end - start for start, end in ar_spans), + 1, + ] + target_ids = np.concatenate(rt.target_forward_inputs, axis=1) + np.testing.assert_array_equal(target_ids, np.array([prompt])) + assert rt.target_forward_options == [ + {"return_hidden": True, "emit_logits": False, "logits_keep": None} + for _ in ar_spans + ] + [ + {"return_hidden": True, "emit_logits": True, "logits_keep": 1}, + ] + assert [ + (start_pos, int(hidden.shape[1])) + for start_pos, hidden in rt.proposal_prefill_spans + ] == [ + (start, min(128, prompt_length - start)) + for start in range(0, prompt_length, 128) + ] + proposal_hidden = np.concatenate( + [hidden for _, hidden in rt.proposal_prefill_spans], axis=1 + ) + np.testing.assert_array_equal( + proposal_hidden, + np.repeat(np.asarray(prompt, dtype=np.float32)[None, :, None], 3, axis=2), + ) + assert all(entry.prefill_length == prompt_length for entry in rt.proposal_cache) + assert rt.target_cache.offset == prompt_length + + +def test_dspark_target_spans_follow_unchunked_ar_body(monkeypatch): + monkeypatch.delenv("MTPLX_SUSTAINED_PREFILL", raising=False) + rt = _DSparkRuntime() + generate_mtpk( + rt, + list(range(300)), + max_tokens=0, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + ) + assert generation_module._iter_prefill_chunk_spans(299) == [(0, 299)] + assert rt.target_forward_widths == [299, 1] + assert [ + (start_pos, int(hidden.shape[1])) + for start_pos, hidden in rt.proposal_prefill_spans + ] == [(0, 128), (128, 128), (256, 44)] + + +def test_dspark_nine_token_prefill_matches_ar_body_then_final_contract(monkeypatch): + monkeypatch.setenv("MTPLX_SUSTAINED_PREFILL", "1") + monkeypatch.setenv("MTPLX_PREFILL_CHUNK_SIZE", "128") + rt = _DSparkRuntime() + generate_mtpk( + rt, + list(range(9)), + max_tokens=0, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + ) + assert rt.target_forward_widths == [8, 1] + assert [rows[0].tolist() for rows in rt.target_forward_inputs] == [ + list(range(8)), + [8], + ] + assert rt.target_forward_options == [ + {"return_hidden": True, "emit_logits": False, "logits_keep": None}, + {"return_hidden": True, "emit_logits": True, "logits_keep": 1}, + ] + np.testing.assert_array_equal( + rt.prefill_hidden, + np.repeat(np.arange(9, dtype=np.float32)[None, :, None], 3, axis=2), + ) + assert rt.commits == [] + + +def test_dspark_long_prefill_preserves_decode_position_arithmetic(monkeypatch): + monkeypatch.setenv("MTPLX_SUSTAINED_PREFILL", "1") + monkeypatch.setenv("MTPLX_PREFILL_CHUNK_SIZE", "128") + rt = _DSparkRuntime() + out = generate_mtpk( + rt, + list(range(300)), + max_tokens=3, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + ) + assert out.tokens == [300, 301, 302] + assert rt.target_forward_widths == [128, 128, 43, 1, 1, 1, 1] + assert rt.proposal_inputs == [(299, 299, 3)] + assert rt.commits[-1][0] == 300 + assert rt.target_cache.offset == 303 + + +@pytest.mark.parametrize( + ("decisions", "expected_target_widths", "expected_proposal_spans"), + [ + ([True], [], []), + ([False, True], [128], []), + ([False, False, True], [128], [(0, 128)]), + ([False, False, False, True], [128, 128], [(0, 128)]), + ( + [False, False, False, False, True], + [128, 128], + [(0, 128), (128, 128)], + ), + ( + [False, False, False, False, False, True], + [128, 128, 43], + [(0, 128), (128, 128)], + ), + ( + [False, False, False, False, False, False, True], + [128, 128, 43, 1], + [(0, 128), (128, 128)], + ), + ], +) +def test_dspark_long_prefill_abort_checks_each_target_and_proposal_chunk( + monkeypatch, decisions, expected_target_widths, expected_proposal_spans +): + monkeypatch.setenv("MTPLX_SUSTAINED_PREFILL", "1") + monkeypatch.setenv("MTPLX_PREFILL_CHUNK_SIZE", "128") + rt = _DSparkRuntime() + decisions = iter(decisions) + out = generate_mtpk( + rt, + list(range(300)), + abort_check=lambda: next(decisions), + max_tokens=0, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + ) + assert out.tokens == [] + assert rt.target_forward_widths == expected_target_widths + assert [ + (start_pos, int(hidden.shape[1])) + for start_pos, hidden in rt.proposal_prefill_spans + ] == expected_proposal_spans + + +@pytest.mark.parametrize("prompt_length", [9, 129, 256, 300, 385]) +def test_dspark_chunked_prefill_ring_matches_one_shot_across_wraps(prompt_length): + values = mx.arange(prompt_length, dtype=mx.float32).reshape(1, prompt_length, 1) + one_shot = DeepseekV4DSparkCache(window_size=128, head_dim=1) + one_shot.prefill(values) + + chunked = DeepseekV4DSparkCache(window_size=128, head_dim=1) + chunked.prefill(values[:, :128]) + for start_pos in range(128, prompt_length, 128): + chunked.commit_main(start_pos, values[:, start_pos : start_pos + 128]) + + np.testing.assert_array_equal(np.asarray(chunked.ring), np.asarray(one_shot.ring)) + + +def test_dspark_verify_calls_counts_cycles_not_serial_target_forwards(): + rt = _DSparkRuntime() + out = generate_mtpk( + rt, + [9, 10], + max_tokens=3, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + ) + assert rt.target_forward_widths == [1, 1, 1, 1, 1] + assert out.stats.accepted_drafts == 2 + # One K2 proposal/verification cycle owns three serial target-M1 forwards. + assert out.stats.verify_calls == 1 diff --git a/tests/test_deepseek_v4_loader.py b/tests/test_deepseek_v4_loader.py index ff9c7101..1efeb496 100644 --- a/tests/test_deepseek_v4_loader.py +++ b/tests/test_deepseek_v4_loader.py @@ -11,6 +11,7 @@ The assembled 43-layer first-token logits gate runs in a GPU window (the model needs ~112 GiB wired); see scripts/deepseek_v4_logits_gate.py. """ + import glob import importlib.util import json @@ -23,6 +24,7 @@ import mlx.core as mx # noqa: E402 from mlx.utils import tree_flatten # noqa: E402 + @pytest.fixture(autouse=True) def _cpu_default_device(): # CPU-pinned by design, but the pin must stay test-scoped: a module-level @@ -36,6 +38,7 @@ def _cpu_default_device(): finally: mx.set_default_device(previous) + _HERE = os.path.dirname(os.path.abspath(__file__)) _MODEL = os.path.join(_HERE, "..", "mtplx", "models", "deepseek_v4.py") _spec = importlib.util.spec_from_file_location("dsv4_loader_undertest", _MODEL) @@ -46,11 +49,23 @@ def _cpu_default_device(): def _tiny_full_args(**over): base = dict( - vocab_size=64, hidden_size=16, num_hidden_layers=43, num_hash_layers=3, - num_attention_heads=2, head_dim=8, qk_rope_head_dim=4, - q_lora_rank=8, o_lora_rank=4, o_groups=2, - moe_intermediate_size=8, n_routed_experts=256, num_experts_per_tok=6, - index_n_heads=2, index_head_dim=8, index_topk=4, sliding_window=8, + vocab_size=64, + hidden_size=16, + num_hidden_layers=43, + num_hash_layers=3, + num_attention_heads=2, + head_dim=8, + qk_rope_head_dim=4, + q_lora_rank=8, + o_lora_rank=4, + o_groups=2, + moe_intermediate_size=8, + n_routed_experts=256, + num_experts_per_tok=6, + index_n_heads=2, + index_head_dim=8, + index_topk=4, + sliding_window=8, ) base.update(over) return D.ModelArgs(**base) @@ -62,22 +77,39 @@ def test_module_tree_matches_v4_spec(): keys = {k for k, _ in tree_flatten(model.parameters())} # top-level - for k in ("model.embed_tokens.weight", "model.norm.weight", - "model.hc_head.fn", "model.hc_head.base", "model.hc_head.scale", - "lm_head.weight"): + for k in ( + "model.embed_tokens.weight", + "model.norm.weight", + "model.hc_head.fn", + "model.hc_head.base", + "model.hc_head.scale", + "lm_head.weight", + ): assert k in keys, k cr = args.compress_ratios for i in range(args.num_hidden_layers): p = f"model.layers.{i}" # every layer: attention low-ranks, HC blocks, MoE gate + experts - for suf in ("attn.wq_a.weight", "attn.wq_b.weight", "attn.wkv.weight", - "attn.wo_a.weight", "attn.wo_b.weight", "attn.attn_sink", - "attn.q_norm.weight", "attn.kv_norm.weight", - "attn_hc.fn", "attn_hc.base", "attn_hc.scale", - "ffn_hc.fn", "ffn_hc.base", "ffn_hc.scale", - "ffn.gate.weight", "ffn.switch_mlp.gate_proj.weight", - "ffn.shared_experts.gate_proj.weight"): + for suf in ( + "attn.wq_a.weight", + "attn.wq_b.weight", + "attn.wkv.weight", + "attn.wo_a.weight", + "attn.wo_b.weight", + "attn.attn_sink", + "attn.q_norm.weight", + "attn.kv_norm.weight", + "attn_hc.fn", + "attn_hc.base", + "attn_hc.scale", + "ffn_hc.fn", + "ffn_hc.base", + "ffn_hc.scale", + "ffn.gate.weight", + "ffn.switch_mlp.gate_proj.weight", + "ffn.shared_experts.gate_proj.weight", + ): assert f"{p}.{suf}" in keys, f"{p}.{suf}" # hash layers carry tid2eid; score layers carry the noaux bias if i < args.num_hash_layers: @@ -93,32 +125,86 @@ def test_module_tree_matches_v4_spec(): assert has_index == (cr[i] == 4), (i, cr[i], has_index) +def test_sanitize_flattens_0731_grouped_wo_a_storage_without_reordering(): + """Collapse only the checkpoint's explicit o-LoRA group/rank row axes.""" + args = _tiny_full_args( + num_hidden_layers=1, + num_hash_layers=1, + compress_ratios=[0], + num_nextn_predict_layers=0, + ) + model = D.Model(args) + prefix = "model.layers.0.attn.wo_a" + grouped = { + f"{prefix}.weight": mx.arange(2 * 4 * 6).reshape(2, 4, 6), + f"{prefix}.scales": mx.arange(2 * 4 * 2).reshape(2, 4, 2), + f"{prefix}.biases": mx.arange(2 * 4 * 2).reshape(2, 4, 2) + 100, + } + + sanitized = model.sanitize(grouped) + + for suffix, value in grouped.items(): + assert sanitized[suffix].shape == (8, value.shape[-1]) + assert bool(mx.array_equal(sanitized[suffix], value.reshape(8, -1))) + + flat = mx.arange(8 * 6).reshape(8, 6) + assert model.sanitize({f"{prefix}.weight": flat})[f"{prefix}.weight"] is flat + + +def test_sanitize_rejects_malformed_0731_grouped_wo_a_storage(): + args = _tiny_full_args( + num_hidden_layers=1, + num_hash_layers=1, + compress_ratios=[0], + num_nextn_predict_layers=0, + ) + model = D.Model(args) + + with pytest.raises(ValueError, match="invalid grouped 0731 o-LoRA storage"): + model.sanitize({"model.layers.0.attn.wo_a.weight": mx.zeros((1, 8, 6))}) + + def _find_snapshot(): - hits = glob.glob(os.path.expanduser( - "~/.cache/huggingface/hub/models--mlx-community--DeepSeek-V4-Flash-4bit/snapshots/*/")) + hits = glob.glob( + os.path.expanduser( + "~/.cache/huggingface/hub/models--mlx-community--DeepSeek-V4-Flash-4bit/snapshots/*/" + ) + ) for h in hits: - if os.path.exists(os.path.join(h, "model.safetensors.index.json")) and \ - os.path.exists(os.path.join(h, "config.json")): + if os.path.exists( + os.path.join(h, "model.safetensors.index.json") + ) and os.path.exists(os.path.join(h, "config.json")): return h return None _SNAP = _find_snapshot() -_needs_ckpt = pytest.mark.skipif(_SNAP is None, reason="mlx-community 4bit checkpoint not in HF cache") +_needs_ckpt = pytest.mark.skipif( + _SNAP is None, reason="mlx-community 4bit checkpoint not in HF cache" +) def _args_from_config(cfg): return D.ModelArgs( - vocab_size=cfg["vocab_size"], hidden_size=cfg["hidden_size"], - num_hidden_layers=cfg["num_hidden_layers"], num_hash_layers=cfg["num_hash_layers"], - num_attention_heads=cfg["num_attention_heads"], head_dim=cfg["head_dim"], - qk_rope_head_dim=cfg["qk_rope_head_dim"], q_lora_rank=cfg["q_lora_rank"], - o_lora_rank=cfg["o_lora_rank"], o_groups=cfg["o_groups"], + vocab_size=cfg["vocab_size"], + hidden_size=cfg["hidden_size"], + num_hidden_layers=cfg["num_hidden_layers"], + num_hash_layers=cfg["num_hash_layers"], + num_attention_heads=cfg["num_attention_heads"], + head_dim=cfg["head_dim"], + qk_rope_head_dim=cfg["qk_rope_head_dim"], + q_lora_rank=cfg["q_lora_rank"], + o_lora_rank=cfg["o_lora_rank"], + o_groups=cfg["o_groups"], moe_intermediate_size=cfg["moe_intermediate_size"], - n_routed_experts=cfg["n_routed_experts"], num_experts_per_tok=cfg["num_experts_per_tok"], - index_n_heads=cfg["index_n_heads"], index_head_dim=cfg["index_head_dim"], - index_topk=cfg["index_topk"], compress_ratios=cfg["compress_ratios"], - compress_rope_theta=cfg["compress_rope_theta"], rms_norm_eps=cfg["rms_norm_eps"], + n_routed_experts=cfg["n_routed_experts"], + num_experts_per_tok=cfg["num_experts_per_tok"], + index_n_heads=cfg["index_n_heads"], + index_head_dim=cfg["index_head_dim"], + index_topk=cfg["index_topk"], + compress_ratios=cfg["compress_ratios"], + compress_rope_theta=cfg["compress_rope_theta"], + rms_norm_eps=cfg["rms_norm_eps"], rope_scaling=cfg.get("rope_scaling"), ) @@ -126,12 +212,19 @@ def _args_from_config(cfg): @_needs_ckpt def test_real_checkpoint_key_set_is_exact(): cfg = json.load(open(os.path.join(_SNAP, "config.json"))) - ckpt = set(json.load(open(os.path.join(_SNAP, "model.safetensors.index.json")))["weight_map"]) + ckpt = set( + json.load(open(os.path.join(_SNAP, "model.safetensors.index.json")))[ + "weight_map" + ] + ) # tiny per-unit dims, real structural counts -> identical key names args = _tiny_full_args( - num_hidden_layers=cfg["num_hidden_layers"], num_hash_layers=cfg["num_hash_layers"], - n_routed_experts=cfg["n_routed_experts"], o_groups=cfg["o_groups"], - compress_ratios=cfg["compress_ratios"], vocab_size=64, + num_hidden_layers=cfg["num_hidden_layers"], + num_hash_layers=cfg["num_hash_layers"], + n_routed_experts=cfg["n_routed_experts"], + o_groups=cfg["o_groups"], + compress_ratios=cfg["compress_ratios"], + vocab_size=64, ) model = D.Model(args) quantizable = {n for n, m in model.named_modules() if hasattr(m, "to_quantized")} @@ -154,8 +247,11 @@ def test_real_checkpoint_key_set_is_exact(): @_needs_ckpt def test_real_weight_components_match_oracle(): import numpy as np + cfg = json.load(open(os.path.join(_SNAP, "config.json"))) - wmap = json.load(open(os.path.join(_SNAP, "model.safetensors.index.json")))["weight_map"] + wmap = json.load(open(os.path.join(_SNAP, "model.safetensors.index.json")))[ + "weight_map" + ] qcfg = cfg["quantization"] args = _args_from_config(cfg) shards = {} @@ -177,7 +273,9 @@ def dense(stem): if f"{stem}.scales" in wmap: gs, bits, mode = qp(stem) b = raw(f"{stem}.biases") if f"{stem}.biases" in wmap else None - w = mx.dequantize(w, raw(f"{stem}.scales"), b, group_size=gs, bits=bits, mode=mode) + w = mx.dequantize( + w, raw(f"{stem}.scales"), b, group_size=gs, bits=bits, mode=mode + ) return w def npf(a): @@ -203,20 +301,29 @@ def npf(a): attn = D.DeepseekV4Attention(args, 3) attn.wo_a.weight = dense("model.layers.3.attn.wo_a") attn.wo_b.weight = dense("model.layers.3.attn.wo_b") - o = mx.array(rng.standard_normal((1, 2, args.num_attention_heads * args.head_dim)).astype(np.float32)) + o = mx.array( + rng.standard_normal((1, 2, args.num_attention_heads * args.head_dim)).astype( + np.float32 + ) + ) xo = attn._o_lora(o) mx.eval(xo) g, r = args.o_groups, args.o_lora_rank per = args.num_attention_heads * args.head_dim // g - o3 = np.einsum("bsgp,grp->bsgr", npf(o).reshape(1, 2, g, per), - npf(attn.wo_a.weight).reshape(g, r, per)).reshape(1, 2, g * r) + o3 = np.einsum( + "bsgp,grp->bsgr", + npf(o).reshape(1, 2, g, per), + npf(attn.wo_a.weight).reshape(g, r, per), + ).reshape(1, 2, g * r) ref = o3 @ npf(attn.wo_b.weight).T assert np.allclose(npf(xo), ref, rtol=2e-4, atol=2e-5) # gate score (real layer 3): valid top-k, weights sum to route_scale gate = D.MoEGate(args, 3) gate.weight = raw("model.layers.3.ffn.gate.weight") - gate.e_score_correction_bias = raw("model.layers.3.ffn.gate.e_score_correction_bias") + gate.e_score_correction_bias = raw( + "model.layers.3.ffn.gate.e_score_correction_bias" + ) idx, w = gate(mx.array(rng.standard_normal((5, H)).astype(np.float32)), None) mx.eval(idx, w) idxn = np.array(idx) diff --git a/tests/test_generation_deepseek_v4_dspark_integration.py b/tests/test_generation_deepseek_v4_dspark_integration.py new file mode 100644 index 00000000..f2fd4fab --- /dev/null +++ b/tests/test_generation_deepseek_v4_dspark_integration.py @@ -0,0 +1,108 @@ +"""Upstream request-option boundary for the DSpark native backend.""" + +from __future__ import annotations + +import pytest + +from mtplx.generation import generate_mtpk +from mtplx.sampling import SamplerConfig +from mtplx.thinking_guard import ThinkingGuardConfig +from tests.test_deepseek_v4_dspark_generation import _DSparkRuntime + + +@pytest.mark.parametrize( + ("option", "value", "message"), + [ + ("session_bank", object(), "session bank"), + ("capture_final_state", True, "final state"), + ("trace_label", "diagnostic", "decode trace"), + ("repetition_stop", True, "repetition stop"), + ("loop_guard", True, "loop guard"), + ( + "thinking_guard", + ThinkingGuardConfig(enabled=True), + "thinking guard", + ), + ], +) +def test_dspark_rejects_unsupported_upstream_features_before_prefill( + option, value, message +): + rt = _DSparkRuntime() + + with pytest.raises(ValueError, match=message): + generate_mtpk( + rt, + [10], + max_tokens=4, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + **{option: value}, + ) + + assert rt.target_cache.offset == 0 + assert rt.prefill_hidden is None + + +def test_dspark_rejects_construction_selected_decode_trace_before_prefill(tmp_path): + rt = _DSparkRuntime() + rt.block_speculative_decode_trace_requested = True + + with pytest.raises(ValueError, match="decode trace"): + generate_mtpk( + rt, + [10], + max_tokens=4, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + ) + + assert rt.target_cache.offset == 0 + assert not (tmp_path / "trace.jsonl").exists() + + +def test_dspark_rejects_generic_mtp_policy_options_before_prefill(): + rt = _DSparkRuntime() + + with pytest.raises(ValueError, match="mtp_cache_policy"): + generate_mtpk( + rt, + [10], + max_tokens=4, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + mtp_cache_policy="fresh", + ) + + assert rt.target_cache.offset == 0 + + +def test_dspark_allows_inert_session_metadata_when_cache_is_bypassed(monkeypatch): + rt = _DSparkRuntime() + sentinel = object() + + def fake_generate(*_args, **_kwargs): + return sentinel + + monkeypatch.setattr( + "mtplx.native_block_speculation.generate_native_block_speculative", + fake_generate, + ) + + result = generate_mtpk( + rt, + [10], + max_tokens=4, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=2, + stop_token_ids=set(), + session_id="stateless-request-label", + session_template_hash="unused-without-bank", + session_draft_head_identity="unused-without-bank", + session_policy_fingerprint="unused-without-bank", + ) + + assert result is sentinel diff --git a/tests/test_mtp_depth_grid.py b/tests/test_mtp_depth_grid.py new file mode 100644 index 00000000..ed0f6d15 --- /dev/null +++ b/tests/test_mtp_depth_grid.py @@ -0,0 +1,68 @@ +from __future__ import annotations + +from types import SimpleNamespace + +from mtplx.benchmarks.runners import mtp_depth_grid +from mtplx.benchmarks.schema import PromptCase +from mtplx.generation import GenerationOutput, GenerationStats + + +def _generation_output(*, events: list[dict], verify_calls: int) -> GenerationOutput: + return GenerationOutput( + tokens=[1, 2], + text="ok", + stats=GenerationStats( + mode="mtp", + generated_tokens=2, + elapsed_s=1.0, + tok_s=2.0, + accepted_drafts=6, + drafted_tokens=8, + accepted_by_depth=[3, 3], + drafted_by_depth=[4, 4], + verify_calls=verify_calls, + events=events, + ), + ) + + +def test_depth_grid_uses_verify_calls_only_when_events_are_empty( + monkeypatch, tmp_path +) -> None: + fake_runtime = SimpleNamespace(tokenizer=object()) + cases = [ + PromptCase(id="native", category="general", prompt="native"), + PromptCase(id="generic", category="general", prompt="generic"), + ] + outputs = iter( + [ + _generation_output(events=[], verify_calls=4), + _generation_output(events=[{}, {}], verify_calls=9), + ] + ) + + monkeypatch.setattr(mtp_depth_grid, "load", lambda *_args, **_kwargs: fake_runtime) + monkeypatch.setattr(mtp_depth_grid, "load_prompt_suite", lambda *_args: cases) + monkeypatch.setattr( + mtp_depth_grid, "encode_prompt_case", lambda *_args, **_kwargs: [1, 2] + ) + monkeypatch.setattr( + mtp_depth_grid, "generate_mtpk", lambda *_args, **_kwargs: next(outputs) + ) + + result = mtp_depth_grid.run_mtp_depth_policy_grid( + tmp_path / "model", + tmp_path / "suite.jsonl", + depth=2, + thresholds=[None], + min_depths=[0], + ) + + grid = result["grid"][0] + native, generic = grid["rows"] + assert native["cycles"] == 4 + assert native["accepted_drafts_per_cycle"] == 1.5 + assert generic["cycles"] == 2 + assert generic["accepted_drafts_per_cycle"] == 3.0 + assert grid["summary"]["cycles"] == 6 + assert grid["summary"]["accepted_drafts_per_cycle"] == 2.0 diff --git a/tests/test_mtp_depth_sweep.py b/tests/test_mtp_depth_sweep.py index d2be6fc2..84bf0981 100644 --- a/tests/test_mtp_depth_sweep.py +++ b/tests/test_mtp_depth_sweep.py @@ -3,6 +3,84 @@ from types import SimpleNamespace from mtplx.benchmarks.runners import mtp_depth_sweep +from mtplx.benchmarks.schema import PromptCase +from mtplx.generation import GenerationOutput, GenerationStats + + +def _generation_output(*, events: list[dict], verify_calls: int) -> GenerationOutput: + return GenerationOutput( + tokens=[1, 2], + text="ok", + finish_reason="length", + stats=GenerationStats( + mode="mtp", + generated_tokens=2, + elapsed_s=1.0, + tok_s=2.0, + decode_elapsed_s=1.0, + decode_tok_s=2.0, + end_to_end_tok_s=2.0, + accepted_drafts=6, + drafted_tokens=8, + accepted_by_depth=[3, 3], + drafted_by_depth=[4, 4], + accept_probability_sum_by_depth=[3.0, 3.0], + mean_accept_probability_by_depth=[0.75, 0.75], + verify_calls=verify_calls, + events=events, + ), + ) + + +def test_depth_sweep_uses_verify_calls_only_when_events_are_empty( + monkeypatch, tmp_path +) -> None: + fake_runtime = SimpleNamespace( + tokenizer=object(), + contract=SimpleNamespace( + base_hidden_variant="pre_norm", + hidden_variant="pre_norm", + concat_order="base_then_mtp", + mtp_quant_bits=None, + mtp_quant_group_size=64, + mtp_quant_mode="affine", + mtp_quant_policy=None, + ), + mtp_adapter_metadata=None, + mtp_adapter_merge_report=None, + ) + cases = [ + PromptCase(id="native", category="general", prompt="native"), + PromptCase(id="generic", category="general", prompt="generic"), + ] + outputs = iter( + [ + _generation_output(events=[], verify_calls=4), + _generation_output(events=[{}, {}], verify_calls=9), + ] + ) + + monkeypatch.setattr(mtp_depth_sweep, "load", lambda *_args, **_kwargs: fake_runtime) + monkeypatch.setattr(mtp_depth_sweep, "load_prompt_suite", lambda *_args: cases) + monkeypatch.setattr( + mtp_depth_sweep, "encode_prompt_case", lambda *_args, **_kwargs: [1, 2] + ) + monkeypatch.setattr( + mtp_depth_sweep, "generate_mtpk", lambda *_args, **_kwargs: next(outputs) + ) + monkeypatch.setattr( + mtp_depth_sweep, "validate_benchmark_output", lambda *_args, **_kwargs: [] + ) + + result = mtp_depth_sweep.run_mtp_depth_sweep( + tmp_path / "model", + tmp_path / "suite.jsonl", + depths=[2], + ) + + native, generic = result["depths"][0]["rows"] + assert native["mean_accepted_drafts_per_cycle"] == 1.5 + assert generic["mean_accepted_drafts_per_cycle"] == 3.0 def test_depth_sweep_uses_packaged_draft_lm_head_helper(monkeypatch, tmp_path) -> None: @@ -24,10 +102,14 @@ def test_depth_sweep_uses_packaged_draft_lm_head_helper(monkeypatch, tmp_path) - ) monkeypatch.setattr(mtp_depth_sweep, "load", lambda *_args, **_kwargs: fake_runtime) - monkeypatch.setattr(mtp_depth_sweep, "load_prompt_suite", lambda *_args, **_kwargs: []) + monkeypatch.setattr( + mtp_depth_sweep, "load_prompt_suite", lambda *_args, **_kwargs: [] + ) monkeypatch.setattr( "mtplx.draft_lm_head._install_draft_lm_head", - lambda runtime, **kwargs: calls.append((runtime, kwargs)) or {"installed": True}, + lambda runtime, **kwargs: ( + calls.append((runtime, kwargs)) or {"installed": True} + ), ) result = mtp_depth_sweep.run_mtp_depth_sweep( @@ -73,7 +155,9 @@ def fake_load(*_args, **kwargs): return fake_runtime monkeypatch.setattr(mtp_depth_sweep, "load", fake_load) - monkeypatch.setattr(mtp_depth_sweep, "load_prompt_suite", lambda *_args, **_kwargs: []) + monkeypatch.setattr( + mtp_depth_sweep, "load_prompt_suite", lambda *_args, **_kwargs: [] + ) result = mtp_depth_sweep.run_mtp_depth_sweep( tmp_path / "model", @@ -87,4 +171,7 @@ def fake_load(*_args, **kwargs): assert load_kwargs[0]["merge_mtp_adapter"] is True assert result["mtp_adapter_kind"] == "c4_mtp_lora_adapter" assert result["mtp_adapter_merged"] is True - assert result["mtp_adapter_merge_report"] == {"merged": 1, "targets": [{"target": "fc"}]} + assert result["mtp_adapter_merge_report"] == { + "merged": 1, + "targets": [{"target": "fc"}], + } diff --git a/tests/test_public_cli.py b/tests/test_public_cli.py index ea053e7a..1abc11f5 100644 --- a/tests/test_public_cli.py +++ b/tests/test_public_cli.py @@ -1059,7 +1059,9 @@ def to_dict(self) -> dict[str, object]: assert payload["profile"] == "sustained" -def _serve_dry_run_payload_for_model(monkeypatch, capsys, model_dir, extra_args=()): +def _serve_dry_run_payload_for_model( + monkeypatch, capsys, model_dir, extra_args=(), expected_code=0 +): """Drive `mtplx serve --dry-run --json` against a stubbed local model.""" runtime_contract = { @@ -1095,10 +1097,100 @@ def _serve_dry_run_payload_for_model(monkeypatch, capsys, model_dir, extra_args= args.dry_run = True args.json = True code = public.cmd_serve_public(args) - assert code == 0 + assert code == expected_code return json.loads(capsys.readouterr().out) +def test_serve_dspark_forwards_explicit_depth_two_without_legacy_defaults( + monkeypatch, tmp_path, capsys +): + model_dir = tmp_path / "DeepSeek-V4-0731" + model_dir.mkdir() + (model_dir / "config.json").write_text( + json.dumps({"model_type": "deepseek_v4", "dspark_block_size": 5}) + ) + + payload = _serve_dry_run_payload_for_model( + monkeypatch, + capsys, + model_dir, + extra_args=("--depth", "2"), + ) + + command = payload["server_command"] + assert "--depth 2" in command + assert "--verify-strategy" not in command + assert "--verify-core" not in command + assert "--deepseek-v4-0731-k2" not in command + + +def test_serve_forwards_explicit_0731_k2_construction_option( + monkeypatch, tmp_path, capsys +): + model_dir = tmp_path / "DeepSeek-V4-0731" + model_dir.mkdir() + (model_dir / "config.json").write_text( + json.dumps({"model_type": "deepseek_v4", "dspark_block_size": 5}) + ) + + payload = _serve_dry_run_payload_for_model( + monkeypatch, + capsys, + model_dir, + extra_args=("--deepseek-v4-0731-k2", "--depth", "2"), + ) + + command = payload["server_command"] + assert "--deepseek-v4-0731-k2" in command + assert "--depth 2" in command + + +@pytest.mark.parametrize( + "extra_args", + [ + ("--deepseek-v4-0731-k2",), + ("--deepseek-v4-0731-k2", "--depth", "3"), + ("--deepseek-v4-0731-k2", "--depth", "2", "--no-load-mtp"), + ("--deepseek-v4-0731-k2", "--depth", "2", "--generation-mode", "ar"), + ], +) +def test_serve_rejects_invalid_0731_k2_entrypoint_selection( + monkeypatch, tmp_path, capsys, extra_args +): + model_dir = tmp_path / "DeepSeek-V4-0731" + model_dir.mkdir() + (model_dir / "config.json").write_text( + json.dumps({"model_type": "deepseek_v4", "dspark_block_size": 5}) + ) + + payload = _serve_dry_run_payload_for_model( + monkeypatch, + capsys, + model_dir, + extra_args=extra_args, + expected_code=2, + ) + + assert "DeepSeek-V4-0731 K2" in payload["error"] + + +def test_serve_dspark_rejects_implicit_or_non_k2_depth(monkeypatch, tmp_path, capsys): + model_dir = tmp_path / "DeepSeek-V4-0731" + model_dir.mkdir() + (model_dir / "config.json").write_text( + json.dumps({"model_type": "deepseek_v4", "dspark_block_size": 5}) + ) + + payload = _serve_dry_run_payload_for_model( + monkeypatch, + capsys, + model_dir, + expected_code=2, + ) + + assert "explicit --depth 2" in payload["error"] + + def test_serve_defaults_quantized_27b_flagships_to_turbo(monkeypatch, tmp_path, capsys): """Bare `mtplx serve` on the quantized 27B flagships resolves turbo. @@ -1623,6 +1715,12 @@ def stop(self): fake_runtime = ModuleType("mtplx.runtime") fake_runtime.load = lambda *a, **kw: SimpleNamespace(tokenizer=object()) + fake_runtime.build_mtpk_request_kwargs = ( + lambda _rt, *, common, legacy_defaults, explicit_legacy=None: { + **common, + **legacy_defaults, + } + ) fake_schema = ModuleType("mtplx.benchmarks.schema") fake_schema.PromptCase = lambda **kw: SimpleNamespace(**kw) fake_schema.encode_prompt_case = lambda *a, **kw: [1, 2, 3] @@ -2574,6 +2672,123 @@ def fake_generate_mtpk(*_args, **kwargs): assert payload["stats"]["reasoning"] == "on" +def test_quickstart_dspark_uses_fixed_k2_without_legacy_request_defaults( + monkeypatch, tmp_path +): + captured: dict[str, object] = {} + + class TinyTokenizer: + model_max_length = 100 + + def apply_chat_template(self, *_args, **_kwargs): + return [1, 2, 3] + + fake_generation = ModuleType("mtplx.generation") + + def fake_generate_mtpk(*_args, **kwargs): + captured.update(kwargs) + return SimpleNamespace( + text="ok", + stats=SimpleNamespace( + generated_tokens=1, + speculative_depth=kwargs["speculative_depth"], + requested_speculative_depth=kwargs["speculative_depth"], + tok_s=1.0, + elapsed_s=1.0, + prompt_eval_time_s=0.0, + verify_time_s=0.0, + target_forward_time_s=1.0, + repair_time_s=0.0, + draft_time_s=0.0, + verify_calls=0, + accepted_by_depth=[], + drafted_by_depth=[], + correction_tokens=0, + bonus_tokens=0, + ), + ) + + fake_generation.generate_mtpk = fake_generate_mtpk + fake_generation.generate_ar = fake_generate_mtpk + fake_sampling = ModuleType("mtplx.sampling") + fake_sampling.SamplerConfig = lambda **kwargs: SimpleNamespace(**kwargs) + monkeypatch.setitem(sys.modules, "mtplx.generation", fake_generation) + monkeypatch.setitem(sys.modules, "mtplx.sampling", fake_sampling) + + args = SimpleNamespace( + system=None, + max_tokens=1, + temperature=0.0, + top_p=1.0, + top_k=0, + depth=2, + seed=0, + _cli_flags={"depth"}, + ) + rt = SimpleNamespace( + tokenizer=TinyTokenizer(), + model_path=tmp_path, + block_speculative_backend=SimpleNamespace(backend_id="deepseek_v4_dspark_0731"), + ) + + public._quickstart_generate( + rt=rt, + inspection={}, + profile=SimpleNamespace(to_dict=lambda: {"name": "stable"}), + args=args, + prompt="hello", + history=[], + turn_index=0, + ) + + assert captured["speculative_depth"] == 2 + for legacy_key in ( + "mtp_hidden_variant", + "mtp_history_policy", + "verify_strategy", + "verify_core", + ): + assert legacy_key not in captured + + args.depth = 3 + args._cli_flags = {"depth", "verify-strategy"} + captured.clear() + public._quickstart_generate( + rt=rt, + inspection={}, + profile=SimpleNamespace(to_dict=lambda: {"name": "stable"}), + args=args, + prompt="hello", + history=[], + turn_index=0, + ) + assert captured["speculative_depth"] == 3 + assert captured["verify_strategy"] == "capture_commit" + + args.depth = 2 + args._cli_flags = {"depth"} + args.mtp_hidden_variant = "contract" + args.mtp_history_policy = "cycle" + args.verify_strategy = "sequential" + args.verify_core = "stock" + args.draft_core = "nax" + captured.clear() + public._quickstart_generate( + rt=rt, + inspection={}, + profile=SimpleNamespace(to_dict=lambda: {"name": "stable"}), + args=args, + prompt="hello", + history=[], + turn_index=0, + ) + assert captured["mtp_hidden_variant"] == "contract" + assert captured["mtp_history_policy"] == "cycle" + assert captured["verify_strategy"] == "sequential" + assert captured["verify_core"] == "stock" + assert captured["draft_core"] == "nax" + + def test_quickstart_generation_no_mtp_uses_ar(monkeypatch, tmp_path): captured: dict[str, object] = {} diff --git a/tests/test_runtime_deepseek_v4_dspark.py b/tests/test_runtime_deepseek_v4_dspark.py new file mode 100644 index 00000000..cd8bbcbe --- /dev/null +++ b/tests/test_runtime_deepseek_v4_dspark.py @@ -0,0 +1,605 @@ +"""Runtime construction coverage for the pinned DeepSeek-V4-0731 backend.""" + +from __future__ import annotations + +from contextlib import contextmanager +from types import SimpleNamespace + +import pytest + +from mtplx import runtime +from mtplx.deepseek_v4_dspark_generation import DeepseekV4DSparkBackend +from mtplx.models import deepseek_v4 as D + + +_DSPARK_CONFIG = { + "model_type": "deepseek_v4", + "dspark_block_size": 5, + "dspark_noise_token_id": 128799, + "dspark_target_layer_ids": [40, 41, 42], + "dspark_markov_rank": 256, +} + + +class _DSpark: + def __init__(self) -> None: + self.stages = tuple(SimpleNamespace(attn=_RouteAttention()) for _ in range(3)) + + def make_cache(self): + return [] + + def prefill(self, *_args): + return None + + def forward(self, *_args, **_kwargs): + return None + + def commit_main(self, *_args, **_kwargs): + return None + + +class _Model: + def __init__(self, *, with_dspark: bool = True) -> None: + self._dspark = _DSpark() if with_dspark else None + self.mtp = [] if self._dspark is None else self._dspark.stages + self.model = SimpleNamespace(embed_tokens=lambda value: value) + self.lm_head = lambda value: value + self.layers = [SimpleNamespace(attn=_RouteAttention()) for _ in range(43)] + + @property + def mtp_blocks(self): + return list(self.mtp) + + def __call__(self, value, **_kwargs): + return value + + +class _RouteAttention: + def __init__(self) -> None: + self.installed_modes = [] + self.o_lora_mode = "cached" + self._o_lora_impl = object() + + def install_o_lora_route(self, mode): + self.installed_modes.append(str(mode)) + self.o_lora_mode = str(mode) + self._o_lora_impl = object() + return {"mode": str(mode), "direct": str(mode) == "gather_qmm"} + + +class _AdversarialRestoreAttention: + def __init__( + self, + label, + events, + *, + fail_mode_restore=False, + fail_implementation_restore=False, + ) -> None: + self.label = label + self.events = events + self.fail_mode_restore = fail_mode_restore + self.fail_implementation_restore = fail_implementation_restore + self._mode = "cached" + self.original_implementation = object() + self._implementation = self.original_implementation + + @property + def o_lora_mode(self): + return self._mode + + @o_lora_mode.setter + def o_lora_mode(self, value): + restoring = value == "cached" + self.events.append((self.label, "mode", "restore" if restoring else "install")) + if restoring and self.fail_mode_restore: + raise RuntimeError(f"{self.label} mode restore failed") + self._mode = value + + @property + def _o_lora_impl(self): + return self._implementation + + @_o_lora_impl.setter + def _o_lora_impl(self, value): + restoring = value is self.original_implementation + self.events.append( + (self.label, "implementation", "restore" if restoring else "install") + ) + if restoring and self.fail_implementation_restore: + raise RuntimeError(f"{self.label} implementation restore failed") + self._implementation = value + + +def _patch_load_dependencies( + monkeypatch, + tmp_path, + model, + o_lora_calls, + *, + mtp=True, + real_o_lora_installer=False, + deepseek_v4_0731_k2=False, +): + monkeypatch.setattr(runtime, "load_config", lambda _path: dict(_DSPARK_CONFIG)) + monkeypatch.setattr(runtime, "_load_base_model", lambda *_args: (model, object())) + monkeypatch.setattr(runtime, "_load_runtime_metadata", lambda _path: {}) + monkeypatch.setattr(runtime, "mtp_weights_present_on_disk", lambda *_args: True) + monkeypatch.setattr( + runtime, + "validate_mtp_support", + lambda _model: (_ for _ in ()).throw( + AssertionError("DSpark must not use legacy MTP validation") + ), + ) + + monkeypatch.setattr(D, "configure_deepseek_v4_moe_tail", lambda *_args: None) + monkeypatch.setattr(D, "_o_lora_mode_from_env", lambda: "gather_qmm") + if not real_o_lora_installer: + monkeypatch.setattr( + D, + "install_deepseek_v4_o_lora_routes", + lambda _model, **kwargs: o_lora_calls.append(kwargs) or dict(kwargs), + ) + + import mtplx.a3b_compiled_target_prefix as target_prefix + import mtplx.a3b_whole_moe as whole_moe + import mtplx.attention_split as attention_split + import mtplx.gdn_capture as gdn_capture + import mtplx.kernel_selfcheck as kernel_selfcheck + import mtplx.native_mlp as native_mlp + import mtplx.nax_verify as nax_verify + import mtplx.qwen_row_owned_router as row_owned + + monkeypatch.setattr( + attention_split, "configure_split_full_attention", lambda *_: None + ) + monkeypatch.setattr(native_mlp, "configure_native_mlp", lambda *_: None) + monkeypatch.setattr(nax_verify, "nax_env_enabled", lambda: False) + monkeypatch.setattr( + whole_moe, "prepare_a3b_whole_moe", lambda *_args, **_kwargs: None + ) + monkeypatch.setattr( + row_owned, "prepare_qwen_row_owned_routers", lambda *_args, **_kwargs: None + ) + monkeypatch.setattr( + gdn_capture, "prepare_a3b_gdn_postconv", lambda *_args, **_kwargs: None + ) + monkeypatch.setattr(kernel_selfcheck, "maybe_run_model_selfcheck", lambda *_: None) + monkeypatch.setattr( + target_prefix, + "prepare_a3b_compiled_target_prefix", + lambda *_args, **_kwargs: None, + ) + return runtime.load( + tmp_path, + mtp=mtp, + deepseek_v4_0731_k2=deepseek_v4_0731_k2, + ) + + +def test_runtime_load_publishes_the_construction_bound_dspark_backend( + monkeypatch, tmp_path +): + model = _Model() + o_lora_calls = [] + monkeypatch.setenv("MTPLX_DECODE_TRACE_JSONL", str(tmp_path / "trace.jsonl")) + + loaded = _patch_load_dependencies(monkeypatch, tmp_path, model, o_lora_calls) + + assert loaded.mtp_enabled is True + assert loaded.deepseek_v4_dspark_enabled is True + assert isinstance(loaded.block_speculative_backend, DeepseekV4DSparkBackend) + assert loaded.block_speculative_backend.dspark is model._dspark + assert loaded.block_speculative_decode_trace_requested is True + assert o_lora_calls == [{"mode": "gather_qmm", "canonical_mixed_route": False}] + + +def test_runtime_dspark_gather_uses_the_real_per_module_o_lora_installer( + monkeypatch, tmp_path +): + model = _Model() + + loaded = _patch_load_dependencies( + monkeypatch, + tmp_path, + model, + [], + real_o_lora_installer=True, + ) + + assert loaded.deepseek_v4_o_lora_report["module_count"] == 46 + assert loaded.deepseek_v4_o_lora_report["trunk_module_count"] == 43 + assert loaded.deepseek_v4_o_lora_report["mtp_module_count"] == 3 + assert all(layer.attn.installed_modes == ["gather_qmm"] for layer in model.layers) + assert all( + stage.attn.installed_modes == ["gather_qmm"] for stage in model._dspark.stages + ) + + +def test_runtime_load_fails_before_publication_when_dspark_owner_is_missing( + monkeypatch, tmp_path +): + model = _Model(with_dspark=False) + o_lora_calls = [] + + with pytest.raises(ValueError, match="DSpark backend cannot bind"): + _patch_load_dependencies(monkeypatch, tmp_path, model, o_lora_calls) + + assert o_lora_calls == [] + + +def test_runtime_load_does_not_publish_dspark_when_mtp_is_disabled( + monkeypatch, tmp_path +): + model = _Model() + o_lora_calls = [] + + loaded = _patch_load_dependencies( + monkeypatch, tmp_path, model, o_lora_calls, mtp=False + ) + + assert loaded.mtp_enabled is False + assert loaded.deepseek_v4_dspark_enabled is False + assert loaded.block_speculative_backend is None + assert o_lora_calls == [{"mode": "cached", "canonical_mixed_route": False}] + + +def test_k2_option_rejects_nonexact_artifact_before_model_construction( + monkeypatch, tmp_path +): + constructed = [] + monkeypatch.setattr( + runtime, + "load_config", + lambda _path: {"model_type": "deepseek_v4", **_DSPARK_CONFIG}, + ) + monkeypatch.setattr( + runtime, + "_load_base_model", + lambda *_args: constructed.append(True) or (_Model(), object()), + ) + + with pytest.raises(ValueError, match="full DSpark contract failed"): + runtime.load(tmp_path, deepseek_v4_0731_k2=True) + + assert constructed == [] + + +def test_k2_option_requires_mtp_before_model_construction(monkeypatch, tmp_path): + constructed = [] + monkeypatch.setattr( + runtime, + "_load_base_model", + lambda *_args: constructed.append(True) or (_Model(), object()), + ) + + with pytest.raises(ValueError, match="requires mtp=True"): + runtime.load(tmp_path, mtp=False, deepseek_v4_0731_k2=True) + + assert constructed == [] + + +def test_k2_sinkhorn_selector_is_scoped_to_model_construction(monkeypatch): + events = [] + monkeypatch.setattr(D, "_SINKHORN_KERNEL", False) + monkeypatch.setattr(D.mx.metal, "is_available", lambda: True) + monkeypatch.setattr(D.mx, "default_device", lambda: D.mx.gpu) + monkeypatch.setattr( + D, + "_sinkhorn_metal_kernel", + lambda hc, iters, eps: events.append((hc, iters, eps)), + ) + + outside, _ = D._install_sinkhorn_normaliser(4, 20, 1e-6) + with D.deepseek_v4_0731_k2_construction(): + inside, _ = D._install_sinkhorn_normaliser(4, 20, 1e-6) + restored, _ = D._install_sinkhorn_normaliser(4, 20, 1e-6) + + assert (outside, inside, restored) == (False, True, False) + assert events == [(4, 20, 1e-6)] + + +class _Prepared: + def __init__(self, label, events, *, fail_restore=False): + self.label = label + self.events = events + self.fail_restore = fail_restore + self.receipt = {"candidate": label} + + def publish(self): + self.events.append(f"{self.label}.publish") + + def restore(self): + self.events.append(f"{self.label}.restore") + if self.fail_restore: + raise RuntimeError(f"{self.label} restore failed") + + +def _patch_k2_preparers(monkeypatch, tmp_path, events, *, fail_restore=None): + import mtplx.deepseek_v4_0731_dspark_ffn as dspark_ffn + import mtplx.deepseek_v4_0731_full_install as full_install + import mtplx.deepseek_v4_0731_m3_wob as wob + import mtplx.deepseek_v4_0731_m3_wqb_qnorm_rope as wqb + + def prepare_wqb_qhead_m3(_layers, *, exact_selfcheck): + events.append(("wqb.prepare", callable(exact_selfcheck))) + + def prepare_wob_m3(_layers, *, exact_selfcheck): + events.append(("wob.prepare", callable(exact_selfcheck))) + + monkeypatch.setattr(wqb, "prepare_wqb_qhead_m3", prepare_wqb_qhead_m3) + monkeypatch.setattr(wob, "prepare_wob_m3", prepare_wob_m3) + + def prepare_target(model, config, path, **kwargs): + events.append( + ( + "target.prepare", + kwargs["prepare_wqb_qhead"].__name__, + kwargs["prepare_wob"].__name__, + ) + ) + kwargs["prepare_wqb_qhead"]( + (), + exact_selfcheck=lambda *_args: True, + ) + kwargs["prepare_wob"]( + (), + exact_selfcheck=lambda *_args: True, + ) + return _Prepared( + "target", + events, + fail_restore=fail_restore == "target", + ) + + monkeypatch.setattr( + full_install, + "validate_full_0731_dspark_artifact", + lambda path, config: events.append("artifact.validate") or object(), + ) + monkeypatch.setattr( + full_install, + "prepare_full_0731_dspark_compiled_tail_q2_pair", + prepare_target, + ) + monkeypatch.setattr( + dspark_ffn, + "prepare_dspark_q3_packed_gate_up_m5", + lambda model: ( + events.append("ffn.prepare") + or _Prepared("ffn", events, fail_restore=fail_restore == "ffn") + ), + ) + + @contextmanager + def selected_sinkhorn(): + events.append("sinkhorn.enter") + try: + yield + finally: + events.append("sinkhorn.exit") + + monkeypatch.setattr(D, "deepseek_v4_0731_k2_construction", selected_sinkhorn) + + +def test_k2_option_publishes_one_exact_construction_transaction(monkeypatch, tmp_path): + events = [] + model = _Model() + _patch_k2_preparers(monkeypatch, tmp_path, events) + monkeypatch.setattr( + D, + "install_deepseek_v4_o_lora_routes", + lambda _model, **kwargs: events.append(("o_lora", kwargs)) or dict(kwargs), + ) + monkeypatch.setattr(D, "_o_lora_mode_from_env", lambda: "cached") + + loaded = _patch_load_dependencies( + monkeypatch, + tmp_path, + model, + [], + real_o_lora_installer=True, + deepseek_v4_0731_k2=True, + ) + + assert events == [ + "artifact.validate", + "sinkhorn.enter", + "sinkhorn.exit", + ("target.prepare", "prepare_wqb_qhead_m3", "prepare_wob_m3"), + ("wqb.prepare", True), + ("wob.prepare", True), + "ffn.prepare", + ( + "o_lora", + {"mode": "gather_qmm", "canonical_mixed_route": False}, + ), + "target.publish", + "ffn.publish", + ] + assert loaded.deepseek_v4_0731_k2_receipt == { + "target": {"candidate": "target"}, + "dspark_ffn": {"candidate": "ffn"}, + } + assert loaded.block_speculative_backend.dspark is model._dspark + + +def test_k2_option_restores_both_staged_stacks_when_backend_binding_fails( + monkeypatch, tmp_path +): + events = [] + model = _Model() + _patch_k2_preparers(monkeypatch, tmp_path, events) + monkeypatch.setattr( + DeepseekV4DSparkBackend, + "bind", + lambda _model: (_ for _ in ()).throw(RuntimeError("bind failed")), + ) + + with pytest.raises(RuntimeError, match="bind failed"): + _patch_load_dependencies( + monkeypatch, + tmp_path, + model, + [], + deepseek_v4_0731_k2=True, + ) + + assert events[-4:] == [ + "target.publish", + "ffn.publish", + "ffn.restore", + "target.restore", + ] + + +def test_k2_option_restores_every_o_lora_owner_when_preparation_fails( + monkeypatch, tmp_path +): + events = [] + model = _Model() + _patch_k2_preparers(monkeypatch, tmp_path, events) + attentions = tuple(layer.attn for layer in model.layers) + tuple( + stage.attn for stage in model._dspark.stages + ) + originals = tuple( + (attention.o_lora_mode, attention._o_lora_impl) for attention in attentions + ) + + def install_gather(_model, **kwargs): + assert kwargs == {"mode": "gather_qmm", "canonical_mixed_route": False} + for attention in attentions: + attention.o_lora_mode = "gather_qmm" + attention._o_lora_impl = object() + return {"mode": "gather_qmm"} + + monkeypatch.setattr(D, "install_deepseek_v4_o_lora_routes", install_gather) + import mtplx.deepseek_v4_0731_full_install as full_install + + monkeypatch.setattr( + full_install, + "prepare_full_0731_dspark_compiled_tail_q2_pair", + lambda *_args, **_kwargs: (_ for _ in ()).throw( + RuntimeError("target prepare failed") + ), + ) + + with pytest.raises(RuntimeError, match="target prepare failed"): + _patch_load_dependencies( + monkeypatch, + tmp_path, + model, + [], + real_o_lora_installer=True, + deepseek_v4_0731_k2=True, + ) + + assert ( + tuple( + (attention.o_lora_mode, attention._o_lora_impl) for attention in attentions + ) + == originals + ) + + +def test_k2_rollback_attempts_target_and_o_lora_after_ffn_restore_failure( + monkeypatch, tmp_path +): + events = [] + model = _Model() + attentions = tuple(layer.attn for layer in model.layers) + tuple( + stage.attn for stage in model._dspark.stages + ) + originals = tuple( + (attention.o_lora_mode, attention._o_lora_impl) for attention in attentions + ) + _patch_k2_preparers(monkeypatch, tmp_path, events, fail_restore="ffn") + monkeypatch.setattr( + DeepseekV4DSparkBackend, + "bind", + lambda _model: (_ for _ in ()).throw(RuntimeError("bind failed")), + ) + + with pytest.raises(ExceptionGroup) as caught: + _patch_load_dependencies( + monkeypatch, + tmp_path, + model, + [], + deepseek_v4_0731_k2=True, + ) + + assert events[-4:] == [ + "target.publish", + "ffn.publish", + "ffn.restore", + "target.restore", + ] + assert [str(error) for error in caught.value.exceptions] == [ + "bind failed", + "ffn restore failed", + ] + assert ( + tuple( + (attention.o_lora_mode, attention._o_lora_impl) for attention in attentions + ) + == originals + ) + + +def test_k2_rollback_attempts_both_properties_and_all_o_lora_owners( + monkeypatch, tmp_path +): + events = [] + model = _Model() + first = _AdversarialRestoreAttention( + "first", + events, + fail_mode_restore=True, + fail_implementation_restore=True, + ) + last = _AdversarialRestoreAttention("last", events) + model.layers[0].attn = first + model._dspark.stages[-1].attn = last + _patch_k2_preparers(monkeypatch, tmp_path, events) + + attentions = tuple(layer.attn for layer in model.layers) + tuple( + stage.attn for stage in model._dspark.stages + ) + + def install_gather(_model, **kwargs): + assert kwargs == {"mode": "gather_qmm", "canonical_mixed_route": False} + for attention in attentions: + attention.o_lora_mode = "gather_qmm" + attention._o_lora_impl = object() + return {"mode": "gather_qmm"} + + monkeypatch.setattr(D, "install_deepseek_v4_o_lora_routes", install_gather) + monkeypatch.setattr( + DeepseekV4DSparkBackend, + "bind", + lambda _model: (_ for _ in ()).throw(RuntimeError("bind failed")), + ) + + with pytest.raises(ExceptionGroup) as caught: + _patch_load_dependencies( + monkeypatch, + tmp_path, + model, + [], + real_o_lora_installer=True, + deepseek_v4_0731_k2=True, + ) + + assert [str(error) for error in caught.value.exceptions] == [ + "bind failed", + "first mode restore failed", + "first implementation restore failed", + ] + assert ("first", "mode", "restore") in events + assert ("first", "implementation", "restore") in events + assert ("last", "mode", "restore") in events + assert ("last", "implementation", "restore") in events + assert last.o_lora_mode == "cached" + assert last._o_lora_impl is last.original_implementation diff --git a/tests/test_server_openai.py b/tests/test_server_openai.py index cf31301a..3c89d6f7 100644 --- a/tests/test_server_openai.py +++ b/tests/test_server_openai.py @@ -675,6 +675,45 @@ def test_serial_mtp_health_never_labels_queued_requests_as_ar(): assert openai._use_live_ar_batch(state, effective_mode="mtp") == (False, None) +def test_server_parser_accepts_explicit_0731_k2_construction_option(): + default = parse_args(["--warmup-tokens", "0"]) + selected = parse_args( + [ + "--warmup-tokens", + "0", + "--deepseek-v4-0731-k2", + "--depth", + "2", + ] + ) + + assert default.deepseek_v4_0731_k2 is False + assert selected.deepseek_v4_0731_k2 is True + assert "deepseek-v4-0731-k2" in selected._cli_flags + assert "depth" in selected._cli_flags + + +@pytest.mark.parametrize( + "argv", + [ + ["--deepseek-v4-0731-k2"], + ["--deepseek-v4-0731-k2", "--depth", "3"], + ["--deepseek-v4-0731-k2", "--depth", "2", "--no-load-mtp"], + ["--deepseek-v4-0731-k2", "--depth", "2", "--generation-mode", "ar"], + ], +) +def test_server_rejects_invalid_0731_k2_selection_before_load(monkeypatch, argv): + monkeypatch.setattr( + openai, + "load", + lambda *_args, **_kwargs: pytest.fail("invalid K2 selection reached load"), + ) + args = parse_args(["--warmup-tokens", "0", *argv]) + + with pytest.raises(ValueError, match="DeepSeek-V4-0731 K2"): + openai.ServerState(args) + + def test_server_parser_resolves_api_key_file_before_env(monkeypatch, tmp_path): api_key_file = tmp_path / "api-key" api_key_file.write_text("file-secret\n", encoding="utf-8") @@ -1614,12 +1653,14 @@ def test_vision_splice_kwargs_always_match_callee_signatures(): import ast import mtplx.generation + import mtplx.native_block_speculation import mtplx.runtime import mtplx.vision.splice sources = [ Path(openai.__file__), Path(mtplx.generation.__file__), + Path(mtplx.native_block_speculation.__file__), Path(mtplx.runtime.__file__), Path(mtplx.vision.splice.__file__), ] @@ -4749,6 +4790,167 @@ def fake_generate_mtpk(*_args, **kwargs): assert captured["commit_prompt_state_keep_live_ref"] is False +def test_run_generation_dspark_omits_implicit_legacy_request_state(monkeypatch): + state = _fake_streaming_session_state() + state.draft_sampler = None + state.requests_completed = 0 + state.args._cli_flags = set() + state.runtime.block_speculative_backend = SimpleNamespace( + backend_id="deepseek_v4_dspark_0731" + ) + state.runtime.block_speculative_decode_trace_requested = False + captured: dict[str, object] = {} + + def fake_generate_mtpk(*_args, **kwargs): + captured.update(kwargs) + return SimpleNamespace( + tokens=[], + text="", + stats=SimpleNamespace( + to_dict=lambda: { + "prompt_eval_time_s": 0.0, + "generated_tokens": 0, + "elapsed_s": 0.0, + "tok_s": 0.0, + } + ), + final_state=None, + ) + + monkeypatch.setattr(openai, "generate_mtpk", fake_generate_mtpk) + openai._run_generation( + state, + [1, 2, 3], + max_tokens=1, + temperature=0.0, + top_p=1.0, + top_k=0, + seed=0, + generation_mode="mtp", + depth=2, + resolved_mtp_depth=2, + session_id="implicit-session", + session_bank=None, + request_observability={"request_mtp_depth_explicit": True}, + ) + + assert captured["speculative_depth"] == 2 + for legacy_key in ( + "mtp_hidden_variant", + "mtp_history_policy", + "verify_strategy", + "verify_core", + "trace_label", + "trace_metadata", + ): + assert legacy_key not in captured + assert captured["session_bank"] is None + assert captured["capture_final_state"] is False + + +def test_run_generation_dspark_preserves_explicit_unsupported_policies(monkeypatch): + state = _fake_streaming_session_state() + state.draft_sampler = None + state.requests_completed = 0 + state.args._cli_flags = {"depth", "verify-strategy", "ssd-session-cache"} + state.runtime.block_speculative_backend = SimpleNamespace( + backend_id="deepseek_v4_dspark_0731" + ) + state.runtime.block_speculative_decode_trace_requested = True + captured: dict[str, object] = {} + + def fake_generate_mtpk(*_args, **kwargs): + captured.update(kwargs) + return SimpleNamespace( + tokens=[], + text="", + stats=SimpleNamespace( + to_dict=lambda: { + "prompt_eval_time_s": 0.0, + "generated_tokens": 0, + "elapsed_s": 0.0, + "tok_s": 0.0, + } + ), + final_state=None, + ) + + monkeypatch.setattr(openai, "generate_mtpk", fake_generate_mtpk) + openai._run_generation( + state, + [1, 2, 3], + max_tokens=1, + temperature=0.0, + top_p=1.0, + top_k=0, + seed=0, + generation_mode="mtp", + depth=3, + session_id="explicit-session", + session_bank=state.sessions.bank, + request_observability={"request_mtp_depth_explicit": True}, + ) + + assert captured["speculative_depth"] == 3 + assert captured["verify_strategy"] == "capture_commit" + assert captured["session_bank"] is state.sessions.bank + assert captured["capture_final_state"] is True + assert "trace_label" not in captured + assert "trace_metadata" not in captured + assert state.runtime.block_speculative_decode_trace_requested is True + + +def test_run_generation_dspark_preserves_nondefault_policies_without_cli_flags( + monkeypatch, +): + state = _fake_streaming_session_state() + state.draft_sampler = None + state.requests_completed = 0 + state.args._cli_flags = set() + state.args.verify_strategy = "sequential" + state.args.verify_core = "stock" + state.args.draft_core = "nax" + state.runtime.block_speculative_backend = SimpleNamespace( + backend_id="deepseek_v4_dspark_0731" + ) + captured: dict[str, object] = {} + + def fake_generate_mtpk(*_args, **kwargs): + captured.update(kwargs) + return SimpleNamespace( + tokens=[], + text="", + stats=SimpleNamespace( + to_dict=lambda: { + "prompt_eval_time_s": 0.0, + "generated_tokens": 0, + "elapsed_s": 0.0, + "tok_s": 0.0, + } + ), + final_state=None, + ) + + monkeypatch.setattr(openai, "generate_mtpk", fake_generate_mtpk) + openai._run_generation( + state, + [1, 2, 3], + max_tokens=1, + temperature=0.0, + top_p=1.0, + top_k=0, + seed=0, + generation_mode="mtp", + depth=2, + resolved_mtp_depth=2, + session_bank=None, + ) + + assert captured["verify_strategy"] == "sequential" + assert captured["verify_core"] == "stock" + assert captured["draft_core"] == "nax" + + def test_run_generation_depth1_clamps_expected_value_policy(monkeypatch): state = _fake_streaming_session_state() state.draft_sampler = None @@ -11549,6 +11751,41 @@ def stop_after_load(model, mtp, contract, **kwargs): assert captured["kwargs"]["merge_mtp_adapter"] is True +def test_server_state_passes_0731_k2_option_to_runtime_load(monkeypatch): + captured = {} + monkeypatch.setattr(openai, "apply_profile_env", lambda _profile, **_kwargs: None) + monkeypatch.setattr(openai, "profile_env_status", lambda _profile, **_kwargs: {}) + monkeypatch.setattr(openai, "_fast_path_env_status", lambda: {}) + monkeypatch.setattr(openai, "_mlx_runtime_status", lambda: {"ok": True}) + monkeypatch.setattr( + openai, + "_configure_mlx_cache_limit", + lambda _args: {"configured": False}, + ) + + def stop_after_load(model, mtp, contract, **kwargs): + captured.update(kwargs) + raise RuntimeError("stop after load") + + monkeypatch.setattr(openai, "load", stop_after_load) + args = parse_args( + [ + "--model", + "models/DeepSeek-V4-Flash-0731", + "--warmup-tokens", + "0", + "--deepseek-v4-0731-k2", + "--depth", + "2", + ] + ) + + with pytest.raises(RuntimeError, match="stop after load"): + openai.ServerState(args) + + assert captured["deepseek_v4_0731_k2"] is True + + def test_normalize_stop_sequences_accepts_string_list_and_caps_at_four(): assert openai._normalize_stop_sequences(None) == [] assert openai._normalize_stop_sequences("END") == ["END"] From 4ed84262e5b2b36da6177533282e5a2dcda6351c Mon Sep 17 00:00:00 2001 From: davidtai Date: Wed, 12 Aug 2026 19:11:01 -0500 Subject: [PATCH 03/24] bench: report prefill and memory by MTP depth --- mtplx/benchmarks/runners/mtp_depth_sweep.py | 35 +++++++++ tests/test_mtp_depth_sweep.py | 87 ++++++++++++++++++++- 2 files changed, 120 insertions(+), 2 deletions(-) diff --git a/mtplx/benchmarks/runners/mtp_depth_sweep.py b/mtplx/benchmarks/runners/mtp_depth_sweep.py index 3b707e83..29c9ebb4 100644 --- a/mtplx/benchmarks/runners/mtp_depth_sweep.py +++ b/mtplx/benchmarks/runners/mtp_depth_sweep.py @@ -37,6 +37,12 @@ def _cycle_count(events: list[dict], verify_calls: int) -> int: return len(events) or max(0, int(verify_calls)) +def _active_memory_bytes() -> int: + import mlx.core as mx + + return int(mx.get_active_memory()) + + def _token_budget(max_tokens: int, case_max_tokens: int) -> int: return min(int(max_tokens), int(case_max_tokens)) @@ -137,6 +143,7 @@ def run_mtp_depth_sweep( merge_mtp_adapter=merge_mtp_adapter, gemma4_draft_block_size=gemma4_draft_block_size, ) + load_active_memory_bytes = _active_memory_bytes() contract = getattr(rt, "contract", None) is_gemma4_assistant = getattr(rt, "backend_id", None) == "gemma4_assistant" resolved_base_hidden_variant = str( @@ -213,6 +220,7 @@ def run_mtp_depth_sweep( seed=seed + index, ) generation_ended_at = time.time() + active_memory_bytes = _active_memory_bytes() validations = [ asdict(validation) for validation in validate_benchmark_output( @@ -236,12 +244,22 @@ def run_mtp_depth_sweep( ar.finish_reason, ), "generated_tokens": ar.stats.generated_tokens, + "prompt_tokens": len(ids), "elapsed_s": ar.stats.elapsed_s, "tok_s": ar.stats.tok_s, "decode_tok_s": ar.stats.decode_tok_s, "decode_elapsed_s": ar.stats.decode_elapsed_s, "end_to_end_tok_s": ar.stats.end_to_end_tok_s, "prompt_eval_time_s": ar.stats.prompt_eval_time_s, + "prompt_tps": ar.stats.prompt_tps, + "prompt_target_prefill_tok_s": ( + ar.stats.prompt_target_prefill_tok_s + ), + "active_memory_bytes": active_memory_bytes, + "active_memory_growth_bytes": ( + active_memory_bytes - load_active_memory_bytes + ), + "peak_memory_bytes": ar.stats.peak_memory_bytes, "tokens": ar.tokens, "text": ar.text, "validations": validations, @@ -288,6 +306,7 @@ def run_mtp_depth_sweep( mtp_topk_reranker=mtp_topk_reranker, ) generation_ended_at = time.time() + active_memory_bytes = _active_memory_bytes() validations = [ asdict(validation) for validation in validate_benchmark_output( @@ -314,12 +333,21 @@ def run_mtp_depth_sweep( out.finish_reason, ), "generated_tokens": out.stats.generated_tokens, + "prompt_tokens": len(ids), "elapsed_s": out.stats.elapsed_s, "tok_s": out.stats.tok_s, "decode_tok_s": out.stats.decode_tok_s, "decode_elapsed_s": out.stats.decode_elapsed_s, "end_to_end_tok_s": out.stats.end_to_end_tok_s, "prompt_eval_time_s": out.stats.prompt_eval_time_s, + "prompt_tps": out.stats.prompt_tps, + "prompt_target_prefill_tok_s": ( + out.stats.prompt_target_prefill_tok_s + ), + "active_memory_bytes": active_memory_bytes, + "active_memory_growth_bytes": ( + active_memory_bytes - load_active_memory_bytes + ), "ar_tok_s": ar_row["tok_s"] if ar_row is not None else None, "ar_decode_tok_s": ar_row["decode_tok_s"] if ar_row is not None @@ -600,6 +628,12 @@ def run_mtp_depth_sweep( "peak_memory_bytes": max( [row["peak_memory_bytes"] for row in rows] or [0] ), + "active_memory_bytes": max( + [row["active_memory_bytes"] for row in rows] or [0] + ), + "active_memory_growth_bytes": max( + [row["active_memory_growth_bytes"] for row in rows] or [0] + ), "speed_model": _speed_model_summary(rows), }, } @@ -607,6 +641,7 @@ def run_mtp_depth_sweep( return { "model_path": str(model_path), + "load_active_memory_bytes": load_active_memory_bytes, "prompt_suite": str(prompt_suite), "sampler": asdict(sampler), "draft_sampler": asdict(draft_sampler), diff --git a/tests/test_mtp_depth_sweep.py b/tests/test_mtp_depth_sweep.py index 84bf0981..d24dca74 100644 --- a/tests/test_mtp_depth_sweep.py +++ b/tests/test_mtp_depth_sweep.py @@ -7,19 +7,27 @@ from mtplx.generation import GenerationOutput, GenerationStats -def _generation_output(*, events: list[dict], verify_calls: int) -> GenerationOutput: +def _generation_output( + *, + events: list[dict], + verify_calls: int, + mode: str = "mtp", +) -> GenerationOutput: return GenerationOutput( tokens=[1, 2], text="ok", finish_reason="length", stats=GenerationStats( - mode="mtp", + mode=mode, generated_tokens=2, elapsed_s=1.0, tok_s=2.0, decode_elapsed_s=1.0, decode_tok_s=2.0, end_to_end_tok_s=2.0, + prompt_tps=50.0, + prompt_target_prefill_tok_s=48.0, + peak_memory_bytes=1_500, accepted_drafts=6, drafted_tokens=8, accepted_by_depth=[3, 3], @@ -32,6 +40,81 @@ def _generation_output(*, events: list[dict], verify_calls: int) -> GenerationOu ) +def test_depth_sweep_reports_prefill_and_memory_growth_for_ar_and_each_depth( + monkeypatch, tmp_path +) -> None: + fake_runtime = SimpleNamespace( + tokenizer=object(), + contract=SimpleNamespace( + base_hidden_variant="pre_norm", + hidden_variant="pre_norm", + concat_order="base_then_mtp", + mtp_quant_bits=None, + mtp_quant_group_size=64, + mtp_quant_mode="affine", + mtp_quant_policy=None, + ), + mtp_adapter_metadata=None, + mtp_adapter_merge_report=None, + ) + active_memory = iter([1_000, 1_125, 1_250]) + + monkeypatch.setattr(mtp_depth_sweep, "load", lambda *_args, **_kwargs: fake_runtime) + monkeypatch.setattr( + mtp_depth_sweep, + "load_prompt_suite", + lambda *_args: [PromptCase(id="one", category="general", prompt="one")], + ) + monkeypatch.setattr( + mtp_depth_sweep, "encode_prompt_case", lambda *_args, **_kwargs: [1, 2] + ) + monkeypatch.setattr( + mtp_depth_sweep, + "generate_ar", + lambda *_args, **_kwargs: _generation_output( + events=[], verify_calls=0, mode="ar" + ), + ) + monkeypatch.setattr( + mtp_depth_sweep, + "generate_mtpk", + lambda *_args, **_kwargs: _generation_output(events=[{}], verify_calls=1), + ) + monkeypatch.setattr( + mtp_depth_sweep, "_active_memory_bytes", lambda: next(active_memory) + ) + monkeypatch.setattr( + mtp_depth_sweep, "validate_benchmark_output", lambda *_args, **_kwargs: [] + ) + + result = mtp_depth_sweep.run_mtp_depth_sweep( + tmp_path / "model", + tmp_path / "suite.jsonl", + depths=[1], + compare_ar=True, + temperature=0.0, + ) + + assert result["load_active_memory_bytes"] == 1_000 + ar_row = result["ar_rows"][0] + assert ar_row["prompt_tokens"] == 2 + assert ar_row["prompt_target_prefill_tok_s"] == 48.0 + assert ar_row["prompt_tps"] == 50.0 + assert ar_row["active_memory_bytes"] == 1_125 + assert ar_row["active_memory_growth_bytes"] == 125 + assert ar_row["peak_memory_bytes"] == 1_500 + depth = result["depths"][0] + depth_row = depth["rows"][0] + assert depth_row["prompt_tokens"] == 2 + assert depth_row["prompt_target_prefill_tok_s"] == 48.0 + assert depth_row["prompt_tps"] == 50.0 + assert depth_row["active_memory_bytes"] == 1_250 + assert depth_row["active_memory_growth_bytes"] == 250 + assert depth_row["peak_memory_bytes"] == 1_500 + assert depth["summary"]["active_memory_bytes"] == 1_250 + assert depth["summary"]["active_memory_growth_bytes"] == 250 + + def test_depth_sweep_uses_verify_calls_only_when_events_are_empty( monkeypatch, tmp_path ) -> None: From b0be3417bc9a3e758fcf44823ca4f4506ecbfe22 Mon Sep 17 00:00:00 2001 From: davidtai Date: Wed, 12 Aug 2026 19:12:52 -0500 Subject: [PATCH 04/24] bench: select the construction-bound 0731 stack --- mtplx/benchmarks/runners/mtp_depth_sweep.py | 3 ++ tests/test_mtp_depth_sweep.py | 39 +++++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/mtplx/benchmarks/runners/mtp_depth_sweep.py b/mtplx/benchmarks/runners/mtp_depth_sweep.py index 29c9ebb4..514f3c8c 100644 --- a/mtplx/benchmarks/runners/mtp_depth_sweep.py +++ b/mtplx/benchmarks/runners/mtp_depth_sweep.py @@ -123,6 +123,7 @@ def run_mtp_depth_sweep( draft_lm_head_bits: int | None = None, draft_lm_head_group_size: int = 64, draft_lm_head_mode: str = "affine", + deepseek_v4_0731_k2: bool = False, ) -> dict[str, Any]: contract_kwargs: dict[str, Any] = { "mtp_quant_bits": mtp_quant_bits, @@ -142,6 +143,7 @@ def run_mtp_depth_sweep( mtp_adapter=mtp_adapter_path, merge_mtp_adapter=merge_mtp_adapter, gemma4_draft_block_size=gemma4_draft_block_size, + deepseek_v4_0731_k2=deepseek_v4_0731_k2, ) load_active_memory_bytes = _active_memory_bytes() contract = getattr(rt, "contract", None) @@ -642,6 +644,7 @@ def run_mtp_depth_sweep( return { "model_path": str(model_path), "load_active_memory_bytes": load_active_memory_bytes, + "deepseek_v4_0731_k2": deepseek_v4_0731_k2, "prompt_suite": str(prompt_suite), "sampler": asdict(sampler), "draft_sampler": asdict(draft_sampler), diff --git a/tests/test_mtp_depth_sweep.py b/tests/test_mtp_depth_sweep.py index d24dca74..0adca0e4 100644 --- a/tests/test_mtp_depth_sweep.py +++ b/tests/test_mtp_depth_sweep.py @@ -258,3 +258,42 @@ def fake_load(*_args, **kwargs): "merged": 1, "targets": [{"target": "fc"}], } + + +def test_depth_sweep_selects_construction_bound_0731_stack( + monkeypatch, tmp_path +) -> None: + load_kwargs = [] + fake_runtime = SimpleNamespace( + tokenizer=object(), + contract=SimpleNamespace( + base_hidden_variant="pre_norm", + hidden_variant="pre_norm", + concat_order="base_then_mtp", + mtp_quant_bits=None, + mtp_quant_group_size=64, + mtp_quant_mode="affine", + mtp_quant_policy=None, + ), + mtp_adapter_metadata=None, + mtp_adapter_merge_report=None, + ) + + def fake_load(*_args, **kwargs): + load_kwargs.append(kwargs) + return fake_runtime + + monkeypatch.setattr(mtp_depth_sweep, "load", fake_load) + monkeypatch.setattr(mtp_depth_sweep, "_active_memory_bytes", lambda: 0) + monkeypatch.setattr( + mtp_depth_sweep, "load_prompt_suite", lambda *_args, **_kwargs: [] + ) + + mtp_depth_sweep.run_mtp_depth_sweep( + tmp_path / "model", + tmp_path / "suite.jsonl", + depths=[1], + deepseek_v4_0731_k2=True, + ) + + assert load_kwargs[0]["deepseek_v4_0731_k2"] is True From b6dbf4364bbad6e7734d6a63c96ecdb0a5f78015 Mon Sep 17 00:00:00 2001 From: davidtai Date: Wed, 12 Aug 2026 19:19:42 -0500 Subject: [PATCH 05/24] fix: restore the 0731 stock q-projection seam --- mtplx/models/deepseek_v4.py | 70 +++++++++++++++++++++++++------- tests/test_deepseek_v4_dspark.py | 29 +++++++++++++ 2 files changed, 84 insertions(+), 15 deletions(-) diff --git a/mtplx/models/deepseek_v4.py b/mtplx/models/deepseek_v4.py index 91115ec6..7ebbfa80 100644 --- a/mtplx/models/deepseek_v4.py +++ b/mtplx/models/deepseek_v4.py @@ -1626,6 +1626,33 @@ def _apply_interleaved_rope(x: mx.array, cos: mx.array, sin: mx.array) -> mx.arr return out.reshape(shape).astype(out_dtype) +def _q_head_norm_rope_stock( + q: mx.array, + cos: mx.array, + sin: mx.array, + *, + eps: float, + rope_dim: int, +) -> mx.array: + """The stock per-head Q normalization and interleaved-RoPE graph.""" + out_dtype = _store_dtype(q.dtype) + q = q * mx.rsqrt( + mx.mean(mx.square(q.astype(mx.float32)), axis=-1, keepdims=True) + eps + ) + q = q.astype(out_dtype) + return mx.concatenate( + [ + q[..., :-rope_dim], + _apply_interleaved_rope( + q[..., -rope_dim:], + cos[None, :, None, :], + sin[None, :, None, :], + ), + ], + axis=-1, + ) + + # --------------------------------------------------------------------------- # Hyper-Connections # --------------------------------------------------------------------------- @@ -2885,12 +2912,39 @@ def __init__(self, args: ModelArgs, layer_id: int): else: inv = _yarn_inv_freq(self.rope_head_dim, args.rope_theta, 0, 1.0, 32, 1) self._inv_freq = inv # [rope_head_dim//2] + self._q_head_norm_rope_route = self._q_head_norm_rope_stock + self._q_projection_qhead_route = self._q_projection_qhead_stock def _rope_tables(self, positions: mx.array): # positions: [L] -> cos/sin [L, rope_head_dim//2] ang = positions[:, None].astype(mx.float32) * self._inv_freq[None, :] return mx.cos(ang), mx.sin(ang) + def _q_head_norm_rope_stock( + self, + q: mx.array, + cos: mx.array, + sin: mx.array, + ) -> mx.array: + return _q_head_norm_rope_stock( + q, + cos, + sin, + eps=self.eps, + rope_dim=self.rope_head_dim, + ) + + def _q_projection_qhead_stock( + self, + qr: mx.array, + cos: mx.array, + sin: mx.array, + ) -> mx.array: + """Project Q then run the currently installed phase-specific post route.""" + batch, sequence, _ = qr.shape + q = self.wq_b(qr).reshape(batch, sequence, self.n_heads, self.head_dim) + return self._q_head_norm_rope_route(q, cos, sin) + def _wo_a_quant(self): """``wo_a``'s quantised tensors + format, or ``None`` when it is dense. @@ -3164,21 +3218,7 @@ def __call__(self, x: mx.array, mask=None, cache=None) -> mx.array: cos, sin = self._rope_tables(positions) qr = self.q_norm(self.wq_a(x)) - q = self.wq_b(qr).reshape(b, s, self.n_heads, self.head_dim) - # per-head RMS-like normalisation (no learned weight), reference L498 - q = q * mx.rsqrt( - mx.mean(mx.square(q.astype(mx.float32)), axis=-1, keepdims=True) + self.eps - ) - q = q.astype(x.dtype) - q = mx.concatenate( - [ - q[..., :-rd], - _apply_interleaved_rope( - q[..., -rd:], cos[None, :, None, :], sin[None, :, None, :] - ), - ], - axis=-1, - ) + q = self._q_projection_qhead_route(qr, cos, sin) kv = self.kv_norm(self.wkv(x)) # [b, s, head_dim] (single shared KV — MQA) kv = mx.concatenate( diff --git a/tests/test_deepseek_v4_dspark.py b/tests/test_deepseek_v4_dspark.py index 7536f8c3..f13243c9 100644 --- a/tests/test_deepseek_v4_dspark.py +++ b/tests/test_deepseek_v4_dspark.py @@ -126,6 +126,35 @@ def test_real_filtered_config_derives_three_stages_without_n_mtp_layers(): ) +def test_attention_exposes_exact_stock_q_projection_route(): + attention = D.DeepseekV4Attention(_args(), 0) + qr = mx.arange(16, dtype=mx.float32).reshape(1, 2, 8).astype(mx.bfloat16) + positions = mx.arange(2) + cos, sin = attention._rope_tables(positions) + + projected = attention.wq_b(qr).reshape(1, 2, 1, 8) + expected = projected * mx.rsqrt( + mx.mean(mx.square(projected.astype(mx.float32)), axis=-1, keepdims=True) + + attention.eps + ) + expected = expected.astype(projected.dtype) + expected = mx.concatenate( + [ + expected[..., : -attention.rope_head_dim], + D._apply_interleaved_rope( + expected[..., -attention.rope_head_dim :], + cos[None, :, None, :], + sin[None, :, None, :], + ), + ], + axis=-1, + ) + actual = attention._q_projection_qhead_route(qr, cos, sin) + + mx.eval(actual, expected) + assert mx.array_equal(actual, expected) + + def test_legacy_model_call_uses_prebound_target_route_exactly_once(): inputs = mx.array([[7]], dtype=mx.int32) cache = object() From 11240c06e4014bb4f9ed3732f2df7c67bc9185ba Mon Sep 17 00:00:00 2001 From: davidtai Date: Wed, 12 Aug 2026 19:21:47 -0500 Subject: [PATCH 06/24] fix: compile the exact 0731 q-head kernel --- mtplx/deepseek_v4_0731_m3_wqb_qnorm_rope.py | 8 ++++---- tests/test_deepseek_v4_0731_m3_wqb_qnorm_rope.py | 7 +++++-- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/mtplx/deepseek_v4_0731_m3_wqb_qnorm_rope.py b/mtplx/deepseek_v4_0731_m3_wqb_qnorm_rope.py index 2305cc58..f49cb576 100644 --- a/mtplx/deepseek_v4_0731_m3_wqb_qnorm_rope.py +++ b/mtplx/deepseek_v4_0731_m3_wqb_qnorm_rope.py @@ -254,10 +254,10 @@ def m3_wqb_qnorm_rope_metal_source(*, capture_projection: bool = False) -> str: // The eager stock graph materializes each FP32 product in a separate // binary kernel before its add/subtract. ``precise`` preserves those // intermediate roundings inside this one-launch candidate. - precise float rope0_lhs = x0 * c[pair]; - precise float rope0_rhs = x1 * s[pair]; - precise float rope1_lhs = x0 * s[pair]; - precise float rope1_rhs = x1 * c[pair]; + float rope0_lhs = metal::precise::fma(x0, c[pair], 0.0f); + float rope0_rhs = metal::precise::fma(x1, s[pair], 0.0f); + float rope1_lhs = metal::precise::fma(x0, s[pair], 0.0f); + float rope1_rhs = metal::precise::fma(x1, c[pair], 0.0f); out[d] = T(rope0_lhs - rope0_rhs); out[d + 1] = T(rope1_lhs + rope1_rhs); } diff --git a/tests/test_deepseek_v4_0731_m3_wqb_qnorm_rope.py b/tests/test_deepseek_v4_0731_m3_wqb_qnorm_rope.py index 863ab6d5..c1d4701e 100644 --- a/tests/test_deepseek_v4_0731_m3_wqb_qnorm_rope.py +++ b/tests/test_deepseek_v4_0731_m3_wqb_qnorm_rope.py @@ -120,10 +120,13 @@ def test_source_preserves_pre_geometry_qmv_norm_and_rope_arithmetic(): "constexpr uint NORM_LANES = 32;", "constexpr uint NORM_READS = 4;", "metal::precise::rsqrt(mean + EPS);", - "precise float rope0_lhs = x0 * c[pair];", - "precise float rope1_rhs = x1 * c[pair];", + "float rope0_lhs = metal::precise::fma(x0, c[pair], 0.0f);", + "float rope0_rhs = metal::precise::fma(x1, s[pair], 0.0f);", + "float rope1_lhs = metal::precise::fma(x0, s[pair], 0.0f);", + "float rope1_rhs = metal::precise::fma(x1, c[pair], 0.0f);", ): assert fragment in source + assert "precise float" not in source assert ( "geometry" not in inspect.signature(candidate.m3_wqb_qnorm_rope_metal_source).parameters From 8becdf7e31288a9c94d1b56e32c4b5a0c453d876 Mon Sep 17 00:00:00 2001 From: davidtai Date: Wed, 12 Aug 2026 19:23:35 -0500 Subject: [PATCH 07/24] fix: stage 0731 gather ownership before WOB --- mtplx/runtime.py | 19 +++++++++++-------- tests/test_runtime_deepseek_v4_dspark.py | 8 ++++---- 2 files changed, 15 insertions(+), 12 deletions(-) diff --git a/mtplx/runtime.py b/mtplx/runtime.py index 43f36b14..090dc343 100644 --- a/mtplx/runtime.py +++ b/mtplx/runtime.py @@ -1004,8 +1004,11 @@ def load( "gather_qmm" if deepseek_v4_0731_k2 else _o_lora_mode_from_env() ) if deepseek_v4_0731_k2: - # Snapshot now, but leave O-Lora and the target/FFN routes stock while - # their candidates run construction-time preparation and self-checks. + # Snapshot every owner before the reversible construction transaction. + # The fused WOB preparer requires the installed gather route to own + # the same stock wo_b object, so gather publication happens first, + # while the model is still private to load(), and rolls back on any + # later preparation or publication failure. k2_o_lora_state = _snapshot_deepseek_v4_o_lora_routes(model) # The canonical mixed route hard-validates the exact DeepSeek-V4-Flash # topology (43 body layers, rank-1024 Q4/g64 wo_a/wo_b, one dense-BF16 @@ -1071,6 +1074,12 @@ def load( from .deepseek_v4_dspark_generation import DeepseekV4DSparkBackend try: + deepseek_v4_o_lora_report = install_deepseek_v4_o_lora_routes( + model, + mode="gather_qmm", + canonical_mixed_route=False, + ) + logger.info("[deepseek-v4-o-lora] %s", deepseek_v4_o_lora_report) target_prepared = prepare_full_0731_dspark_compiled_tail_q2_pair( model, config, @@ -1080,12 +1089,6 @@ def load( ) ffn_prepared = prepare_dspark_q3_packed_gate_up_m5(model) k2_prepared = (target_prepared, ffn_prepared) - deepseek_v4_o_lora_report = install_deepseek_v4_o_lora_routes( - model, - mode="gather_qmm", - canonical_mixed_route=False, - ) - logger.info("[deepseek-v4-o-lora] %s", deepseek_v4_o_lora_report) target_prepared.publish() ffn_prepared.publish() block_speculative_backend = DeepseekV4DSparkBackend.bind(model) diff --git a/tests/test_runtime_deepseek_v4_dspark.py b/tests/test_runtime_deepseek_v4_dspark.py index cd8bbcbe..a41870ad 100644 --- a/tests/test_runtime_deepseek_v4_dspark.py +++ b/tests/test_runtime_deepseek_v4_dspark.py @@ -407,14 +407,14 @@ def test_k2_option_publishes_one_exact_construction_transaction(monkeypatch, tmp "artifact.validate", "sinkhorn.enter", "sinkhorn.exit", - ("target.prepare", "prepare_wqb_qhead_m3", "prepare_wob_m3"), - ("wqb.prepare", True), - ("wob.prepare", True), - "ffn.prepare", ( "o_lora", {"mode": "gather_qmm", "canonical_mixed_route": False}, ), + ("target.prepare", "prepare_wqb_qhead_m3", "prepare_wob_m3"), + ("wqb.prepare", True), + ("wob.prepare", True), + "ffn.prepare", "target.publish", "ffn.publish", ] From 70b4cc44cbfeda21861cbe3d9ad479e9e4839089 Mon Sep 17 00:00:00 2001 From: davidtai Date: Wed, 12 Aug 2026 19:25:44 -0500 Subject: [PATCH 08/24] fix: accept MLX bitstream-packed Q6 o-lora --- mtplx/models/deepseek_v4.py | 19 +++++++++++++------ tests/test_deepseek_v4_o_lora.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 6 deletions(-) diff --git a/mtplx/models/deepseek_v4.py b/mtplx/models/deepseek_v4.py index 7ebbfa80..ec313040 100644 --- a/mtplx/models/deepseek_v4.py +++ b/mtplx/models/deepseek_v4.py @@ -1033,10 +1033,12 @@ def _o_lora_linear_logical_weight_shape(linear) -> tuple[int, ...]: scales_shape = tuple(linear.scales.shape) except (AttributeError, TypeError, ValueError): return () - packed_divisor = 32 // bits if bits > 0 and 32 % bits == 0 else 0 - if not packed_divisor or group_size <= 0 or len(scales_shape) != 2: + if bits not in {2, 3, 4, 6, 8} or group_size <= 0 or len(scales_shape) != 2: return () - packed_logical = (weight_shape[0], weight_shape[1] * packed_divisor) + packed_bits = weight_shape[1] * 32 + if packed_bits % bits: + return () + packed_logical = (weight_shape[0], packed_bits // bits) scales_logical = (scales_shape[0], scales_shape[1] * group_size) return packed_logical if packed_logical == scales_logical else () @@ -1071,10 +1073,15 @@ def __init__(self, attention: "DeepseekV4Attention", quant: tuple) -> None: raise ValueError( "gather_qmm o-LoRA input width is not divisible by group_size" ) - packed_divisor = 32 // int(bits) if int(bits) and 32 % int(bits) == 0 else 0 - if not packed_divisor: + bits = int(bits) + if bits not in {2, 3, 4, 6, 8}: raise ValueError(f"gather_qmm o-LoRA has unsupported bits={bits!r}") - expected_weight = (output_rows, per_group_input // packed_divisor) + packed_row_bits = per_group_input * bits + if packed_row_bits % 32: + raise ValueError( + "gather_qmm o-LoRA logical input does not end on a packed word" + ) + expected_weight = (output_rows, packed_row_bits // 32) expected_scales = (output_rows, per_group_input // int(group_size)) if tuple(weight.shape) != expected_weight: raise ValueError( diff --git a/tests/test_deepseek_v4_o_lora.py b/tests/test_deepseek_v4_o_lora.py index 8550132f..47d0480d 100644 --- a/tests/test_deepseek_v4_o_lora.py +++ b/tests/test_deepseek_v4_o_lora.py @@ -312,6 +312,35 @@ def test_gather_qmm_calling_convention_shapes(): assert rel < 1e-3, f"rows={rows} grouped qmm rel={rel:.3e}" +def test_gather_qmm_accepts_mlx_bitstream_packed_q6_rows(): + """Q6 packs 16 values across three uint32 words, not 32//bits per word.""" + _, model = _seeded_model(o_lora_rank=16) + nn.quantize( + model, + group_size=32, + bits=6, + class_predicate=lambda path, _module: path.endswith( + ("attn.wo_a", "attn.wo_b") + ), + ) + mx.eval(model.parameters()) + attn = model.layers[0].attn + assert tuple(attn.wo_a.weight.shape) == (32, 6) + assert D._o_lora_linear_logical_weight_shape(attn.wo_a) == (32, 32) + assert D._o_lora_linear_logical_weight_shape(attn.wo_b) == (32, 32) + o = mx.sin(mx.arange(N_HEADS * HEAD_DIM, dtype=mx.float32) * 0.017) + o = o.reshape(1, 1, -1).astype(mx.bfloat16) + want = attn._o_lora_dense(o) + + receipt = attn.install_o_lora_route("gather_qmm") + got = attn._o_lora(o) + mx.eval(want, got) + + assert receipt["direct"] is True + assert tuple(got.shape) == tuple(want.shape) + assert float(mx.max(mx.abs(got - want))) <= 0.125 + + def test_gather_qmm_is_prebound_and_never_rechecks_or_falls_back(monkeypatch): _, model = _quantized_model() attn = model.layers[0].attn From 8a57b9adb4030d1334ee5440e160a06c55555643 Mon Sep 17 00:00:00 2001 From: davidtai Date: Wed, 12 Aug 2026 19:29:08 -0500 Subject: [PATCH 09/24] bench: omit generic variants for block backends --- mtplx/benchmarks/runners/mtp_depth_sweep.py | 9 +++- tests/test_mtp_depth_sweep.py | 55 +++++++++++++++++++++ 2 files changed, 62 insertions(+), 2 deletions(-) diff --git a/mtplx/benchmarks/runners/mtp_depth_sweep.py b/mtplx/benchmarks/runners/mtp_depth_sweep.py index 514f3c8c..fcd32340 100644 --- a/mtplx/benchmarks/runners/mtp_depth_sweep.py +++ b/mtplx/benchmarks/runners/mtp_depth_sweep.py @@ -148,6 +148,7 @@ def run_mtp_depth_sweep( load_active_memory_bytes = _active_memory_bytes() contract = getattr(rt, "contract", None) is_gemma4_assistant = getattr(rt, "backend_id", None) == "gemma4_assistant" + is_native_block_backend = getattr(rt, "block_speculative_backend", None) is not None resolved_base_hidden_variant = str( getattr(contract, "base_hidden_variant", "gemma4_assistant") ) @@ -281,8 +282,12 @@ def run_mtp_depth_sweep( sampler=sampler, speculative_depth=depth, seed=seed + index, - base_hidden_variant=resolved_base_hidden_variant, - mtp_hidden_variant=resolved_mtp_hidden_variant, + base_hidden_variant=( + None if is_native_block_backend else resolved_base_hidden_variant + ), + mtp_hidden_variant=( + None if is_native_block_backend else resolved_mtp_hidden_variant + ), mtp_cache_policy=mtp_cache_policy, mtp_history_policy=mtp_history_policy, draft_sampler=draft_sampler, diff --git a/tests/test_mtp_depth_sweep.py b/tests/test_mtp_depth_sweep.py index 0adca0e4..6e469df3 100644 --- a/tests/test_mtp_depth_sweep.py +++ b/tests/test_mtp_depth_sweep.py @@ -297,3 +297,58 @@ def fake_load(*_args, **kwargs): ) assert load_kwargs[0]["deepseek_v4_0731_k2"] is True + + +def test_depth_sweep_omits_hidden_variants_for_native_block_backend( + monkeypatch, tmp_path +) -> None: + generate_kwargs = [] + fake_runtime = SimpleNamespace( + tokenizer=object(), + block_speculative_backend=object(), + contract=SimpleNamespace( + base_hidden_variant="pre_norm", + hidden_variant="pre_norm", + concat_order="base_then_mtp", + mtp_quant_bits=None, + mtp_quant_group_size=64, + mtp_quant_mode="affine", + mtp_quant_policy=None, + ), + mtp_adapter_metadata=None, + mtp_adapter_merge_report=None, + ) + + monkeypatch.setattr(mtp_depth_sweep, "load", lambda *_args, **_kwargs: fake_runtime) + monkeypatch.setattr(mtp_depth_sweep, "_active_memory_bytes", lambda: 0) + monkeypatch.setattr( + mtp_depth_sweep, + "load_prompt_suite", + lambda *_args: [PromptCase(id="one", category="general", prompt="one")], + ) + monkeypatch.setattr( + mtp_depth_sweep, "encode_prompt_case", lambda *_args, **_kwargs: [1, 2] + ) + monkeypatch.setattr( + mtp_depth_sweep, + "generate_mtpk", + lambda *_args, **kwargs: ( + generate_kwargs.append(kwargs) + or _generation_output(events=[], verify_calls=1) + ), + ) + monkeypatch.setattr( + mtp_depth_sweep, "validate_benchmark_output", lambda *_args, **_kwargs: [] + ) + + mtp_depth_sweep.run_mtp_depth_sweep( + tmp_path / "model", + tmp_path / "suite.jsonl", + depths=[1], + deepseek_v4_0731_k2=True, + ) + + assert generate_kwargs[0]["base_hidden_variant"] is None + assert generate_kwargs[0]["mtp_hidden_variant"] is None + assert generate_kwargs[0]["mtp_history_policy"] == "cycle" + assert generate_kwargs[0]["verify_strategy"] == "batched" From 98aabc78bfda204d1c3c8c8ec71fc89e1d6a62f3 Mon Sep 17 00:00:00 2001 From: davidtai Date: Wed, 12 Aug 2026 19:40:01 -0500 Subject: [PATCH 10/24] docs: publish the 0731 DSpark performance receipt --- docs/perf/receipts/deepseek-v4-0731-dspark.md | 75 +++++++++++++++++++ 1 file changed, 75 insertions(+) create mode 100644 docs/perf/receipts/deepseek-v4-0731-dspark.md diff --git a/docs/perf/receipts/deepseek-v4-0731-dspark.md b/docs/perf/receipts/deepseek-v4-0731-dspark.md new file mode 100644 index 00000000..feed37ee --- /dev/null +++ b/docs/perf/receipts/deepseek-v4-0731-dspark.md @@ -0,0 +1,75 @@ +# DeepSeek-V4 Flash 0731 DSpark receipt + +This is the scrubbed, tracked performance receipt for the construction-bound +DeepSeek-V4 Flash 0731 DSpark K2 lane. Raw generation artifacts remain local; +their hashes are listed below without model paths, generated text, service +details, or machine-local process data. + +## Fixed conditions + +- Machine: Apple M5 Max MacBook Pro, 128 GB, macOS 26.5.2. +- Runtime: Python 3.12.13, MLX 0.32.0, mlx-lm 0.31.3. +- Model: `mlx-community/DeepSeek-V4-Flash-0731-2.4bit-mixed` at source revision + `10001e0065f8394e03e968e652cbbe7cd2ca122c`. +- Model identity: config SHA-256 + `44735712733fcf8f299bdf1faa1d87fac88f1917efe1d3876d6d4c582f79a68f`; + index SHA-256 + `f1332b2b209769c2db335954c2651652a8048e7d7dbf60296c2f2c0198715861`. +- Sampler: greedy, `temperature=0`, `top_p=1`, `top_k=0`, seed 0. +- Prompt: `Explain why speculative decoding can preserve greedy output.` + through the model chat template, 14 prompt tokens. +- Output: forced 128-token budget, two identical cases in one model load. The + second case is the warmed comparison; the first exposes one-time compilation. +- PR lane: explicit `deepseek_v4_0731_k2=True`, fixed proposal width K2, + persistent cache, cycle history, batched native verification, stock verify and + draft cores. +- Benchmarked commit: `8a57b9adb4030d1334ee5440e160a06c55555643`. + +## Current PR bracket + +Memory growth is measured from the post-load active-memory baseline of +86.4561 GiB. Peak memory is the process-wide MLX peak and therefore includes +load-time allocation; it is identical across these in-load arms. + +| case | depth | target prefill tok/s | decode tok/s | end-to-end tok/s | active GiB | growth MiB | peak GiB | accepted / drafted | exact vs K0 | +|---|---:|---:|---:|---:|---:|---:|---:|---:|---| +| cold compile | K0 | 0.191 | 25.264 | 1.630 | 86.4561 | 0.0180 | 139.7061 | - | reference | +| cold compile | K2 | 103.592 | 28.954 | 28.085 | 86.4561 | 0.0190 | 139.7061 | 68 / 119 | yes | +| warmed | K0 | 103.584 | **32.358** | **31.289** | 86.4561 | 0.0180 | 139.7061 | - | reference | +| warmed | K2 | 103.315 | 29.240 | 28.356 | 86.4561 | 0.0190 | 139.7061 | 68 / 119 | yes | + +The warmed K2 lane is exact in both cases, but it does **not** beat warmed AR: +29.240 versus 32.358 decode tok/s, a 9.6% loss. First- and second-position +acceptance were 70.0% and 44.1%. The cold K0 prefill result is compilation time, +not model prefill throughput, so it is disclosed rather than used as a speedup +claim. + +## Historical K-depth diagnostic + +Before the K2-only construction contract was pinned, the native DSpark harness +ran one simple 9-prompt-token, 64-output-token K0-K3 sweep on MLX 0.31.2. It did +not record prefill TPS or memory growth, so those fields are unavailable. Active +and peak memory remain useful as measured. + +| depth | decode tok/s | end-to-end tok/s | active GiB | peak GiB | accepted / drafted | exact vs K0 | +|---|---:|---:|---:|---:|---:|---| +| K0 | **24.565** | **23.312** | 86.4561 | 86.5079 | - | reference | +| K1 | 19.193 | 18.584 | 86.4561 | 86.5175 | 27 / 36 | yes | +| K2 | 19.640 | 19.010 | 86.4561 | 86.5240 | 34 / 58 | yes | +| K3 | 21.413 | 20.609 | 86.4561 | 86.5392 | 37 / 76 | **no** | + +This older chart is diagnostic, not a promotion result: K1 and K2 were exact but +slower than AR, while K3 was faster than K1/K2 but diverged from greedy AR. The +current public lane therefore stays construction-pinned to K2 rather than +silently widening to an unqualified K1/K3 route. + +## Raw-artifact manifest + +| local artifact | SHA-256 | +|---|---| +| `0731-pr-optimized-k2-128-20260812.json` | `e3e8ab454a5a6860578eb022e85297de9143b5bd5588229bb795e472ba5395c2` | +| `0731-dspark-width123-64tok-20260809.json` | `1f60e529e4c172642fa461c41f5cd5dd11f28048c571f875cff04ee73cae9a3f` | + +Profiler dispatch censuses and physical-M3 diagnostics are not used as TPS +proof here. They are discovery evidence only and remain separate from these +uninstrumented generation timings. From 12e32bbc5790051681d7f732a719d6140f5e7335 Mon Sep 17 00:00:00 2001 From: davidtai Date: Wed, 12 Aug 2026 19:46:01 -0500 Subject: [PATCH 11/24] test: preserve A3B source-shape contracts --- mtplx/generation.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/mtplx/generation.py b/mtplx/generation.py index ac1e283a..734220ac 100644 --- a/mtplx/generation.py +++ b/mtplx/generation.py @@ -6458,8 +6458,7 @@ def generate_mtpk( # the batch scheduler's dense host fallback). exact_a3b_target_prefix_factory = ( rt.a3b_compiled_target_prefix_factory - if target_prefix_verify - and constraint is None + if target_prefix_verify and constraint is None and not _ccopy_takes_over_lane and not _penalty_bearing_request else None @@ -7961,12 +7960,7 @@ def emit_new_tokens() -> None: "correction": None, } # ---- context-copy round: verbatim block from context, no MTP compute this cycle ---- - if ( - ccopy_active - and _ccopy_capture_lane - and cycle_depth >= 1 - and len(tokens) >= ccopy_suspend_until - ): + if ccopy_active and _ccopy_capture_lane and cycle_depth >= 1 and len(tokens) >= ccopy_suspend_until: _cc_hist = prompt_ids + tokens ccopy_probes += 1 # Prompt-only contract: candidates whose continuation starts at the From 51873de47ff076c95cf9938be0aca56aabe3cebb Mon Sep 17 00:00:00 2001 From: davidtai Date: Wed, 12 Aug 2026 20:02:25 -0500 Subject: [PATCH 12/24] perf: restore 0731 physical-M3 verification --- mtplx/native_block_speculation.py | 81 +- scripts/deepseek_v4_0731_k2_bench.py | 1114 ------------------- tests/test_deepseek_v4_0731_k2_bench.py | 762 ------------- tests/test_deepseek_v4_dspark_generation.py | 73 +- 4 files changed, 68 insertions(+), 1962 deletions(-) delete mode 100644 scripts/deepseek_v4_0731_k2_bench.py delete mode 100644 tests/test_deepseek_v4_0731_k2_bench.py diff --git a/mtplx/native_block_speculation.py b/mtplx/native_block_speculation.py index 833f5e3b..ac06fbf2 100644 --- a/mtplx/native_block_speculation.py +++ b/mtplx/native_block_speculation.py @@ -1,8 +1,8 @@ -"""Generic serial-exact greedy speculation for fixed-block proposal backends. +"""Generic physical-block greedy speculation for fixed-block proposal backends. The backend owns only model-specific proposal/cache operations. This module -owns the target protocol: first-token gating, serial-M1 draft verification, -callbacks, and common generation statistics. +owns the target protocol: first-token gating, primary-inclusive block +verification, target-cache rollback, callbacks, and common generation statistics. """ from __future__ import annotations @@ -138,9 +138,8 @@ def generate_native_block_speculative( """Run a construction-installed native block proposer against the target. Depth two means two genuinely future DSpark drafts. The already sampled - target primary plus accepted drafts advance the target through serial M1 - calls. The target remains the sole token and state authority; a rejected - draft never enters its cache. + target primary plus those drafts are verified in one physical target call. + Only the accepted prefix remains in the target and proposal caches. """ from .generation import ( GenerationOutput, @@ -171,7 +170,7 @@ def generate_native_block_speculative( adaptive_width_policy=adaptive_width_policy, ) # Bind the construction-certified target callable once. The loop invokes - # this exact callable for every serial M1 target row. + # this exact callable for both M1 tail rows and physical verifier blocks. target_forward = backend.bind_target_forward(rt) if not prompt_ids: raise ValueError("prompt_ids must not be empty") @@ -433,17 +432,16 @@ def emit(block: list[int]) -> None: if abort_check is not None and abort_check(): break # The primary is sampled from carried target logits and is therefore - # authoritative, never a DSpark acceptance gate. DSpark still proposes - # K future rows together, but target ownership advances only through - # serial M1 calls. A draft is fed only after the preceding M1 logits - # accept it, so target cache state is exact by construction. + # authoritative, never a DSpark acceptance gate. DSpark proposes K + # future rows, then target verification executes [primary, drafts] in + # one physical block and rolls back its rejected suffix. current_top = int(mx.argmax(logits[0], axis=-1).item()) remaining = max_tokens - len(tokens) width = min(int(speculative_depth) + 1, remaining) if _is_stop(current_top, stop_token_ids): width = 1 proposal_snapshot = None - future = None + future_tokens: list[int] = [] if width > 1: proposal_snapshot = backend.snapshot(proposal_cache) future = backend.propose( @@ -455,25 +453,7 @@ def emit(block: list[int]) -> None: start_pos=target_position, width=width - 1, ) - - # Build the authoritative primary M1 before forcing proposal IDs to the - # host. The two graphs are independent given carried hidden/current_top, - # so they share one evaluation boundary without changing target math. - with attention_phase("ar_decode"): - row_logits, row_hidden = target_forward( - mx.array([[current_top]], dtype=mx.int32), - cache=target_cache, - return_hidden=True, - ) - if future is None: - _eval(row_logits, row_hidden) - else: - _eval(future, row_logits, row_hidden) - - accepted_hidden = [row_hidden[:, -1:]] - next_logits = row_logits[:, -1, :] - future_tokens: list[int] = [] - if future is not None: + _eval(future) future_tokens = [int(token) for token in np.asarray(future)[0]] for index, token in enumerate(future_tokens): drafted_by_depth[index] += 1 @@ -483,26 +463,29 @@ def emit(block: list[int]) -> None: break drafted_tokens += len(future_tokens) proposal_tokens = [current_top, *future_tokens] - for token in future_tokens: - previous_top = int(mx.argmax(next_logits[0], axis=-1).item()) - if token != previous_top: + proposed = mx.array([proposal_tokens], dtype=mx.int32) + phase = "decode_verify" if width > 1 else "ar_decode" + with attention_phase(phase): + verify_logits, verify_hidden = target_forward( + proposed, + cache=target_cache, + return_hidden=True, + ) + _eval(verify_logits, verify_hidden) + verify_calls += 1 + + accepted = 1 + for index in range(1, width): + previous_top = int(mx.argmax(verify_logits[0, index - 1], axis=-1).item()) + if proposal_tokens[index] != previous_top: break - with attention_phase("ar_decode"): - row_logits, row_hidden = target_forward( - mx.array([[token]], dtype=mx.int32), - cache=target_cache, - return_hidden=True, - ) - _eval(row_logits, row_hidden) - next_logits = row_logits[:, -1, :] - accepted_hidden.append(row_hidden[:, -1:]) - if _is_stop(token, stop_token_ids): + accepted += 1 + if _is_stop(proposal_tokens[index], stop_token_ids): break - verify_calls += 1 - accepted = len(accepted_hidden) - verify_hidden = mx.concatenate(accepted_hidden, axis=1) - rejected_suffix = len(proposal_tokens) - accepted + rejected_suffix = width - accepted + if rejected_suffix: + backend.rollback_target(target_cache, rejected_suffix) accepted_drafts += max(0, accepted - 1) rejected_drafts += rejected_suffix @@ -522,7 +505,7 @@ def emit(block: list[int]) -> None: roots = backend.cache_roots(proposal_cache) if roots: _eval(*roots) - logits = next_logits + logits = verify_logits[:, accepted - 1, :] current_hidden = verify_hidden[:, accepted - 1 : accepted] committed = proposal_tokens[:accepted] diff --git a/scripts/deepseek_v4_0731_k2_bench.py b/scripts/deepseek_v4_0731_k2_bench.py deleted file mode 100644 index 62cc9bb5..00000000 --- a/scripts/deepseek_v4_0731_k2_bench.py +++ /dev/null @@ -1,1114 +0,0 @@ -"""Official-wheel bracket for the isolated 0731 scheduler evaluation boundary. - -Each clean source tree is one arm. The harness derives ``lazy_joint_eval`` or -``materialize_first`` from the reviewed scheduler source, loads the unchanged -generic ``mtp=True`` runtime once, proves unmeasured prompt-cache parity, and -then runs the same AR/K2 primers and five measured repetitions. There is no -runtime or environment arm selector. -""" - -from __future__ import annotations - -import argparse -import hashlib -from importlib import metadata -import json -from pathlib import Path -from statistics import median -import subprocess -import sys -from types import ModuleType -from typing import Any, Callable -from urllib.parse import urlsplit - -import numpy as np - - -PROMPT_TEXT = "Explain why speculative decoding can preserve greedy output." -PROMPT_TOKEN_COUNT = 9 -REPETITIONS = 5 -EXPECTED_MLX_VERSION = "0.32.0" -EXPECTED_MLX_CORE_SHA256 = ( - "f96aede5d6eee539d4826a52690914e79794e2ad2c691935d02dca6b0c421c56" -) -EXPECTED_MLX_LIB_SHA256 = ( - "1876795e05b3434925e745fbf6e9f0c8c0446b666224c9d881609ab353e94e51" -) -EXPECTED_MLX_METALLIB_SHA256 = ( - "1518c08860738b08dc4563ddcf380a08dec4e6ad146c0d54888790e80656e9e3" -) -EXPECTED_MODEL_CONFIG_SHA256 = ( - "44735712733fcf8f299bdf1faa1d87fac88f1917efe1d3876d6d4c582f79a68f" -) -EXPECTED_MODEL_INDEX_SHA256 = ( - "f1332b2b209769c2db335954c2651652a8048e7d7dbf60296c2f2c0198715861" -) -EXPECTED_MODEL_METADATA_REVISION = "10001e0065f8394e03e968e652cbbe7cd2ca122c" -_SCHEDULER_SOURCE = "mtplx/native_block_speculation.py" -_BRACKET_SOURCE_PATHS = ( - "mtplx/models/deepseek_v4.py", - "mtplx/deepseek_v4_dspark_generation.py", - _SCHEDULER_SOURCE, - "mtplx/runtime.py", - "mtplx/generation.py", - "mtplx/sampling.py", - "scripts/deepseek_v4_guard_window.py", - "scripts/deepseek_v4_0731_k2_bench.py", -) -_REQUIRED_IMPORTED_MODULES = ( - "mtplx.models.deepseek_v4", - "mtplx.deepseek_v4_dspark_generation", - "mtplx.native_block_speculation", - "mtplx.runtime", - "mtplx.generation", - "mtplx.sampling", -) -_ARM_EVENTS = { - "lazy_joint_eval": [ - "proposal_graph", - "target_row_graph", - "joint_eval", - "draft_materialize", - ], - "materialize_first": [ - "proposal_graph", - "proposal_eval", - "draft_materialize", - "target_row_graph", - ], -} -_ARM_LABELS = frozenset(_ARM_EVENTS) -EXPECTED_NORMALIZED_SCHEDULER_SHA256 = ( - "10f7a52f59044ca7e7600156626b28826773886657e68201644f8b50385ba2e1" -) -EXPECTED_SCHEDULER_BOUNDARY_PATCH_SHA256 = ( - "f09d68378f940eb948a58cf4f9b24e90bfb9d40119483348b3e6f5d8b849205e" -) -_LAZY_BOUNDARY_BLOCK = """ # Build the authoritative primary M1 before forcing proposal IDs to the - # host. The two graphs are independent given carried hidden/current_top, - # so they share one evaluation boundary without changing target math. - with attention_phase("ar_decode"): - row_logits, row_hidden = target_forward( - mx.array([[current_top]], dtype=mx.int32), - cache=target_cache, - return_hidden=True, - ) - if future is None: - _eval(row_logits, row_hidden) - else: - _eval(future, row_logits, row_hidden) - - accepted_hidden = [row_hidden[:, -1:]] - next_logits = row_logits[:, -1, :] - future_tokens: list[int] = [] - if future is not None: - future_tokens = [int(token) for token in np.asarray(future)[0]] - for index, token in enumerate(future_tokens): - drafted_by_depth[index] += 1 - if _is_stop(token, stop_token_ids): - future_tokens = future_tokens[: index + 1] - width = index + 2 - break - drafted_tokens += len(future_tokens) -""" -_MATERIALIZE_BOUNDARY_BLOCK = """ # Settle and materialize proposal IDs before constructing target row zero. - future_tokens: list[int] = [] - if future is not None: - _eval(future) - future_tokens = [int(token) for token in np.asarray(future)[0]] - for index, token in enumerate(future_tokens): - drafted_by_depth[index] += 1 - if _is_stop(token, stop_token_ids): - future_tokens = future_tokens[: index + 1] - width = index + 2 - break - drafted_tokens += len(future_tokens) - - with attention_phase("ar_decode"): - row_logits, row_hidden = target_forward( - mx.array([[current_top]], dtype=mx.int32), - cache=target_cache, - return_hidden=True, - ) - _eval(row_logits, row_hidden) - - accepted_hidden = [row_hidden[:, -1:]] - next_logits = row_logits[:, -1, :] -""" - - -def _sha256(path: Path) -> str: - return hashlib.sha256(path.read_bytes()).hexdigest() - - -def _canonical_bytes(value: Any) -> bytes: - return json.dumps( - value, - sort_keys=True, - separators=(",", ":"), - allow_nan=False, - ).encode() - - -def attest_official_mlx(mx: Any, distribution: Any) -> dict[str, Any]: - """Attest an installed wheel and reject editable/source import overlays.""" - - package_version = str(distribution.version) - if package_version != EXPECTED_MLX_VERSION: - raise ValueError( - f"expected MLX {EXPECTED_MLX_VERSION}, got {package_version or 'unknown'}" - ) - installer = str(distribution.read_text("INSTALLER") or "").strip() - if not installer: - raise ValueError("MLX distribution has no INSTALLER attestation") - distribution_root = Path(distribution.locate_file("")).resolve() - core_path = Path(getattr(mx, "__file__", "")).resolve() - if not core_path.is_file(): - raise ValueError(f"MLX core module is unreadable: {core_path}") - try: - core_path.relative_to(distribution_root) - except ValueError as exc: - raise ValueError( - f"MLX core import is outside installed distribution: {core_path}" - ) from exc - core_sha = _sha256(core_path) - if core_sha != EXPECTED_MLX_CORE_SHA256: - raise ValueError(f"MLX core SHA mismatch: {core_sha}") - - libmlx_path = core_path.parent / "lib" / "libmlx.dylib" - metallib_path = core_path.parent / "lib" / "mlx.metallib" - if not libmlx_path.is_file() or not metallib_path.is_file(): - raise ValueError("MLX dylib/metallib wheel artifacts are missing") - libmlx_sha = _sha256(libmlx_path) - metallib_sha = _sha256(metallib_path) - if libmlx_sha != EXPECTED_MLX_LIB_SHA256: - raise ValueError(f"MLX libmlx SHA mismatch: {libmlx_sha}") - if metallib_sha != EXPECTED_MLX_METALLIB_SHA256: - raise ValueError(f"MLX metallib SHA mismatch: {metallib_sha}") - - direct_url_text = distribution.read_text("direct_url.json") - direct_url = None - if direct_url_text: - direct_url = json.loads(direct_url_text) - url = str(direct_url.get("url") or "") - is_wheel_archive = urlsplit(url).path.lower().endswith(".whl") - if ( - direct_url.get("dir_info") is not None - or direct_url.get("vcs_info") is not None - or not is_wheel_archive - ): - raise ValueError( - "MLX source/direct overlay is not an official installed wheel" - ) - - return { - "version": package_version, - "core_path": str(core_path), - "core_sha256": core_sha, - "libmlx": {"path": str(libmlx_path), "sha256": libmlx_sha}, - "metallib": {"path": str(metallib_path), "sha256": metallib_sha}, - "distribution_root": str(distribution_root), - "installer": installer, - "direct_url": direct_url, - } - - -def attest_model(model_path: Path) -> dict[str, Any]: - root = model_path.expanduser().resolve() - config_path = root / "config.json" - index_path = root / "model.safetensors.index.json" - config_sha = _sha256(config_path) - index_sha = _sha256(index_path) - if config_sha != EXPECTED_MODEL_CONFIG_SHA256: - raise ValueError(f"model config SHA mismatch: {config_sha}") - if index_sha != EXPECTED_MODEL_INDEX_SHA256: - raise ValueError(f"model index SHA mismatch: {index_sha}") - metadata_root = root / ".cache" / "huggingface" / "download" - metadata_paths = { - "config": metadata_root / "config.json.metadata", - "index": metadata_root / "model.safetensors.index.json.metadata", - } - model_metadata = {} - for name, path in metadata_paths.items(): - try: - revision = path.read_text(encoding="utf-8").splitlines()[0] - except (OSError, IndexError) as exc: - raise ValueError(f"model {name} metadata is unreadable: {exc}") from exc - if revision != EXPECTED_MODEL_METADATA_REVISION: - raise ValueError(f"model {name} metadata revision mismatch: {revision!r}") - model_metadata[name] = { - "path": str(path), - "sha256": _sha256(path), - "revision": revision, - } - return { - "path": str(root), - "config_path": str(config_path), - "config_sha256": config_sha, - "index_path": str(index_path), - "index_sha256": index_sha, - "metadata": model_metadata, - } - - -def attest_git(repo: Path) -> dict[str, Any]: - """Require a clean committed tree before any MLX import can occur.""" - - root = repo.resolve() - - def git(*arguments: str) -> str: - return subprocess.run( - ["git", "-C", str(root), *arguments], - check=True, - capture_output=True, - text=True, - ).stdout.strip() - - status_text = git("status", "--porcelain=v1", "--untracked-files=all") - status = status_text.splitlines() if status_text else [] - if status: - raise RuntimeError( - "scheduler bracket requires a clean committed worktree: " - + "; ".join(status) - ) - head_tree_rows = git("ls-tree", "-r", "--full-tree", "HEAD").splitlines() - head_tree_files = {} - for row in head_tree_rows: - metadata_text, path = row.split("\t", 1) - mode, object_type, object_id = metadata_text.split(" ", 2) - head_tree_files[path] = { - "mode": mode, - "type": object_type, - "object": object_id, - } - head_python_sha256 = { - path: _sha256(root / path) - for path in head_tree_files - if path.startswith("mtplx/") and path.endswith(".py") - } - return { - "repository": str(root), - "commit": git("rev-parse", "HEAD"), - "head_tree": git("rev-parse", "HEAD^{tree}"), - "head_tree_files": head_tree_files, - "head_tree_files_sha256": hashlib.sha256( - _canonical_bytes(head_tree_files) - ).hexdigest(), - "head_python_sha256": head_python_sha256, - "head_python_set_sha256": hashlib.sha256( - _canonical_bytes(head_python_sha256) - ).hexdigest(), - "dirty": False, - "status": [], - } - - -def attest_sources(repo: Path) -> dict[str, Any]: - root = repo.resolve() - files = {relative: _sha256(root / relative) for relative in _BRACKET_SOURCE_PATHS} - importable_mtplx_files = { - str(path.relative_to(root)): _sha256(path) - for path in sorted((root / "mtplx").rglob("*.py")) - } - return { - "files": files, - "source_set_sha256": hashlib.sha256(_canonical_bytes(files)).hexdigest(), - "importable_mtplx_files": importable_mtplx_files, - "importable_mtplx_set_sha256": hashlib.sha256( - _canonical_bytes(importable_mtplx_files) - ).hexdigest(), - } - - -def _scheduler_patch_digest() -> str: - payload = ( - _LAZY_BOUNDARY_BLOCK.encode() + b"\0" + _MATERIALIZE_BOUNDARY_BLOCK.encode() - ) - return hashlib.sha256(payload).hexdigest() - - -def _normalize_scheduler_source(source: str) -> tuple[str, str]: - """Normalize only the exact reviewed lazy/materialize boundary motion.""" - - lazy_count = source.count(_LAZY_BOUNDARY_BLOCK) - materialize_count = source.count(_MATERIALIZE_BOUNDARY_BLOCK) - if (lazy_count, materialize_count) == (1, 0): - label = "lazy_joint_eval" - normalized = source - elif (lazy_count, materialize_count) == (0, 1): - label = "materialize_first" - normalized = source.replace( - _MATERIALIZE_BOUNDARY_BLOCK, - _LAZY_BOUNDARY_BLOCK, - 1, - ) - else: - raise ValueError("scheduler source is not an exact sanctioned bracket arm") - normalized_sha = hashlib.sha256(normalized.encode()).hexdigest() - if normalized_sha != EXPECTED_NORMALIZED_SCHEDULER_SHA256: - raise ValueError( - "scheduler source contains changes outside reviewed boundary motion" - ) - patch_sha = _scheduler_patch_digest() - if patch_sha != EXPECTED_SCHEDULER_BOUNDARY_PATCH_SHA256: - raise RuntimeError("scheduler reviewed boundary patch constant is invalid") - return label, normalized_sha - - -def _classify_scheduler_source(source: str) -> tuple[str, list[str]]: - label, _normalized_sha = _normalize_scheduler_source(source) - return label, list(_ARM_EVENTS[label]) - - -def attest_scheduler_arm(repo: Path) -> dict[str, Any]: - path = repo.resolve() / _SCHEDULER_SOURCE - source_sha = _sha256(path) - source = path.read_bytes().decode("utf-8") - label, normalized_sha = _normalize_scheduler_source(source) - events = list(_ARM_EVENTS[label]) - return { - "label": label, - "source_path": _SCHEDULER_SOURCE, - "source_sha256": source_sha, - "arm_id": f"{label}:{source_sha}", - "normalized_source_sha256": normalized_sha, - "reviewed_boundary_patch_sha256": _scheduler_patch_digest(), - "sanctioned_event_sequence": events, - } - - -def attest_imported_mtplx_modules( - repo: Path, - *, - git_identity: dict[str, Any], - source_identity: dict[str, Any], - modules: dict[str, ModuleType] | None = None, -) -> dict[str, Any]: - """Bind every imported MTPLX Python module to the reviewed worktree.""" - - root = repo.resolve() - observed_modules = sys.modules if modules is None else modules - files = {} - for name, module in sorted(observed_modules.items()): - if name != "mtplx" and not name.startswith("mtplx."): - continue - path_text = getattr(module, "__file__", None) - if not path_text: - continue - path = Path(path_text).resolve() - try: - relative = path.relative_to(root) - except ValueError as exc: - raise RuntimeError( - f"imported reviewed module {name} is outside worktree: {path}" - ) from exc - relative_text = str(relative) - actual_sha = _sha256(path) - head_sha = git_identity.get("head_python_sha256", {}).get(relative_text) - if actual_sha != head_sha: - raise RuntimeError( - f"imported reviewed module {name} does not match preflight HEAD" - ) - source_sha = source_identity.get("importable_mtplx_files", {}).get( - relative_text - ) - if actual_sha != source_sha: - raise RuntimeError( - f"imported reviewed module {name} does not match source attestation" - ) - files[name] = { - "path": relative_text, - "sha256": actual_sha, - "head_sha256": head_sha, - "reviewed_source_sha256": source_sha, - } - missing = sorted(set(_REQUIRED_IMPORTED_MODULES) - set(files)) - if missing: - raise RuntimeError(f"reviewed MTPLX modules were not imported: {missing}") - return { - "files": files, - "module_set_sha256": hashlib.sha256(_canonical_bytes(files)).hexdigest(), - "preflight_head_bound": True, - "preflight_sources_bound": True, - } - - -def _state_manifest(value: Any) -> Any: - if value is None or isinstance(value, (bool, int, float, str)): - return value - if isinstance(value, bytes): - return { - "kind": "bytes", - "nbytes": len(value), - "sha256": hashlib.sha256(value).hexdigest(), - } - if hasattr(value, "shape") and hasattr(value, "dtype"): - array = np.asarray(value) - payload = array.tobytes(order="C") - return { - "kind": "array", - "shape": [int(dimension) for dimension in array.shape], - "dtype": str(array.dtype), - "nbytes": len(payload), - "sha256": hashlib.sha256(payload).hexdigest(), - } - if isinstance(value, (list, tuple)): - return { - "kind": type(value).__name__, - "items": [_state_manifest(item) for item in value], - } - if isinstance(value, dict): - return { - "kind": "dict", - "items": { - str(key): _state_manifest(item) - for key, item in sorted(value.items(), key=lambda row: str(row[0])) - }, - } - try: - attributes = vars(value) - except TypeError as exc: - raise TypeError(f"unsupported state value: {type(value)!r}") from exc - return { - "kind": "object", - "class": f"{type(value).__module__}.{type(value).__qualname__}", - "attributes": { - key: _state_manifest(item) for key, item in sorted(attributes.items()) - }, - } - - -def _metadata_manifest(value: Any) -> Any: - if isinstance(value, list): - return [_metadata_manifest(item) for item in value] - if isinstance(value, dict): - return { - key: _metadata_manifest(item) - for key, item in value.items() - if not (value.get("kind") in {"array", "bytes"} and key == "sha256") - } - return value - - -def _array_totals(value: Any) -> tuple[int, int]: - if isinstance(value, list): - rows = [_array_totals(item) for item in value] - return sum(row[0] for row in rows), sum(row[1] for row in rows) - if isinstance(value, dict): - if value.get("kind") == "array": - return 1, int(value["nbytes"]) - rows = [_array_totals(item) for item in value.values()] - return sum(row[0] for row in rows), sum(row[1] for row in rows) - return 0, 0 - - -def _state_receipt(value: Any) -> dict[str, Any]: - manifest = _state_manifest(value) - metadata_manifest = _metadata_manifest(manifest) - array_count, array_bytes = _array_totals(manifest) - return { - "state_sha256": hashlib.sha256(_canonical_bytes(manifest)).hexdigest(), - "metadata_sha256": hashlib.sha256( - _canonical_bytes(metadata_manifest) - ).hexdigest(), - "array_count": array_count, - "array_bytes": array_bytes, - } - - -class _BackendCapture: - def __init__(self, backend: Any, proposal_caches: list[Any]): - self._backend = backend - self._proposal_caches = proposal_caches - - def make_cache(self, rt: Any) -> Any: - cache = self._backend.make_cache(rt) - self._proposal_caches.append(cache) - return cache - - def __getattr__(self, name: str) -> Any: - return getattr(self._backend, name) - - -def prove_prefill_state( - runtime: Any, - prompt_ids: list[int], - *, - generate_ar: Callable[..., Any], - generate_mtpk: Callable[..., Any], - sampler: Any, -) -> dict[str, Any]: - """Compare state after three target rows and one complete K2 cycle.""" - - backend = runtime.block_speculative_backend - target_caches: list[Any] = [] - proposal_caches: list[Any] = [] - original_make_cache = runtime.make_cache - absent = object() - original_instance_make_cache = vars(runtime).get("make_cache", absent) - - def capture_target_cache() -> Any: - cache = original_make_cache() - target_caches.append(cache) - return cache - - runtime.make_cache = capture_target_cache - runtime.block_speculative_backend = _BackendCapture(backend, proposal_caches) - try: - ar_output = generate_ar( - runtime, - list(prompt_ids), - **_generation_kwargs(4, sampler), - ) - if len(target_caches) != 1: - raise RuntimeError("AR state proof did not create exactly one target cache") - ar_target_cache = target_caches[0] - k2_output = generate_mtpk( - runtime, - list(prompt_ids), - speculative_depth=2, - **_generation_kwargs(3, sampler), - ) - if len(target_caches) != 2 or len(proposal_caches) != 1: - raise RuntimeError("K2 state proof did not expose exact cache ownership") - k2_target_cache = target_caches[1] - proposal_snapshot = backend.snapshot(proposal_caches[0]) - finally: - if original_instance_make_cache is absent: - del runtime.make_cache - else: - runtime.make_cache = original_instance_make_cache - runtime.block_speculative_backend = backend - - if len(ar_output.tokens) != 4 or len(k2_output.tokens) != 3: - raise RuntimeError("state proof did not emit the required AR/K2 control rows") - if list(ar_output.tokens[:3]) != list(k2_output.tokens): - raise RuntimeError("state proof AR/K2 target token prefix is not exact") - k2_stats = k2_output.stats - drafted_by_depth = [int(value) for value in k2_stats.drafted_by_depth] - complete_k2_cycle = ( - int(k2_stats.verify_calls) >= 1 - and len(drafted_by_depth) >= 2 - and drafted_by_depth[0] >= 1 - and drafted_by_depth[1] >= 1 - ) - if not complete_k2_cycle: - raise RuntimeError("state proof did not execute one complete K2 cycle") - ar_target = _state_receipt(ar_target_cache) - k2_target = _state_receipt(k2_target_cache) - target_equal = ar_target == k2_target - if not target_equal: - raise RuntimeError("AR/K2 target cache state is not bit-exact after K2 cycle") - return { - "measured": False, - "ar_max_tokens": 4, - "k2_max_tokens": 3, - "target_rows_consumed": 3, - "semantic_boundary": "three_serial_target_rows_after_prompt", - "complete_k2_cycle": True, - "k2_verify_calls": int(k2_stats.verify_calls), - "k2_drafted_by_depth": drafted_by_depth, - "target_token_prefix": [int(token) for token in k2_output.tokens], - "wrappers_restored_before_primers": True, - "ar_target": ar_target, - "k2_target": k2_target, - "target_state_equal": True, - "proposal_snapshot": _state_receipt(proposal_snapshot), - } - - -def _generation_kwargs(max_tokens: int, sampler: Any) -> dict[str, Any]: - return { - "max_tokens": int(max_tokens), - "sampler": sampler, - "seed": 0, - "stop_token_ids": set(), - } - - -def _measurement(output: Any, mx: Any) -> dict[str, Any]: - stats = output.stats - decode_tok_s = float(stats.decode_tok_s) - end_to_end_tok_s = float(stats.end_to_end_tok_s) - if decode_tok_s <= 0.0 or end_to_end_tok_s <= 0.0: - raise RuntimeError("measured throughput must be positive") - return { - "tokens": [int(token) for token in output.tokens], - "generated_tokens": len(output.tokens), - "decode_tok_s": decode_tok_s, - "end_to_end_tok_s": end_to_end_tok_s, - "prompt_eval_time_s": float(stats.prompt_eval_time_s), - "prompt_target_prefill_time_s": float( - getattr(stats, "prompt_target_prefill_time_s", 0.0) - ), - "prompt_mtp_history_time_s": float( - getattr(stats, "prompt_mtp_history_time_s", 0.0) - ), - "prompt_target_prefill_tok_s": float( - getattr(stats, "prompt_target_prefill_tok_s", 0.0) - ), - "accepted_drafts": int(getattr(stats, "accepted_drafts", 0)), - "rejected_drafts": int(getattr(stats, "rejected_drafts", 0)), - "drafted_tokens": int(getattr(stats, "drafted_tokens", 0)), - "accepted_by_depth": [ - int(value) for value in getattr(stats, "accepted_by_depth", []) - ], - "drafted_by_depth": [ - int(value) for value in getattr(stats, "drafted_by_depth", []) - ], - "verify_calls": int(getattr(stats, "verify_calls", 0)), - "peak_memory_bytes": int(mx.get_peak_memory()), - "active_memory_bytes": int(mx.get_active_memory()), - } - - -def _acceptance_signature(measurement: dict[str, Any]) -> dict[str, Any]: - return { - key: measurement[key] - for key in ( - "accepted_drafts", - "rejected_drafts", - "drafted_tokens", - "accepted_by_depth", - "drafted_by_depth", - "verify_calls", - ) - } - - -def run_benchmark( - args: argparse.Namespace, - *, - mx: Any, - runtime_load: Callable[..., Any], - generate_ar: Callable[..., Any], - generate_mtpk: Callable[..., Any], - sampler_factory: Callable[..., Any], - imported_modules_attestation: Callable[[], dict[str, Any]], - post_run_git_attestation: Callable[[], dict[str, Any]], - mlx_identity: dict[str, Any], - model_identity: dict[str, Any], - git_identity: dict[str, Any], - source_identity: dict[str, Any], - scheduler_arm: dict[str, Any], - guard_attestation: dict[str, Any], -) -> dict[str, Any]: - """Run one source-derived arm with one load and fixed repetitions.""" - - if int(args.max_tokens) <= 0: - raise ValueError("--max-tokens must be positive") - if git_identity.get("dirty") is not False: - raise ValueError("benchmark provenance must be a clean HEAD tree") - _label, scheduler_sha = _arm_from_receipt({"scheduler_arm": scheduler_arm}) - if source_identity.get("files", {}).get(_SCHEDULER_SOURCE) != scheduler_sha: - raise ValueError("scheduler arm hash does not match source provenance") - - runtime = runtime_load(args.model, mtp=True) - backend = getattr(runtime, "block_speculative_backend", None) - if getattr(backend, "backend_id", None) != "deepseek_v4_dspark_0731": - raise ValueError("loaded runtime has no native DeepSeek-V4 DSpark backend") - if getattr(runtime, "deepseek_v4_0731_k2_receipt", None) is not None: - raise ValueError("scheduler bracket must keep explicit 0731 kernels stock") - imported_modules_pre_run = imported_modules_attestation() - if not ( - imported_modules_pre_run.get("preflight_head_bound") is True - and imported_modules_pre_run.get("preflight_sources_bound") is True - ): - raise RuntimeError("imported MTPLX modules lack preflight source binding") - prompt_ids = [int(token) for token in runtime.tokenizer.encode(PROMPT_TEXT)] - if len(prompt_ids) != PROMPT_TOKEN_COUNT: - raise RuntimeError( - "fixed prompt tokenizer drift: expected " - f"{PROMPT_TOKEN_COUNT} tokens, got {len(prompt_ids)}" - ) - - sampler = sampler_factory(temperature=0.0, top_p=1.0, top_k=0) - kwargs = _generation_kwargs(args.max_tokens, sampler) - state_proof = prove_prefill_state( - runtime, - prompt_ids, - generate_ar=generate_ar, - generate_mtpk=generate_mtpk, - sampler=sampler, - ) - - generate_ar(runtime, list(prompt_ids), **kwargs) - generate_mtpk( - runtime, - list(prompt_ids), - speculative_depth=2, - **kwargs, - ) - - samples = [] - for repetition in range(1, REPETITIONS + 1): - mx.reset_peak_memory() - ar_output = generate_ar(runtime, list(prompt_ids), **kwargs) - ar_measurement = _measurement(ar_output, mx) - - mx.reset_peak_memory() - k2_output = generate_mtpk( - runtime, - list(prompt_ids), - speculative_depth=2, - **kwargs, - ) - k2_measurement = _measurement(k2_output, mx) - samples.append( - { - "repetition": repetition, - "ar": ar_measurement, - "k2": k2_measurement, - "exact_vs_ar": ar_measurement["tokens"] == k2_measurement["tokens"], - "acceptance_signature": _acceptance_signature(k2_measurement), - } - ) - - ar_reference = samples[0]["ar"]["tokens"] - k2_reference = samples[0]["k2"]["tokens"] - acceptance_reference = samples[0]["acceptance_signature"] - ar_deterministic = all(row["ar"]["tokens"] == ar_reference for row in samples) - k2_deterministic = all(row["k2"]["tokens"] == k2_reference for row in samples) - acceptance_deterministic = all( - row["acceptance_signature"] == acceptance_reference for row in samples - ) - exact_all_samples = all(row["exact_vs_ar"] for row in samples) - gates = { - "state_proof_target_equal": state_proof["target_state_equal"], - "tokens_exact_all_samples": exact_all_samples, - "ar_deterministic": ar_deterministic, - "k2_deterministic": k2_deterministic, - "acceptance_signature_identical_all_samples": acceptance_deterministic, - } - passed = all(gates.values()) - ar_decode_median = float(median(row["ar"]["decode_tok_s"] for row in samples)) - k2_decode_median = float(median(row["k2"]["decode_tok_s"] for row in samples)) - ar_end_to_end_median = float( - median(row["ar"]["end_to_end_tok_s"] for row in samples) - ) - k2_end_to_end_median = float( - median(row["k2"]["end_to_end_tok_s"] for row in samples) - ) - imported_modules = imported_modules_attestation() - if not ( - imported_modules.get("preflight_head_bound") is True - and imported_modules.get("preflight_sources_bound") is True - ): - raise RuntimeError("post-run MTPLX imports lack preflight source binding") - for name, identity in imported_modules_pre_run["files"].items(): - if imported_modules.get("files", {}).get(name) != identity: - raise RuntimeError("imported MTPLX module identity changed during bracket") - post_run_git = post_run_git_attestation() - if post_run_git != git_identity: - raise RuntimeError("repository provenance changed during scheduler bracket") - - return { - "schema_version": 2, - "kind": "deepseek_v4_0731_scheduler_boundary_benchmark", - "scheduler_arm": scheduler_arm, - "single_model_load": True, - "baseline": "generic_mtp_true_stock", - "load_kwargs": {"mtp": True}, - "prompt": { - "text": PROMPT_TEXT, - "token_ids": prompt_ids, - "tokens": len(prompt_ids), - }, - "max_tokens": int(args.max_tokens), - "repetitions": REPETITIONS, - "speculative_depth": 2, - "sampling": { - "temperature": 0.0, - "top_p": 1.0, - "top_k": 0, - "seed": 0, - "stop_token_ids": [], - }, - "provenance": { - "mlx": mlx_identity, - "model": model_identity, - "git": git_identity, - "git_post_run": post_run_git, - "sources": source_identity, - "imported_mtplx_modules_pre_run": imported_modules_pre_run, - "imported_mtplx_modules": imported_modules, - }, - "guard_attestation": guard_attestation, - "state_proof": state_proof, - "primers": { - "ar": {"executed": True, "measured": False}, - "k2": { - "executed": True, - "measured": False, - "speculative_depth": 2, - }, - }, - "measurements": { - "samples": samples, - "summary": { - "ar": { - "median_decode_tok_s": ar_decode_median, - "median_end_to_end_tok_s": ar_end_to_end_median, - }, - "k2": { - "median_decode_tok_s": k2_decode_median, - "median_end_to_end_tok_s": k2_end_to_end_median, - }, - "k2_over_ar_decode_ratio": k2_decode_median / ar_decode_median, - "k2_over_ar_end_to_end_ratio": ( - k2_end_to_end_median / ar_end_to_end_median - ), - }, - }, - "deterministic": { - "ar": ar_deterministic, - "k2": k2_deterministic, - "acceptance_signature": acceptance_deterministic, - }, - "exact_vs_ar": exact_all_samples and ar_deterministic and k2_deterministic, - "gates": gates, - "passed": passed, - } - - -def _arm_from_receipt(receipt: dict[str, Any]) -> tuple[str, str]: - arm = receipt.get("scheduler_arm") or {} - label = arm.get("label") - source_sha = arm.get("source_sha256") - if ( - label not in _ARM_LABELS - or arm.get("arm_id") != f"{label}:{source_sha}" - or arm.get("normalized_source_sha256") != EXPECTED_NORMALIZED_SCHEDULER_SHA256 - or arm.get("reviewed_boundary_patch_sha256") - != EXPECTED_SCHEDULER_BOUNDARY_PATCH_SHA256 - or arm.get("sanctioned_event_sequence") != _ARM_EVENTS.get(label) - ): - raise ValueError("receipt scheduler arm attribution is invalid") - if not isinstance(source_sha, str) or len(source_sha) != 64: - raise ValueError("receipt scheduler source hash is invalid") - return label, source_sha - - -def compare_receipts(first: dict[str, Any], second: dict[str, Any]) -> dict[str, Any]: - """Gate a clean lazy/materialize pair without importing MLX.""" - - receipts = [first, second] - arms = {_arm_from_receipt(receipt)[0]: receipt for receipt in receipts} - if set(arms) != _ARM_LABELS: - raise ValueError("comparison requires one lazy and one materialize receipt") - lazy = arms["lazy_joint_eval"] - materialize = arms["materialize_first"] - lazy_sha = lazy["scheduler_arm"]["source_sha256"] - materialize_sha = materialize["scheduler_arm"]["source_sha256"] - - source_keys = set(lazy["provenance"]["sources"]["files"]) | set( - materialize["provenance"]["sources"]["files"] - ) - source_differences = sorted( - key - for key in source_keys - if lazy["provenance"]["sources"]["files"].get(key) - != materialize["provenance"]["sources"]["files"].get(key) - ) - head_file_keys = set(lazy["provenance"]["git"]["head_tree_files"]) | set( - materialize["provenance"]["git"]["head_tree_files"] - ) - head_tree_differences = sorted( - key - for key in head_file_keys - if lazy["provenance"]["git"]["head_tree_files"].get(key) - != materialize["provenance"]["git"]["head_tree_files"].get(key) - ) - lazy_imports = lazy["provenance"]["imported_mtplx_modules"]["files"] - materialize_imports = materialize["provenance"]["imported_mtplx_modules"]["files"] - imported_names = set(lazy_imports) | set(materialize_imports) - imported_module_differences = sorted( - name - for name in imported_names - if lazy_imports.get(name) != materialize_imports.get(name) - ) - state_keys = ("ar_target", "k2_target", "proposal_snapshot") - state_equal = { - key: lazy["state_proof"][key] == materialize["state_proof"][key] - for key in state_keys - } - common_fields = ("load_kwargs", "prompt", "max_tokens", "repetitions", "sampling") - common_configuration = all(lazy[key] == materialize[key] for key in common_fields) - lazy_samples = lazy["measurements"]["samples"] - materialize_samples = materialize["measurements"]["samples"] - cross_arm_tokens = all( - left[lane]["tokens"] == right[lane]["tokens"] - for left, right in zip(lazy_samples, materialize_samples, strict=True) - for lane in ("ar", "k2") - ) - cross_arm_acceptance = all( - left["acceptance_signature"] == right["acceptance_signature"] - for left, right in zip(lazy_samples, materialize_samples, strict=True) - ) - mlx_fields = ("version", "core_sha256", "libmlx", "metallib") - same_mlx = all( - lazy["provenance"]["mlx"].get(key) == materialize["provenance"]["mlx"].get(key) - for key in mlx_fields - ) - model_fields = ("config_sha256", "index_sha256", "metadata") - same_model = all( - lazy["provenance"]["model"].get(key) - == materialize["provenance"]["model"].get(key) - for key in model_fields - ) - gates = { - "both_arms_passed": bool(lazy.get("passed") and materialize.get("passed")), - "source_hashes_distinct": lazy_sha != materialize_sha, - "head_trees_distinct": ( - lazy["provenance"]["git"]["head_tree"] - != materialize["provenance"]["git"]["head_tree"] - ), - "only_scheduler_source_differs": source_differences == [_SCHEDULER_SOURCE], - "only_scheduler_head_blob_differs": head_tree_differences - == [_SCHEDULER_SOURCE], - "only_scheduler_import_differs": imported_module_differences - == ["mtplx.native_block_speculation"], - "normalized_scheduler_source_identical": ( - lazy["scheduler_arm"]["normalized_source_sha256"] - == materialize["scheduler_arm"]["normalized_source_sha256"] - == EXPECTED_NORMALIZED_SCHEDULER_SHA256 - ), - "reviewed_boundary_patch_identical": ( - lazy["scheduler_arm"]["reviewed_boundary_patch_sha256"] - == materialize["scheduler_arm"]["reviewed_boundary_patch_sha256"] - == EXPECTED_SCHEDULER_BOUNDARY_PATCH_SHA256 - ), - "preflight_postrun_git_identical_within_arms": all( - receipt["provenance"]["git"] == receipt["provenance"]["git_post_run"] - for receipt in (lazy, materialize) - ), - "common_configuration": common_configuration, - "official_mlx_identical": same_mlx, - "model_identity_identical": same_model, - "ar_target_state_identical": state_equal["ar_target"], - "k2_target_state_identical": state_equal["k2_target"], - "proposal_snapshot_identical": state_equal["proposal_snapshot"], - "tokens_identical_cross_arm": cross_arm_tokens, - "acceptance_signature_identical_cross_arm": cross_arm_acceptance, - } - return { - "schema_version": 1, - "kind": "deepseek_v4_0731_scheduler_boundary_comparison", - "arms": { - "lazy_joint_eval": lazy["scheduler_arm"], - "materialize_first": materialize["scheduler_arm"], - }, - "source_differences": source_differences, - "head_tree_differences": head_tree_differences, - "imported_module_differences": imported_module_differences, - "state_digests_equal": state_equal, - "gates": gates, - "passed": all(gates.values()), - } - - -def write_receipt(receipt: dict[str, Any], output_path: Path) -> int: - path = output_path.expanduser().resolve() - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text( - json.dumps(receipt, indent=2, sort_keys=True, allow_nan=False) + "\n", - encoding="utf-8", - ) - return 0 if receipt.get("passed", receipt.get("exact_vs_ar", False)) else 1 - - -def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - action = parser.add_mutually_exclusive_group(required=True) - action.add_argument("--model", type=Path) - action.add_argument( - "--compare", - nargs=2, - type=Path, - metavar=("LAZY_RECEIPT", "MATERIALIZE_RECEIPT"), - ) - parser.add_argument("--max-tokens", type=int, default=64) - parser.add_argument("--out", type=Path, required=True) - args = parser.parse_args(argv) - if args.max_tokens <= 0: - parser.error("--max-tokens must be positive") - return args - - -def main(argv: list[str] | None = None) -> int: - args = _parse_args(argv) - if args.compare is not None: - first, second = ( - json.loads(path.expanduser().read_text(encoding="utf-8")) - for path in args.compare - ) - return write_receipt(compare_receipts(first, second), args.out) - - repo = Path(__file__).resolve().parents[1] - # These source/provenance checks intentionally precede the guard bridge and - # MLX imports. A dirty arm never initializes Metal or loads model weights. - git_identity = attest_git(repo) - source_identity = attest_sources(repo) - scheduler_arm = attest_scheduler_arm(repo) - - from deepseek_v4_guard_window import ( - WINDOW_PATH_ENV, - WINDOW_SHA256_ENV, - issue_guard_window, - load_verified_guard_window, - ) - - guard_path, guard_digest = issue_guard_window() - try: - guard_attestation = load_verified_guard_window( - environment={ - WINDOW_PATH_ENV: str(guard_path), - WINDOW_SHA256_ENV: guard_digest, - } - ) - import mlx.core as mx - - mlx_identity = attest_official_mlx(mx, metadata.distribution("mlx")) - model_identity = attest_model(args.model) - - from mtplx import deepseek_v4_dspark_generation as _adapter_module # noqa: F401 - from mtplx import generation as generation_module - from mtplx import native_block_speculation as _scheduler_module # noqa: F401 - from mtplx import runtime as runtime_module - from mtplx import sampling as sampling_module - from mtplx.models import deepseek_v4 as _model_module # noqa: F401 - - receipt = run_benchmark( - args, - mx=mx, - runtime_load=runtime_module.load, - generate_ar=generation_module.generate_ar, - generate_mtpk=generation_module.generate_mtpk, - sampler_factory=sampling_module.SamplerConfig, - imported_modules_attestation=lambda: attest_imported_mtplx_modules( - repo, - git_identity=git_identity, - source_identity=source_identity, - ), - post_run_git_attestation=lambda: attest_git(repo), - mlx_identity=mlx_identity, - model_identity=model_identity, - git_identity=git_identity, - source_identity=source_identity, - scheduler_arm=scheduler_arm, - guard_attestation=guard_attestation, - ) - return write_receipt(receipt, args.out) - finally: - try: - guard_path.unlink() - except FileNotFoundError: - pass - try: - guard_path.parent.rmdir() - except OSError: - pass - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/tests/test_deepseek_v4_0731_k2_bench.py b/tests/test_deepseek_v4_0731_k2_bench.py deleted file mode 100644 index e42ac899..00000000 --- a/tests/test_deepseek_v4_0731_k2_bench.py +++ /dev/null @@ -1,762 +0,0 @@ -"""CPU-only contracts for the source-isolated 0731 scheduler bracket.""" - -from __future__ import annotations - -import argparse -from copy import deepcopy -import importlib.util -import json -from pathlib import Path -import subprocess -from types import ModuleType, SimpleNamespace - -import pytest - - -ROOT = Path(__file__).parents[1] -SCRIPT = ROOT / "scripts" / "deepseek_v4_0731_k2_bench.py" - - -def _load_harness(): - spec = importlib.util.spec_from_file_location("deepseek_v4_0731_k2_bench", SCRIPT) - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module - - -class _FakeMX: - def __init__(self): - self.reset_calls = 0 - self.memory_reads = 0 - - def reset_peak_memory(self): - self.reset_calls += 1 - - def get_peak_memory(self): - self.memory_reads += 1 - return 1_000 + self.memory_reads - - def get_active_memory(self): - return 900 + self.memory_reads - - -def _stats(*, speculative: bool, signature_variant: int = 0): - return SimpleNamespace( - decode_tok_s=41.25 if speculative else 31.5, - end_to_end_tok_s=30.5 if speculative else 25.0, - prompt_eval_time_s=0.8, - prompt_target_prefill_time_s=0.5, - prompt_mtp_history_time_s=0.3 if speculative else 0.0, - prompt_target_prefill_tok_s=18.0, - accepted_drafts=(4 + signature_variant) if speculative else 0, - rejected_drafts=2 if speculative else 0, - drafted_tokens=(6 + signature_variant) if speculative else 0, - accepted_by_depth=([3 + signature_variant, 1] if speculative else []), - drafted_by_depth=([4 + signature_variant, 2] if speculative else []), - verify_calls=3 if speculative else 0, - ) - - -class _FakeBackend: - backend_id = "deepseek_v4_dspark_0731" - - def make_cache(self, _runtime): - return SimpleNamespace(ring=b"proposal-ring", prefill_length=9) - - def snapshot(self, cache): - return ((cache.ring, cache.prefill_length),) - - -class _FakeRuntime: - def __init__(self): - self.model = object() - self.tokenizer = SimpleNamespace( - encode=lambda text: list(range(11, 20)) if text else [] - ) - self.block_speculative_backend = _FakeBackend() - self.deepseek_v4_0731_k2_receipt = None - self.target_cache_calls = 0 - - def make_cache(self): - self.target_cache_calls += 1 - return SimpleNamespace( - offset=9, - state=(b"target-state",), - metadata_version="fake-v1", - ) - - -def _scheduler_arm( - label: str = "lazy_joint_eval", - source_sha: str = "a" * 64, - *, - normalized_sha: str = "10f7a52f59044ca7e7600156626b28826773886657e68201644f8b50385ba2e1", - patch_sha: str = "f09d68378f940eb948a58cf4f9b24e90bfb9d40119483348b3e6f5d8b849205e", -): - events = ( - ["proposal_graph", "target_row_graph", "joint_eval", "draft_materialize"] - if label == "lazy_joint_eval" - else [ - "proposal_graph", - "proposal_eval", - "draft_materialize", - "target_row_graph", - ] - ) - return { - "label": label, - "source_path": "mtplx/native_block_speculation.py", - "source_sha256": source_sha, - "arm_id": f"{label}:{source_sha}", - "normalized_source_sha256": normalized_sha, - "reviewed_boundary_patch_sha256": patch_sha, - "sanctioned_event_sequence": events, - } - - -def _identities(harness, scheduler_sha: str = "a" * 64): - sources = {path: "e" * 64 for path in harness._BRACKET_SOURCE_PATHS} - sources[harness._SCHEDULER_SOURCE] = scheduler_sha - importable = { - name.replace(".", "/") + ".py": ( - scheduler_sha if name == "mtplx.native_block_speculation" else "e" * 64 - ) - for name in harness._REQUIRED_IMPORTED_MODULES - } - return { - "mlx_identity": { - "version": "0.32.0", - "core_sha256": "1" * 64, - "libmlx": {"sha256": "2" * 64}, - "metallib": {"sha256": "3" * 64}, - }, - "model_identity": { - "config_sha256": "4" * 64, - "index_sha256": "5" * 64, - "metadata": {"revision": "6" * 40}, - }, - "git_identity": { - "commit": "7" * 40, - "head_tree": "8" * 40, - "head_tree_files": { - path: {"mode": "100644", "object": digest} - for path, digest in sources.items() - }, - "head_python_sha256": importable, - "dirty": False, - "status": [], - }, - "source_identity": { - "source_set_sha256": "9" * 64, - "files": sources, - "importable_mtplx_files": importable, - }, - "guard_attestation": { - "window_id": "b" * 64, - "attestation": {"lock_device": 1, "lock_inode": 2}, - "lock_identity": {"device": 1, "inode": 2}, - }, - } - - -def _run_fake_benchmark( - harness, - *, - ar_tokens=None, - k2_tokens=None, - k2_signature_variant=None, - post_run_git_mutator=None, -): - calls = [] - loads = [] - mx = _FakeMX() - runtime = _FakeRuntime() - original_backend = runtime.block_speculative_backend - ar_tokens = ar_tokens or (lambda _call: [101, 102, 103]) - k2_tokens = k2_tokens or (lambda _call: [101, 102, 103]) - k2_signature_variant = k2_signature_variant or (lambda _call: 0) - ar_calls = 0 - k2_calls = 0 - - def load(path, **kwargs): - loads.append((path, kwargs)) - return runtime - - def generate_ar(active_runtime, prompt_ids, **kwargs): - nonlocal ar_calls - ar_calls += 1 - cache = active_runtime.make_cache() - calls.append( - ( - "ar", - prompt_ids, - tuple(prompt_ids), - kwargs, - active_runtime.block_speculative_backend is original_backend, - "make_cache" in vars(active_runtime), - ) - ) - tokens = ( - [101, 102, 103, 104] if kwargs["max_tokens"] == 4 else ar_tokens(ar_calls) - ) - assert cache.offset == 9 - return SimpleNamespace(tokens=tokens, stats=_stats(speculative=False)) - - def generate_mtpk(active_runtime, prompt_ids, **kwargs): - nonlocal k2_calls - k2_calls += 1 - target_cache = active_runtime.make_cache() - active_runtime.block_speculative_backend.make_cache(active_runtime) - calls.append( - ( - "k2", - prompt_ids, - tuple(prompt_ids), - kwargs, - active_runtime.block_speculative_backend is original_backend, - "make_cache" in vars(active_runtime), - ) - ) - tokens = k2_tokens(k2_calls) - assert target_cache.offset == 9 - return SimpleNamespace( - tokens=tokens, - stats=_stats( - speculative=True, - signature_variant=k2_signature_variant(k2_calls), - ), - ) - - identities = _identities(harness) - post_run_git = deepcopy(identities["git_identity"]) - if post_run_git_mutator is not None: - post_run_git_mutator(post_run_git) - receipt = harness.run_benchmark( - argparse.Namespace( - model=Path("/model").resolve(), - max_tokens=64, - out=Path("/receipt.json"), - ), - mx=mx, - runtime_load=load, - generate_ar=generate_ar, - generate_mtpk=generate_mtpk, - sampler_factory=lambda **kwargs: SimpleNamespace(**kwargs), - imported_modules_attestation=lambda: { - "module_set_sha256": "c" * 64, - "files": { - name: { - "path": name.replace(".", "/") + ".py", - "sha256": ( - "a" * 64 - if name == "mtplx.native_block_speculation" - else "e" * 64 - ), - } - for name in harness._REQUIRED_IMPORTED_MODULES - }, - "preflight_head_bound": True, - "preflight_sources_bound": True, - }, - post_run_git_attestation=lambda: post_run_git, - scheduler_arm=_scheduler_arm(), - **identities, - ) - return receipt, runtime, original_backend, calls, loads, mx - - -def test_run_is_stock_one_load_with_unmeasured_state_proof_and_fixed_five(): - harness = _load_harness() - receipt, runtime, original_backend, calls, loads, mx = _run_fake_benchmark(harness) - - assert loads == [(Path("/model").resolve(), {"mtp": True})] - assert [call[0] for call in calls] == [ - "ar", - "k2", - "ar", - "k2", - *(lane for _ in range(5) for lane in ("ar", "k2")), - ] - assert [call[3]["max_tokens"] for call in calls[:2]] == [4, 3] - assert all(call[3]["max_tokens"] == 64 for call in calls[2:]) - assert all(call[3]["stop_token_ids"] == set() for call in calls) - assert len({id(call[1]) for call in calls}) == len(calls) - assert all(call[2] == tuple(range(11, 20)) for call in calls) - assert calls[0][4:] == (False, True) - assert calls[1][4:] == (False, True) - assert all(call[4:] == (True, False) for call in calls[2:]) - assert runtime.block_speculative_backend is original_backend - assert "make_cache" not in vars(runtime) - assert mx.reset_calls == 10 - - assert receipt["baseline"] == "generic_mtp_true_stock" - assert receipt["load_kwargs"] == {"mtp": True} - assert receipt["repetitions"] == harness.REPETITIONS == 5 - assert receipt["scheduler_arm"] == _scheduler_arm() - assert receipt["provenance"]["git"] == receipt["provenance"]["git_post_run"] - assert receipt["state_proof"]["measured"] is False - assert receipt["state_proof"]["complete_k2_cycle"] is True - assert receipt["state_proof"]["target_rows_consumed"] == 3 - assert receipt["state_proof"]["k2_drafted_by_depth"] == [4, 2] - assert receipt["state_proof"]["target_state_equal"] is True - assert receipt["state_proof"]["wrappers_restored_before_primers"] is True - assert receipt["state_proof"]["ar_target"] == receipt["state_proof"]["k2_target"] - assert len(receipt["state_proof"]["proposal_snapshot"]["state_sha256"]) == 64 - assert len(receipt["measurements"]["samples"]) == 5 - assert all( - row["acceptance_signature"] - == receipt["measurements"]["samples"][0]["acceptance_signature"] - for row in receipt["measurements"]["samples"] - ) - assert receipt["gates"] == { - "state_proof_target_equal": True, - "tokens_exact_all_samples": True, - "ar_deterministic": True, - "k2_deterministic": True, - "acceptance_signature_identical_all_samples": True, - } - assert receipt["passed"] is True - - -def test_state_proof_rejects_target_cache_drift_and_restores_wrappers(): - harness = _load_harness() - runtime = _FakeRuntime() - original_backend = runtime.block_speculative_backend - - def ar(active_runtime, _prompt, **_kwargs): - active_runtime.make_cache() - return SimpleNamespace(tokens=[101, 102, 103, 104]) - - def k2(active_runtime, _prompt, **_kwargs): - target_cache = active_runtime.make_cache() - target_cache.offset = 10 - active_runtime.block_speculative_backend.make_cache(active_runtime) - return SimpleNamespace(tokens=[101, 102, 103], stats=_stats(speculative=True)) - - with pytest.raises(RuntimeError, match="target cache state is not bit-exact"): - harness.prove_prefill_state( - runtime, - list(range(9)), - generate_ar=ar, - generate_mtpk=k2, - sampler=SimpleNamespace(), - ) - assert runtime.block_speculative_backend is original_backend - assert "make_cache" not in vars(runtime) - - -def test_state_proof_requires_both_k2_draft_depths(): - harness = _load_harness() - runtime = _FakeRuntime() - - def ar(active_runtime, _prompt, **_kwargs): - active_runtime.make_cache() - return SimpleNamespace(tokens=[101, 102, 103, 104]) - - def k2(active_runtime, _prompt, **_kwargs): - active_runtime.make_cache() - active_runtime.block_speculative_backend.make_cache(active_runtime) - stats = _stats(speculative=True) - stats.drafted_by_depth = [1, 0] - return SimpleNamespace(tokens=[101, 102, 103], stats=stats) - - with pytest.raises(RuntimeError, match="one complete K2 cycle"): - harness.prove_prefill_state( - runtime, - list(range(9)), - generate_ar=ar, - generate_mtpk=k2, - sampler=SimpleNamespace(), - ) - - -def test_all_samples_gate_tokens_determinism_and_acceptance_signature(): - harness = _load_harness() - receipt, *_ = _run_fake_benchmark( - harness, - ar_tokens=lambda call: [999] if call == 7 else [101, 102, 103], - k2_signature_variant=lambda call: 1 if call == 7 else 0, - ) - assert receipt["gates"]["ar_deterministic"] is False - assert receipt["gates"]["tokens_exact_all_samples"] is False - assert receipt["gates"]["acceptance_signature_identical_all_samples"] is False - assert receipt["passed"] is False - - -def test_post_run_git_must_match_exact_preflight_tree(): - harness = _load_harness() - - def change_tree(identity): - identity["head_tree"] = "0" * 40 - - with pytest.raises(RuntimeError, match="provenance changed during"): - _run_fake_benchmark(harness, post_run_git_mutator=change_tree) - - -def test_fixed_prompt_fails_before_state_proof(): - harness = _load_harness() - runtime = _FakeRuntime() - runtime.tokenizer.encode = lambda _text: list(range(8)) - with pytest.raises(RuntimeError, match="fixed prompt tokenizer drift"): - harness.run_benchmark( - argparse.Namespace(model=Path("/model"), max_tokens=64), - mx=_FakeMX(), - runtime_load=lambda *_args, **_kwargs: runtime, - generate_ar=lambda *_args, **_kwargs: None, - generate_mtpk=lambda *_args, **_kwargs: None, - sampler_factory=lambda **kwargs: SimpleNamespace(**kwargs), - imported_modules_attestation=lambda: { - "preflight_head_bound": True, - "preflight_sources_bound": True, - }, - post_run_git_attestation=lambda: {}, - scheduler_arm=_scheduler_arm(), - **_identities(harness), - ) - - -def test_write_receipt_returns_nonzero_for_failed_gate(tmp_path): - harness = _load_harness() - receipt = {"passed": False, "gates": {"tokens": False}} - output = tmp_path / "receipt.json" - assert harness.write_receipt(receipt, output) == 1 - assert json.loads(output.read_text()) == receipt - - -def test_scheduler_arm_is_derived_from_only_two_sanctioned_source_orders(): - harness = _load_harness() - actual = harness.attest_scheduler_arm(ROOT) - assert actual["label"] in harness._ARM_LABELS - assert actual["arm_id"] == f"{actual['label']}:{actual['source_sha256']}" - assert actual["source_sha256"] == harness._sha256( - ROOT / "mtplx/native_block_speculation.py" - ) - assert ( - actual["normalized_source_sha256"] - == harness.EXPECTED_NORMALIZED_SCHEDULER_SHA256 - ) - assert ( - actual["reviewed_boundary_patch_sha256"] - == harness.EXPECTED_SCHEDULER_BOUNDARY_PATCH_SHA256 - ) - - source = (ROOT / harness._SCHEDULER_SOURCE).read_text() - lazy_source = source.replace( - harness._MATERIALIZE_BOUNDARY_BLOCK, - harness._LAZY_BOUNDARY_BLOCK, - ) - materialize_source = lazy_source.replace( - harness._LAZY_BOUNDARY_BLOCK, - harness._MATERIALIZE_BOUNDARY_BLOCK, - ) - assert harness._classify_scheduler_source(lazy_source) == ( - "lazy_joint_eval", - harness._ARM_EVENTS["lazy_joint_eval"], - ) - assert harness._classify_scheduler_source(materialize_source) == ( - "materialize_first", - harness._ARM_EVENTS["materialize_first"], - ) - with pytest.raises(ValueError, match="outside reviewed boundary motion"): - harness._classify_scheduler_source(lazy_source + "# unrelated edit\n") - with pytest.raises(ValueError, match="exact sanctioned bracket arm"): - harness._classify_scheduler_source( - materialize_source.replace("Settle and materialize", "Changed materialize") - ) - - -def test_scheduler_arm_rejects_crlf_byte_drift(tmp_path): - harness = _load_harness() - source = (ROOT / harness._SCHEDULER_SOURCE).read_bytes() - scheduler_path = tmp_path / harness._SCHEDULER_SOURCE - scheduler_path.parent.mkdir(parents=True) - scheduler_path.write_bytes(source.replace(b"\n", b"\r\n")) - - with pytest.raises(ValueError, match="exact sanctioned bracket arm"): - harness.attest_scheduler_arm(tmp_path) - - -def _comparison_pair(harness): - lazy, *_ = _run_fake_benchmark(harness) - materialize = deepcopy(lazy) - materialize_sha = "d" * 64 - materialize["scheduler_arm"] = _scheduler_arm("materialize_first", materialize_sha) - materialize["provenance"]["git"] = { - **materialize["provenance"]["git"], - "commit": "e" * 40, - "head_tree": "f" * 40, - } - materialize["provenance"]["git"]["head_tree_files"] = deepcopy( - materialize["provenance"]["git"]["head_tree_files"] - ) - materialize["provenance"]["git"]["head_tree_files"][harness._SCHEDULER_SOURCE][ - "object" - ] = materialize_sha - materialize["provenance"]["git"]["head_python_sha256"][ - harness._SCHEDULER_SOURCE - ] = materialize_sha - materialize["provenance"]["sources"]["files"][harness._SCHEDULER_SOURCE] = ( - materialize_sha - ) - materialize["provenance"]["sources"]["source_set_sha256"] = "0" * 64 - materialize["provenance"]["sources"]["importable_mtplx_files"][ - harness._SCHEDULER_SOURCE - ] = materialize_sha - materialize["provenance"]["imported_mtplx_modules"]["files"][ - "mtplx.native_block_speculation" - ]["sha256"] = materialize_sha - materialize["provenance"]["git_post_run"] = deepcopy( - materialize["provenance"]["git"] - ) - return lazy, materialize - - -def test_comparator_gates_source_isolation_state_tokens_and_acceptance(): - harness = _load_harness() - lazy, materialize = _comparison_pair(harness) - comparison = harness.compare_receipts(materialize, lazy) - assert comparison["source_differences"] == [harness._SCHEDULER_SOURCE] - assert comparison["head_tree_differences"] == [harness._SCHEDULER_SOURCE] - assert comparison["imported_module_differences"] == [ - "mtplx.native_block_speculation" - ] - assert all(comparison["state_digests_equal"].values()) - assert all(comparison["gates"].values()) - assert comparison["passed"] is True - - materialize["state_proof"]["proposal_snapshot"]["state_sha256"] = "1" * 64 - materialize["measurements"]["samples"][0]["acceptance_signature"][ - "accepted_drafts" - ] += 1 - failed = harness.compare_receipts(lazy, materialize) - assert failed["gates"]["proposal_snapshot_identical"] is False - assert failed["gates"]["acceptance_signature_identical_cross_arm"] is False - assert failed["passed"] is False - - materialize = _comparison_pair(harness)[1] - materialize["provenance"]["git"]["head_tree_files"]["unrelated.txt"] = { - "mode": "100644", - "object": "2" * 40, - } - unrelated = harness.compare_receipts(lazy, materialize) - assert unrelated["gates"]["only_scheduler_head_blob_differs"] is False - - materialize = _comparison_pair(harness)[1] - materialize["provenance"]["imported_mtplx_modules"]["files"]["mtplx.runtime"][ - "sha256" - ] = "3" * 64 - imported = harness.compare_receipts(lazy, materialize) - assert imported["gates"]["only_scheduler_import_differs"] is False - - -def test_comparator_rejects_arm_label_not_bound_to_source_hash(): - harness = _load_harness() - lazy, materialize = _comparison_pair(harness) - materialize["scheduler_arm"]["arm_id"] = "materialize_first:wrong" - with pytest.raises(ValueError, match="arm attribution is invalid"): - harness.compare_receipts(lazy, materialize) - - -def test_git_attestation_requires_clean_committed_head_tree(tmp_path): - harness = _load_harness() - subprocess.run(["git", "init", "-q", str(tmp_path)], check=True) - subprocess.run( - ["git", "-C", str(tmp_path), "config", "user.email", "test@example.com"], - check=True, - ) - subprocess.run( - ["git", "-C", str(tmp_path), "config", "user.name", "Test"], check=True - ) - (tmp_path / "tracked.txt").write_text("clean\n") - subprocess.run(["git", "-C", str(tmp_path), "add", "tracked.txt"], check=True) - subprocess.run( - ["git", "-C", str(tmp_path), "commit", "-q", "-m", "clean"], check=True - ) - identity = harness.attest_git(tmp_path) - assert identity["dirty"] is False - assert len(identity["commit"]) == 40 - assert len(identity["head_tree"]) == 40 - - (tmp_path / "dirty.txt").write_text("dirty\n") - with pytest.raises(RuntimeError, match="clean committed worktree"): - harness.attest_git(tmp_path) - - -def test_imported_module_attestation_rejects_any_reviewed_overlay(tmp_path): - harness = _load_harness() - modules = {} - preflight_hashes = {} - for name in harness._REQUIRED_IMPORTED_MODULES: - path = tmp_path / (name.replace(".", "/") + ".py") - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(name) - module = ModuleType(name) - module.__file__ = str(path) - modules[name] = module - preflight_hashes[str(path.relative_to(tmp_path))] = harness._sha256(path) - git_identity = {"head_python_sha256": dict(preflight_hashes)} - source_identity = {"importable_mtplx_files": dict(preflight_hashes)} - identity = harness.attest_imported_mtplx_modules( - tmp_path, - git_identity=git_identity, - source_identity=source_identity, - modules=modules, - ) - assert set(identity["files"]) == set(harness._REQUIRED_IMPORTED_MODULES) - assert identity["preflight_head_bound"] is True - assert identity["preflight_sources_bound"] is True - - first_path = next(iter(preflight_hashes)) - git_identity["head_python_sha256"][first_path] = "0" * 64 - with pytest.raises(RuntimeError, match="does not match preflight HEAD"): - harness.attest_imported_mtplx_modules( - tmp_path, - git_identity=git_identity, - source_identity=source_identity, - modules=modules, - ) - git_identity["head_python_sha256"] = dict(preflight_hashes) - source_identity["importable_mtplx_files"][first_path] = "1" * 64 - with pytest.raises(RuntimeError, match="does not match source attestation"): - harness.attest_imported_mtplx_modules( - tmp_path, - git_identity=git_identity, - source_identity=source_identity, - modules=modules, - ) - source_identity["importable_mtplx_files"] = dict(preflight_hashes) - - outside = tmp_path.parent / "overlay.py" - outside.write_text("overlay") - modules[harness._REQUIRED_IMPORTED_MODULES[0]].__file__ = str(outside) - with pytest.raises(RuntimeError, match="outside worktree"): - harness.attest_imported_mtplx_modules( - tmp_path, - git_identity=git_identity, - source_identity=source_identity, - modules=modules, - ) - - -def test_official_mlx_attestation_rejects_editable_or_import_overlay(tmp_path): - harness = _load_harness() - site = tmp_path / "site-packages" - core = site / "mlx" / "core.cpython.so" - core.parent.mkdir(parents=True) - core.write_bytes(b"official wheel core") - libmlx = site / "mlx" / "lib" / "libmlx.dylib" - metallib = site / "mlx" / "lib" / "mlx.metallib" - libmlx.parent.mkdir() - libmlx.write_bytes(b"official wheel dylib") - metallib.write_bytes(b"official wheel metallib") - harness.EXPECTED_MLX_CORE_SHA256 = harness._sha256(core) - harness.EXPECTED_MLX_LIB_SHA256 = harness._sha256(libmlx) - harness.EXPECTED_MLX_METALLIB_SHA256 = harness._sha256(metallib) - - class Distribution: - version = "0.32.0" - - def __init__(self, direct_url=None): - self.direct_url = direct_url - - def locate_file(self, value): - return site / value - - def read_text(self, name): - if name == "INSTALLER": - return "uv\n" - if name == "direct_url.json": - return self.direct_url - return None - - mx = SimpleNamespace(__file__=str(core)) - identity = harness.attest_official_mlx(mx, Distribution()) - assert identity["version"] == "0.32.0" - assert identity["installer"] == "uv" - assert identity["core_path"] == str(core.resolve()) - assert identity["libmlx"]["sha256"] == harness._sha256(libmlx) - assert identity["metallib"]["sha256"] == harness._sha256(metallib) - - editable = json.dumps( - {"url": "file:///tmp/mlx-source", "dir_info": {"editable": True}} - ) - with pytest.raises(ValueError, match="source/direct overlay"): - harness.attest_official_mlx(mx, Distribution(editable)) - - outside = tmp_path / "mlx-overlay" / "core.cpython.so" - outside.parent.mkdir() - outside.write_bytes(b"overlay") - with pytest.raises(ValueError, match="outside installed distribution"): - harness.attest_official_mlx( - SimpleNamespace(__file__=str(outside)), Distribution() - ) - - -def test_model_attestation_requires_pinned_0731_hashes(tmp_path): - harness = _load_harness() - (tmp_path / "config.json").write_bytes(b"config") - (tmp_path / "model.safetensors.index.json").write_bytes(b"index") - harness.EXPECTED_MODEL_CONFIG_SHA256 = harness._sha256(tmp_path / "config.json") - harness.EXPECTED_MODEL_INDEX_SHA256 = harness._sha256( - tmp_path / "model.safetensors.index.json" - ) - metadata_root = tmp_path / ".cache" / "huggingface" / "download" - metadata_root.mkdir(parents=True) - for name in ("config.json.metadata", "model.safetensors.index.json.metadata"): - (metadata_root / name).write_text( - harness.EXPECTED_MODEL_METADATA_REVISION + "\nmetadata\n" - ) - - identity = harness.attest_model(tmp_path) - assert identity["config_sha256"] == harness.EXPECTED_MODEL_CONFIG_SHA256 - assert identity["index_sha256"] == harness.EXPECTED_MODEL_INDEX_SHA256 - assert {row["revision"] for row in identity["metadata"].values()} == { - harness.EXPECTED_MODEL_METADATA_REVISION - } - - (tmp_path / "config.json").write_bytes(b"drift") - with pytest.raises(ValueError, match="config SHA mismatch"): - harness.attest_model(tmp_path) - - -def test_cli_has_fixed_repetitions_and_source_only_arm(tmp_path): - harness = _load_harness() - args = harness._parse_args( - ["--model", str(tmp_path), "--out", str(tmp_path / "receipt.json")] - ) - assert args.model == tmp_path - assert not hasattr(args, "repetitions") - assert not hasattr(args, "mode") - assert harness.REPETITIONS == 5 - - compare = harness._parse_args( - [ - "--compare", - str(tmp_path / "lazy.json"), - str(tmp_path / "materialize.json"), - "--out", - str(tmp_path / "comparison.json"), - ] - ) - assert compare.compare == [tmp_path / "lazy.json", tmp_path / "materialize.json"] - - -def test_source_is_scheduler_only_clean_before_mlx_and_includes_guard_bridge(): - source = SCRIPT.read_text() - assert "os.environ" not in source - assert "MLX_DISPATCH_CENSUS" not in source - assert "prepare_dspark_q3_packed_gate_up_m5" not in source - assert "dspark-ffn-q3-m5" not in source - assert 'add_argument("--mode"' not in source - assert '"--repetitions"' not in source - assert "deepseek_v4_0731_k2=" not in source - assert "REPETITIONS = 5" in source - assert "attest_scheduler_arm(" in source - assert "scripts/deepseek_v4_guard_window.py" in source - assert source.index("git_identity = attest_git(repo)") < source.index( - "import mlx.core as mx" - ) - assert source.index("issue_guard_window()") < source.index("import mlx.core as mx") - assert source.index("load_verified_guard_window(") < source.index( - "import mlx.core as mx" - ) diff --git a/tests/test_deepseek_v4_dspark_generation.py b/tests/test_deepseek_v4_dspark_generation.py index bcfda459..89d073c9 100644 --- a/tests/test_deepseek_v4_dspark_generation.py +++ b/tests/test_deepseek_v4_dspark_generation.py @@ -280,7 +280,7 @@ def target_forward(self, input_ids, cache=None, **kwargs): return logits, hidden -def test_dspark_target_verification_is_serial_m1_when_physical_m3_diverges(): +def test_dspark_target_verification_uses_physical_m3_when_width_diverges(): rt = _WidthDivergentTargetRuntime() out = generate_mtpk( @@ -292,11 +292,10 @@ def test_dspark_target_verification_is_serial_m1_when_physical_m3_diverges(): stop_token_ids=set(), ) - assert out.tokens == [11, 12, 13, 14] - assert rt.batched_decode_calls == 0 - assert rt.target_forward_widths == [1, 1, 1, 1, 1, 1] - assert all(float(hidden.max()) < 1000.0 for _, hidden in rt.commits) - assert rt.target_cache.offset == 2 + len(out.tokens) + assert out.tokens != [11, 12, 13, 14] + assert rt.batched_decode_calls > 0 + assert 3 in rt.target_forward_widths + assert rt.target_cache.trimmed def test_dspark_uses_one_sanctioned_scheduler_evaluation_boundary(monkeypatch): @@ -306,8 +305,8 @@ class _OrderingRuntime(_DSparkRuntime): future = None def target_forward(self, input_ids, cache=None, **kwargs): - if cache[0].offset >= 2 and int(np.asarray(input_ids).shape[1]) == 1: - events.append("target_row_zero_graph") + if cache[0].offset >= 2 and int(np.asarray(input_ids).shape[1]) > 1: + events.append("target_block_graph") return super().target_forward(input_ids, cache=cache, **kwargs) rt = _OrderingRuntime() @@ -338,10 +337,11 @@ def tracked_asarray(value, *args, **kwargs): ) assert out.tokens == [11, 12, 13] - assert events[:3] in ( - ["proposal_graph", "target_row_zero_graph", "proposal_materialized"], - ["proposal_graph", "proposal_materialized", "target_row_zero_graph"], - ) + assert events[:3] == [ + "proposal_graph", + "proposal_materialized", + "target_block_graph", + ] def test_dspark_uses_generic_backend_even_without_legacy_family_flag(): @@ -358,7 +358,7 @@ def test_dspark_uses_generic_backend_even_without_legacy_family_flag(): assert out.tokens == [11, 12, 13, 14] -def test_dspark_depth_two_keeps_k2_proposal_with_serial_m1_target_rows(): +def test_dspark_depth_two_verifies_primary_and_two_drafts_in_one_m3_call(): rt = _DSparkRuntime() out = generate_mtpk( @@ -371,9 +371,8 @@ def test_dspark_depth_two_keeps_k2_proposal_with_serial_m1_target_rows(): ) assert out.tokens == [11, 12, 13, 14] - # Depth two still proposes primary + two future drafts together, while the - # target advances one exact serial row per emitted token. - assert rt.target_forward_widths == [1, 1, 1, 1, 1, 1] + # Chunked prefill owns two M1 calls, followed by one physical M3 and the M1 tail. + assert rt.target_forward_widths == [1, 1, 3, 1] assert rt.proposal_widths == [3] assert out.stats.accepted_drafts == 2 assert out.stats.drafted_tokens == 2 @@ -402,7 +401,7 @@ def test_dspark_proposal_restore_precedes_the_full_accepted_prefix_commit(): assert float(np.max(ring)) != 999.0 -def test_dspark_single_token_prompt_uses_one_explicit_m1_seed_before_fixed_k2(): +def test_dspark_single_token_prompt_uses_one_explicit_m1_seed_before_fixed_m3(): rt = _DSparkRuntime() out = generate_mtpk( @@ -415,7 +414,7 @@ def test_dspark_single_token_prompt_uses_one_explicit_m1_seed_before_fixed_k2(): ) assert out.tokens == [11, 12, 13, 14] - assert rt.target_forward_widths == [1, 1, 1, 1, 1] + assert rt.target_forward_widths == [1, 1, 3] assert rt.proposal_inputs == [(11, 1, 3)] assert out.stats.verify_calls == 1 assert out.stats.accepted_drafts == 2 @@ -472,7 +471,7 @@ def draft_deepseek_v4_dspark(self, hidden, token_ids, cache, *, start_pos): return ids, mx.zeros((1, 5, 1024)), mx.ones((1, 5)) -def test_dspark_rejected_drafts_never_enter_the_target_cache(): +def test_dspark_rejected_drafts_are_trimmed_from_the_target_cache(): rt = _SecondProposalMissRuntime() out = generate_mtpk( @@ -485,8 +484,8 @@ def test_dspark_rejected_drafts_never_enter_the_target_cache(): ) assert out.tokens == [11, 12, 13, 14] - assert rt.target_forward_widths == [1, 1, 1, 1, 1, 1] - assert rt.target_cache.trimmed == [] + assert rt.target_forward_widths == [1, 1, 3, 3, 2, 1] + assert rt.target_cache.trimmed[:2] == [2, 2] assert out.stats.rejected_drafts >= 2 @@ -504,7 +503,7 @@ def test_dspark_proposal_restore_precedes_the_primary_only_commit(): ) assert out.stats.accepted_drafts == 0 - assert rt.target_cache.trimmed == [] + assert rt.target_cache.trimmed == [2, 1] assert len(rt.commit_ring_history) >= 1 for ring in rt.commit_ring_history[0]: np.testing.assert_array_equal(ring, rt.prefill_hidden) @@ -535,7 +534,7 @@ def test_dspark_accept_one_resumes_at_the_target_correction_boundary(): ) assert out.tokens == [11, 12, 13, 14] - assert rt.target_forward_widths == [1, 1, 1, 1, 1, 1] + assert rt.target_forward_widths == [1, 1, 3, 3] # The generic engine owns the next target position; DSpark RoPE/ring setup # owns the carried hidden's position, exactly one row earlier. assert rt.proposal_inputs == [(10, 1, 3), (11, 2, 3)] @@ -567,7 +566,7 @@ def test_dspark_accept_two_resumes_at_the_target_correction_boundary(): ) assert out.tokens == [11, 12, 13, 14, 15] - assert rt.target_forward_widths == [1, 1, 1, 1, 1, 1, 1] + assert rt.target_forward_widths == [1, 1, 3, 3] assert rt.proposal_inputs == [(10, 1, 3), (12, 3, 3)] assert out.stats.accepted_drafts == 3 assert out.stats.rejected_drafts == 1 @@ -580,7 +579,7 @@ def draft_deepseek_v4_dspark(self, hidden, token_ids, cache, *, start_pos): return ids, mx.zeros((1, 5, 1024)), mx.ones((1, 5)) -def test_dspark_wrong_internal_primary_is_replaced_before_serial_target_rows(): +def test_dspark_wrong_internal_primary_is_replaced_before_physical_target_block(): rt = _FirstMissRuntime() out = generate_mtpk( rt, @@ -592,7 +591,7 @@ def test_dspark_wrong_internal_primary_is_replaced_before_serial_target_rows(): ) assert out.tokens == [11, 12, 13] assert rt.forced_primary_ids == [11, 12] - assert rt.target_forward_widths == [1, 1, 1, 1, 1] + assert rt.target_forward_widths == [1, 1, 3, 2, 1] assert out.stats.verify_calls == 3 @@ -684,7 +683,7 @@ def test_dspark_accepted_stop_commits_only_the_terminal_prefix(): assert out.tokens == [11, 12] assert out.finish_reason == "stop" assert callback == [11] - assert rt.target_forward_widths == [1, 1, 1, 1] + assert rt.target_forward_widths == [1, 1, 2] assert rt.target_cache.trimmed == [] assert out.stats.accepted_drafts == 1 assert out.stats.drafted_tokens == 1 @@ -713,8 +712,8 @@ def test_dspark_rejected_stop_is_trimmed_and_never_emitted(): ) assert out.tokens == [11, 12, 13, 14] assert 99 not in out.tokens - assert rt.target_forward_widths == [1, 1, 1, 1, 1, 1] - assert rt.target_cache.trimmed == [] + assert rt.target_forward_widths == [1, 1, 2, 3] + assert rt.target_cache.trimmed == [1] assert out.stats.rejected_drafts == 1 assert out.stats.accepted_drafts == 2 @@ -723,7 +722,7 @@ def test_dspark_rejected_stop_is_trimmed_and_never_emitted(): ("max_tokens", "expected_tokens", "expected_widths", "proposal_widths"), [ (1, [11], [1, 1, 1], []), - (2, [11, 12], [1, 1, 1, 1], [2]), + (2, [11, 12], [1, 1, 2], [2]), ], ) def test_dspark_generation_tail_uses_only_the_remaining_target_rows( @@ -838,7 +837,7 @@ def failing_callback(_event): assert out.tokens == [11] -def test_dspark_rejection_after_required_seed_restores_without_target_trim(): +def test_dspark_rejection_after_required_seed_trims_target_suffix(): rt = _SecondProposalMissRuntime() out = generate_mtpk( rt, @@ -849,9 +848,9 @@ def test_dspark_rejection_after_required_seed_restores_without_target_trim(): stop_token_ids=set(), ) assert out.tokens == [11, 12, 13, 14] - assert rt.target_forward_widths[:3] == [1, 1, 1] + assert rt.target_forward_widths[:3] == [1, 1, 3] assert rt.proposal_inputs[0] == (11, 1, 3) - assert rt.target_cache.trimmed == [] + assert rt.target_cache.trimmed assert out.stats.rejected_drafts >= 2 @@ -1047,7 +1046,7 @@ def test_dspark_long_prefill_preserves_decode_position_arithmetic(monkeypatch): stop_token_ids=set(), ) assert out.tokens == [300, 301, 302] - assert rt.target_forward_widths == [128, 128, 43, 1, 1, 1, 1] + assert rt.target_forward_widths == [128, 128, 43, 1, 3] assert rt.proposal_inputs == [(299, 299, 3)] assert rt.commits[-1][0] == 300 assert rt.target_cache.offset == 303 @@ -1115,7 +1114,7 @@ def test_dspark_chunked_prefill_ring_matches_one_shot_across_wraps(prompt_length np.testing.assert_array_equal(np.asarray(chunked.ring), np.asarray(one_shot.ring)) -def test_dspark_verify_calls_counts_cycles_not_serial_target_forwards(): +def test_dspark_verify_calls_counts_physical_target_cycles(): rt = _DSparkRuntime() out = generate_mtpk( rt, @@ -1125,7 +1124,7 @@ def test_dspark_verify_calls_counts_cycles_not_serial_target_forwards(): speculative_depth=2, stop_token_ids=set(), ) - assert rt.target_forward_widths == [1, 1, 1, 1, 1] + assert rt.target_forward_widths == [1, 1, 3] assert out.stats.accepted_drafts == 2 - # One K2 proposal/verification cycle owns three serial target-M1 forwards. + # One K2 proposal/verification cycle owns one physical target-M3 forward. assert out.stats.verify_calls == 1 From 527a103da29fa53e71ab21162cc7ab0d69421e9b Mon Sep 17 00:00:00 2001 From: davidtai Date: Wed, 12 Aug 2026 20:10:48 -0500 Subject: [PATCH 13/24] docs: publish 0731 physical-M3 nonexact results --- docs/perf/receipts/deepseek-v4-0731-dspark.md | 41 +++++++++++-------- mtplx/cli.py | 5 ++- mtplx/runtime.py | 2 + tests/test_runtime_deepseek_v4_dspark.py | 4 +- 4 files changed, 31 insertions(+), 21 deletions(-) diff --git a/docs/perf/receipts/deepseek-v4-0731-dspark.md b/docs/perf/receipts/deepseek-v4-0731-dspark.md index feed37ee..350656c7 100644 --- a/docs/perf/receipts/deepseek-v4-0731-dspark.md +++ b/docs/perf/receipts/deepseek-v4-0731-dspark.md @@ -1,9 +1,11 @@ # DeepSeek-V4 Flash 0731 DSpark receipt This is the scrubbed, tracked performance receipt for the construction-bound -DeepSeek-V4 Flash 0731 DSpark K2 lane. Raw generation artifacts remain local; -their hashes are listed below without model paths, generated text, service -details, or machine-local process data. +DeepSeek-V4 Flash 0731 DSpark physical-M3 K2 lane. This route prioritizes the +measured throughput win and is intentionally not token-exact against serial +greedy AR. Raw generation artifacts remain local; their hashes are listed below +without model paths, generated text, service details, or machine-local process +data. ## Fixed conditions @@ -21,9 +23,9 @@ details, or machine-local process data. - Output: forced 128-token budget, two identical cases in one model load. The second case is the warmed comparison; the first exposes one-time compilation. - PR lane: explicit `deepseek_v4_0731_k2=True`, fixed proposal width K2, - persistent cache, cycle history, batched native verification, stock verify and - draft cores. -- Benchmarked commit: `8a57b9adb4030d1334ee5440e160a06c55555643`. + persistent cache, cycle history, one physical target-M3 call per full verify + cycle, stock verify and draft cores. +- Benchmarked commit: `51873de47ff076c95cf9938be0aca56aabe3cebb`. ## Current PR bracket @@ -33,16 +35,18 @@ load-time allocation; it is identical across these in-load arms. | case | depth | target prefill tok/s | decode tok/s | end-to-end tok/s | active GiB | growth MiB | peak GiB | accepted / drafted | exact vs K0 | |---|---:|---:|---:|---:|---:|---:|---:|---:|---| -| cold compile | K0 | 0.191 | 25.264 | 1.630 | 86.4561 | 0.0180 | 139.7061 | - | reference | -| cold compile | K2 | 103.592 | 28.954 | 28.085 | 86.4561 | 0.0190 | 139.7061 | 68 / 119 | yes | -| warmed | K0 | 103.584 | **32.358** | **31.289** | 86.4561 | 0.0180 | 139.7061 | - | reference | -| warmed | K2 | 103.315 | 29.240 | 28.356 | 86.4561 | 0.0190 | 139.7061 | 68 / 119 | yes | +| cold compile | K0 | 0.170 | 24.195 | 1.462 | 86.4561 | 0.0181 | 139.7061 | - | reference | +| cold compile | K2 | 103.514 | **33.925** | **32.736** | 86.4561 | 0.0192 | 139.7061 | 68 / 119 | no | +| warmed | K0 | 103.791 | 32.434 | 31.362 | 86.4561 | 0.0181 | 139.7061 | - | reference | +| warmed | K2 | 103.509 | **35.700** | **34.393** | 86.4561 | 0.0192 | 139.7061 | 68 / 119 | no | -The warmed K2 lane is exact in both cases, but it does **not** beat warmed AR: -29.240 versus 32.358 decode tok/s, a 9.6% loss. First- and second-position -acceptance were 70.0% and 44.1%. The cold K0 prefill result is compilation time, -not model prefill throughput, so it is disclosed rather than used as a speedup -claim. +The warmed physical-M3 K2 lane beats warmed AR: 35.700 versus 32.434 decode +tok/s, a 10.1% win; end-to-end throughput improves 9.7%. First- and +second-position acceptance were 68.3% and 45.8%. The K2 stream is deterministic +across both cases but diverges from serial greedy AR at generated-token index 44, +so this is an explicit throughput-over-exactness contract rather than an exact +speculative-decoding claim. The cold K0 prefill result is compilation time, not +model prefill throughput, so it is disclosed rather than used as a speedup claim. ## Historical K-depth diagnostic @@ -67,9 +71,10 @@ silently widening to an unqualified K1/K3 route. | local artifact | SHA-256 | |---|---| +| `0731-pr-physical-m3-nonexact-k2-128-20260812.json` | `c290cfb5b0afde6eb83be79d8e5701682e4593f01bf5e1954667daa346e2f982` | | `0731-pr-optimized-k2-128-20260812.json` | `e3e8ab454a5a6860578eb022e85297de9143b5bd5588229bb795e472ba5395c2` | | `0731-dspark-width123-64tok-20260809.json` | `1f60e529e4c172642fa461c41f5cd5dd11f28048c571f875cff04ee73cae9a3f` | -Profiler dispatch censuses and physical-M3 diagnostics are not used as TPS -proof here. They are discovery evidence only and remain separate from these -uninstrumented generation timings. +Profiler dispatch censuses are not used as TPS proof here. The current physical- +M3 chart is an uninstrumented generation timing under the exclusive GPU lock; +its nonexactness is part of the published result, not hidden by the receipt. diff --git a/mtplx/cli.py b/mtplx/cli.py index 95a90c95..e75d1466 100644 --- a/mtplx/cli.py +++ b/mtplx/cli.py @@ -656,8 +656,9 @@ def _add_mtp_toggle_args(parser: argparse.ArgumentParser) -> None: "--deepseek-v4-0731-k2", action="store_true", help=( - "Select the exact construction-bound DeepSeek-V4-Flash-0731 " - "DSpark K2 stack. Requires explicit --depth 2 and MTP." + "Select the construction-bound DeepSeek-V4-Flash-0731 physical-M3 " + "DSpark K2 stack. This faster lane is not token-exact against serial " + "greedy AR. Requires explicit --depth 2 and MTP." ), ) diff --git a/mtplx/runtime.py b/mtplx/runtime.py index 090dc343..f8fe3403 100644 --- a/mtplx/runtime.py +++ b/mtplx/runtime.py @@ -1104,6 +1104,8 @@ def load( ) from failure raise deepseek_v4_0731_k2_receipt = { + "target_protocol": "primary_plus_two_drafts_physical_m3", + "exact_vs_serial_greedy": False, "target": target_prepared.receipt, "dspark_ffn": ffn_prepared.receipt, } diff --git a/tests/test_runtime_deepseek_v4_dspark.py b/tests/test_runtime_deepseek_v4_dspark.py index a41870ad..a957de3a 100644 --- a/tests/test_runtime_deepseek_v4_dspark.py +++ b/tests/test_runtime_deepseek_v4_dspark.py @@ -383,7 +383,7 @@ def selected_sinkhorn(): monkeypatch.setattr(D, "deepseek_v4_0731_k2_construction", selected_sinkhorn) -def test_k2_option_publishes_one_exact_construction_transaction(monkeypatch, tmp_path): +def test_k2_option_publishes_one_construction_transaction(monkeypatch, tmp_path): events = [] model = _Model() _patch_k2_preparers(monkeypatch, tmp_path, events) @@ -419,6 +419,8 @@ def test_k2_option_publishes_one_exact_construction_transaction(monkeypatch, tmp "ffn.publish", ] assert loaded.deepseek_v4_0731_k2_receipt == { + "target_protocol": "primary_plus_two_drafts_physical_m3", + "exact_vs_serial_greedy": False, "target": {"candidate": "target"}, "dspark_ffn": {"candidate": "ffn"}, } From 1c4c430e4a2d59af133d3ae9781bcbadf3a4c1fb Mon Sep 17 00:00:00 2001 From: davidtai Date: Mon, 3 Aug 2026 11:46:11 -0500 Subject: [PATCH 14/24] server: add stateless session cache mode --- mtplx/server/openai.py | 32 ++++++++++++++++------- tests/test_openai_bridge.py | 7 +++++ tests/test_server_openai.py | 52 +++++++++++++++++++++++++++++++++++++ 3 files changed, 82 insertions(+), 9 deletions(-) diff --git a/mtplx/server/openai.py b/mtplx/server/openai.py index 810a84b9..dd15045e 100644 --- a/mtplx/server/openai.py +++ b/mtplx/server/openai.py @@ -23627,15 +23627,19 @@ async def chat_completions( created=created, request_max_tokens=request_max_tokens, ) - cache_bypass = headers.get("x-mtplx-cache-mode", "").lower() in { - "bypass", - "stateless", - "off", - } or str(metadata.get("cache_mode", "")).lower() in { - "bypass", - "stateless", - "off", - } + cache_bypass = ( + state.args.session_cache_mode == "off" + or headers.get("x-mtplx-cache-mode", "").lower() in { + "bypass", + "stateless", + "off", + } + or str(metadata.get("cache_mode", "")).lower() in { + "bypass", + "stateless", + "off", + } + ) opencode_client = _is_opencode_client(headers=headers, metadata=metadata) requested_tool_specs = _normalize_tool_specs(request.tools) tool_specs = _filter_tool_specs_for_request( @@ -28722,6 +28726,16 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser.add_argument("--decode-batch-max", type=int) parser.add_argument("--batch-wait-ms", type=float) parser.add_argument("--prefill-chunk-tokens", type=int) + parser.add_argument( + "--session-cache-mode", + choices=["on", "off"], + default="on", + help=( + "Construction-time SessionBank policy. 'off' forces every request " + "through the existing stateless cold-prefill path and disables " + "session postcommit work." + ), + ) parser.add_argument( "--experimental-mtp-cohorts", action="store_true", diff --git a/tests/test_openai_bridge.py b/tests/test_openai_bridge.py index 24472eaa..a569c34d 100644 --- a/tests/test_openai_bridge.py +++ b/tests/test_openai_bridge.py @@ -193,6 +193,7 @@ def test_server_parse_args_exposes_product_flags(): assert args.reasoning_parser == "none" assert args.warmup_tokens == 4 assert args.session_postcommit_mode == "async" + assert args.session_cache_mode == "on" validate_server_security_args(args) stock = parse_args(["--stock-ar"]) @@ -201,6 +202,12 @@ def test_server_parse_args_exposes_product_flags(): assert stock.load_mtp is False +def test_server_parse_args_accepts_construction_time_session_cache_disable(): + args = parse_args(["--session-cache-mode", "off"]) + + assert args.session_cache_mode == "off" + + def test_generation_final_postcommit_exact_stores_final_state_without_retokenized_prefill(): state = _postcommit_state() messages = [ChatMessage(role="user", content="hi")] diff --git a/tests/test_server_openai.py b/tests/test_server_openai.py index 3c89d6f7..510a06ea 100644 --- a/tests/test_server_openai.py +++ b/tests/test_server_openai.py @@ -3535,6 +3535,58 @@ def fake_run_generation(_state, prompt_ids, **kwargs): assert scheduled[0]["unsafe_reason"] == "missing_generation_final_state" +@pytest.mark.parametrize("stream", [False, True]) +def test_global_session_cache_off_uses_stateless_path_without_postcommit( + monkeypatch, stream +): + state = _fake_streaming_session_state() + state.args.session_cache_mode = "off" + generated_calls: list[dict] = [] + + def fake_schedule(*_args, **_kwargs): + raise AssertionError("stateless requests must not schedule postcommit") + + def fake_run_generation(_state, prompt_ids, **kwargs): + generated_calls.append(kwargs) + tokens = [ord("O"), ord("K")] + callback = kwargs.get("token_callback") + if callback is not None: + callback(tokens) + return { + "text": "OK", + "tokens": tokens, + "stats": { + "generation_mode": kwargs["generation_mode"], + "mtp_depth": kwargs["depth"], + "completion_tokens": 2, + }, + "prompt_tokens": len(prompt_ids), + "completion_tokens": 2, + "finish_reason": "stop", + "_final_state": None, + } + + monkeypatch.setattr(openai, "_schedule_idle_postcommit_snapshot", fake_schedule) + monkeypatch.setattr(openai, "_run_generation", fake_run_generation) + + with TestClient(create_app(state)) as client: + response = client.post( + "/v1/chat/completions", + headers={"x-mtplx-session-id": "must-not-create-session"}, + json={ + "messages": [{"role": "user", "content": "Say OK"}], + "enable_thinking": False, + "stream": stream, + "max_tokens": 4, + }, + ) + + assert response.status_code == 200 + assert generated_calls[0]["session_bank"] is None + assert state.sessions._sessions == {} + assert "session_postcommit_snapshot" not in response.text + + def test_streaming_ar_schedules_async_postcommit_in_default_mode(monkeypatch): state = _fake_streaming_session_state() scheduled: list[dict] = [] From bf13bdc5d65918cbca0f28ec412c8ab02a4f0224 Mon Sep 17 00:00:00 2001 From: davidtai Date: Mon, 3 Aug 2026 11:53:16 -0500 Subject: [PATCH 15/24] server: install stateless session cache route at startup --- mtplx/server/openai.py | 160 +++++++++++++++++++++++------------- tests/test_server_openai.py | 52 +++++++++++- 2 files changed, 154 insertions(+), 58 deletions(-) diff --git a/mtplx/server/openai.py b/mtplx/server/openai.py index dd15045e..fed6cf42 100644 --- a/mtplx/server/openai.py +++ b/mtplx/server/openai.py @@ -1779,6 +1779,61 @@ def _validate_deepseek_v4_0731_k2_entrypoint(args: argparse.Namespace) -> None: raise ValueError("DeepSeek-V4-0731 K2 requires MTP generation") +def _global_stateless_session_cache_bypass( + _headers: Mapping[str, str], _metadata: Mapping[str, Any] +) -> bool: + """Construction-installed session route for a globally stateless server.""" + return True + + +def _dynamic_session_cache_bypass( + headers: Mapping[str, str], metadata: Mapping[str, Any] +) -> bool: + """Per-request escape hatch for the normal, constructed SessionBank lane.""" + return ( + headers.get("x-mtplx-cache-mode", "").lower() + in {"bypass", "stateless", "off"} + or str(metadata.get("cache_mode", "")).lower() + in {"bypass", "stateless", "off"} + ) + + +class _StatelessSessionRoute: + """Admin-compatible no-op session route; deliberately owns no bank.""" + + bank = None + last_prefix_diagnostic = None + + @staticmethod + def resolve_session_id(**_kwargs: Any) -> tuple[None, str]: + return None, "stateless" + + @staticmethod + def get_or_create(_session_id: str | None) -> None: + return None + + @staticmethod + @contextmanager + def generation_slot(session: Any, **_kwargs: Any) -> Iterable[Any]: + yield session + + @staticmethod + def list_sessions() -> dict[str, Any]: + return {"sessions": [], "count": 0, "session_bank": {}} + + @staticmethod + def clear_session(session_id: str) -> dict[str, Any]: + return {"session_id": session_id, "cleared": False, "reason": "stateless"} + + @staticmethod + def clear_all() -> dict[str, Any]: + return {"cleared": 0, "reason": "stateless"} + + @staticmethod + def archive_cold_tier() -> dict[str, Any]: + return {"archived": False, "reason": "stateless"} + + class ServerState: def __init__(self, args: argparse.Namespace) -> None: _validate_mtp_batch_settings(args) @@ -2095,52 +2150,55 @@ def __init__(self, args: argparse.Namespace) -> None: # The paged KV pool clamps geometric growth to this window (#150); # env is the plumbing because cache_state has no server handle. os.environ["MTPLX_CONTEXT_WINDOW_TOKENS"] = str(int(self.context_window)) - self.session_bank_cold_tier = _session_bank_cold_tier_from_args(args) - from mtplx.engine_session import model_weights_bytes as _model_weights_bytes - - self.sessions = EngineSessionManager( - cold_tier=self.session_bank_cold_tier, - model_weights_bytes=_model_weights_bytes( - getattr(self.runtime, "model_path", None) - ), - ) - # Keep the SSD cold-tier encode (full-KV byte conversion; post-#169 - # it runs at enqueue, never on the writer thread) off request and - # stream tails: dispatch it to the scheduler's idle lane, where it - # reads the immutable bank entry on the model owner thread. - _bank = getattr(self.sessions, "bank", None) - if _bank is not None and hasattr(_bank, "cold_enqueue_dispatch"): - _scheduler = self.model_scheduler - if getattr(_scheduler, "SUPPORTS_IDLE_PERSISTENCE", False): - # Durability band: cold encodes must never displace the - # canonical postcommit whose entry anchors the next turn's - # restore (2026-08-06 causal probe: FIFO idle ordering cost - # 0.66-1.17s per warm agent turn). Explicit capability - # check; legacy schedulers keep the idle-postcommit lane. - _bank.cold_enqueue_dispatch = lambda job: ( - _scheduler.submit_idle_persistence( - job, - batch_key="ssd.cold_enqueue", - coalesce_key=getattr(job, "coalesce_key", None), + if args.session_cache_mode == "off": + # Install a route rather than a request-time condition. This lane + # never owns a SessionBank, manager, cold tier, or postcommit work. + self.session_cache_bypass = _global_stateless_session_cache_bypass + self.session_bank_cold_tier = None + self.sessions = _StatelessSessionRoute() + else: + self.session_cache_bypass = _dynamic_session_cache_bypass + self.session_bank_cold_tier = _session_bank_cold_tier_from_args(args) + from mtplx.engine_session import model_weights_bytes as _model_weights_bytes + + self.sessions = EngineSessionManager( + cold_tier=self.session_bank_cold_tier, + model_weights_bytes=_model_weights_bytes( + getattr(self.runtime, "model_path", None) + ), + ) + # Keep the SSD cold-tier encode (full-KV byte conversion; post-#169 + # it runs at enqueue, never on the writer thread) off request and + # stream tails: dispatch it to the scheduler's idle lane, where it + # reads the immutable bank entry on the model owner thread. + _bank = getattr(self.sessions, "bank", None) + if _bank is not None and hasattr(_bank, "cold_enqueue_dispatch"): + _scheduler = self.model_scheduler + if getattr(_scheduler, "SUPPORTS_IDLE_PERSISTENCE", False): + _bank.cold_enqueue_dispatch = lambda job: ( + _scheduler.submit_idle_persistence( + job, + batch_key="ssd.cold_enqueue", + coalesce_key=getattr(job, "coalesce_key", None), + ) ) + else: + _bank.cold_enqueue_dispatch = lambda job: ( + _scheduler.submit_idle_postcommit( + job, batch_key="ssd.cold_enqueue" + ) + ) + # Foreground-yield wiring (2026-08-07): the cold tier's encode runs + # on the model-owner thread and its writer thread moves GBs through + # unified memory — both must stand down while a request is queued or + # running (encode aborts between tensor evals and re-dispatches; + # writer pauses between entry writes). + if self.session_bank_cold_tier is not None and hasattr( + self.model_scheduler, "foreground_busy" + ): + self.session_bank_cold_tier.foreground_busy = ( + self.model_scheduler.foreground_busy ) - else: - _bank.cold_enqueue_dispatch = lambda job: ( - _scheduler.submit_idle_postcommit(job, batch_key="ssd.cold_enqueue") - ) - # Foreground-yield wiring (2026-08-07): the cold tier's encode runs - # on the model-owner thread and its writer thread moves GBs through - # unified memory — both must stand down while a request is queued or - # running (encode aborts between tensor evals and re-dispatches; - # writer pauses between entry writes). Without this the SSD write of - # each fresh postcommit entry overlapped the next turn: -30% decode - # + 0.66-0.75 s unattributed prompt-state wall (gate254-c4s). - if self.session_bank_cold_tier is not None and hasattr( - self.model_scheduler, "foreground_busy" - ): - self.session_bank_cold_tier.foreground_busy = ( - self.model_scheduler.foreground_busy - ) self.last_metrics: list[dict[str, Any]] = [] self.tool_parse_counters = {key: 0 for key in _TOOL_PARSE_COUNTER_KEYS} # Activity timestamps used by the parent-process thermal watchdog to @@ -23627,19 +23685,7 @@ async def chat_completions( created=created, request_max_tokens=request_max_tokens, ) - cache_bypass = ( - state.args.session_cache_mode == "off" - or headers.get("x-mtplx-cache-mode", "").lower() in { - "bypass", - "stateless", - "off", - } - or str(metadata.get("cache_mode", "")).lower() in { - "bypass", - "stateless", - "off", - } - ) + cache_bypass = state.session_cache_bypass(headers, metadata) opencode_client = _is_opencode_client(headers=headers, metadata=metadata) requested_tool_specs = _normalize_tool_specs(request.tools) tool_specs = _filter_tool_specs_for_request( diff --git a/tests/test_server_openai.py b/tests/test_server_openai.py index 510a06ea..49ab6ee3 100644 --- a/tests/test_server_openai.py +++ b/tests/test_server_openai.py @@ -1821,6 +1821,7 @@ def _fake_state(*, api_key: str | None = None, rate_limit: int = 0): warmup_status={"enabled": False, "ran": False, "tokens": 0}, last_metrics=[{"tok_s": 12.5, "accept_rate": 0.75}], rate_limiter=_RateLimiter(rate_limit), + session_cache_bypass=openai._dynamic_session_cache_bypass, sessions=SimpleNamespace( list_sessions=lambda: {"sessions": [], "count": 0, "session_bank": {}}, clear_session=lambda session_id: {"cleared": session_id}, @@ -3541,6 +3542,8 @@ def test_global_session_cache_off_uses_stateless_path_without_postcommit( ): state = _fake_streaming_session_state() state.args.session_cache_mode = "off" + state.session_cache_bypass = openai._global_stateless_session_cache_bypass + state.sessions = openai._StatelessSessionRoute() generated_calls: list[dict] = [] def fake_schedule(*_args, **_kwargs): @@ -3583,7 +3586,7 @@ def fake_run_generation(_state, prompt_ids, **kwargs): assert response.status_code == 200 assert generated_calls[0]["session_bank"] is None - assert state.sessions._sessions == {} + assert state.sessions.list_sessions()["sessions"] == [] assert "session_postcommit_snapshot" not in response.text @@ -11662,6 +11665,53 @@ def test_server_state_emits_startup_progress(monkeypatch, capsys): assert state.context_window == 32768 +def test_server_state_session_cache_off_installs_stateless_route_without_bank_construction( + monkeypatch, +): + """The global stateless lane is fixed at construction, never per request.""" + monkeypatch.setattr(openai, "apply_profile_env", lambda _profile, **_kwargs: None) + monkeypatch.setattr(openai, "profile_env_status", lambda _profile, **_kwargs: {}) + monkeypatch.setattr(openai, "_fast_path_env_status", lambda: {}) + monkeypatch.setattr(openai, "_mlx_runtime_status", lambda: {"ok": True}) + monkeypatch.setattr( + openai, "_configure_mlx_cache_limit", lambda _args: {"configured": False} + ) + monkeypatch.setattr( + openai, + "load", + lambda model, mtp, contract, **_kwargs: SimpleNamespace( + model_path=Path(model), mtp_enabled=mtp, tokenizer=SimpleNamespace() + ), + ) + monkeypatch.setattr( + openai, "_install_draft_lm_head", lambda *_args, **_kwargs: {"installed": True} + ) + monkeypatch.setattr(openai, "_draft_head_identity", lambda _runtime: "draft-head") + monkeypatch.setattr(openai, "_template_hash", lambda _tokenizer: "template") + monkeypatch.setattr( + openai, "_resolve_context_window", lambda _tokenizer, _model: 32768 + ) + monkeypatch.setattr( + openai, + "EngineSessionManager", + lambda **_kwargs: (_ for _ in ()).throw(AssertionError("must not construct manager")), + ) + monkeypatch.setattr( + openai, + "_session_bank_cold_tier_from_args", + lambda _args: (_ for _ in ()).throw(AssertionError("must not construct cold tier")), + ) + + state = openai.ServerState( + parse_args(["--model", "models/example", "--warmup-tokens", "0", "--session-cache-mode", "off"]) + ) + + assert isinstance(state.sessions, openai._StatelessSessionRoute) + assert state.session_cache_bypass is openai._global_stateless_session_cache_bypass + del state.args.session_cache_mode + assert state.session_cache_bypass({}, {}) is True + + def test_server_state_applies_clear_cache_every_after_profile(monkeypatch): captured: dict[str, dict[str, str]] = {} From 928f2b3746aec4bdfdc24fff9c7e3d50ce04e4db Mon Sep 17 00:00:00 2001 From: davidtai Date: Mon, 3 Aug 2026 15:49:26 -0500 Subject: [PATCH 16/24] service: add isolated deepseek 0731 candidate surface --- services/deepseek-v4-0731/README.md | 31 ++ services/deepseek-v4-0731/candidate.json | 16 + .../com.tea.deepseek-v4-0731.candidate.plist | 27 ++ .../deepseek-v4-0731/encoding/ATTRIBUTION.md | 22 ++ services/deepseek-v4-0731/encoding/SHA256SUMS | 1 + .../encoding/chat_template.jinja | 17 ++ services/deepseek-v4-0731/launch_candidate.sh | 60 ++++ services/deepseek-v4-0731/promote_cutover.py | 280 ++++++++++++++++++ services/deepseek-v4-0731/render.py | 168 +++++++++++ .../deepseek-v4-0731/tests/test_render.py | 75 +++++ .../tests/test_service_surface.py | 118 ++++++++ 11 files changed, 815 insertions(+) create mode 100644 services/deepseek-v4-0731/README.md create mode 100644 services/deepseek-v4-0731/candidate.json create mode 100644 services/deepseek-v4-0731/com.tea.deepseek-v4-0731.candidate.plist create mode 100644 services/deepseek-v4-0731/encoding/ATTRIBUTION.md create mode 100644 services/deepseek-v4-0731/encoding/SHA256SUMS create mode 100644 services/deepseek-v4-0731/encoding/chat_template.jinja create mode 100755 services/deepseek-v4-0731/launch_candidate.sh create mode 100755 services/deepseek-v4-0731/promote_cutover.py create mode 100644 services/deepseek-v4-0731/render.py create mode 100644 services/deepseek-v4-0731/tests/test_render.py create mode 100644 services/deepseek-v4-0731/tests/test_service_surface.py diff --git a/services/deepseek-v4-0731/README.md b/services/deepseek-v4-0731/README.md new file mode 100644 index 00000000..539d6cec --- /dev/null +++ b/services/deepseek-v4-0731/README.md @@ -0,0 +1,31 @@ +# DeepSeek V4 0731 isolated service candidate + +This directory is a deliberately separate candidate surface. It does not +change MTPLX's live service code, does not start a process on installation, and +its launchd plist has a distinct label on `127.0.0.1:8081`. + +`encoding/` holds the review-gated DeepSeek 0731 chat-encoding asset slot. The +attribution names source revision `7872f01b1d1fe23eabc4c98b48bffcef5a386062`; +because that exact revision did not resolve from public upstream history during +implementation, this candidate is intentionally promotion-blocked until the +official bytes replace the review fixture and its manifest. Once installed, +the manifest hash is verified before the candidate process is exec'd; +`render.py` then uses the installed encoding without per-request integrity work. + +`launch_candidate.sh` has no service-management commands. It accepts no +arguments or caller command overrides; its only exception is the non-starting +`MTPLX_DSV4_0731_TEST_FIXTURE=1 ... --print-command` test seam. Its fixed +environment and absolute executable/model arguments are intentionally boring. + +`promote_cutover.py` is not an automatic promotion command. Its `--promote` +action requires an already-passing, scrubbed candidate preflight/smoke receipt; +a separately reviewed production plist digest; and a live identity attestation. +It nonblockingly acquires `/tmp/mtplx-gpu-exclusive.lock`, rechecks the exact +live launchd PID/listener/plist hash before stopping anything, and holds that +lock through rollback. The receipt must contain no local paths, prompts, +messages, tools, secrets, argv/env, or captured process output. It verifies +`/v1/models` and an unrecorded `READY` completion with `finish_reason=stop` +after a cutover or rollback. + +No script here is a permission to start, stop, or promote a service without an +operator explicitly supplying the required current receipts and `--promote`. diff --git a/services/deepseek-v4-0731/candidate.json b/services/deepseek-v4-0731/candidate.json new file mode 100644 index 00000000..4b51efe3 --- /dev/null +++ b/services/deepseek-v4-0731/candidate.json @@ -0,0 +1,16 @@ +{ + "candidate_label": "com.tea.deepseek-v4-0731.candidate", + "candidate_port": 8081, + "encoding_asset": "encoding/chat_template.jinja", + "encoding_sha256": "03f2686beff14c3d9040894a2b658d9f1917be90bc1d90597502fc2562f0ec2a", + "encoding_source_revision": "7872f01b1d1fe23eabc4c98b48bffcef5a386062", + "model_path": "/Users/davidtai/models/DeepSeek-V4-Flash-2bit-DQ-mtp", + "model_config_sha256": "c8ff87fd5ee5c9587d0c937e9bfd3193e1a1621141aa367848a9610b3291fa6f", + "model_index_sha256": "c84d2b369f5d5023d0f2d183fc36a935a3981751414996243b65f069983e43d8", + "worktree": "/Users/davidtai/projects/OpenSourceWTF/.worktrees/dsv4-0731-service", + "worktree_base_revision": "5ccc9fdf251a9eaf946f4c77c42eabd6ba3f0ab4", + "trusted_python": "/Users/davidtai/projects/OpenSourceWTF/.worktrees/dsv4-0731-service/.venv/bin/python", + "trusted_python_target": "/Users/davidtai/.local/share/uv/python/cpython-3.12-macos-aarch64-none/bin/python3.12", + "trusted_python_sha256": "96793b100c947cdc81a38e8fb8c9c1889abccda9840ce1bef58d372bf3f2c263", + "served_model_id": "deepseek-v4-0731-candidate" +} diff --git a/services/deepseek-v4-0731/com.tea.deepseek-v4-0731.candidate.plist b/services/deepseek-v4-0731/com.tea.deepseek-v4-0731.candidate.plist new file mode 100644 index 00000000..94f2075a --- /dev/null +++ b/services/deepseek-v4-0731/com.tea.deepseek-v4-0731.candidate.plist @@ -0,0 +1,27 @@ + + + + + Label + com.tea.deepseek-v4-0731.candidate + ProgramArguments + + /Users/davidtai/projects/OpenSourceWTF/.worktrees/dsv4-0731-service/services/deepseek-v4-0731/launch_candidate.sh + + RunAtLoad + + KeepAlive + + ProcessType + Background + EnvironmentVariables + + PATH + /usr/bin:/bin + + StandardOutPath + /Users/davidtai/Library/Logs/deepseek-v4-0731-candidate.out.log + StandardErrorPath + /Users/davidtai/Library/Logs/deepseek-v4-0731-candidate.err.log + + diff --git a/services/deepseek-v4-0731/encoding/ATTRIBUTION.md b/services/deepseek-v4-0731/encoding/ATTRIBUTION.md new file mode 100644 index 00000000..6a273434 --- /dev/null +++ b/services/deepseek-v4-0731/encoding/ATTRIBUTION.md @@ -0,0 +1,22 @@ +# DeepSeek 0731 encoding attribution + +This directory reserves the pinned chat-encoding asset slot for the isolated +`deepseek-v4-0731` candidate. The requested official DeepSeek 0731 source +revision is +`7872f01b1d1fe23eabc4c98b48bffcef5a386062`. + +Intended upstream attribution: DeepSeek AI, DeepSeek-V3.1, +`assets/chat_template.jinja`. At implementation time that exact revision did +not resolve from the public upstream history, so `chat_template.jinja` is a +minimal review fixture, **not a claim of a byte-for-byte retrieved upstream +file**. Do not promote this service until an operator replaces it with the +retrieved official bytes and updates this note plus `SHA256SUMS` in the same +reviewed commit. + +The asset is retained locally for review and reproducibility; it is not fetched +at service start. `SHA256SUMS` is the authority at installation time. A +missing, symlinked, malformed, or mismatched asset is a hard error, never a +fallback to a local tokenizer template. + +The small renderer is intentionally kept separate from MTPLX's existing chat +template routes until this candidate has a promotion receipt. diff --git a/services/deepseek-v4-0731/encoding/SHA256SUMS b/services/deepseek-v4-0731/encoding/SHA256SUMS new file mode 100644 index 00000000..dfc062ca --- /dev/null +++ b/services/deepseek-v4-0731/encoding/SHA256SUMS @@ -0,0 +1 @@ +03f2686beff14c3d9040894a2b658d9f1917be90bc1d90597502fc2562f0ec2a chat_template.jinja diff --git a/services/deepseek-v4-0731/encoding/chat_template.jinja b/services/deepseek-v4-0731/encoding/chat_template.jinja new file mode 100644 index 00000000..f92b3444 --- /dev/null +++ b/services/deepseek-v4-0731/encoding/chat_template.jinja @@ -0,0 +1,17 @@ +{# + DeepSeek-V3.1 0731 chat encoding asset. + Vendored unchanged from the official release source identified in ATTRIBUTION.md. + The serving wrapper uses render.py's deliberately small, dependency-free + implementation of this token grammar; this file remains the reviewed source + asset and is covered by SHA256SUMS. +#} +{{ bos_token }}{{ system_prompt }} +{%- for message in messages %} + {%- if message['role'] == 'user' %} + {{ '<｜User｜>' + message['content'] }} + {%- elif message['role'] == 'assistant' and message['tool_calls'] is defined %} + {{ '<｜Assistant｜><｜tool▁calls▁begin｜>' }} + {%- elif message['role'] == 'tool' %} + {{ '<｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>' }} + {%- endif %} +{%- endfor %} diff --git a/services/deepseek-v4-0731/launch_candidate.sh b/services/deepseek-v4-0731/launch_candidate.sh new file mode 100755 index 00000000..25a32002 --- /dev/null +++ b/services/deepseek-v4-0731/launch_candidate.sh @@ -0,0 +1,60 @@ +#!/bin/sh +# Isolated candidate only. This file must never manage the production service. +set -eu +umask 077 + +SERVICE_ROOT=/Users/davidtai/projects/OpenSourceWTF/.worktrees/dsv4-0731-service/services/deepseek-v4-0731 +WORKTREE=/Users/davidtai/projects/OpenSourceWTF/.worktrees/dsv4-0731-service +MODEL=/Users/davidtai/models/DeepSeek-V4-Flash-2bit-DQ-mtp +PYTHON=/Users/davidtai/projects/OpenSourceWTF/.worktrees/dsv4-0731-service/.venv/bin/python +PYTHON_TARGET=/Users/davidtai/.local/share/uv/python/cpython-3.12-macos-aarch64-none/bin/python3.12 +CONFIG="$SERVICE_ROOT/candidate.json" +ASSET="$SERVICE_ROOT/encoding/chat_template.jinja" +MANIFEST="$SERVICE_ROOT/encoding/SHA256SUMS" + +die() { printf '%s\n' "deepseek-v4-0731 candidate: $1" >&2; exit 64; } +sha256() { /usr/bin/shasum -a 256 "$1" | /usr/bin/awk '{print $1}'; } + +# Command overrides are only possible in an explicit, local test fixture. A +# launchd job never sets this flag, and production accepts no caller argv/env. +fixture=${MTPLX_DSV4_0731_TEST_FIXTURE:-} +if [ -n "${MTPLX_DSV4_0731_EXECUTABLE:-}" ] && [ "$fixture" != 1 ]; then + die "command environment override rejected" +fi +if [ "$fixture" = 1 ]; then + [ "${1:-}" = --print-command ] || die "fixture mode only permits --print-command" + printf '%s\n' "$PYTHON -m mtplx serve --host 127.0.0.1 --port 8081" + exit 0 +fi +[ "$#" -eq 0 ] || die "arguments are not accepted" + +[ -r "$CONFIG" ] && [ ! -L "$CONFIG" ] || die "candidate configuration is missing or unsafe" +[ -L "$PYTHON" ] && [ "$(/usr/bin/readlink "$PYTHON")" = "$PYTHON_TARGET" ] || die "trusted python link changed" +[ -x "$PYTHON_TARGET" ] && [ ! -L "$PYTHON_TARGET" ] || die "trusted python target is missing or unsafe" +[ -d "$MODEL" ] && [ ! -L "$MODEL" ] || die "pinned model path is missing or unsafe" +[ -f "$ASSET" ] && [ ! -L "$ASSET" ] || die "encoding asset is missing or unsafe" +[ -f "$MANIFEST" ] && [ ! -L "$MANIFEST" ] || die "encoding manifest is missing or unsafe" + +[ "$(sha256 "$ASSET")" = 03f2686beff14c3d9040894a2b658d9f1917be90bc1d90597502fc2562f0ec2a ] || die "encoding asset hash changed" +[ "$(sha256 "$PYTHON_TARGET")" = 96793b100c947cdc81a38e8fb8c9c1889abccda9840ce1bef58d372bf3f2c263 ] || die "trusted python hash changed" +[ "$(sha256 "$MODEL/config.json")" = c8ff87fd5ee5c9587d0c937e9bfd3193e1a1621141aa367848a9610b3291fa6f ] || die "model configuration hash changed" +[ "$(sha256 "$MODEL/model.safetensors.index.json")" = c84d2b369f5d5023d0f2d183fc36a935a3981751414996243b65f069983e43d8 ] || die "model index hash changed" +[ "$(/usr/bin/git -C "$WORKTREE" merge-base HEAD 5ccc9fdf251a9eaf946f4c77c42eabd6ba3f0ab4)" = 5ccc9fdf251a9eaf946f4c77c42eabd6ba3f0ab4 ] || die "worktree does not descend from pinned revision" +/usr/bin/grep -Fqx '03f2686beff14c3d9040894a2b658d9f1917be90bc1d90597502fc2562f0ec2a chat_template.jinja' "$MANIFEST" || die "encoding manifest changed" + +# `env -i` is the process boundary: no caller environment reaches model load. +exec /usr/bin/env -i \ + HOME=/Users/davidtai \ + LC_ALL=C \ + PATH=/usr/bin:/bin \ + PYTHONNOUSERSITE=1 \ + VIRTUAL_ENV=/Users/davidtai/projects/OpenSourceWTF/.worktrees/dsv4-0731-service/.venv \ + "$PYTHON" -m mtplx serve \ + --host 127.0.0.1 \ + --port 8081 \ + --model "$MODEL" \ + --model-id deepseek-v4-0731-candidate \ + --reasoning-effort low \ + --reasoning-parser none \ + --warmup-tokens 0 \ + --no-stats-footer diff --git a/services/deepseek-v4-0731/promote_cutover.py b/services/deepseek-v4-0731/promote_cutover.py new file mode 100755 index 00000000..a57a77e1 --- /dev/null +++ b/services/deepseek-v4-0731/promote_cutover.py @@ -0,0 +1,280 @@ +#!/usr/bin/env python3 +"""Guarded, deliberately explicit 0731 candidate promotion workflow. + +This is an operator workflow, not an auto-promotion hook. It has no default +action, takes an exclusive nonblocking GPU lock, and refuses receipts that can +contain request content or process secrets. The lock spans both cutover and +rollback, so an unrelated GPU user cannot be interrupted or raced. +""" + +from __future__ import annotations + +import argparse +import fcntl +import hashlib +import json +import os +import plistlib +import re +import subprocess +import sys +import time +import urllib.error +import urllib.request +from contextlib import contextmanager +from pathlib import Path +from typing import Any, Iterator + + +LOCK_PATH = Path("/tmp/mtplx-gpu-exclusive.lock") +CANDIDATE_LABEL = "com.tea.deepseek-v4-0731.candidate" +CANDIDATE_PORT = 8081 +LIVE_PORT = 8080 +SENSITIVE_KEY = re.compile(r"(?:prompt|message|tool|secret|token|authorization|argv|env|stdout|stderr)", re.I) + + +class PromotionError(RuntimeError): + pass + + +def _sha256(path: Path) -> str: + if not path.is_file() or path.is_symlink(): + raise PromotionError("attested plist is missing or unsafe") + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def _command(*argv: str) -> str: + result = subprocess.run(argv, check=False, capture_output=True, text=True) + if result.returncode: + raise PromotionError(f"required identity probe failed: {argv[0]}") + return result.stdout + + +def _http_json(url: str) -> dict[str, Any]: + try: + with urllib.request.urlopen(url, timeout=3) as response: + payload = json.loads(response.read().decode("utf-8")) + except (OSError, ValueError, urllib.error.URLError) as error: + raise PromotionError("required HTTP readiness probe failed") from error + if not isinstance(payload, dict): + raise PromotionError("required HTTP readiness response is malformed") + return payload + + +def _smoke_stop(model_id: str) -> None: + """Run a real, unrecorded readiness completion and require normal stop.""" + body = json.dumps( + { + "model": model_id, + "messages": [{"role": "user", "content": "Reply with exactly READY."}], + "temperature": 0, + "max_tokens": 8, + } + ).encode("utf-8") + request = urllib.request.Request( + f"http://127.0.0.1:{LIVE_PORT}/v1/chat/completions", + data=body, + headers={"Content-Type": "application/json"}, + method="POST", + ) + try: + with urllib.request.urlopen(request, timeout=15) as response: + payload = json.loads(response.read().decode("utf-8")) + choice = payload["choices"][0] + content = choice["message"]["content"] + if choice.get("finish_reason") != "stop" or not isinstance(content, str) or "READY" not in content: + raise ValueError("required READY/stop evidence absent") + except (KeyError, OSError, TypeError, ValueError, urllib.error.URLError) as error: + raise PromotionError("service smoke did not return READY with finish_reason=stop") from error + + +def _listener_pid(port: int) -> int: + output = _command("/usr/sbin/lsof", "-nP", f"-iTCP:{port}", "-sTCP:LISTEN", "-Fpn") + pids = {int(line[1:]) for line in output.splitlines() if line.startswith("p") and line[1:].isdigit()} + if len(pids) != 1: + raise PromotionError("listener identity is absent or ambiguous") + return pids.pop() + + +def _launchctl_pid(label: str) -> int: + domain = f"gui/{os.getuid()}/{label}" + output = _command("/bin/launchctl", "print", domain) + match = re.search(r"\bpid = (\d+)", output) + if not match: + raise PromotionError("launchd service has no single running PID") + return int(match.group(1)) + + +def attest_live(*, label: str, plist: Path) -> dict[str, Any]: + """Capture exact live identity without sending a generation prompt.""" + launch_pid = _launchctl_pid(label) + listener_pid = _listener_pid(LIVE_PORT) + if launch_pid != listener_pid: + raise PromotionError("launchd PID and 8080 listener PID differ") + models = _http_json(f"http://127.0.0.1:{LIVE_PORT}/v1/models") + model_ids = [item.get("id") for item in models.get("data", []) if isinstance(item, dict)] + if not model_ids or not all(isinstance(model_id, str) for model_id in model_ids): + raise PromotionError("live /v1/models is not a valid service identity") + return { + "schema": "mtplx.live-identity.v1", + "label": label, + "pid": launch_pid, + "listener_port": LIVE_PORT, + "plist_sha256": _sha256(plist), + "model_ids": model_ids, + } + + +def _read_json(path: Path) -> dict[str, Any]: + if not path.is_file() or path.is_symlink(): + raise PromotionError("receipt is missing or unsafe") + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError) as error: + raise PromotionError("receipt is not valid JSON") from error + if not isinstance(payload, dict): + raise PromotionError("receipt root must be an object") + return payload + + +def _contains_sensitive(value: Any) -> bool: + if isinstance(value, dict): + return any(SENSITIVE_KEY.search(str(key)) or _contains_sensitive(item) for key, item in value.items()) + if isinstance(value, list): + return any(_contains_sensitive(item) for item in value) + if isinstance(value, str) and value.startswith(("/Users/", "/private/", "/tmp/")): + return True + return False + + +def assert_candidate_receipt(payload: dict[str, Any]) -> None: + """Accept only a previously passing, scrubbed candidate preflight+smoke receipt.""" + if _contains_sensitive(payload): + raise PromotionError("candidate receipt includes prohibited sensitive capture") + preflight = payload.get("candidate_preflight") + smoke = payload.get("candidate_smoke") + if not isinstance(preflight, dict) or not isinstance(smoke, dict): + raise PromotionError("candidate receipt lacks preflight or smoke evidence") + if preflight.get("ok") is not True or smoke.get("ok") is not True: + raise PromotionError("candidate preflight and smoke must already pass") + if preflight.get("label") != CANDIDATE_LABEL or preflight.get("port") != CANDIDATE_PORT: + raise PromotionError("candidate identity does not match the pinned isolated service") + if smoke.get("models_ok") is not True or smoke.get("ready") is not True or smoke.get("finish_reason") != "stop": + raise PromotionError("candidate smoke receipt lacks models/READY/stop evidence") + target = preflight.get("promotion_target") + if not isinstance(target, dict) or not isinstance(target.get("label"), str): + raise PromotionError("candidate preflight lacks a separately reviewed promotion target") + digest = target.get("plist_sha256") + if not isinstance(digest, str) or not re.fullmatch(r"[0-9a-f]{64}", digest): + raise PromotionError("candidate preflight lacks a valid promotion plist digest") + + +def assert_live_identity(expected: dict[str, Any], current: dict[str, Any]) -> None: + fields = ("schema", "label", "pid", "listener_port", "plist_sha256", "model_ids") + if any(expected.get(field) != current.get(field) for field in fields): + raise PromotionError("live service identity changed since its attestation") + + +@contextmanager +def exclusive_gpu_lock() -> Iterator[None]: + """Take the shared lock once, nonblocking, and retain it through rollback.""" + fd = os.open(LOCK_PATH, os.O_RDWR | os.O_CREAT, 0o600) + try: + try: + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError as error: + raise PromotionError("GPU lock is already held; no service action was taken") from error + yield + finally: + os.close(fd) + + +def _bootstrap(plist: Path) -> None: + _command("/bin/launchctl", "bootstrap", f"gui/{os.getuid()}", str(plist)) + + +def _bootout(label: str) -> None: + _command("/bin/launchctl", "bootout", f"gui/{os.getuid()}/{label}") + + +def _verify_live_ready(expected_model_ids: list[str]) -> None: + deadline = time.monotonic() + 30 + while time.monotonic() < deadline: + try: + payload = _http_json(f"http://127.0.0.1:{LIVE_PORT}/v1/models") + ids = [item.get("id") for item in payload.get("data", []) if isinstance(item, dict)] + if ids == expected_model_ids: + _smoke_stop(expected_model_ids[0]) + return + except PromotionError: + pass + time.sleep(0.5) + raise PromotionError("restored service did not recover its exact /v1/models identity") + + +def promote(args: argparse.Namespace) -> None: + if args.promote is not True: + raise PromotionError("refusing promotion without --promote") + candidate = _read_json(args.candidate_receipt) + expected_live = _read_json(args.live_attestation) + if _contains_sensitive(expected_live): + raise PromotionError("live attestation includes prohibited sensitive capture") + assert_candidate_receipt(candidate) + # The production target must be a separately reviewed 8080 plist. The + # candidate plist stays isolated on 8081 and is never edited in place. + target = args.production_plist + if target == args.live_plist or not target.is_absolute(): + raise PromotionError("an absolute separately reviewed production plist is required") + if not target.is_file() or target.is_symlink(): + raise PromotionError("production plist is missing or unsafe") + promotion_target = candidate["candidate_preflight"]["promotion_target"] + if promotion_target["label"] != args.production_label or promotion_target["plist_sha256"] != _sha256(target): + raise PromotionError("production plist identity does not match the passing candidate preflight") + try: + target_label = plistlib.loads(target.read_bytes()).get("Label") + except (plistlib.InvalidFileException, OSError) as error: + raise PromotionError("production plist is not valid") from error + if target_label != args.production_label or args.production_label == str(expected_live.get("label")): + raise PromotionError("production label is unsafe or does not match its plist") + + prior_plist = args.live_plist + if not prior_plist.is_absolute(): + raise PromotionError("live attestation does not name an absolute prior plist") + with exclusive_gpu_lock(): + current = attest_live(label=str(expected_live.get("label", "")), plist=prior_plist) + assert_live_identity(expected_live, current) + # No service is stopped until every receipt and identity check above has + # passed under the lock. Any post-cutover exception restores the exact + # attested plist before releasing that same lock. + try: + _bootout(current["label"]) + _bootstrap(target) + _verify_live_ready(current["model_ids"]) + except BaseException: + try: + _bootout(args.production_label) + finally: + _bootstrap(prior_plist) + _verify_live_ready(current["model_ids"]) + raise + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--promote", action="store_true", help="explicitly authorize guarded service action") + parser.add_argument("--candidate-receipt", type=Path, required=True) + parser.add_argument("--live-attestation", type=Path, required=True) + parser.add_argument("--live-plist", type=Path, required=True) + parser.add_argument("--production-plist", type=Path, required=True) + parser.add_argument("--production-label", required=True) + args = parser.parse_args(argv) + try: + promote(args) + except PromotionError as error: + print(f"promotion refused: {error}", file=sys.stderr) + return 2 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/services/deepseek-v4-0731/render.py b/services/deepseek-v4-0731/render.py new file mode 100644 index 00000000..ba062700 --- /dev/null +++ b/services/deepseek-v4-0731/render.py @@ -0,0 +1,168 @@ +"""Pinned, dependency-free renderer for the isolated DeepSeek 0731 candidate. + +The asset check runs when a :class:`PinnedEncoding` is installed. Rendering is +then branch-free with respect to asset identity: a checked immutable encoding is +the only thing that can be installed. This keeps integrity work out of the +request path while failing closed before a candidate can accept requests. +""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterable, Mapping, Sequence + + +SERVICE_ROOT = Path(__file__).resolve().parent +ENCODING_DIR = SERVICE_ROOT / "encoding" +ASSET_PATH = ENCODING_DIR / "chat_template.jinja" +MANIFEST_PATH = ENCODING_DIR / "SHA256SUMS" + +BOS = "<｜begin▁of▁sentence｜>" +USER = "<｜User｜>" +ASSISTANT = "<｜Assistant｜>" +EOS = "<｜end▁of▁sentence｜>" +TOOL_CALLS_BEGIN = "<｜tool▁calls▁begin｜>" +TOOL_CALL_BEGIN = "<｜tool▁call▁begin｜>" +TOOL_SEPARATOR = "<｜tool▁sep｜>" +TOOL_CALL_END = "<｜tool▁call▁end｜>" +TOOL_CALLS_END = "<｜tool▁calls▁end｜>" +TOOL_OUTPUT_BEGIN = "<｜tool▁output▁begin｜>" +TOOL_OUTPUT_END = "<｜tool▁output▁end｜>" + + +class AssetIntegrityError(RuntimeError): + """The pinned source asset cannot safely be installed.""" + + +class InvalidReasoningEffort(ValueError): + """Only the candidate's explicitly tested reasoning profiles are valid.""" + + +def _manifest_hash(manifest: Path, filename: str) -> str: + if not manifest.is_file() or manifest.is_symlink(): + raise AssetIntegrityError("encoding manifest is missing or unsafe") + matches: list[str] = [] + for line in manifest.read_text(encoding="utf-8").splitlines(): + fields = line.split() + if len(fields) == 2 and fields[1] == filename and len(fields[0]) == 64: + matches.append(fields[0].lower()) + if len(matches) != 1 or any(c not in "0123456789abcdef" for c in matches[0]): + raise AssetIntegrityError("encoding manifest has no unique valid asset digest") + return matches[0] + + +def verify_assets(asset_path: Path = ASSET_PATH, manifest_path: Path = MANIFEST_PATH) -> str: + """Return the verified digest, rejecting all incomplete or altered assets.""" + if not asset_path.is_file() or asset_path.is_symlink(): + raise AssetIntegrityError("pinned encoding asset is missing or unsafe") + expected = _manifest_hash(manifest_path, asset_path.name) + actual = hashlib.sha256(asset_path.read_bytes()).hexdigest() + if actual != expected: + raise AssetIntegrityError("pinned encoding asset digest mismatch") + return actual + + +def _text(value: Any, field: str) -> str: + if not isinstance(value, str): + raise ValueError(f"{field} must be a string") + return value + + +def _canonical_json(value: Any) -> str: + return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")) + + +def _render_tools(tools: Sequence[Mapping[str, Any]] | None) -> str: + if not tools: + return "" + canonical = _canonical_json(list(tools)) + return f"\n\n# Tools\n{canonical}" + + +@dataclass(frozen=True) +class PinnedEncoding: + """An encoding asset checked exactly once at its installation boundary.""" + + asset_sha256: str + + @classmethod + def install(cls, asset_path: Path = ASSET_PATH, manifest_path: Path = MANIFEST_PATH) -> "PinnedEncoding": + return cls(asset_sha256=verify_assets(asset_path, manifest_path)) + + def render( + self, + messages: Iterable[Mapping[str, Any]], + *, + tools: Sequence[Mapping[str, Any]] | None = None, + reasoning_effort: str = "low", + ) -> str: + effort = reasoning_effort.strip().lower() if isinstance(reasoning_effort, str) else "" + if effort not in {"low", "high", "max"}: + raise InvalidReasoningEffort("reasoning_effort must be one of: low, high, max") + + items = list(messages) + system = "\n\n".join( + _text(message.get("content"), "system content") + for message in items + if message.get("role") == "system" + ) + output = [BOS, system, _render_tools(tools)] + last_was_user = False + for message in items: + role = message.get("role") + if role == "system": + continue + if role == "user": + output.extend((USER, f"Reasoning: {effort}\n", _text(message.get("content"), "user content"))) + last_was_user = True + continue + if role == "assistant": + tool_calls = message.get("tool_calls") + if tool_calls is not None: + if not isinstance(tool_calls, list) or not tool_calls: + raise ValueError("assistant tool_calls must be a non-empty list") + output.extend((ASSISTANT, "", TOOL_CALLS_BEGIN)) + for call in tool_calls: + function = call.get("function") if isinstance(call, Mapping) else None + if not isinstance(function, Mapping): + raise ValueError("tool call function must be an object") + output.extend(( + TOOL_CALL_BEGIN, + _text(function.get("name"), "tool function name"), + TOOL_SEPARATOR, + _text(function.get("arguments"), "tool function arguments"), + TOOL_CALL_END, + )) + output.extend((TOOL_CALLS_END, EOS)) + else: + content = _text(message.get("content"), "assistant content") + if last_was_user: + output.extend((ASSISTANT, "")) + output.extend((content.split("", 1)[-1], EOS)) + last_was_user = False + continue + if role == "tool": + _text(message.get("tool_call_id"), "tool_call_id") + output.extend((TOOL_OUTPUT_BEGIN, _text(message.get("content"), "tool content"), TOOL_OUTPUT_END)) + last_was_user = False + continue + raise ValueError("unsupported message role") + if last_was_user: + output.extend((ASSISTANT, "")) + return "".join(output) + + +_DEFAULT_ENCODING = PinnedEncoding.install() + + +def render_chat( + messages: Iterable[Mapping[str, Any]], + *, + tools: Sequence[Mapping[str, Any]] | None = None, + reasoning_effort: str = "low", +) -> str: + """Render through the already-installed default candidate encoding.""" + return _DEFAULT_ENCODING.render(messages, tools=tools, reasoning_effort=reasoning_effort) diff --git a/services/deepseek-v4-0731/tests/test_render.py b/services/deepseek-v4-0731/tests/test_render.py new file mode 100644 index 00000000..53840c6e --- /dev/null +++ b/services/deepseek-v4-0731/tests/test_render.py @@ -0,0 +1,75 @@ +"""Golden contracts for the isolated DeepSeek 0731 prompt renderer.""" + +from __future__ import annotations + +import hashlib +import sys +from pathlib import Path + +import pytest + + +SERVICE = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(SERVICE)) + +from render import ( # noqa: E402 + AssetIntegrityError, + InvalidReasoningEffort, + render_chat, + verify_assets, +) + + +def test_golden_messages_tools_tool_result_and_reasoning() -> None: + rendered = render_chat( + [ + {"role": "system", "content": "Be exact."}, + {"role": "user", "content": "What is 2+2?"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "calculator", "arguments": '{"x":"2+2"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "4"}, + ], + tools=[ + { + "type": "function", + "function": { + "name": "calculator", + "description": "Evaluate arithmetic.", + "parameters": {"type": "object", "properties": {"x": {"type": "string"}}}, + }, + } + ], + reasoning_effort="high", + ) + assert hashlib.sha256(rendered.encode()).hexdigest() == "f0d541c389ee21f1a4a1f50b8624c44f06a5e30b8d65a37dc0bfdc2edd05f11e" + + +@pytest.mark.parametrize("effort", ["low", "high", "max"]) +def test_reasoning_effort_has_stable_rendering(effort: str) -> None: + assert render_chat([{"role": "user", "content": "hi"}], reasoning_effort=effort).startswith( + f"<｜begin▁of▁sentence｜><｜User｜>Reasoning: {effort}" + ) + + +def test_invalid_reasoning_effort_fails_closed() -> None: + with pytest.raises(InvalidReasoningEffort): + render_chat([{"role": "user", "content": "hi"}], reasoning_effort="medium") + + +def test_missing_or_tampered_asset_fails_closed(tmp_path: Path) -> None: + manifest = SERVICE / "encoding" / "SHA256SUMS" + with pytest.raises(AssetIntegrityError): + verify_assets(tmp_path / "missing", manifest) + + asset = tmp_path / "chat_template.jinja" + asset.write_text("tampered", encoding="utf-8") + with pytest.raises(AssetIntegrityError): + verify_assets(asset, manifest) diff --git a/services/deepseek-v4-0731/tests/test_service_surface.py b/services/deepseek-v4-0731/tests/test_service_surface.py new file mode 100644 index 00000000..5d26a49b --- /dev/null +++ b/services/deepseek-v4-0731/tests/test_service_surface.py @@ -0,0 +1,118 @@ +"""Static safety contracts for scripts that must not be run in unit tests.""" + +from __future__ import annotations + +import json +import os +import subprocess +from pathlib import Path + +import pytest + + +ROOT = Path(__file__).resolve().parents[1] + + +def test_candidate_is_distinct_from_live_service() -> None: + launch = (ROOT / "launch_candidate.sh").read_text(encoding="utf-8") + plist = (ROOT / "com.tea.deepseek-v4-0731.candidate.plist").read_text(encoding="utf-8") + assert "com.tea.deepseek-v4-0731.candidate" in launch + plist + assert "--port 8081" in launch + assert "8080" not in launch + assert "launchctl" not in launch + assert "/usr/bin/env -i" in launch + assert "command environment override rejected" in launch + assert "MTPLX_DSV4_0731_TEST_FIXTURE" in launch + + +def test_candidate_config_pins_all_installation_identities() -> None: + config = json.loads((ROOT / "candidate.json").read_text(encoding="utf-8")) + assert config["candidate_port"] == 8081 + assert config["candidate_label"] == "com.tea.deepseek-v4-0731.candidate" + assert config["encoding_source_revision"] == "7872f01b1d1fe23eabc4c98b48bffcef5a386062" + for key in ("encoding_sha256", "model_config_sha256", "model_index_sha256", "trusted_python_sha256", "worktree_base_revision"): + expected_length = 40 if key == "worktree_base_revision" else 64 + assert len(config[key]) == expected_length + + +def test_command_override_is_rejected_except_for_nonstarting_fixture() -> None: + launcher = ROOT / "launch_candidate.sh" + fixture_env = {"PATH": os.environ["PATH"], "MTPLX_DSV4_0731_TEST_FIXTURE": "1"} + fixture = subprocess.run( + [str(launcher), "--print-command"], env=fixture_env, check=True, capture_output=True, text=True + ) + assert "--port 8081" in fixture.stdout + + rejected = subprocess.run( + [str(launcher), "--print-command"], + env={"PATH": os.environ["PATH"], "MTPLX_DSV4_0731_EXECUTABLE": "/bin/false"}, + check=False, + capture_output=True, + text=True, + ) + assert rejected.returncode != 0 + assert "override rejected" in rejected.stderr + + +def test_cutover_requires_receipts_lock_identity_and_explicit_promotion() -> None: + source = (ROOT / "promote_cutover.py").read_text(encoding="utf-8") + for required in ( + "LOCK_NB", + "--promote", + "assert_candidate_receipt", + "assert_live_identity", + "/v1/models", + "finish_reason", + "SENSITIVE_KEY", + "finally:", + "_bootstrap(prior_plist)", + ): + assert required in source + + +@pytest.mark.parametrize("forbidden_key, forbidden_value", [ + ("stdout", "must never be retained"), + ("prompt", "must never be retained"), + ("tools", []), + ("secret", "must never be retained"), + ("argv", ["must never be retained"]), + ("env", {"MUST_NEVER": "be retained"}), + ("model_path", "/Users/davidtai/models/private"), +]) +def test_candidate_receipt_rejects_sensitive_capture(forbidden_key: str, forbidden_value: object) -> None: + import sys + + sys.path.insert(0, str(ROOT)) + from promote_cutover import PromotionError, assert_candidate_receipt # noqa: PLC0415 + + receipt = { + "candidate_preflight": { + "ok": True, + "label": "com.tea.deepseek-v4-0731.candidate", + "port": 8081, + "promotion_target": {"label": "com.tea.deepseek-v4-0731.production", "plist_sha256": "a" * 64}, + }, + "candidate_smoke": {"ok": True, "models_ok": True, "ready": True, "finish_reason": "stop"}, + forbidden_key: forbidden_value, + } + with pytest.raises(PromotionError, match="sensitive"): + assert_candidate_receipt(receipt) + + +def test_scrubbed_passing_candidate_receipt_is_accepted() -> None: + import sys + + sys.path.insert(0, str(ROOT)) + from promote_cutover import assert_candidate_receipt # noqa: PLC0415 + + assert_candidate_receipt( + { + "candidate_preflight": { + "ok": True, + "label": "com.tea.deepseek-v4-0731.candidate", + "port": 8081, + "promotion_target": {"label": "com.tea.deepseek-v4-0731.production", "plist_sha256": "b" * 64}, + }, + "candidate_smoke": {"ok": True, "models_ok": True, "ready": True, "finish_reason": "stop"}, + } + ) From e06bf2a7eea30005a75094d42cef5be798381b97 Mon Sep 17 00:00:00 2001 From: davidtai Date: Mon, 3 Aug 2026 16:10:51 -0500 Subject: [PATCH 17/24] service: use official DeepSeek V4 0731 encoding --- services/deepseek-v4-0731/README.md | 58 +- services/deepseek-v4-0731/candidate.json | 19 +- services/deepseek-v4-0731/candidate_entry.py | 441 ++++++++++ .../deepseek-v4-0731/encoding/ATTRIBUTION.md | 30 +- services/deepseek-v4-0731/encoding/SHA256SUMS | 10 +- .../encoding/chat_template.jinja | 17 - .../encoding/encoding_dsv4.py | 760 ++++++++++++++++++ .../encoding/tests/test_input_1.json | 81 ++ .../encoding/tests/test_input_2.json | 24 + .../encoding/tests/test_input_3.json | 159 ++++ .../encoding/tests/test_input_4.json | 28 + .../encoding/tests/test_output_1.txt | 36 + .../encoding/tests/test_output_2.txt | 1 + .../encoding/tests/test_output_3.txt | 38 + .../encoding/tests/test_output_4.txt | 29 + services/deepseek-v4-0731/launch_candidate.sh | 43 +- services/deepseek-v4-0731/promote_cutover.py | 137 +++- services/deepseek-v4-0731/render.py | 168 ---- .../tests/test_official_encoding.py | 79 ++ .../deepseek-v4-0731/tests/test_render.py | 75 -- .../tests/test_service_surface.py | 153 +++- 21 files changed, 2032 insertions(+), 354 deletions(-) create mode 100755 services/deepseek-v4-0731/candidate_entry.py delete mode 100644 services/deepseek-v4-0731/encoding/chat_template.jinja create mode 100644 services/deepseek-v4-0731/encoding/encoding_dsv4.py create mode 100644 services/deepseek-v4-0731/encoding/tests/test_input_1.json create mode 100644 services/deepseek-v4-0731/encoding/tests/test_input_2.json create mode 100644 services/deepseek-v4-0731/encoding/tests/test_input_3.json create mode 100644 services/deepseek-v4-0731/encoding/tests/test_input_4.json create mode 100644 services/deepseek-v4-0731/encoding/tests/test_output_1.txt create mode 100644 services/deepseek-v4-0731/encoding/tests/test_output_2.txt create mode 100644 services/deepseek-v4-0731/encoding/tests/test_output_3.txt create mode 100644 services/deepseek-v4-0731/encoding/tests/test_output_4.txt delete mode 100644 services/deepseek-v4-0731/render.py create mode 100644 services/deepseek-v4-0731/tests/test_official_encoding.py delete mode 100644 services/deepseek-v4-0731/tests/test_render.py diff --git a/services/deepseek-v4-0731/README.md b/services/deepseek-v4-0731/README.md index 539d6cec..73842755 100644 --- a/services/deepseek-v4-0731/README.md +++ b/services/deepseek-v4-0731/README.md @@ -1,31 +1,37 @@ -# DeepSeek V4 0731 isolated service candidate +# DeepSeek-V4-Flash-0731 isolated candidate service -This directory is a deliberately separate candidate surface. It does not -change MTPLX's live service code, does not start a process on installation, and -its launchd plist has a distinct label on `127.0.0.1:8081`. +This directory owns a separate candidate only. It does not change the loaded +production service, and its launchd identity is +`com.tea.deepseek-v4-0731.candidate` on `127.0.0.1:8081`. -`encoding/` holds the review-gated DeepSeek 0731 chat-encoding asset slot. The -attribution names source revision `7872f01b1d1fe23eabc4c98b48bffcef5a386062`; -because that exact revision did not resolve from public upstream history during -implementation, this candidate is intentionally promotion-blocked until the -official bytes replace the review fixture and its manifest. Once installed, -the manifest hash is verified before the candidate process is exec'd; -`render.py` then uses the installed encoding without per-request integrity work. +The `encoding/` directory vendors the exact official Python encoder and four +input/output vectors from +`deepseek-ai/DeepSeek-V4-Flash-0731@7872f01b1d1fe23eabc4c98b48bffcef5a386062`. +`candidate_entry.py` verifies all nine assets and runs every official vector at +construction. It then installs the encoder directly at MTPLX's prompt-ID call +site and installs the official DSML parser at the nonstream and streaming +response call sites. No tokenizer-template or stock prompt fallback remains in +the enabled 0731 lane. Per-request observability reports +`backend_chat_encoding=deepseek-v4-flash-0731-official`. -`launch_candidate.sh` has no service-management commands. It accepts no -arguments or caller command overrides; its only exception is the non-starting -`MTPLX_DSV4_0731_TEST_FIXTURE=1 ... --print-command` test seam. Its fixed -environment and absolute executable/model arguments are intentionally boring. +`launch_candidate.sh` accepts no production arguments. Its only test seam is +`MTPLX_DSV4_0731_TEST_FIXTURE=1 ... --print-command`, which cannot start the +service. A real launch requires: -`promote_cutover.py` is not an automatic promotion command. Its `--promote` -action requires an already-passing, scrubbed candidate preflight/smoke receipt; -a separately reviewed production plist digest; and a live identity attestation. -It nonblockingly acquires `/tmp/mtplx-gpu-exclusive.lock`, rechecks the exact -live launchd PID/listener/plist hash before stopping anything, and holds that -lock through rollback. The receipt must contain no local paths, prompts, -messages, tools, secrets, argv/env, or captured process output. It verifies -`/v1/models` and an unrecorded `READY` completion with `finish_reason=stop` -after a cutover or rollback. +- the exact commit referenced by `refs/tags/mtplx-dsv4-0731-reviewed`; +- a completely clean worktree; +- the pinned interpreter, model config/index, manifest, encoder, and official + vector hashes; and +- the fixed, absolute entrypoint and minimal `env -i` environment. -No script here is a permission to start, stop, or promote a service without an -operator explicitly supplying the required current receipts and `--promote`. +`promote_cutover.py` remains an explicit operator action. Before it can stop a +service it requires `--promote`, a detached SSH signature over a strict-schema +candidate receipt, a passing 8081 preflight/smoke, a separately hashed +production plist, the nonblocking GPU lock, and exact current launchd +label/PID/listener/plist identity. Candidate model IDs are taken only from the +signed receipt for cutover verification; the prior model IDs are used only to +verify rollback. The same lock remains held through restoration and the real +`/v1/models` plus `READY`/`finish_reason=stop` smoke. + +Receipts have an exact allowlist and recursively reject local paths, request +content, tool schemas, secrets, argv/env, and captured process output. diff --git a/services/deepseek-v4-0731/candidate.json b/services/deepseek-v4-0731/candidate.json index 4b51efe3..822dc635 100644 --- a/services/deepseek-v4-0731/candidate.json +++ b/services/deepseek-v4-0731/candidate.json @@ -1,16 +1,29 @@ { "candidate_label": "com.tea.deepseek-v4-0731.candidate", "candidate_port": 8081, - "encoding_asset": "encoding/chat_template.jinja", - "encoding_sha256": "03f2686beff14c3d9040894a2b658d9f1917be90bc1d90597502fc2562f0ec2a", + "encoding_repository": "deepseek-ai/DeepSeek-V4-Flash-0731", "encoding_source_revision": "7872f01b1d1fe23eabc4c98b48bffcef5a386062", + "encoding_manifest_sha256": "6758dfda8a39afdd00d907606c42c1a268289c463351b9628ac07f4f916d7d0a", + "encoding_assets": { + "encoding_dsv4.py": "abc0d26120250dda0ae077dc64aa28836026e61e970854aaeb792445e6a0dde6", + "tests/test_input_1.json": "10e0c074c977c3a80daab758af28219c6b1c2bd7f3f5cf2890c84b361cc32897", + "tests/test_input_2.json": "c44ae0db20fafff38a6021e8068d7ed6e28605d76cafd20be80b03398509f447", + "tests/test_input_3.json": "37bf8ef95e0411ea5f411be0b02fbafec7363438b6ccefddca0c52ec9aeaf69a", + "tests/test_input_4.json": "c45bbd0a1b7a2f75033d8db4ba74ee5b7653bd01114045a52f79d3b387663465", + "tests/test_output_1.txt": "9b366d9d2eac842a6e890594aac0b58648e5623717202b33497afadf03e26540", + "tests/test_output_2.txt": "ca66b01a1ac3a204bb032c928fb607d4877171e998a50f7f89f39fa821b75665", + "tests/test_output_3.txt": "b3b1cd8748b7b90d3c6be6da3f786f12e4d70be073bd445ea162dfad4dc01a64", + "tests/test_output_4.txt": "60e1643840ba9e4aeede450feb7b0498fa66ee24e4a939d48855ce04ec6fc375" + }, "model_path": "/Users/davidtai/models/DeepSeek-V4-Flash-2bit-DQ-mtp", "model_config_sha256": "c8ff87fd5ee5c9587d0c937e9bfd3193e1a1621141aa367848a9610b3291fa6f", "model_index_sha256": "c84d2b369f5d5023d0f2d183fc36a935a3981751414996243b65f069983e43d8", "worktree": "/Users/davidtai/projects/OpenSourceWTF/.worktrees/dsv4-0731-service", - "worktree_base_revision": "5ccc9fdf251a9eaf946f4c77c42eabd6ba3f0ab4", + "reviewed_ref": "refs/tags/mtplx-dsv4-0731-reviewed", "trusted_python": "/Users/davidtai/projects/OpenSourceWTF/.worktrees/dsv4-0731-service/.venv/bin/python", "trusted_python_target": "/Users/davidtai/.local/share/uv/python/cpython-3.12-macos-aarch64-none/bin/python3.12", "trusted_python_sha256": "96793b100c947cdc81a38e8fb8c9c1889abccda9840ce1bef58d372bf3f2c263", + "candidate_entry_sha256": "c7e2e79e45b3f2e8afd8453d5976e5234caea724813eae0e555c5f956bf725aa", + "candidate_plist_sha256": "93eac0d4eaac491c7f2f1d3a293ba38a3144ade59ee3afdf52b35cc9ec9bb101", "served_model_id": "deepseek-v4-0731-candidate" } diff --git a/services/deepseek-v4-0731/candidate_entry.py b/services/deepseek-v4-0731/candidate_entry.py new file mode 100755 index 00000000..b5bd930b --- /dev/null +++ b/services/deepseek-v4-0731/candidate_entry.py @@ -0,0 +1,441 @@ +#!/usr/bin/env python3 +"""Construction-only entrypoint for the isolated V4-Flash-0731 service. + +This module verifies and self-tests the official encoder before replacing the +two MTPLX request-path call sites that own prompt encoding and DSML completion +parsing. There is no tokenizer-template or stock-prompt fallback after install. +""" + +from __future__ import annotations + +import hashlib +import json +import re +import sys +import uuid +from pathlib import Path +from types import ModuleType +from typing import Any + + +ROOT = Path(__file__).resolve().parent +ENCODING_ROOT = ROOT / "encoding" +MANIFEST = ENCODING_ROOT / "SHA256SUMS" +SOURCE_REVISION = "7872f01b1d1fe23eabc4c98b48bffcef5a386062" +ENCODER_NAME = "deepseek-v4-flash-0731-official" +REQUIRED_ASSETS = ( + "encoding_dsv4.py", + "tests/test_input_1.json", + "tests/test_input_2.json", + "tests/test_input_3.json", + "tests/test_input_4.json", + "tests/test_output_1.txt", + "tests/test_output_2.txt", + "tests/test_output_3.txt", + "tests/test_output_4.txt", +) + + +class CandidateConstructionError(RuntimeError): + """The isolated service cannot install its reviewed request surface.""" + + +def _manifest_entries() -> dict[str, str]: + if not MANIFEST.is_file() or MANIFEST.is_symlink(): + raise CandidateConstructionError("official encoding manifest is missing or unsafe") + entries: dict[str, str] = {} + for line in MANIFEST.read_text(encoding="utf-8").splitlines(): + match = re.fullmatch(r"([0-9a-f]{64}) ([A-Za-z0-9_./-]+)", line) + if match is None or match.group(2) in entries: + raise CandidateConstructionError("official encoding manifest is malformed") + entries[match.group(2)] = match.group(1) + if tuple(entries) != REQUIRED_ASSETS: + raise CandidateConstructionError("official encoding manifest asset set changed") + return entries + + +def verify_official_assets() -> dict[str, str]: + """Verify the exact official source and vector set once at construction.""" + entries = _manifest_entries() + root = ENCODING_ROOT.resolve() + for relative, expected in entries.items(): + path = ENCODING_ROOT / relative + if not path.is_file() or path.is_symlink() or path.resolve().parent != (root / relative).parent: + raise CandidateConstructionError("official encoding asset is missing or unsafe") + actual = hashlib.sha256(path.read_bytes()).hexdigest() + if actual != expected: + raise CandidateConstructionError("official encoding asset digest mismatch") + return entries + + +def _load_official_encoder() -> ModuleType: + verify_official_assets() + path = ENCODING_ROOT / "encoding_dsv4.py" + module = ModuleType("mtplx_dsv4_0731_official_encoding") + module.__file__ = str(path) + # Compile the already-verified source bytes directly. This cannot select an + # ignored or stale __pycache__ artifact in place of the reviewed encoder. + code = compile(path.read_bytes(), str(path), "exec") + exec(code, module.__dict__) + return module + + +def _self_test(encoding: ModuleType) -> None: + """Run all four official byte vectors before installing the request path.""" + for case in range(1, 5): + payload = json.loads((ENCODING_ROOT / f"tests/test_input_{case}.json").read_text(encoding="utf-8")) + if case == 1: + messages = payload["messages"] + messages[0]["tools"] = payload["tools"] + else: + messages = payload + mode = "chat" if case == 4 else "thinking" + expected = (ENCODING_ROOT / f"tests/test_output_{case}.txt").read_text(encoding="utf-8") + if encoding.encode_messages(messages, thinking_mode=mode) != expected: + raise CandidateConstructionError(f"official encoding vector {case} failed") + + +def _message_dict(message: Any) -> dict[str, Any]: + if isinstance(message, dict): + return dict(message) + dump = getattr(message, "model_dump", None) + if callable(dump): + value = dump(exclude_none=True) + if isinstance(value, dict): + return value + value: dict[str, Any] = {} + for field in ("role", "content", "name", "tool_call_id", "tool_calls", "reasoning_content"): + if hasattr(message, field): + item = getattr(message, field) + if item is not None: + value[field] = item + if "role" not in value: + raise CandidateConstructionError("request message has no role") + return value + + +def _install_encoder(server: ModuleType, encoding: ModuleType): + encode_text = server._encode_rendered_chat_text + + def encode_messages( + tokenizer: Any, + messages: list[Any], + *, + enable_thinking: bool, + reasoning_effort: str | None = None, + strip_assistant_reasoning_history: bool = False, + scoped_reasoning_history: bool = False, + add_generation_prompt: bool = True, + tools: list[dict[str, Any]] | None = None, + tool_choice: Any = None, + tool_prompt_mode: str = "native", + template_observability: dict[str, Any] | None = None, + ) -> list[int]: + del strip_assistant_reasoning_history, scoped_reasoning_history, tool_prompt_mode + if tool_choice not in (None, "auto"): + raise CandidateConstructionError("V4-0731 candidate does not support forced tool_choice") + effort = reasoning_effort or "low" + if effort not in encoding.REASONING_EFFORT_PROMPTS: + raise CandidateConstructionError("reasoning_effort must be one of: low, high, max") + prepared = [_message_dict(message) for message in messages] + if not prepared: + prepared = [{"role": "user", "content": ""}] + if tools: + if prepared[0].get("role") != "system": + prepared.insert(0, {"role": "system", "content": "", "tools": tools}) + else: + prepared[0] = {**prepared[0], "tools": tools} + mode = "thinking" if enable_thinking else "chat" + rendered = encoding.encode_messages( + prepared, + thinking_mode=mode, + reasoning_effort=effort, + ) + if not add_generation_prompt: + suffix = encoding.ASSISTANT_SP_TOKEN + ( + encoding.thinking_start_token if enable_thinking else encoding.thinking_end_token + ) + if rendered.endswith(suffix): + rendered = rendered[: -len(suffix)] + if template_observability is not None: + template_observability.update( + { + "backend_chat_encoding": ENCODER_NAME, + "encoding_source_revision": SOURCE_REVISION, + } + ) + return encode_text(tokenizer, rendered) + + return encode_messages + + +def _install_completion_parser(server: ModuleType, encoding: ModuleType): + stock = server._parse_generated_tool_calls_or_content + dsml_marker = f"<{encoding.dsml_token}{encoding.tool_calls_block_name}>" + + def parse_generated_tool_calls_or_content( + text: str, + *, + tools: list[dict[str, Any]], + tokenizer: Any | None = None, + state: Any | None = None, + response_id: str | None = None, + stream: bool = False, + ): + if dsml_marker not in text: + return stock( + text, + tools=tools, + tokenizer=tokenizer, + state=state, + response_id=response_id, + stream=stream, + ) + completion = text if text.endswith(encoding.eos_token) else text + encoding.eos_token + mode = "thinking" if encoding.thinking_end_token in completion.split(dsml_marker, 1)[0] else "chat" + try: + parsed = encoding.parse_message_from_completion_text(completion, thinking_mode=mode) + except (AssertionError, ValueError) as error: + raise CandidateConstructionError("malformed V4-0731 DSML completion") from error + calls = parsed.get("tool_calls") or None + return calls, None + + return parse_generated_tool_calls_or_content + + +def _install_actual_tool_extractor(server: ModuleType, encoding: ModuleType) -> None: + """Install official DSML completion parsing at the live response call site.""" + from mtplx.server.omlx_bridge import ToolCallExtraction + + stock = server.omlx_extract_tool_calls_with_thinking + dsml_marker = f"<{encoding.dsml_token}{encoding.tool_calls_block_name}>" + + def extract( + thinking_content: str, + regular_content: str, + tokenizer: Any | None, + tools: list[dict[str, Any]] | None = None, + ) -> ToolCallExtraction: + combined = thinking_content + regular_content + if dsml_marker not in combined: + return stock(thinking_content, regular_content, tokenizer, tools) + mode = "thinking" if thinking_content else "chat" + completion = ( + thinking_content + encoding.thinking_end_token + regular_content + if mode == "thinking" + else regular_content + ) + if not completion.endswith(encoding.eos_token): + completion += encoding.eos_token + try: + parsed = encoding.parse_message_from_completion_text(completion, thinking_mode=mode) + except (AssertionError, ValueError) as error: + raise CandidateConstructionError("malformed V4-0731 DSML completion") from error + calls = parsed.get("tool_calls") or None + if calls: + calls = [ + {**call, "id": str(call.get("id") or f"call_{uuid.uuid4().hex[:24]}")} + for call in calls + ] + return ToolCallExtraction( + cleaned_text=str(parsed.get("content") or ""), + tool_calls=calls, + cleaned_thinking=str(parsed.get("reasoning_content") or ""), + parser_source="deepseek_v4_0731_official", + status="parsed" if calls else "no_tool", + raw_tool_markup_suppressed=True, + ) + + server.omlx_extract_tool_calls_with_thinking = extract + + +def _install_stream_translator(server: ModuleType, encoding: ModuleType) -> None: + """Buffer the official DSML envelope so streaming never leaks it as text.""" + stock_class = server._ToolAwareContentStreamTranslator + dsml_marker = f"<{encoding.dsml_token}{encoding.tool_calls_block_name}>" + + class DSV40731StreamTranslator: + def __init__(self, *, tools, argument_chunk_chars, tokenizer=None, **kwargs) -> None: + self._tools = tools + self._argument_chunk_chars = argument_chunk_chars + self._tokenizer = tokenizer + self._stock = stock_class( + tools=tools, + argument_chunk_chars=argument_chunk_chars, + tokenizer=tokenizer, + **kwargs, + ) + self._pending = "" + self._mode = "undecided" + self.tool_calls = None + self.fallback_reason = None + self.tool_parser_dialect = "deepseek_v4_0731_official" + self._suppressed = False + + @property + def has_tool_calls(self): + return bool(self.tool_calls) if self._mode == "dsml" else self._stock.has_tool_calls + + @property + def has_emitted_tool_deltas(self): + return False if self._mode == "dsml" else self._stock.has_emitted_tool_deltas + + @property + def suppressed_tool_markup(self): + return self._suppressed or self._stock.suppressed_tool_markup + + @property + def buffering_tool_call(self): + return self._mode == "dsml" or self._stock.buffering_tool_call + + @property + def tool_argument_in_progress(self): + return self._mode == "dsml" or self._stock.tool_argument_in_progress + + @property + def ready_to_finish_tool_turn(self): + return False if self._mode == "dsml" else self._stock.ready_to_finish_tool_turn + + @property + def invalid_trailing_after_tool_call(self): + return False if self._mode == "dsml" else self._stock.invalid_trailing_after_tool_call + + def feed(self, field: str, text: str): + if self._mode == "stock": + return self._stock.feed(field, text) + if field != "content": + return self._stock.feed(field, text) + self._pending += text + stripped = self._pending.lstrip() + if dsml_marker in stripped: + self._mode = "dsml" + self._suppressed = True + return [] + if dsml_marker.startswith(stripped): + return [] + self._mode = "stock" + pending, self._pending = self._pending, "" + return self._stock.feed(field, pending) + + def finish(self, *, defer_content_resolution: bool = False): + if self._mode != "dsml": + if self._pending: + self._stock.feed("content", self._pending) + self._pending = "" + return self._stock.finish(defer_content_resolution=defer_content_resolution) + extraction = server.omlx_extract_tool_calls_with_thinking( + "", self._pending, self._tokenizer, self._tools + ) + self.tool_calls = extraction.tool_calls + self._pending = "" + if not self.tool_calls: + raise CandidateConstructionError("official DSML stream ended without tool calls") + return list( + server._stream_tool_call_deltas( + self.tool_calls, + argument_chunk_chars=self._argument_chunk_chars, + ) + ) + + def resolve_deferred_content(self, *, has_tool_calls: bool): + if self._mode == "dsml": + return [] + return self._stock.resolve_deferred_content(has_tool_calls=has_tool_calls) + + server._ToolAwareContentStreamTranslator = DSV40731StreamTranslator + + +def _install_reasoning_policy(server: ModuleType) -> None: + def normalize(value: Any, *, default: str = "low") -> str: + effort = str(value or default).strip().lower() + if effort not in {"auto", "low", "high", "max"}: + raise ValueError("reasoning_effort must be one of: auto, low, high, max") + return effort + + def for_state( + state: Any, + *, + thinking_enabled: bool, + request_effort: str | None = None, + allow_client_controls: bool = True, + ) -> str | None: + if not thinking_enabled: + return None + raw = request_effort if request_effort is not None and allow_client_controls else state.args.reasoning_effort + effort = normalize(raw, default="low") + return "low" if effort == "auto" else effort + + server._normalize_reasoning_effort = normalize + server._reasoning_effort_for_state = for_state + + +def _install_construction_identity(server: ModuleType) -> None: + manifest_digest = hashlib.sha256(MANIFEST.read_bytes()).hexdigest() + + def apply_profile(_tokenizer: Any, _args: Any) -> dict[str, Any]: + return { + "profile": ENCODER_NAME, + "source": "official_python_encoder", + "path": None, + "applied": True, + "sha256": manifest_digest, + } + + server._apply_chat_template_profile = apply_profile + server._template_hash = lambda _tokenizer: f"{ENCODER_NAME}:{manifest_digest}" + server._template_supports_scoped_reasoning = lambda _tokenizer: True + + +def install_candidate_surface(server: ModuleType) -> dict[str, str]: + """Install the verified encoder/parser directly into the imported server.""" + encoding = _load_official_encoder() + _self_test(encoding) + if not hasattr(server, "_encode_rendered_chat_text"): + # Unit fixture: retain the same strict no-special-token encoding contract. + server._encode_rendered_chat_text = lambda tokenizer, text: list( + tokenizer.encode(text, add_special_tokens=False) + ) + server._encode_messages = _install_encoder(server, encoding) + server._parse_generated_tool_calls_or_content = _install_completion_parser(server, encoding) + if hasattr(server, "omlx_extract_tool_calls_with_thinking"): + _install_actual_tool_extractor(server, encoding) + if hasattr(server, "_ToolAwareContentStreamTranslator"): + _install_stream_translator(server, encoding) + _install_reasoning_policy(server) + _install_construction_identity(server) + server._DSV4_0731_ENCODER_INSTALLED = True + manifest_digest = hashlib.sha256(MANIFEST.read_bytes()).hexdigest() + return { + "encoder": ENCODER_NAME, + "source_revision": SOURCE_REVISION, + "asset_set_sha256": manifest_digest, + } + + +def main() -> int: + if sys.argv[1:]: + raise CandidateConstructionError("candidate entrypoint accepts no arguments") + from mtplx.server import openai as server + from mtplx.cli import main as mtplx_main + + install_candidate_surface(server) + return mtplx_main( + [ + "serve", + "--host", "127.0.0.1", + "--port", "8081", + "--model", "/Users/davidtai/models/DeepSeek-V4-Flash-2bit-DQ-mtp", + "--model-id", "deepseek-v4-0731-candidate", + "--reasoning", "on", + "--reasoning-effort", "low", + "--reasoning-parser", "qwen3", + "--tool-prompt-mode", "native", + "--chat-template-profile", "tokenizer", + "--warmup-tokens", "0", + "--no-stats-footer", + ] + ) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/services/deepseek-v4-0731/encoding/ATTRIBUTION.md b/services/deepseek-v4-0731/encoding/ATTRIBUTION.md index 6a273434..695d0ce9 100644 --- a/services/deepseek-v4-0731/encoding/ATTRIBUTION.md +++ b/services/deepseek-v4-0731/encoding/ATTRIBUTION.md @@ -1,22 +1,14 @@ -# DeepSeek 0731 encoding attribution +# DeepSeek-V4-Flash-0731 encoding attribution -This directory reserves the pinned chat-encoding asset slot for the isolated -`deepseek-v4-0731` candidate. The requested official DeepSeek 0731 source -revision is -`7872f01b1d1fe23eabc4c98b48bffcef5a386062`. +The files listed in `SHA256SUMS` are vendored byte-for-byte from the official +Hugging Face repository: -Intended upstream attribution: DeepSeek AI, DeepSeek-V3.1, -`assets/chat_template.jinja`. At implementation time that exact revision did -not resolve from the public upstream history, so `chat_template.jinja` is a -minimal review fixture, **not a claim of a byte-for-byte retrieved upstream -file**. Do not promote this service until an operator replaces it with the -retrieved official bytes and updates this note plus `SHA256SUMS` in the same -reviewed commit. +- Repository: `deepseek-ai/DeepSeek-V4-Flash-0731` +- Source revision: `7872f01b1d1fe23eabc4c98b48bffcef5a386062` +- Encoder: `encoding/encoding_dsv4.py` +- Vectors: `encoding/tests/test_input_{1..4}.json` and + `encoding/tests/test_output_{1..4}.txt` +- Upstream owner: DeepSeek AI -The asset is retained locally for review and reproducibility; it is not fetched -at service start. `SHA256SUMS` is the authority at installation time. A -missing, symlinked, malformed, or mismatched asset is a hard error, never a -fallback to a local tokenizer template. - -The small renderer is intentionally kept separate from MTPLX's existing chat -template routes until this candidate has a promotion receipt. +The service verifies all nine files at its construction boundary. It never +downloads encoding code or falls back to a tokenizer chat template at runtime. diff --git a/services/deepseek-v4-0731/encoding/SHA256SUMS b/services/deepseek-v4-0731/encoding/SHA256SUMS index dfc062ca..5e84ab36 100644 --- a/services/deepseek-v4-0731/encoding/SHA256SUMS +++ b/services/deepseek-v4-0731/encoding/SHA256SUMS @@ -1 +1,9 @@ -03f2686beff14c3d9040894a2b658d9f1917be90bc1d90597502fc2562f0ec2a chat_template.jinja +abc0d26120250dda0ae077dc64aa28836026e61e970854aaeb792445e6a0dde6 encoding_dsv4.py +10e0c074c977c3a80daab758af28219c6b1c2bd7f3f5cf2890c84b361cc32897 tests/test_input_1.json +c44ae0db20fafff38a6021e8068d7ed6e28605d76cafd20be80b03398509f447 tests/test_input_2.json +37bf8ef95e0411ea5f411be0b02fbafec7363438b6ccefddca0c52ec9aeaf69a tests/test_input_3.json +c45bbd0a1b7a2f75033d8db4ba74ee5b7653bd01114045a52f79d3b387663465 tests/test_input_4.json +9b366d9d2eac842a6e890594aac0b58648e5623717202b33497afadf03e26540 tests/test_output_1.txt +ca66b01a1ac3a204bb032c928fb607d4877171e998a50f7f89f39fa821b75665 tests/test_output_2.txt +b3b1cd8748b7b90d3c6be6da3f786f12e4d70be073bd445ea162dfad4dc01a64 tests/test_output_3.txt +60e1643840ba9e4aeede450feb7b0498fa66ee24e4a939d48855ce04ec6fc375 tests/test_output_4.txt diff --git a/services/deepseek-v4-0731/encoding/chat_template.jinja b/services/deepseek-v4-0731/encoding/chat_template.jinja deleted file mode 100644 index f92b3444..00000000 --- a/services/deepseek-v4-0731/encoding/chat_template.jinja +++ /dev/null @@ -1,17 +0,0 @@ -{# - DeepSeek-V3.1 0731 chat encoding asset. - Vendored unchanged from the official release source identified in ATTRIBUTION.md. - The serving wrapper uses render.py's deliberately small, dependency-free - implementation of this token grammar; this file remains the reviewed source - asset and is covered by SHA256SUMS. -#} -{{ bos_token }}{{ system_prompt }} -{%- for message in messages %} - {%- if message['role'] == 'user' %} - {{ '<｜User｜>' + message['content'] }} - {%- elif message['role'] == 'assistant' and message['tool_calls'] is defined %} - {{ '<｜Assistant｜><｜tool▁calls▁begin｜>' }} - {%- elif message['role'] == 'tool' %} - {{ '<｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>' }} - {%- endif %} -{%- endfor %} diff --git a/services/deepseek-v4-0731/encoding/encoding_dsv4.py b/services/deepseek-v4-0731/encoding/encoding_dsv4.py new file mode 100644 index 00000000..66f2d30e --- /dev/null +++ b/services/deepseek-v4-0731/encoding/encoding_dsv4.py @@ -0,0 +1,760 @@ +""" +DeepSeek-V4 Encoding + +A self-contained implementation for encoding/decoding DeepSeek-V4 chat messages +with tool calling, thinking mode, and quick instruction task support. +""" + +from typing import Any, Dict, List, Union, Optional, Tuple +import copy +import json +import re + +# ============================================================ +# Special Tokens +# ============================================================ + +bos_token: str = "<｜begin▁of▁sentence｜>" +eos_token: str = "<｜end▁of▁sentence｜>" +thinking_start_token: str = "" +thinking_end_token: str = "" +dsml_token: str = "｜DSML｜" + +USER_SP_TOKEN = "<｜User｜>" +ASSISTANT_SP_TOKEN = "<｜Assistant｜>" +LATEST_REMINDER_SP_TOKEN = "<｜latest_reminder｜>" + +# Task special tokens for internal classification tasks +DS_TASK_SP_TOKENS = { + "action": "<｜action｜>", + "query": "<｜query｜>", + "authority": "<｜authority｜>", + "domain": "<｜domain｜>", + "title": "<｜title｜>", + "read_url": "<｜read_url｜>", +} +VALID_TASKS = set(DS_TASK_SP_TOKENS.keys()) + +# ============================================================ +# Templates +# ============================================================ + +system_msg_template: str = "{content}" +user_msg_template: str = "{content}" +latest_reminder_msg_template: str = "{content}" +assistant_msg_template: str = "{reasoning}{content}{tool_calls}" + eos_token +assistant_msg_wo_eos_template: str = "{reasoning}{content}{tool_calls}" +thinking_template: str = "{reasoning_content}" + +response_format_template: str = ( + "## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{schema}" +) +tool_call_template: str = ( + "<{dsml_token}invoke name=\"{name}\">\n{arguments}\n" +) +tool_calls_template = ( + "<{dsml_token}{tc_block_name}>\n{tool_calls}\n" +) +tool_calls_block_name: str = "tool_calls" + +tool_output_template: str = ( + "{content}" +) + +# Reasoning effort levels. In thinking mode, the prompt for the selected level is +# prepended at the very beginning of the conversation. `low` is the default and +# adds nothing. +REASONING_EFFORT_PROMPTS: Dict[str, str] = { + "low": "", + "high": ( + "Reasoning Effort: Absolute maximum with no shortcuts permitted.\n" + "You MUST be very thorough in your thinking and comprehensively decompose the problem to resolve the root cause, rigorously stress-testing your logic against all potential paths, edge cases, and adversarial scenarios.\n" + "Explicitly write out your entire deliberation process, documenting every intermediate step, considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.\n\n" + ), + "max": ( + "Reasoning Effort: Beyond maximum — exhaustive, relentless, and uncompromising.\n" + "You MUST reason with the utmost depth and rigor, leaving absolutely nothing to chance: exhaustively decompose the problem into its most fundamental components, trace every causal chain to its root, and resolve the underlying cause rather than any surface symptom.\n" + "Do not stop reasoning until you have independently verified the solution from multiple angles and are certain that no assumption remains unchecked and no error remains undiscovered.\n\n" + ), +} +DEFAULT_REASONING_EFFORT = "low" + +TOOLS_TEMPLATE = """## Tools + +You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<{dsml_token}tool_calls>" block like the following: + +<{dsml_token}tool_calls> +<{dsml_token}invoke name="$TOOL_NAME"> +<{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE +... + +<{dsml_token}invoke name="$TOOL_NAME2"> +... + + + +String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`. + +If thinking_mode is enabled (triggered by {thinking_start_token}), you MUST output your complete reasoning inside {thinking_start_token}...{thinking_end_token} BEFORE any tool calls or final response. + +Otherwise, output directly after {thinking_end_token} with tool calls or final response. + +### Available Tool Schemas + +{tool_schemas} + +You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls. +""" + +# ============================================================ +# Utility Functions +# ============================================================ + +def to_json(value: Any) -> str: + """Serialize a value to JSON string.""" + try: + return json.dumps(value, ensure_ascii=False) + except: + return json.dumps(value, ensure_ascii=True) + + +def tools_from_openai_format(tools): + """Extract function definitions from OpenAI-format tool list.""" + return [tool["function"] for tool in tools] + + +def tool_calls_from_openai_format(tool_calls): + """Convert OpenAI-format tool calls to internal format.""" + return [ + { + "name": tool_call["function"]["name"], + "arguments": tool_call["function"]["arguments"], + } + for tool_call in tool_calls + ] + + +def tool_calls_to_openai_format(tool_calls): + """Convert internal tool calls to OpenAI format.""" + return [ + { + "type": "function", + "function": { + "name": tool_call["name"], + "arguments": tool_call["arguments"], + } + } + for tool_call in tool_calls + ] + + +def encode_arguments_to_dsml(tool_call: Dict[str, str]) -> str: + """ + Encode tool call arguments into DSML parameter format. + + Args: + tool_call: Dict with "name" and "arguments" (JSON string) keys. + + Returns: + DSML-formatted parameter string. + """ + p_dsml_template = '<{dsml_token}parameter name="{key}" string="{is_str}">{value}' + P_dsml_strs = [] + + try: + arguments = json.loads(tool_call["arguments"]) + except Exception as err: + arguments = {"arguments": tool_call["arguments"]} + + for k, v in arguments.items(): + p_dsml_str = p_dsml_template.format( + dsml_token=dsml_token, + key=k, + is_str="true" if isinstance(v, str) else "false", + value=v if isinstance(v, str) else to_json(v), + ) + P_dsml_strs.append(p_dsml_str) + + return "\n".join(P_dsml_strs) + + +def decode_dsml_to_arguments(tool_name: str, tool_args: Dict[str, Tuple[str, str]]) -> Dict[str, str]: + """ + Decode DSML parameters back to a tool call dict. + + Args: + tool_name: Name of the tool. + tool_args: Dict mapping param_name -> (value, is_string_flag). + + Returns: + Dict with "name" and "arguments" (JSON string) keys. + """ + def _decode_value(key: str, value: str, string: str): + if string == "true": + value = to_json(value) + return f"{to_json(key)}: {value}" + + tool_args_json = "{" + ", ".join([_decode_value(k, v, string=is_str) for k, (v, is_str) in tool_args.items()]) + "}" + return dict(name=tool_name, arguments=tool_args_json) + + +def render_tools(tools: List[Dict[str, Union[str, Dict[str, Any]]]]) -> str: + """ + Render tool schemas into the system prompt format. + + Args: + tools: List of tool schema dicts (each with name, description, parameters). + + Returns: + Formatted tools section string. + """ + tools_json = [to_json(t) for t in tools] + + return TOOLS_TEMPLATE.format( + tool_schemas="\n".join(tools_json), + dsml_token=dsml_token, + thinking_start_token=thinking_start_token, + thinking_end_token=thinking_end_token, + ) + + +def find_last_user_index(messages: List[Dict[str, Any]]) -> int: + """Find the index of the last user/developer message.""" + last_user_index = -1 + for idx in range(len(messages) - 1, -1, -1): + if messages[idx].get("role") in ["user", "developer"]: + last_user_index = idx + break + return last_user_index + + +# ============================================================ +# Message Rendering +# ============================================================ + +def render_message(index: int, messages: List[Dict[str, Any]], thinking_mode: str, drop_thinking: bool = True, reasoning_effort: Optional[str] = None) -> str: + """ + Render a single message at the given index into its encoded string form. + + This is the core function that converts each message in the conversation + into the DeepSeek-V4 format. + + Args: + index: Index of the message to render. + messages: Full list of messages in the conversation. + thinking_mode: Either "chat" or "thinking". + drop_thinking: Whether to drop reasoning content from earlier turns. + reasoning_effort: Reasoning effort level, one of "low", "high", "max". + None is treated as "low". + + Returns: + Encoded string for this message. + """ + assert 0 <= index < len(messages) + assert thinking_mode in ["chat", "thinking"], f"Invalid thinking_mode `{thinking_mode}`" + + prompt = "" + msg = messages[index] + last_user_idx = find_last_user_index(messages) + + role = msg.get("role") + content = msg.get("content") + tools = msg.get("tools") + response_format = msg.get("response_format") + tool_calls = msg.get("tool_calls") + reasoning_content = msg.get("reasoning_content") + wo_eos = msg.get("wo_eos", False) + + if tools: + tools = tools_from_openai_format(tools) + if tool_calls: + tool_calls = tool_calls_from_openai_format(tool_calls) + + # Reasoning effort prefix (only at index 0 in thinking mode; "low" adds nothing) + reasoning_effort = reasoning_effort or DEFAULT_REASONING_EFFORT + assert reasoning_effort in REASONING_EFFORT_PROMPTS, \ + f"Invalid reasoning effort: {reasoning_effort}, expected one of {list(REASONING_EFFORT_PROMPTS)}" + if index == 0 and thinking_mode == "thinking": + prompt += REASONING_EFFORT_PROMPTS[reasoning_effort] + + if role == "system": + prompt += system_msg_template.format(content=content or "") + if tools: + prompt += "\n\n" + render_tools(tools) + if response_format: + prompt += "\n\n" + response_format_template.format(schema=to_json(response_format)) + + elif role == "developer": + assert content, f"Invalid message for role `{role}`: {msg}" + + content_developer = USER_SP_TOKEN + content_developer += content + + if tools: + content_developer += "\n\n" + render_tools(tools) + if response_format: + content_developer += "\n\n" + response_format_template.format(schema=to_json(response_format)) + + prompt += user_msg_template.format(content=content_developer) + + elif role == "user": + prompt += USER_SP_TOKEN + + # Handle content blocks (tool results mixed with text) + content_blocks = msg.get("content_blocks") + if content_blocks: + parts = [] + for block in content_blocks: + block_type = block.get("type") + if block_type == "text": + parts.append(block.get("text", "")) + elif block_type == "tool_result": + tool_content = block.get("content", "") + if isinstance(tool_content, list): + text_parts = [] + for b in tool_content: + if b.get("type") == "text": + text_parts.append(b.get("text", "")) + else: + text_parts.append(f"[Unsupported {b.get('type')}]") + tool_content = "\n\n".join(text_parts) + parts.append(tool_output_template.format(content=tool_content)) + else: + parts.append(f"[Unsupported {block_type}]") + prompt += "\n\n".join(parts) + else: + prompt += content or "" + + elif role == "latest_reminder": + prompt += LATEST_REMINDER_SP_TOKEN + latest_reminder_msg_template.format(content=content) + + elif role == "tool": + raise NotImplementedError("deepseek_v4 merges tool messages into user; please preprocess with merge_tool_messages()") + + elif role == "assistant": + thinking_part = "" + tc_content = "" + + if tool_calls: + tc_list = [ + tool_call_template.format( + dsml_token=dsml_token, + name=tc.get("name"), + arguments=encode_arguments_to_dsml(tc) + ) + for tc in tool_calls + ] + tc_content += '\n\n' + tool_calls_template.format( + dsml_token=dsml_token, + tool_calls="\n".join(tc_list), + tc_block_name=tool_calls_block_name, + ) + + summary_content = content or "" + rc = reasoning_content or "" + + # Check if previous message has a task - if so, this is a task output (no thinking) + prev_has_task = index - 1 >= 0 and messages[index - 1].get("task") is not None + + if thinking_mode == "thinking" and not prev_has_task: + if not drop_thinking or index > last_user_idx: + thinking_part = thinking_template.format(reasoning_content=rc) + thinking_end_token + else: + thinking_part = "" + + if wo_eos: + prompt += assistant_msg_wo_eos_template.format( + reasoning=thinking_part, + content=summary_content, + tool_calls=tc_content, + ) + else: + prompt += assistant_msg_template.format( + reasoning=thinking_part, + content=summary_content, + tool_calls=tc_content, + ) + else: + raise NotImplementedError(f"Unknown role: {role}") + + # Append transition tokens based on what follows + if index + 1 < len(messages) and messages[index + 1].get("role") not in ["assistant", "latest_reminder"]: + return prompt + + task = messages[index].get("task") + if task is not None: + # Task special token for internal classification tasks + assert task in VALID_TASKS, f"Invalid task: '{task}'. Valid tasks are: {list(VALID_TASKS)}" + task_sp_token = DS_TASK_SP_TOKENS[task] + + if task != "action": + # Non-action tasks: append task sp token directly after the message + prompt += task_sp_token + else: + # Action task: append Assistant + thinking token + action sp token + prompt += ASSISTANT_SP_TOKEN + prompt += thinking_end_token if thinking_mode != "thinking" else thinking_start_token + prompt += task_sp_token + + elif messages[index].get("role") in ["user", "developer"]: + # Normal generation: append Assistant + thinking token + prompt += ASSISTANT_SP_TOKEN + if not drop_thinking and thinking_mode == "thinking": + prompt += thinking_start_token + elif drop_thinking and thinking_mode == "thinking" and index >= last_user_idx: + prompt += thinking_start_token + else: + prompt += thinking_end_token + + return prompt + + +# ============================================================ +# Preprocessing +# ============================================================ + +def merge_tool_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Merge tool messages into the preceding user message using content_blocks format. + + DeepSeek-V4 does not have a standalone "tool" role; instead, tool results + are encoded as blocks within user messages. + + This function converts a standard OpenAI-format conversation (with separate + "tool" role messages) into V4 format where tool results are merged into + user messages. + + Args: + messages: List of message dicts in OpenAI format. + + Returns: + Processed message list with tool messages merged into user messages. + """ + merged: List[Dict[str, Any]] = [] + + for msg in messages: + msg = copy.deepcopy(msg) + role = msg.get("role") + + if role == "tool": + # Convert tool message to a user message with tool_result block + tool_block = { + "type": "tool_result", + "tool_use_id": msg.get("tool_call_id", ""), + "content": msg.get("content", ""), + } + # Merge into previous message if it's already a user (merged tool) + if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1]: + merged[-1]["content_blocks"].append(tool_block) + else: + merged.append({ + "role": "user", + "content_blocks": [tool_block], + }) + elif role == "user": + text_block = {"type": "text", "text": msg.get("content", "")} + if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1] and merged[-1].get("task") is None: + merged[-1]["content_blocks"].append(text_block) + else: + new_msg = { + "role": "user", + "content": msg.get("content", ""), + "content_blocks": [text_block], + } + # Preserve extra fields (task, wo_eos, mask, etc.) + for key in ("task", "wo_eos", "mask"): + if key in msg: + new_msg[key] = msg[key] + merged.append(new_msg) + else: + merged.append(msg) + + return merged + + +def sort_tool_results_by_call_order(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Sort tool_result blocks within user messages by the order of tool_calls + in the preceding assistant message. + + Args: + messages: Preprocessed message list (after merge_tool_messages). + + Returns: + Message list with sorted tool result blocks. + """ + last_tool_call_order: Dict[str, int] = {} + + for msg in messages: + role = msg.get("role") + if role == "assistant" and msg.get("tool_calls"): + last_tool_call_order = {} + for idx, tc in enumerate(msg["tool_calls"]): + tc_id = tc.get("id") or tc.get("function", {}).get("id", "") + if tc_id: + last_tool_call_order[tc_id] = idx + + elif role == "user" and msg.get("content_blocks"): + tool_blocks = [b for b in msg["content_blocks"] if b.get("type") == "tool_result"] + if len(tool_blocks) > 1 and last_tool_call_order: + sorted_blocks = sorted( + tool_blocks, + key=lambda b: last_tool_call_order.get(b.get("tool_use_id", ""), 0) + ) + sorted_idx = 0 + new_blocks = [] + for block in msg["content_blocks"]: + if block.get("type") == "tool_result": + new_blocks.append(sorted_blocks[sorted_idx]) + sorted_idx += 1 + else: + new_blocks.append(block) + msg["content_blocks"] = new_blocks + + return messages + + +# ============================================================ +# Main Encoding Function +# ============================================================ + +def encode_messages( + messages: List[Dict[str, Any]], + thinking_mode: str, + context: Optional[List[Dict[str, Any]]] = None, + drop_thinking: bool = True, + add_default_bos_token: bool = True, + reasoning_effort: Optional[str] = None, +) -> str: + """ + Encode a list of messages into the DeepSeek-V4 prompt format. + + This is the main entry point for encoding conversations. It handles: + - BOS token insertion + - Thinking mode with optional reasoning content dropping + - Tool message merging into user messages + - Multi-turn conversation context + + Args: + messages: List of message dicts to encode. + thinking_mode: Either "chat" or "thinking". + context: Optional preceding context messages (already encoded prefix). + drop_thinking: If True, drop reasoning_content from earlier assistant turns + (only keep reasoning for messages after the last user message). + add_default_bos_token: Whether to prepend BOS token at conversation start. + reasoning_effort: Reasoning effort level, one of "low", "high", "max". + Only takes effect in thinking mode. None is treated as "low". + + Returns: + The encoded prompt string. + """ + context = context if context else [] + + # Preprocess: merge tool messages and sort tool results + messages = merge_tool_messages(messages) + messages = sort_tool_results_by_call_order(context + messages)[len(context):] + if context: + context = merge_tool_messages(context) + context = sort_tool_results_by_call_order(context) + + full_messages = context + messages + + prompt = bos_token if add_default_bos_token and len(context) == 0 else "" + + # Resolve drop_thinking: if any message has tools defined, don't drop thinking + effective_drop_thinking = drop_thinking + if any(m.get("tools") for m in full_messages): + effective_drop_thinking = False + + if thinking_mode == "thinking" and effective_drop_thinking: + full_messages = _drop_thinking_messages(full_messages) + # After dropping, recalculate how many messages to render + # (context may have shrunk too) + num_to_render = len(full_messages) - len(_drop_thinking_messages(context)) + context_len = len(full_messages) - num_to_render + else: + num_to_render = len(messages) + context_len = len(context) + + for idx in range(num_to_render): + prompt += render_message( + idx + context_len, + full_messages, + thinking_mode=thinking_mode, + drop_thinking=effective_drop_thinking, + reasoning_effort=reasoning_effort, + ) + + return prompt + + +def _drop_thinking_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Drop reasoning_content and non-essential messages before the last user message. + + Behavior: + - Messages with role in ["user", "system", "tool", "latest_reminder"] are always kept. + - Messages at or after the last user index are always kept. + - Assistant messages before the last user get reasoning_content removed. + - Developer messages before the last user are dropped entirely. + """ + last_user_idx = find_last_user_index(messages) + result = [] + keep_roles = {"user", "system", "tool", "latest_reminder", "direct_search_results"} + + for idx, msg in enumerate(messages): + role = msg.get("role") + if role in keep_roles or idx >= last_user_idx: + result.append(msg) + elif role == "assistant": + msg = copy.copy(msg) + msg.pop("reasoning_content", None) + result.append(msg) + # developer and other roles before last_user_idx are dropped + + return result + + +# ============================================================ +# Parsing (Decoding model output) +# ============================================================ + +def _read_until_stop(index: int, text: str, stop: List[str]) -> Tuple[int, str, Optional[str]]: + """ + Read text from index until one of the stop strings is found. + + Returns: + Tuple of (new_index, content_before_stop, matched_stop_string_or_None). + """ + min_pos = len(text) + matched_stop = None + + for s in stop: + pos = text.find(s, index) + if pos != -1 and pos < min_pos: + min_pos = pos + matched_stop = s + + if matched_stop: + content = text[index:min_pos] + return min_pos + len(matched_stop), content, matched_stop + else: + content = text[index:] + return len(text), content, None + + +def parse_tool_calls(index: int, text: str) -> Tuple[int, Optional[str], List[Dict[str, str]]]: + """ + Parse DSML tool calls from text starting at the given index. + + Args: + index: Starting position in text. + text: The full text to parse. + + Returns: + Tuple of (new_index, last_stop_token, list_of_tool_call_dicts). + Each tool call dict has "name" and "arguments" keys. + """ + tool_calls: List[Dict[str, Any]] = [] + stop_token = None + tool_calls_end_token = f"" + + while index < len(text): + index, _, stop_token = _read_until_stop(index, text, [f"<{dsml_token}invoke", tool_calls_end_token]) + if _ != ">\n": + raise ValueError(f"Tool call format error: expected '>\\n' but got '{_}'") + + if stop_token == tool_calls_end_token: + break + + if stop_token is None: + raise ValueError("Missing special token in tool calls") + + index, tool_name_content, stop_token = _read_until_stop(index, text, [f"<{dsml_token}parameter", f"\n$', tool_name_content, flags=re.DOTALL) + if len(p_tool_name) != 1: + raise ValueError(f"Tool name format error: '{tool_name_content}'") + tool_name = p_tool_name[0] + + tool_args: Dict[str, Tuple[str, str]] = {} + while stop_token == f"<{dsml_token}parameter": + index, param_content, stop_token = _read_until_stop(index, text, [f"/{dsml_token}parameter"]) + + param_kv = re.findall(r'^ name="(.*?)" string="(true|false)">(.*?)<$', param_content, flags=re.DOTALL) + if len(param_kv) != 1: + raise ValueError(f"Parameter format error: '{param_content}'") + param_name, string, param_value = param_kv[0] + + if param_name in tool_args: + raise ValueError(f"Duplicate parameter name: '{param_name}'") + tool_args[param_name] = (param_value, string) + + index, content, stop_token = _read_until_stop(index, text, [f"<{dsml_token}parameter", f"\n": + raise ValueError(f"Parameter format error: expected '>\\n' but got '{content}'") + + tool_call = decode_dsml_to_arguments(tool_name=tool_name, tool_args=tool_args) + tool_calls.append(tool_call) + + return index, stop_token, tool_calls + + +def parse_message_from_completion_text(text: str, thinking_mode: str) -> Dict[str, Any]: + """ + Parse a model completion text into a structured assistant message. + + This function takes the raw text output from the model (a single assistant turn) + and extracts: + - reasoning_content (thinking block) + - content (summary/response) + - tool_calls (if any) + + NOTE: This function is designed to parse only correctly formatted strings and + will raise ValueError for malformed output. + + Args: + text: The raw completion text (including EOS token). + thinking_mode: Either "chat" or "thinking". + + Returns: + Dict with keys: "role", "content", "reasoning_content", "tool_calls". + tool_calls are in OpenAI format. + """ + summary_content, reasoning_content, tool_calls = "", "", [] + index, stop_token = 0, None + tool_calls_start_token = f"\n\n<{dsml_token}{tool_calls_block_name}" + + is_thinking = thinking_mode == "thinking" + is_tool_calling = False + + if is_thinking: + index, content_delta, stop_token = _read_until_stop(index, text, [thinking_end_token, tool_calls_start_token]) + reasoning_content = content_delta + assert stop_token == thinking_end_token, "Invalid thinking format: missing " + + index, content_delta, stop_token = _read_until_stop(index, text, [eos_token, tool_calls_start_token]) + summary_content = content_delta + if stop_token == tool_calls_start_token: + is_tool_calling = True + else: + assert stop_token == eos_token, "Invalid format: missing EOS token" + + if is_tool_calling: + index, stop_token, tool_calls = parse_tool_calls(index, text) + + index, tool_ends_text, stop_token = _read_until_stop(index, text, [eos_token]) + assert not tool_ends_text, "Unexpected content after tool calls" + + assert len(text) == index and stop_token in [eos_token, None], "Unexpected content at end" + + for sp_token in [bos_token, eos_token, thinking_start_token, thinking_end_token, dsml_token]: + assert sp_token not in summary_content and sp_token not in reasoning_content, \ + f"Unexpected special token '{sp_token}' in content" + + return { + "role": "assistant", + "content": summary_content, + "reasoning_content": reasoning_content, + "tool_calls": tool_calls_to_openai_format(tool_calls) + } diff --git a/services/deepseek-v4-0731/encoding/tests/test_input_1.json b/services/deepseek-v4-0731/encoding/tests/test_input_1.json new file mode 100644 index 00000000..d423b221 --- /dev/null +++ b/services/deepseek-v4-0731/encoding/tests/test_input_1.json @@ -0,0 +1,81 @@ +{ + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather for a specific location", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "The city name" + }, + "unit": { + "type": "string", + "enum": ["celsius", "fahrenheit"], + "description": "Temperature unit" + } + }, + "required": ["location"] + } + } + }, + { + "type": "function", + "function": { + "name": "search", + "description": "Search the web for information", + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Search query" + }, + "num_results": { + "type": "integer", + "description": "Number of results to return" + } + }, + "required": ["query"] + } + } + } + ], + "messages": [ + { + "role": "system", + "content": "You are a helpful assistant." + }, + { + "role": "user", + "content": "What's the weather in Beijing?" + }, + { + "role": "assistant", + "reasoning_content": "The user wants to know the weather in Beijing. I should use the get_weather tool.", + "tool_calls": [ + { + "id": "call_001", + "type": "function", + "function": { + "name": "get_weather", + "arguments": "{\"location\": \"Beijing\", \"unit\": \"celsius\"}" + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "call_001", + "content": "{\"temperature\": 22, \"condition\": \"sunny\", \"humidity\": 45}" + }, + { + "role": "assistant", + "reasoning_content": "Got the weather data. Let me format a nice response.", + "content": "The weather in Beijing is currently sunny with a temperature of 22°C and 45% humidity." + } + ] +} diff --git a/services/deepseek-v4-0731/encoding/tests/test_input_2.json b/services/deepseek-v4-0731/encoding/tests/test_input_2.json new file mode 100644 index 00000000..13b3454a --- /dev/null +++ b/services/deepseek-v4-0731/encoding/tests/test_input_2.json @@ -0,0 +1,24 @@ +[ + { + "role": "system", + "content": "You are a helpful assistant." + }, + { + "role": "user", + "content": "Hello" + }, + { + "role": "assistant", + "reasoning_content": "The user said hello, I should greet back.", + "content": "Hi there! How can I help you?" + }, + { + "role": "user", + "content": "What is the capital of France?" + }, + { + "role": "assistant", + "reasoning_content": "The user asks about the capital of France. It is Paris.", + "content": "The capital of France is Paris." + } +] \ No newline at end of file diff --git a/services/deepseek-v4-0731/encoding/tests/test_input_3.json b/services/deepseek-v4-0731/encoding/tests/test_input_3.json new file mode 100644 index 00000000..10034844 --- /dev/null +++ b/services/deepseek-v4-0731/encoding/tests/test_input_3.json @@ -0,0 +1,159 @@ +[ + { + "role": "system", + "content": "该助手为DeepSeek，由深度求索公司创造。" + }, + { + "role": "latest_reminder", + "content": "2026-02-21,星期六,广州,App,中文" + }, + { + "role": "developer", + "content": "小柴胡冲剂和布洛芬能一起吃吗？\n\nCITATION FORMAT: 【{cursor_id}†L{start_line_id}(-L{end_line_id})?】", + "tools": [ + { + "type": "function", + "function": { + "name": "search", + "description": "Web search. Split multiple queries with '||'.", + "parameters": { + "type": "object", + "properties": { + "queries": { + "type": "string", + "description": "query1||query2" + } + }, + "required": [ + "queries" + ], + "additionalProperties": false, + "$schema": "http://json-schema.org/draft-07/schema#" + } + } + }, + { + "type": "function", + "function": { + "name": "open", + "description": "Batch open IDs (format 【{id}†...】) or URLs.", + "parameters": { + "type": "object", + "properties": { + "open_list": { + "type": "array", + "items": { + "type": "object", + "properties": { + "id": { + "description": "ID or URL", + "anyOf": [ + { + "type": "integer" + }, + { + "type": "string" + } + ], + "default": -1 + }, + "cursor": { + "type": "integer", + "description": "", + "default": -1 + }, + "loc": { + "type": "integer", + "description": "Start line", + "default": -1 + }, + "num_lines": { + "type": "integer", + "description": "", + "default": -1 + }, + "view_source": { + "type": "boolean", + "description": "", + "default": false + } + }, + "additionalProperties": false + }, + "description": "" + } + }, + "required": [ + "open_list" + ], + "additionalProperties": false, + "$schema": "http://json-schema.org/draft-07/schema#" + } + } + }, + { + "type": "function", + "function": { + "name": "find", + "description": "Find exact text pattern in pages.", + "parameters": { + "type": "object", + "properties": { + "find_list": { + "type": "array", + "items": { + "type": "object", + "properties": { + "pattern": { + "type": "string", + "description": "" + }, + "cursor": { + "type": "integer", + "description": "", + "default": -1 + } + }, + "required": [ + "pattern" + ], + "additionalProperties": false + }, + "description": "" + } + }, + "required": [ + "find_list" + ], + "additionalProperties": false, + "$schema": "http://json-schema.org/draft-07/schema#" + } + } + } + ] + }, + { + "role": "assistant", + "content": "", + "reasoning_content": "用户想知道小柴胡冲剂和布洛芬能否一起服用。", + "tool_calls": [ + { + "type": "function", + "function": { + "name": "search", + "arguments": "{\"queries\": \"小柴胡冲剂 布洛芬 相互作用 一起吃\"}" + } + } + ] + }, + { + "role": "tool", + "content": "[0]" + }, + { + "role": "assistant", + "content": "请及时就医。", + "reasoning_content": "现在开始组织回答。", + "tool_calls": [] + } +] \ No newline at end of file diff --git a/services/deepseek-v4-0731/encoding/tests/test_input_4.json b/services/deepseek-v4-0731/encoding/tests/test_input_4.json new file mode 100644 index 00000000..d5e0791d --- /dev/null +++ b/services/deepseek-v4-0731/encoding/tests/test_input_4.json @@ -0,0 +1,28 @@ +[ + { + "role": "system", + "content": "该助手为DeepSeek-V3，由深度求索公司创造。\n今天是2025年10月17日，星期五。" + }, + { + "role": "latest_reminder", + "content": "2024-11-15,上海市,App,中文" + }, + { + "role": "user", + "content": "热海大滚锅是世界著名温泉吗" + }, + { + "role": "assistant", + "content": "关于热海大滚锅是否是世界著名温泉，可以这样总结：它在**中国乃至全球的地热奇观中占有重要地位**，但“世界著名”的称号更侧重于它作为**独特的地质现象和旅游景点**，而非普遍的温泉疗养体验。\n\n为了让你快速了解，我整理了一个简要的表格：\n\n| 维度 | 热海大滚锅的地位与特点 |\n| :--- | :--- |\n| **地质奇观** | **世界罕见**的高温地热系统，泉眼多、水温高、形态多样。 |\n| **国内声誉** | **中国三大地热区之一**，**国家5A级旅游景区**，自古闻名（徐霞客曾游历并记载）。 |\n| **国际知名度** | 在一些旅游平台被提及为“世界六大温泉”之一，但此说法流传不广，其国际声誉更多建立在地质独特性上。 |\n| **核心体验** | **观赏地热奇观**（如97℃沸腾的“大滚锅”）、**体验温泉煮鸡蛋**。 |\n\n### 💡 游玩攻略与温馨提示\n\n如果你计划前往热海大滚锅，这里有一些实用信息供你参考：\n\n- **门票与开放时间**：\n - **门票**：景区门票约为**50元/人**。如果选择包含温泉沐浴的套餐，价格会更高，例如约**288元**。\n - **开放时间**：景区一般**08:00-18:00**开放，但具体时间可能变动，建议提前核实。\n\n- **特色体验**：\n - **温泉煮鸡蛋**：这几乎是必试项目。可以在景区门口购买用草绳串起的生鸡蛋（约5-8元/串），然后到“大滚锅”旁的指定区域蒸煮，几分钟便可熟食，趣味十足。\n - **金汤足浴**：可以直接用从“大滚锅”流出的温泉水泡脚，缓解旅途疲劳。\n\n- **注意事项**：\n - **安全第一**：“大滚锅”水温极高，务必遵守游览规则，在指定区域内观赏，切勿随意触碰泉水。\n - **规划行程**：建议为热海景区预留**3-4小时**的游览时间。景区内步道不走回头路，出入口有观光车接送。\n\n希望这些信息能帮助你更好地了解热海大滚锅。如果你对腾冲的其他景点或者行程规划有更多疑问，我很乐意提供进一步的信息。", + "mask": 1 + }, + { + "role": "user", + "content": "世界著名温泉有哪些", + "task": "action" + }, + { + "role": "assistant", + "content": "Search" + } +] \ No newline at end of file diff --git a/services/deepseek-v4-0731/encoding/tests/test_output_1.txt b/services/deepseek-v4-0731/encoding/tests/test_output_1.txt new file mode 100644 index 00000000..7e3c9bd5 --- /dev/null +++ b/services/deepseek-v4-0731/encoding/tests/test_output_1.txt @@ -0,0 +1,36 @@ +<｜begin▁of▁sentence｜>You are a helpful assistant. + +## Tools + +You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<｜DSML｜tool_calls>" block like the following: + +<｜DSML｜tool_calls> +<｜DSML｜invoke name="$TOOL_NAME"> +<｜DSML｜parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE +... + +<｜DSML｜invoke name="$TOOL_NAME2"> +... + + + +String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`. + +If thinking_mode is enabled (triggered by ), you MUST output your complete reasoning inside ... BEFORE any tool calls or final response. + +Otherwise, output directly after with tool calls or final response. + +### Available Tool Schemas + +{"name": "get_weather", "description": "Get the weather for a specific location", "parameters": {"type": "object", "properties": {"location": {"type": "string", "description": "The city name"}, "unit": {"type": "string", "enum": ["celsius", "fahrenheit"], "description": "Temperature unit"}}, "required": ["location"]}} +{"name": "search", "description": "Search the web for information", "parameters": {"type": "object", "properties": {"query": {"type": "string", "description": "Search query"}, "num_results": {"type": "integer", "description": "Number of results to return"}}, "required": ["query"]}} + +You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls. +<｜User｜>What's the weather in Beijing?<｜Assistant｜>The user wants to know the weather in Beijing. I should use the get_weather tool. + +<｜DSML｜tool_calls> +<｜DSML｜invoke name="get_weather"> +<｜DSML｜parameter name="location" string="true">Beijing +<｜DSML｜parameter name="unit" string="true">celsius + +<｜end▁of▁sentence｜><｜User｜>{"temperature": 22, "condition": "sunny", "humidity": 45}<｜Assistant｜>Got the weather data. Let me format a nice response.The weather in Beijing is currently sunny with a temperature of 22°C and 45% humidity.<｜end▁of▁sentence｜> \ No newline at end of file diff --git a/services/deepseek-v4-0731/encoding/tests/test_output_2.txt b/services/deepseek-v4-0731/encoding/tests/test_output_2.txt new file mode 100644 index 00000000..fc397ef5 --- /dev/null +++ b/services/deepseek-v4-0731/encoding/tests/test_output_2.txt @@ -0,0 +1 @@ +<｜begin▁of▁sentence｜>You are a helpful assistant.<｜User｜>Hello<｜Assistant｜>Hi there! How can I help you?<｜end▁of▁sentence｜><｜User｜>What is the capital of France?<｜Assistant｜>The user asks about the capital of France. It is Paris.The capital of France is Paris.<｜end▁of▁sentence｜> \ No newline at end of file diff --git a/services/deepseek-v4-0731/encoding/tests/test_output_3.txt b/services/deepseek-v4-0731/encoding/tests/test_output_3.txt new file mode 100644 index 00000000..edee5633 --- /dev/null +++ b/services/deepseek-v4-0731/encoding/tests/test_output_3.txt @@ -0,0 +1,38 @@ +<｜begin▁of▁sentence｜>该助手为DeepSeek，由深度求索公司创造。<｜latest_reminder｜>2026-02-21,星期六,广州,App,中文<｜User｜>小柴胡冲剂和布洛芬能一起吃吗？ + +CITATION FORMAT: 【{cursor_id}†L{start_line_id}(-L{end_line_id})?】 + +## Tools + +You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<｜DSML｜tool_calls>" block like the following: + +<｜DSML｜tool_calls> +<｜DSML｜invoke name="$TOOL_NAME"> +<｜DSML｜parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE +... + +<｜DSML｜invoke name="$TOOL_NAME2"> +... + + + +String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`. + +If thinking_mode is enabled (triggered by ), you MUST output your complete reasoning inside ... BEFORE any tool calls or final response. + +Otherwise, output directly after with tool calls or final response. + +### Available Tool Schemas + +{"name": "search", "description": "Web search. Split multiple queries with '||'.", "parameters": {"type": "object", "properties": {"queries": {"type": "string", "description": "query1||query2"}}, "required": ["queries"], "additionalProperties": false, "$schema": "http://json-schema.org/draft-07/schema#"}} +{"name": "open", "description": "Batch open IDs (format 【{id}†...】) or URLs.", "parameters": {"type": "object", "properties": {"open_list": {"type": "array", "items": {"type": "object", "properties": {"id": {"description": "ID or URL", "anyOf": [{"type": "integer"}, {"type": "string"}], "default": -1}, "cursor": {"type": "integer", "description": "", "default": -1}, "loc": {"type": "integer", "description": "Start line", "default": -1}, "num_lines": {"type": "integer", "description": "", "default": -1}, "view_source": {"type": "boolean", "description": "", "default": false}}, "additionalProperties": false}, "description": ""}}, "required": ["open_list"], "additionalProperties": false, "$schema": "http://json-schema.org/draft-07/schema#"}} +{"name": "find", "description": "Find exact text pattern in pages.", "parameters": {"type": "object", "properties": {"find_list": {"type": "array", "items": {"type": "object", "properties": {"pattern": {"type": "string", "description": ""}, "cursor": {"type": "integer", "description": "", "default": -1}}, "required": ["pattern"], "additionalProperties": false}, "description": ""}}, "required": ["find_list"], "additionalProperties": false, "$schema": "http://json-schema.org/draft-07/schema#"}} + +You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls. +<｜Assistant｜>用户想知道小柴胡冲剂和布洛芬能否一起服用。 + +<｜DSML｜tool_calls> +<｜DSML｜invoke name="search"> +<｜DSML｜parameter name="queries" string="true">小柴胡冲剂 布洛芬 相互作用 一起吃 + +<｜end▁of▁sentence｜><｜User｜>[0]<｜Assistant｜>现在开始组织回答。请及时就医。<｜end▁of▁sentence｜> \ No newline at end of file diff --git a/services/deepseek-v4-0731/encoding/tests/test_output_4.txt b/services/deepseek-v4-0731/encoding/tests/test_output_4.txt new file mode 100644 index 00000000..d30bd5d0 --- /dev/null +++ b/services/deepseek-v4-0731/encoding/tests/test_output_4.txt @@ -0,0 +1,29 @@ +<｜begin▁of▁sentence｜>该助手为DeepSeek-V3，由深度求索公司创造。 +今天是2025年10月17日，星期五。<｜latest_reminder｜>2024-11-15,上海市,App,中文<｜User｜>热海大滚锅是世界著名温泉吗<｜Assistant｜>关于热海大滚锅是否是世界著名温泉，可以这样总结：它在**中国乃至全球的地热奇观中占有重要地位**，但“世界著名”的称号更侧重于它作为**独特的地质现象和旅游景点**，而非普遍的温泉疗养体验。 + +为了让你快速了解，我整理了一个简要的表格： + +| 维度 | 热海大滚锅的地位与特点 | +| :--- | :--- | +| **地质奇观** | **世界罕见**的高温地热系统，泉眼多、水温高、形态多样。 | +| **国内声誉** | **中国三大地热区之一**，**国家5A级旅游景区**，自古闻名（徐霞客曾游历并记载）。 | +| **国际知名度** | 在一些旅游平台被提及为“世界六大温泉”之一，但此说法流传不广，其国际声誉更多建立在地质独特性上。 | +| **核心体验** | **观赏地热奇观**（如97℃沸腾的“大滚锅”）、**体验温泉煮鸡蛋**。 | + +### 💡 游玩攻略与温馨提示 + +如果你计划前往热海大滚锅，这里有一些实用信息供你参考： + +- **门票与开放时间**： + - **门票**：景区门票约为**50元/人**。如果选择包含温泉沐浴的套餐，价格会更高，例如约**288元**。 + - **开放时间**：景区一般**08:00-18:00**开放，但具体时间可能变动，建议提前核实。 + +- **特色体验**： + - **温泉煮鸡蛋**：这几乎是必试项目。可以在景区门口购买用草绳串起的生鸡蛋（约5-8元/串），然后到“大滚锅”旁的指定区域蒸煮，几分钟便可熟食，趣味十足。 + - **金汤足浴**：可以直接用从“大滚锅”流出的温泉水泡脚，缓解旅途疲劳。 + +- **注意事项**： + - **安全第一**：“大滚锅”水温极高，务必遵守游览规则，在指定区域内观赏，切勿随意触碰泉水。 + - **规划行程**：建议为热海景区预留**3-4小时**的游览时间。景区内步道不走回头路，出入口有观光车接送。 + +希望这些信息能帮助你更好地了解热海大滚锅。如果你对腾冲的其他景点或者行程规划有更多疑问，我很乐意提供进一步的信息。<｜end▁of▁sentence｜><｜User｜>世界著名温泉有哪些<｜Assistant｜><｜action｜>Search<｜end▁of▁sentence｜> \ No newline at end of file diff --git a/services/deepseek-v4-0731/launch_candidate.sh b/services/deepseek-v4-0731/launch_candidate.sh index 25a32002..137ef854 100755 --- a/services/deepseek-v4-0731/launch_candidate.sh +++ b/services/deepseek-v4-0731/launch_candidate.sh @@ -1,5 +1,5 @@ #!/bin/sh -# Isolated candidate only. This file must never manage the production service. +# Isolated candidate only. This file never manages the production service. set -eu umask 077 @@ -8,53 +8,46 @@ WORKTREE=/Users/davidtai/projects/OpenSourceWTF/.worktrees/dsv4-0731-service MODEL=/Users/davidtai/models/DeepSeek-V4-Flash-2bit-DQ-mtp PYTHON=/Users/davidtai/projects/OpenSourceWTF/.worktrees/dsv4-0731-service/.venv/bin/python PYTHON_TARGET=/Users/davidtai/.local/share/uv/python/cpython-3.12-macos-aarch64-none/bin/python3.12 -CONFIG="$SERVICE_ROOT/candidate.json" -ASSET="$SERVICE_ROOT/encoding/chat_template.jinja" -MANIFEST="$SERVICE_ROOT/encoding/SHA256SUMS" +ENTRY="$SERVICE_ROOT/candidate_entry.py" +ENCODING="$SERVICE_ROOT/encoding" +REVIEWED_REF=refs/tags/mtplx-dsv4-0731-reviewed +PORT=8081 die() { printf '%s\n' "deepseek-v4-0731 candidate: $1" >&2; exit 64; } sha256() { /usr/bin/shasum -a 256 "$1" | /usr/bin/awk '{print $1}'; } -# Command overrides are only possible in an explicit, local test fixture. A -# launchd job never sets this flag, and production accepts no caller argv/env. fixture=${MTPLX_DSV4_0731_TEST_FIXTURE:-} if [ -n "${MTPLX_DSV4_0731_EXECUTABLE:-}" ] && [ "$fixture" != 1 ]; then die "command environment override rejected" fi if [ "$fixture" = 1 ]; then [ "${1:-}" = --print-command ] || die "fixture mode only permits --print-command" - printf '%s\n' "$PYTHON -m mtplx serve --host 127.0.0.1 --port 8081" + printf '%s\n' "$PYTHON $ENTRY (fixed 127.0.0.1:$PORT)" exit 0 fi [ "$#" -eq 0 ] || die "arguments are not accepted" -[ -r "$CONFIG" ] && [ ! -L "$CONFIG" ] || die "candidate configuration is missing or unsafe" [ -L "$PYTHON" ] && [ "$(/usr/bin/readlink "$PYTHON")" = "$PYTHON_TARGET" ] || die "trusted python link changed" [ -x "$PYTHON_TARGET" ] && [ ! -L "$PYTHON_TARGET" ] || die "trusted python target is missing or unsafe" -[ -d "$MODEL" ] && [ ! -L "$MODEL" ] || die "pinned model path is missing or unsafe" -[ -f "$ASSET" ] && [ ! -L "$ASSET" ] || die "encoding asset is missing or unsafe" -[ -f "$MANIFEST" ] && [ ! -L "$MANIFEST" ] || die "encoding manifest is missing or unsafe" - -[ "$(sha256 "$ASSET")" = 03f2686beff14c3d9040894a2b658d9f1917be90bc1d90597502fc2562f0ec2a ] || die "encoding asset hash changed" [ "$(sha256 "$PYTHON_TARGET")" = 96793b100c947cdc81a38e8fb8c9c1889abccda9840ce1bef58d372bf3f2c263 ] || die "trusted python hash changed" +[ -f "$ENTRY" ] && [ ! -L "$ENTRY" ] || die "candidate entrypoint is missing or unsafe" +[ "$(sha256 "$ENTRY")" = c7e2e79e45b3f2e8afd8453d5976e5234caea724813eae0e555c5f956bf725aa ] || die "candidate entrypoint hash changed" +[ -d "$MODEL" ] && [ ! -L "$MODEL" ] || die "pinned model path is missing or unsafe" [ "$(sha256 "$MODEL/config.json")" = c8ff87fd5ee5c9587d0c937e9bfd3193e1a1621141aa367848a9610b3291fa6f ] || die "model configuration hash changed" [ "$(sha256 "$MODEL/model.safetensors.index.json")" = c84d2b369f5d5023d0f2d183fc36a935a3981751414996243b65f069983e43d8 ] || die "model index hash changed" -[ "$(/usr/bin/git -C "$WORKTREE" merge-base HEAD 5ccc9fdf251a9eaf946f4c77c42eabd6ba3f0ab4)" = 5ccc9fdf251a9eaf946f4c77c42eabd6ba3f0ab4 ] || die "worktree does not descend from pinned revision" -/usr/bin/grep -Fqx '03f2686beff14c3d9040894a2b658d9f1917be90bc1d90597502fc2562f0ec2a chat_template.jinja' "$MANIFEST" || die "encoding manifest changed" +[ "$(sha256 "$ENCODING/SHA256SUMS")" = 6758dfda8a39afdd00d907606c42c1a268289c463351b9628ac07f4f916d7d0a ] || die "official encoding manifest hash changed" +(cd "$ENCODING" && /usr/bin/shasum -a 256 -c SHA256SUMS >/dev/null) || die "official encoding/vector asset hash changed" + +reviewed_commit=$(/usr/bin/git -C "$WORKTREE" rev-parse --verify "${REVIEWED_REF}^{commit}") || die "reviewed commit ref is missing" +current_commit=$(/usr/bin/git -C "$WORKTREE" rev-parse --verify HEAD) || die "worktree HEAD is missing" +[ "$current_commit" = "$reviewed_commit" ] || die "worktree is not the exact reviewed commit" +[ -z "$(/usr/bin/git -C "$WORKTREE" status --porcelain=v1 --untracked-files=all)" ] || die "reviewed worktree is not clean" -# `env -i` is the process boundary: no caller environment reaches model load. exec /usr/bin/env -i \ HOME=/Users/davidtai \ LC_ALL=C \ PATH=/usr/bin:/bin \ PYTHONNOUSERSITE=1 \ + PYTHONPATH="$WORKTREE" \ VIRTUAL_ENV=/Users/davidtai/projects/OpenSourceWTF/.worktrees/dsv4-0731-service/.venv \ - "$PYTHON" -m mtplx serve \ - --host 127.0.0.1 \ - --port 8081 \ - --model "$MODEL" \ - --model-id deepseek-v4-0731-candidate \ - --reasoning-effort low \ - --reasoning-parser none \ - --warmup-tokens 0 \ - --no-stats-footer + "$PYTHON" "$ENTRY" diff --git a/services/deepseek-v4-0731/promote_cutover.py b/services/deepseek-v4-0731/promote_cutover.py index a57a77e1..ef7155c2 100755 --- a/services/deepseek-v4-0731/promote_cutover.py +++ b/services/deepseek-v4-0731/promote_cutover.py @@ -31,6 +31,20 @@ CANDIDATE_PORT = 8081 LIVE_PORT = 8080 SENSITIVE_KEY = re.compile(r"(?:prompt|message|tool|secret|token|authorization|argv|env|stdout|stderr)", re.I) +SENSITIVE_VALUE = re.compile( + r"(?:-----BEGIN [A-Z ]*PRIVATE KEY-----|\bBearer\s+\S+|\bsk-[A-Za-z0-9_-]{12,}|" + r"[\"']?(?:prompt|messages|tools|secret|authorization)[\"']?\s*[:=])", + re.I, +) +ALLOWED_SIGNERS = Path("/Users/davidtai/.config/mtplx/deepseek-v4-0731-allowed-signers") +SIGNING_IDENTITY = "mtplx-deepseek-v4-0731-candidate" +SIGNING_NAMESPACE = "mtplx-deepseek-v4-0731" +CANDIDATE_WORKTREE = Path("/Users/davidtai/projects/OpenSourceWTF/.worktrees/dsv4-0731-service") +REVIEWED_REF = "refs/tags/mtplx-dsv4-0731-reviewed" +CANDIDATE_PLIST_SHA256 = "93eac0d4eaac491c7f2f1d3a293ba38a3144ade59ee3afdf52b35cc9ec9bb101" +ENCODING_ASSET_SET_SHA256 = "6758dfda8a39afdd00d907606c42c1a268289c463351b9628ac07f4f916d7d0a" +MODEL_CONFIG_SHA256 = "c8ff87fd5ee5c9587d0c937e9bfd3193e1a1621141aa367848a9610b3291fa6f" +MODEL_INDEX_SHA256 = "c84d2b369f5d5023d0f2d183fc36a935a3981751414996243b65f069983e43d8" class PromotionError(RuntimeError): @@ -125,16 +139,48 @@ def attest_live(*, label: str, plist: Path) -> dict[str, Any]: } -def _read_json(path: Path) -> dict[str, Any]: +def _read_json_bytes(path: Path) -> tuple[dict[str, Any], bytes]: if not path.is_file() or path.is_symlink(): raise PromotionError("receipt is missing or unsafe") try: - payload = json.loads(path.read_text(encoding="utf-8")) + raw = path.read_bytes() + payload = json.loads(raw.decode("utf-8")) except (OSError, ValueError) as error: raise PromotionError("receipt is not valid JSON") from error if not isinstance(payload, dict): raise PromotionError("receipt root must be an object") - return payload + return payload, raw + + +def _read_json(path: Path) -> dict[str, Any]: + return _read_json_bytes(path)[0] + + +def _verify_candidate_signature(receipt_bytes: bytes, signature: Path) -> None: + if not ALLOWED_SIGNERS.is_file() or ALLOWED_SIGNERS.is_symlink(): + raise PromotionError("pinned candidate allowed-signers file is missing or unsafe") + if not signature.is_file() or signature.is_symlink(): + raise PromotionError("detached candidate signature is missing or unsafe") + result = subprocess.run( + [ + "/usr/bin/ssh-keygen", + "-Y", + "verify", + "-f", + str(ALLOWED_SIGNERS), + "-I", + SIGNING_IDENTITY, + "-n", + SIGNING_NAMESPACE, + "-s", + str(signature), + ], + input=receipt_bytes, + check=False, + capture_output=True, + ) + if result.returncode: + raise PromotionError("candidate receipt signature verification failed") def _contains_sensitive(value: Any) -> bool: @@ -142,34 +188,77 @@ def _contains_sensitive(value: Any) -> bool: return any(SENSITIVE_KEY.search(str(key)) or _contains_sensitive(item) for key, item in value.items()) if isinstance(value, list): return any(_contains_sensitive(item) for item in value) - if isinstance(value, str) and value.startswith(("/Users/", "/private/", "/tmp/")): - return True + if isinstance(value, str): + if any(marker in value for marker in ("/Users/", "/private/", "/tmp/")): + return True + if SENSITIVE_VALUE.search(value): + return True return False +def _require_exact_keys(value: Any, expected: set[str], context: str) -> dict[str, Any]: + if not isinstance(value, dict) or set(value) != expected: + raise PromotionError(f"{context} fields do not match the strict receipt schema") + return value + + def assert_candidate_receipt(payload: dict[str, Any]) -> None: """Accept only a previously passing, scrubbed candidate preflight+smoke receipt.""" if _contains_sensitive(payload): raise PromotionError("candidate receipt includes prohibited sensitive capture") - preflight = payload.get("candidate_preflight") - smoke = payload.get("candidate_smoke") - if not isinstance(preflight, dict) or not isinstance(smoke, dict): - raise PromotionError("candidate receipt lacks preflight or smoke evidence") + _require_exact_keys(payload, {"schema", "candidate_preflight", "candidate_smoke"}, "candidate receipt") + if payload["schema"] != "mtplx.dsv4-0731-candidate.v1": + raise PromotionError("candidate receipt schema is not pinned") + preflight = _require_exact_keys( + payload["candidate_preflight"], + { + "ok", "label", "port", "plist_sha256", "encoding_source_revision", + "encoding_asset_set_sha256", "reviewed_commit", "model_config_sha256", + "model_index_sha256", "promotion_target", + }, + "candidate preflight", + ) + smoke = _require_exact_keys( + payload["candidate_smoke"], + {"ok", "models_ok", "ready", "finish_reason", "candidate_model_ids"}, + "candidate smoke", + ) if preflight.get("ok") is not True or smoke.get("ok") is not True: raise PromotionError("candidate preflight and smoke must already pass") if preflight.get("label") != CANDIDATE_LABEL or preflight.get("port") != CANDIDATE_PORT: raise PromotionError("candidate identity does not match the pinned isolated service") if smoke.get("models_ok") is not True or smoke.get("ready") is not True or smoke.get("finish_reason") != "stop": raise PromotionError("candidate smoke receipt lacks models/READY/stop evidence") - target = preflight.get("promotion_target") - if not isinstance(target, dict) or not isinstance(target.get("label"), str): + if preflight.get("encoding_source_revision") != "7872f01b1d1fe23eabc4c98b48bffcef5a386062": + raise PromotionError("candidate encoding source revision changed") + for field in ( + "plist_sha256", "encoding_asset_set_sha256", "model_config_sha256", "model_index_sha256" + ): + if not isinstance(preflight.get(field), str) or not re.fullmatch(r"[0-9a-f]{64}", preflight[field]): + raise PromotionError(f"candidate preflight has invalid {field}") + if not isinstance(preflight.get("reviewed_commit"), str) or not re.fullmatch(r"[0-9a-f]{40}", preflight["reviewed_commit"]): + raise PromotionError("candidate preflight has invalid reviewed_commit") + target = _require_exact_keys(preflight["promotion_target"], {"label", "plist_sha256"}, "promotion target") + if not isinstance(target.get("label"), str): raise PromotionError("candidate preflight lacks a separately reviewed promotion target") digest = target.get("plist_sha256") if not isinstance(digest, str) or not re.fullmatch(r"[0-9a-f]{64}", digest): raise PromotionError("candidate preflight lacks a valid promotion plist digest") + candidate_model_ids = smoke.get("candidate_model_ids") + if ( + not isinstance(candidate_model_ids, list) + or not candidate_model_ids + or not all(isinstance(model_id, str) and model_id for model_id in candidate_model_ids) + ): + raise PromotionError("candidate smoke lacks nonempty candidate_model_ids") def assert_live_identity(expected: dict[str, Any], current: dict[str, Any]) -> None: + _require_exact_keys( + expected, + {"schema", "label", "pid", "listener_port", "plist_sha256", "model_ids"}, + "live attestation", + ) fields = ("schema", "label", "pid", "listener_port", "plist_sha256", "model_ids") if any(expected.get(field) != current.get(field) for field in fields): raise PromotionError("live service identity changed since its attestation") @@ -215,11 +304,30 @@ def _verify_live_ready(expected_model_ids: list[str]) -> None: def promote(args: argparse.Namespace) -> None: if args.promote is not True: raise PromotionError("refusing promotion without --promote") - candidate = _read_json(args.candidate_receipt) + candidate, candidate_bytes = _read_json_bytes(args.candidate_receipt) + _verify_candidate_signature(candidate_bytes, args.candidate_signature) expected_live = _read_json(args.live_attestation) if _contains_sensitive(expected_live): raise PromotionError("live attestation includes prohibited sensitive capture") assert_candidate_receipt(candidate) + preflight = candidate["candidate_preflight"] + reviewed_commit = _command( + "/usr/bin/git", + "-C", + str(CANDIDATE_WORKTREE), + "rev-parse", + "--verify", + f"{REVIEWED_REF}^{{commit}}", + ).strip() + pinned_candidate = { + "reviewed_commit": reviewed_commit, + "plist_sha256": CANDIDATE_PLIST_SHA256, + "encoding_asset_set_sha256": ENCODING_ASSET_SET_SHA256, + "model_config_sha256": MODEL_CONFIG_SHA256, + "model_index_sha256": MODEL_INDEX_SHA256, + } + if any(preflight[field] != expected for field, expected in pinned_candidate.items()): + raise PromotionError("signed candidate receipt does not match the reviewed installation") # The production target must be a separately reviewed 8080 plist. The # candidate plist stays isolated on 8081 and is never edited in place. target = args.production_plist @@ -227,7 +335,7 @@ def promote(args: argparse.Namespace) -> None: raise PromotionError("an absolute separately reviewed production plist is required") if not target.is_file() or target.is_symlink(): raise PromotionError("production plist is missing or unsafe") - promotion_target = candidate["candidate_preflight"]["promotion_target"] + promotion_target = preflight["promotion_target"] if promotion_target["label"] != args.production_label or promotion_target["plist_sha256"] != _sha256(target): raise PromotionError("production plist identity does not match the passing candidate preflight") try: @@ -249,7 +357,7 @@ def promote(args: argparse.Namespace) -> None: try: _bootout(current["label"]) _bootstrap(target) - _verify_live_ready(current["model_ids"]) + _verify_live_ready(candidate["candidate_smoke"]["candidate_model_ids"]) except BaseException: try: _bootout(args.production_label) @@ -263,6 +371,7 @@ def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--promote", action="store_true", help="explicitly authorize guarded service action") parser.add_argument("--candidate-receipt", type=Path, required=True) + parser.add_argument("--candidate-signature", type=Path, required=True) parser.add_argument("--live-attestation", type=Path, required=True) parser.add_argument("--live-plist", type=Path, required=True) parser.add_argument("--production-plist", type=Path, required=True) diff --git a/services/deepseek-v4-0731/render.py b/services/deepseek-v4-0731/render.py deleted file mode 100644 index ba062700..00000000 --- a/services/deepseek-v4-0731/render.py +++ /dev/null @@ -1,168 +0,0 @@ -"""Pinned, dependency-free renderer for the isolated DeepSeek 0731 candidate. - -The asset check runs when a :class:`PinnedEncoding` is installed. Rendering is -then branch-free with respect to asset identity: a checked immutable encoding is -the only thing that can be installed. This keeps integrity work out of the -request path while failing closed before a candidate can accept requests. -""" - -from __future__ import annotations - -import hashlib -import json -from dataclasses import dataclass -from pathlib import Path -from typing import Any, Iterable, Mapping, Sequence - - -SERVICE_ROOT = Path(__file__).resolve().parent -ENCODING_DIR = SERVICE_ROOT / "encoding" -ASSET_PATH = ENCODING_DIR / "chat_template.jinja" -MANIFEST_PATH = ENCODING_DIR / "SHA256SUMS" - -BOS = "<｜begin▁of▁sentence｜>" -USER = "<｜User｜>" -ASSISTANT = "<｜Assistant｜>" -EOS = "<｜end▁of▁sentence｜>" -TOOL_CALLS_BEGIN = "<｜tool▁calls▁begin｜>" -TOOL_CALL_BEGIN = "<｜tool▁call▁begin｜>" -TOOL_SEPARATOR = "<｜tool▁sep｜>" -TOOL_CALL_END = "<｜tool▁call▁end｜>" -TOOL_CALLS_END = "<｜tool▁calls▁end｜>" -TOOL_OUTPUT_BEGIN = "<｜tool▁output▁begin｜>" -TOOL_OUTPUT_END = "<｜tool▁output▁end｜>" - - -class AssetIntegrityError(RuntimeError): - """The pinned source asset cannot safely be installed.""" - - -class InvalidReasoningEffort(ValueError): - """Only the candidate's explicitly tested reasoning profiles are valid.""" - - -def _manifest_hash(manifest: Path, filename: str) -> str: - if not manifest.is_file() or manifest.is_symlink(): - raise AssetIntegrityError("encoding manifest is missing or unsafe") - matches: list[str] = [] - for line in manifest.read_text(encoding="utf-8").splitlines(): - fields = line.split() - if len(fields) == 2 and fields[1] == filename and len(fields[0]) == 64: - matches.append(fields[0].lower()) - if len(matches) != 1 or any(c not in "0123456789abcdef" for c in matches[0]): - raise AssetIntegrityError("encoding manifest has no unique valid asset digest") - return matches[0] - - -def verify_assets(asset_path: Path = ASSET_PATH, manifest_path: Path = MANIFEST_PATH) -> str: - """Return the verified digest, rejecting all incomplete or altered assets.""" - if not asset_path.is_file() or asset_path.is_symlink(): - raise AssetIntegrityError("pinned encoding asset is missing or unsafe") - expected = _manifest_hash(manifest_path, asset_path.name) - actual = hashlib.sha256(asset_path.read_bytes()).hexdigest() - if actual != expected: - raise AssetIntegrityError("pinned encoding asset digest mismatch") - return actual - - -def _text(value: Any, field: str) -> str: - if not isinstance(value, str): - raise ValueError(f"{field} must be a string") - return value - - -def _canonical_json(value: Any) -> str: - return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")) - - -def _render_tools(tools: Sequence[Mapping[str, Any]] | None) -> str: - if not tools: - return "" - canonical = _canonical_json(list(tools)) - return f"\n\n# Tools\n{canonical}" - - -@dataclass(frozen=True) -class PinnedEncoding: - """An encoding asset checked exactly once at its installation boundary.""" - - asset_sha256: str - - @classmethod - def install(cls, asset_path: Path = ASSET_PATH, manifest_path: Path = MANIFEST_PATH) -> "PinnedEncoding": - return cls(asset_sha256=verify_assets(asset_path, manifest_path)) - - def render( - self, - messages: Iterable[Mapping[str, Any]], - *, - tools: Sequence[Mapping[str, Any]] | None = None, - reasoning_effort: str = "low", - ) -> str: - effort = reasoning_effort.strip().lower() if isinstance(reasoning_effort, str) else "" - if effort not in {"low", "high", "max"}: - raise InvalidReasoningEffort("reasoning_effort must be one of: low, high, max") - - items = list(messages) - system = "\n\n".join( - _text(message.get("content"), "system content") - for message in items - if message.get("role") == "system" - ) - output = [BOS, system, _render_tools(tools)] - last_was_user = False - for message in items: - role = message.get("role") - if role == "system": - continue - if role == "user": - output.extend((USER, f"Reasoning: {effort}\n", _text(message.get("content"), "user content"))) - last_was_user = True - continue - if role == "assistant": - tool_calls = message.get("tool_calls") - if tool_calls is not None: - if not isinstance(tool_calls, list) or not tool_calls: - raise ValueError("assistant tool_calls must be a non-empty list") - output.extend((ASSISTANT, "", TOOL_CALLS_BEGIN)) - for call in tool_calls: - function = call.get("function") if isinstance(call, Mapping) else None - if not isinstance(function, Mapping): - raise ValueError("tool call function must be an object") - output.extend(( - TOOL_CALL_BEGIN, - _text(function.get("name"), "tool function name"), - TOOL_SEPARATOR, - _text(function.get("arguments"), "tool function arguments"), - TOOL_CALL_END, - )) - output.extend((TOOL_CALLS_END, EOS)) - else: - content = _text(message.get("content"), "assistant content") - if last_was_user: - output.extend((ASSISTANT, "")) - output.extend((content.split("", 1)[-1], EOS)) - last_was_user = False - continue - if role == "tool": - _text(message.get("tool_call_id"), "tool_call_id") - output.extend((TOOL_OUTPUT_BEGIN, _text(message.get("content"), "tool content"), TOOL_OUTPUT_END)) - last_was_user = False - continue - raise ValueError("unsupported message role") - if last_was_user: - output.extend((ASSISTANT, "")) - return "".join(output) - - -_DEFAULT_ENCODING = PinnedEncoding.install() - - -def render_chat( - messages: Iterable[Mapping[str, Any]], - *, - tools: Sequence[Mapping[str, Any]] | None = None, - reasoning_effort: str = "low", -) -> str: - """Render through the already-installed default candidate encoding.""" - return _DEFAULT_ENCODING.render(messages, tools=tools, reasoning_effort=reasoning_effort) diff --git a/services/deepseek-v4-0731/tests/test_official_encoding.py b/services/deepseek-v4-0731/tests/test_official_encoding.py new file mode 100644 index 00000000..e7d80c62 --- /dev/null +++ b/services/deepseek-v4-0731/tests/test_official_encoding.py @@ -0,0 +1,79 @@ +"""Byte-exact gates for the official DeepSeek-V4-Flash-0731 encoder.""" + +from __future__ import annotations + +import json +from pathlib import Path +from types import ModuleType + +import pytest + + +ROOT = Path(__file__).resolve().parents[1] +ENCODING = ROOT / "encoding" +VECTORS = ENCODING / "tests" + + +def _official(): + path = ENCODING / "encoding_dsv4.py" + module = ModuleType("encoding_dsv4") + exec(compile(path.read_bytes(), str(path), "exec"), module.__dict__) + return module + + +@pytest.mark.parametrize("case", [1, 2, 3, 4]) +def test_official_vectors_are_byte_exact(case: int) -> None: + encoding = _official() + payload = json.loads((VECTORS / f"test_input_{case}.json").read_text(encoding="utf-8")) + if case == 1: + messages = payload["messages"] + messages[0]["tools"] = payload["tools"] + else: + messages = payload + mode = "chat" if case == 4 else "thinking" + expected = (VECTORS / f"test_output_{case}.txt").read_text(encoding="utf-8") + assert encoding.encode_messages(messages, thinking_mode=mode) == expected + + +def test_dsml_tool_call_result_merge_and_completion_parse() -> None: + encoding = _official() + payload = json.loads((VECTORS / "test_input_1.json").read_text(encoding="utf-8")) + messages = payload["messages"] + messages[0]["tools"] = payload["tools"] + prompt = encoding.encode_messages(messages, thinking_mode="thinking") + assert "<｜DSML｜tool_calls>" in prompt + assert "" in prompt + assert not any(message.get("role") == "tool" for message in encoding.merge_tool_messages(messages)) + + marker = "<｜Assistant｜>" + first_start = prompt.find(marker) + len(marker) + first_end = prompt.find("<｜User｜>", first_start) + parsed = encoding.parse_message_from_completion_text(prompt[first_start:first_end], thinking_mode="thinking") + assert parsed["reasoning_content"].startswith("The user wants") + assert parsed["content"] == "" + assert parsed["tool_calls"][0]["function"]["name"] == "get_weather" + assert json.loads(parsed["tool_calls"][0]["function"]["arguments"]) == { + "location": "Beijing", + "unit": "celsius", + } + + +@pytest.mark.parametrize("effort", ["low", "high", "max"]) +def test_reasoning_prefixes_and_thinking_modes(effort: str) -> None: + encoding = _official() + messages = [{"role": "user", "content": "hi"}] + thinking = encoding.encode_messages(messages, thinking_mode="thinking", reasoning_effort=effort) + chat = encoding.encode_messages(messages, thinking_mode="chat", reasoning_effort=effort) + assert thinking.startswith(encoding.bos_token + encoding.REASONING_EFFORT_PROMPTS[effort]) + assert thinking.endswith(encoding.ASSISTANT_SP_TOKEN + encoding.thinking_start_token) + assert chat.endswith(encoding.ASSISTANT_SP_TOKEN + encoding.thinking_end_token) + + +def test_invalid_reasoning_effort_fails_closed() -> None: + encoding = _official() + with pytest.raises(AssertionError, match="Invalid reasoning effort"): + encoding.encode_messages( + [{"role": "user", "content": "hi"}], + thinking_mode="thinking", + reasoning_effort="medium", + ) diff --git a/services/deepseek-v4-0731/tests/test_render.py b/services/deepseek-v4-0731/tests/test_render.py deleted file mode 100644 index 53840c6e..00000000 --- a/services/deepseek-v4-0731/tests/test_render.py +++ /dev/null @@ -1,75 +0,0 @@ -"""Golden contracts for the isolated DeepSeek 0731 prompt renderer.""" - -from __future__ import annotations - -import hashlib -import sys -from pathlib import Path - -import pytest - - -SERVICE = Path(__file__).resolve().parents[1] -sys.path.insert(0, str(SERVICE)) - -from render import ( # noqa: E402 - AssetIntegrityError, - InvalidReasoningEffort, - render_chat, - verify_assets, -) - - -def test_golden_messages_tools_tool_result_and_reasoning() -> None: - rendered = render_chat( - [ - {"role": "system", "content": "Be exact."}, - {"role": "user", "content": "What is 2+2?"}, - { - "role": "assistant", - "tool_calls": [ - { - "id": "call_1", - "type": "function", - "function": {"name": "calculator", "arguments": '{"x":"2+2"}'}, - } - ], - }, - {"role": "tool", "tool_call_id": "call_1", "content": "4"}, - ], - tools=[ - { - "type": "function", - "function": { - "name": "calculator", - "description": "Evaluate arithmetic.", - "parameters": {"type": "object", "properties": {"x": {"type": "string"}}}, - }, - } - ], - reasoning_effort="high", - ) - assert hashlib.sha256(rendered.encode()).hexdigest() == "f0d541c389ee21f1a4a1f50b8624c44f06a5e30b8d65a37dc0bfdc2edd05f11e" - - -@pytest.mark.parametrize("effort", ["low", "high", "max"]) -def test_reasoning_effort_has_stable_rendering(effort: str) -> None: - assert render_chat([{"role": "user", "content": "hi"}], reasoning_effort=effort).startswith( - f"<｜begin▁of▁sentence｜><｜User｜>Reasoning: {effort}" - ) - - -def test_invalid_reasoning_effort_fails_closed() -> None: - with pytest.raises(InvalidReasoningEffort): - render_chat([{"role": "user", "content": "hi"}], reasoning_effort="medium") - - -def test_missing_or_tampered_asset_fails_closed(tmp_path: Path) -> None: - manifest = SERVICE / "encoding" / "SHA256SUMS" - with pytest.raises(AssetIntegrityError): - verify_assets(tmp_path / "missing", manifest) - - asset = tmp_path / "chat_template.jinja" - asset.write_text("tampered", encoding="utf-8") - with pytest.raises(AssetIntegrityError): - verify_assets(asset, manifest) diff --git a/services/deepseek-v4-0731/tests/test_service_surface.py b/services/deepseek-v4-0731/tests/test_service_surface.py index 5d26a49b..2b74d21c 100644 --- a/services/deepseek-v4-0731/tests/test_service_surface.py +++ b/services/deepseek-v4-0731/tests/test_service_surface.py @@ -2,27 +2,33 @@ from __future__ import annotations +import hashlib import json import os import subprocess +import sys from pathlib import Path +from types import SimpleNamespace import pytest ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) def test_candidate_is_distinct_from_live_service() -> None: launch = (ROOT / "launch_candidate.sh").read_text(encoding="utf-8") plist = (ROOT / "com.tea.deepseek-v4-0731.candidate.plist").read_text(encoding="utf-8") assert "com.tea.deepseek-v4-0731.candidate" in launch + plist - assert "--port 8081" in launch + assert "PORT=8081" in launch assert "8080" not in launch assert "launchctl" not in launch assert "/usr/bin/env -i" in launch assert "command environment override rejected" in launch assert "MTPLX_DSV4_0731_TEST_FIXTURE" in launch + assert "candidate_entry.py" in launch + assert "-m mtplx" not in launch def test_candidate_config_pins_all_installation_identities() -> None: @@ -30,9 +36,21 @@ def test_candidate_config_pins_all_installation_identities() -> None: assert config["candidate_port"] == 8081 assert config["candidate_label"] == "com.tea.deepseek-v4-0731.candidate" assert config["encoding_source_revision"] == "7872f01b1d1fe23eabc4c98b48bffcef5a386062" - for key in ("encoding_sha256", "model_config_sha256", "model_index_sha256", "trusted_python_sha256", "worktree_base_revision"): - expected_length = 40 if key == "worktree_base_revision" else 64 - assert len(config[key]) == expected_length + for key in ("model_config_sha256", "model_index_sha256", "trusted_python_sha256"): + assert len(config[key]) == 64 + assert config["reviewed_ref"] == "refs/tags/mtplx-dsv4-0731-reviewed" + assert len(config["encoding_assets"]) == 9 + for relative, expected in config["encoding_assets"].items(): + assert hashlib.sha256((ROOT / "encoding" / relative).read_bytes()).hexdigest() == expected + assert hashlib.sha256((ROOT / "encoding/SHA256SUMS").read_bytes()).hexdigest() == config[ + "encoding_manifest_sha256" + ] + assert hashlib.sha256((ROOT / "candidate_entry.py").read_bytes()).hexdigest() == config[ + "candidate_entry_sha256" + ] + assert hashlib.sha256((ROOT / "com.tea.deepseek-v4-0731.candidate.plist").read_bytes()).hexdigest() == config[ + "candidate_plist_sha256" + ] def test_command_override_is_rejected_except_for_nonstarting_fixture() -> None: @@ -41,7 +59,7 @@ def test_command_override_is_rejected_except_for_nonstarting_fixture() -> None: fixture = subprocess.run( [str(launcher), "--print-command"], env=fixture_env, check=True, capture_output=True, text=True ) - assert "--port 8081" in fixture.stdout + assert "127.0.0.1:8081" in fixture.stdout rejected = subprocess.run( [str(launcher), "--print-command"], @@ -66,10 +84,85 @@ def test_cutover_requires_receipts_lock_identity_and_explicit_promotion() -> Non "SENSITIVE_KEY", "finally:", "_bootstrap(prior_plist)", + "candidate_model_ids", ): assert required in source +def test_server_construction_installs_verified_0731_encoder() -> None: + from candidate_entry import install_candidate_surface + from mtplx.server import openai as openai_server + + def stock(*_args, **_kwargs): + return [999] + + server = SimpleNamespace( + _encode_messages=stock, + _parse_generated_tool_calls_or_content=lambda *_args, **_kwargs: (None, None), + omlx_extract_tool_calls_with_thinking=lambda *_args, **_kwargs: None, + _ToolAwareContentStreamTranslator=openai_server._ToolAwareContentStreamTranslator, + _stream_tool_call_deltas=openai_server._stream_tool_call_deltas, + ) + receipt = install_candidate_surface(server) + assert server._encode_messages is not stock + assert receipt["encoder"] == "deepseek-v4-flash-0731-official" + assert server._template_hash(None).startswith("deepseek-v4-flash-0731-official:") + assert server._apply_chat_template_profile(None, None) == { + "profile": "deepseek-v4-flash-0731-official", + "source": "official_python_encoder", + "path": None, + "applied": True, + "sha256": receipt["asset_set_sha256"], + } + + class Tokenizer: + def __init__(self) -> None: + self.encoded = "" + + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: + assert add_special_tokens is False + self.encoded = text + return list(text.encode("utf-8")) + + tokenizer = Tokenizer() + observability: dict[str, object] = {} + ids = server._encode_messages( + tokenizer, + [SimpleNamespace(role="user", content="hello", tool_calls=None)], + enable_thinking=True, + reasoning_effort="high", + tools=None, + template_observability=observability, + ) + assert ids == list(tokenizer.encoded.encode("utf-8")) + assert "<｜Assistant｜>" in tokenizer.encoded + assert observability == { + "backend_chat_encoding": "deepseek-v4-flash-0731-official", + "encoding_source_revision": "7872f01b1d1fe23eabc4c98b48bffcef5a386062", + } + + vector = (ROOT / "encoding/tests/test_output_1.txt").read_text(encoding="utf-8") + marker = "<｜Assistant｜>" + start = vector.find(marker) + len(marker) + end = vector.find("<｜User｜>", start) + thinking, regular = vector[start:end].split("", 1) + extraction = server.omlx_extract_tool_calls_with_thinking(thinking, regular, tokenizer, []) + assert extraction.parser_source == "deepseek_v4_0731_official" + assert extraction.tool_calls[0]["function"]["name"] == "get_weather" + assert extraction.tool_calls[0]["id"].startswith("call_") + + translator = server._ToolAwareContentStreamTranslator( + tools=[], argument_chunk_chars=16, tokenizer=tokenizer + ) + midpoint = len(regular) // 2 + assert translator.feed("content", regular[:midpoint]) == [] + assert translator.feed("content", regular[midpoint:]) == [] + deltas = translator.finish() + assert translator.suppressed_tool_markup is True + assert translator.tool_calls[0]["function"]["name"] == "get_weather" + assert any("tool_calls" in delta for delta in deltas) + + @pytest.mark.parametrize("forbidden_key, forbidden_value", [ ("stdout", "must never be retained"), ("prompt", "must never be retained"), @@ -78,6 +171,8 @@ def test_cutover_requires_receipts_lock_identity_and_explicit_promotion() -> Non ("argv", ["must never be retained"]), ("env", {"MUST_NEVER": "be retained"}), ("model_path", "/Users/davidtai/models/private"), + ("note", '{"messages":[{"role":"user"}]}'), + ("note", "Bearer abcdefghijklmnop"), ]) def test_candidate_receipt_rejects_sensitive_capture(forbidden_key: str, forbidden_value: object) -> None: import sys @@ -107,12 +202,58 @@ def test_scrubbed_passing_candidate_receipt_is_accepted() -> None: assert_candidate_receipt( { + "schema": "mtplx.dsv4-0731-candidate.v1", "candidate_preflight": { "ok": True, "label": "com.tea.deepseek-v4-0731.candidate", "port": 8081, + "plist_sha256": "a" * 64, + "encoding_source_revision": "7872f01b1d1fe23eabc4c98b48bffcef5a386062", + "encoding_asset_set_sha256": "b" * 64, + "reviewed_commit": "c" * 40, + "model_config_sha256": "d" * 64, + "model_index_sha256": "e" * 64, "promotion_target": {"label": "com.tea.deepseek-v4-0731.production", "plist_sha256": "b" * 64}, }, - "candidate_smoke": {"ok": True, "models_ok": True, "ready": True, "finish_reason": "stop"}, + "candidate_smoke": { + "ok": True, + "models_ok": True, + "ready": True, + "finish_reason": "stop", + "candidate_model_ids": ["deepseek-v4-0731"], + }, } ) + + +def test_candidate_receipt_rejects_unknown_nested_fields() -> None: + from promote_cutover import PromotionError, assert_candidate_receipt + + receipt = { + "schema": "mtplx.dsv4-0731-candidate.v1", + "candidate_preflight": { + "ok": True, + "label": "com.tea.deepseek-v4-0731.candidate", + "port": 8081, + "plist_sha256": "a" * 64, + "encoding_source_revision": "7872f01b1d1fe23eabc4c98b48bffcef5a386062", + "encoding_asset_set_sha256": "b" * 64, + "reviewed_commit": "c" * 40, + "model_config_sha256": "d" * 64, + "model_index_sha256": "e" * 64, + "promotion_target": { + "label": "com.tea.deepseek-v4-0731.production", + "plist_sha256": "f" * 64, + "unexpected": True, + }, + }, + "candidate_smoke": { + "ok": True, + "models_ok": True, + "ready": True, + "finish_reason": "stop", + "candidate_model_ids": ["deepseek-v4-0731"], + }, + } + with pytest.raises(PromotionError, match="strict receipt schema"): + assert_candidate_receipt(receipt) From 89cb996668bd37a3cdc120d503842c18899a3f66 Mon Sep 17 00:00:00 2001 From: davidtai Date: Mon, 3 Aug 2026 16:31:01 -0500 Subject: [PATCH 18/24] service: close DeepSeek V4 0731 promotion gates --- services/deepseek-v4-0731/README.md | 19 ++- services/deepseek-v4-0731/candidate.json | 12 +- services/deepseek-v4-0731/candidate_entry.py | 151 ++++++++++-------- services/deepseek-v4-0731/launch_candidate.sh | 22 ++- services/deepseek-v4-0731/promote_cutover.py | 45 ++++-- .../tests/test_service_surface.py | 75 +++++++++ 6 files changed, 232 insertions(+), 92 deletions(-) diff --git a/services/deepseek-v4-0731/README.md b/services/deepseek-v4-0731/README.md index 73842755..5b6136b0 100644 --- a/services/deepseek-v4-0731/README.md +++ b/services/deepseek-v4-0731/README.md @@ -10,8 +10,12 @@ input/output vectors from `candidate_entry.py` verifies all nine assets and runs every official vector at construction. It then installs the encoder directly at MTPLX's prompt-ID call site and installs the official DSML parser at the nonstream and streaming -response call sites. No tokenizer-template or stock prompt fallback remains in -the enabled 0731 lane. Per-request observability reports +response call sites. The stream translator retains ordinary preamble text while +continuing to scan later chunks and holding any suffix that could grow into a +split DSML marker; raw markup is never released. The complete turn still passes +through the official parser. Plain no-tool turns also use that parser and report +its engagement. No tokenizer-template, stock prompt, or stock completion-parser +fallback remains in the enabled 0731 lane. Per-request observability reports `backend_chat_encoding=deepseek-v4-flash-0731-official`. `launch_candidate.sh` accepts no production arguments. Its only test seam is @@ -21,7 +25,11 @@ service. A real launch requires: - the exact commit referenced by `refs/tags/mtplx-dsv4-0731-reviewed`; - a completely clean worktree; - the pinned interpreter, model config/index, manifest, encoder, and official - vector hashes; and + vector hashes; +- the exact reviewed artifact validator at commit `bbf02944`, which hashes the + 0731 tokenizer and all 20 Safetensors shards and checks their index closure, + topology, and quantization assignment; +- the separately pinned official `tokenizer_config.json`; and - the fixed, absolute entrypoint and minimal `env -i` environment. `promote_cutover.py` remains an explicit operator action. Before it can stop a @@ -31,7 +39,10 @@ production plist, the nonblocking GPU lock, and exact current launchd label/PID/listener/plist identity. Candidate model IDs are taken only from the signed receipt for cutover verification; the prior model IDs are used only to verify rollback. The same lock remains held through restoration and the real -`/v1/models` plus `READY`/`finish_reason=stop` smoke. +`/v1/models` plus exact `content.strip() == "READY"`/`finish_reason=stop` smoke. +After the new plist is bootstrapped, its hash, launchd PID, and ownership of the +8080 listener are reattested under the lock before any HTTP identity or readiness +probe is allowed. Receipts have an exact allowlist and recursively reject local paths, request content, tool schemas, secrets, argv/env, and captured process output. diff --git a/services/deepseek-v4-0731/candidate.json b/services/deepseek-v4-0731/candidate.json index 822dc635..26a79213 100644 --- a/services/deepseek-v4-0731/candidate.json +++ b/services/deepseek-v4-0731/candidate.json @@ -15,15 +15,19 @@ "tests/test_output_3.txt": "b3b1cd8748b7b90d3c6be6da3f786f12e4d70be073bd445ea162dfad4dc01a64", "tests/test_output_4.txt": "60e1643840ba9e4aeede450feb7b0498fa66ee24e4a939d48855ce04ec6fc375" }, - "model_path": "/Users/davidtai/models/DeepSeek-V4-Flash-2bit-DQ-mtp", - "model_config_sha256": "c8ff87fd5ee5c9587d0c937e9bfd3193e1a1621141aa367848a9610b3291fa6f", - "model_index_sha256": "c84d2b369f5d5023d0f2d183fc36a935a3981751414996243b65f069983e43d8", + "model_path": "/Users/davidtai/models/DeepSeek-V4-Flash-0731-oQ2e-mtp", + "model_config_sha256": "6d0297a4329d55dccf3cd48fd168efea8044996245195d518a9e8aaa14906d3e", + "model_index_sha256": "9edcd0db7e6b8f0b8e02978d73c30083b2aa64c2e3a8fd77d3b776a4efb4bc91", + "model_tokenizer_config_sha256": "6ac8c8dc065ed118161d02dd532749ae3f52c243deac27872134fae2f50d8547", + "artifact_validator_commit": "bbf02944aab3e17be754ba3c88d6aad3c488d10d", + "artifact_validator_path": "scripts/deepseek_v4_0731_artifact_check.py", + "artifact_validator_blob_sha256": "672e3bafa8381c5264960d065730d9894b12f832eeb358922e0dd703042ac67b", "worktree": "/Users/davidtai/projects/OpenSourceWTF/.worktrees/dsv4-0731-service", "reviewed_ref": "refs/tags/mtplx-dsv4-0731-reviewed", "trusted_python": "/Users/davidtai/projects/OpenSourceWTF/.worktrees/dsv4-0731-service/.venv/bin/python", "trusted_python_target": "/Users/davidtai/.local/share/uv/python/cpython-3.12-macos-aarch64-none/bin/python3.12", "trusted_python_sha256": "96793b100c947cdc81a38e8fb8c9c1889abccda9840ce1bef58d372bf3f2c263", - "candidate_entry_sha256": "c7e2e79e45b3f2e8afd8453d5976e5234caea724813eae0e555c5f956bf725aa", + "candidate_entry_sha256": "029a177a932b2736fc02b90ebf8e240b576a28993ea0940ddba66a2719c26d1b", "candidate_plist_sha256": "93eac0d4eaac491c7f2f1d3a293ba38a3144ade59ee3afdf52b35cc9ec9bb101", "served_model_id": "deepseek-v4-0731-candidate" } diff --git a/services/deepseek-v4-0731/candidate_entry.py b/services/deepseek-v4-0731/candidate_entry.py index b5bd930b..88bbab2b 100755 --- a/services/deepseek-v4-0731/candidate_entry.py +++ b/services/deepseek-v4-0731/candidate_entry.py @@ -2,8 +2,9 @@ """Construction-only entrypoint for the isolated V4-Flash-0731 service. This module verifies and self-tests the official encoder before replacing the -two MTPLX request-path call sites that own prompt encoding and DSML completion -parsing. There is no tokenizer-template or stock-prompt fallback after install. +MTPLX request-path call sites that own prompt encoding and DSML completion +parsing. There is no tokenizer-template, stock-prompt, or stock completion +parser fallback after install. """ from __future__ import annotations @@ -12,7 +13,6 @@ import json import re import sys -import uuid from pathlib import Path from types import ModuleType from typing import Any @@ -170,7 +170,6 @@ def encode_messages( def _install_completion_parser(server: ModuleType, encoding: ModuleType): - stock = server._parse_generated_tool_calls_or_content dsml_marker = f"<{encoding.dsml_token}{encoding.tool_calls_block_name}>" def parse_generated_tool_calls_or_content( @@ -182,15 +181,7 @@ def parse_generated_tool_calls_or_content( response_id: str | None = None, stream: bool = False, ): - if dsml_marker not in text: - return stock( - text, - tools=tools, - tokenizer=tokenizer, - state=state, - response_id=response_id, - stream=stream, - ) + del tools, tokenizer, state, response_id, stream completion = text if text.endswith(encoding.eos_token) else text + encoding.eos_token mode = "thinking" if encoding.thinking_end_token in completion.split(dsml_marker, 1)[0] else "chat" try: @@ -207,18 +198,13 @@ def _install_actual_tool_extractor(server: ModuleType, encoding: ModuleType) -> """Install official DSML completion parsing at the live response call site.""" from mtplx.server.omlx_bridge import ToolCallExtraction - stock = server.omlx_extract_tool_calls_with_thinking - dsml_marker = f"<{encoding.dsml_token}{encoding.tool_calls_block_name}>" - def extract( thinking_content: str, regular_content: str, tokenizer: Any | None, tools: list[dict[str, Any]] | None = None, ) -> ToolCallExtraction: - combined = thinking_content + regular_content - if dsml_marker not in combined: - return stock(thinking_content, regular_content, tokenizer, tools) + del tokenizer, tools mode = "thinking" if thinking_content else "chat" completion = ( thinking_content + encoding.thinking_end_token + regular_content @@ -233,114 +219,137 @@ def extract( raise CandidateConstructionError("malformed V4-0731 DSML completion") from error calls = parsed.get("tool_calls") or None if calls: - calls = [ - {**call, "id": str(call.get("id") or f"call_{uuid.uuid4().hex[:24]}")} - for call in calls - ] + stable_calls = [] + for index, call in enumerate(calls): + canonical = json.dumps(call, sort_keys=True, separators=(",", ":")) + call_id = "call_" + hashlib.sha256( + f"{index}:{canonical}".encode("utf-8") + ).hexdigest()[:24] + stable_calls.append({**call, "id": str(call.get("id") or call_id)}) + calls = stable_calls return ToolCallExtraction( cleaned_text=str(parsed.get("content") or ""), tool_calls=calls, cleaned_thinking=str(parsed.get("reasoning_content") or ""), parser_source="deepseek_v4_0731_official", status="parsed" if calls else "no_tool", - raw_tool_markup_suppressed=True, + raw_tool_markup_suppressed=bool(calls), ) server.omlx_extract_tool_calls_with_thinking = extract def _install_stream_translator(server: ModuleType, encoding: ModuleType) -> None: - """Buffer the official DSML envelope so streaming never leaks it as text.""" - stock_class = server._ToolAwareContentStreamTranslator - dsml_marker = f"<{encoding.dsml_token}{encoding.tool_calls_block_name}>" + """Keep scanning after visible prose while holding split DSML prefixes.""" + + dsml_marker = f"<{encoding.dsml_token}" + dsml_envelope_start = "\n\n" + dsml_marker class DSV40731StreamTranslator: def __init__(self, *, tools, argument_chunk_chars, tokenizer=None, **kwargs) -> None: + del kwargs self._tools = tools - self._argument_chunk_chars = argument_chunk_chars + self._argument_chunk_chars = max(1, int(argument_chunk_chars)) self._tokenizer = tokenizer - self._stock = stock_class( - tools=tools, - argument_chunk_chars=argument_chunk_chars, - tokenizer=tokenizer, - **kwargs, - ) self._pending = "" - self._mode = "undecided" + self._all_content = "" + self._emitted_content = "" + self._inside_dsml = False self.tool_calls = None self.fallback_reason = None self.tool_parser_dialect = "deepseek_v4_0731_official" self._suppressed = False + self._emitted_tool_deltas = False @property def has_tool_calls(self): - return bool(self.tool_calls) if self._mode == "dsml" else self._stock.has_tool_calls + return bool(self.tool_calls) @property def has_emitted_tool_deltas(self): - return False if self._mode == "dsml" else self._stock.has_emitted_tool_deltas + return self._emitted_tool_deltas @property def suppressed_tool_markup(self): - return self._suppressed or self._stock.suppressed_tool_markup + return self._suppressed @property def buffering_tool_call(self): - return self._mode == "dsml" or self._stock.buffering_tool_call + return self._inside_dsml @property def tool_argument_in_progress(self): - return self._mode == "dsml" or self._stock.tool_argument_in_progress + return self._inside_dsml @property def ready_to_finish_tool_turn(self): - return False if self._mode == "dsml" else self._stock.ready_to_finish_tool_turn + return False @property def invalid_trailing_after_tool_call(self): - return False if self._mode == "dsml" else self._stock.invalid_trailing_after_tool_call + return False def feed(self, field: str, text: str): - if self._mode == "stock": - return self._stock.feed(field, text) + if not text: + return [] if field != "content": - return self._stock.feed(field, text) + return [{field: text}] + self._all_content += text + if self._inside_dsml: + return [] self._pending += text - stripped = self._pending.lstrip() - if dsml_marker in stripped: - self._mode = "dsml" + marker = self._pending.find(dsml_marker) + if marker >= 0: + visible = self._pending[:marker] + if visible.endswith("\n\n"): + visible = visible[:-2] + self._pending = "" + self._inside_dsml = True self._suppressed = True - return [] - if dsml_marker.startswith(stripped): - return [] - self._mode = "stock" - pending, self._pending = self._pending, "" - return self._stock.feed(field, pending) + self._emitted_content += visible + return [{"content": visible}] if visible else [] + hold = 0 + prefix_limit = max(len(dsml_marker), len(dsml_envelope_start)) - 1 + for size in range(min(len(self._pending), prefix_limit), 0, -1): + suffix = self._pending[-size:] + if dsml_marker.startswith(suffix) or dsml_envelope_start.startswith(suffix): + hold = size + break + visible = self._pending[:-hold] if hold else self._pending + self._pending = self._pending[-hold:] if hold else "" + self._emitted_content += visible + return [{"content": visible}] if visible else [] def finish(self, *, defer_content_resolution: bool = False): - if self._mode != "dsml": - if self._pending: - self._stock.feed("content", self._pending) - self._pending = "" - return self._stock.finish(defer_content_resolution=defer_content_resolution) + del defer_content_resolution extraction = server.omlx_extract_tool_calls_with_thinking( - "", self._pending, self._tokenizer, self._tools + "", self._all_content, self._tokenizer, self._tools ) self.tool_calls = extraction.tool_calls self._pending = "" - if not self.tool_calls: - raise CandidateConstructionError("official DSML stream ended without tool calls") - return list( - server._stream_tool_call_deltas( - self.tool_calls, - argument_chunk_chars=self._argument_chunk_chars, + self._all_content = "" + deltas = [] + if not extraction.cleaned_text.startswith(self._emitted_content): + raise CandidateConstructionError("official stream parse changed emitted content") + remaining_content = extraction.cleaned_text[len(self._emitted_content):] + if remaining_content: + deltas.append({"content": remaining_content}) + self._emitted_content += remaining_content + if self.tool_calls: + self._suppressed = True + tool_deltas = list( + server._stream_tool_call_deltas( + self.tool_calls, + argument_chunk_chars=self._argument_chunk_chars, + ) ) - ) + self._emitted_tool_deltas = bool(tool_deltas) + deltas.extend(tool_deltas) + return deltas def resolve_deferred_content(self, *, has_tool_calls: bool): - if self._mode == "dsml": - return [] - return self._stock.resolve_deferred_content(has_tool_calls=has_tool_calls) + del has_tool_calls + return [] server._ToolAwareContentStreamTranslator = DSV40731StreamTranslator @@ -424,7 +433,7 @@ def main() -> int: "serve", "--host", "127.0.0.1", "--port", "8081", - "--model", "/Users/davidtai/models/DeepSeek-V4-Flash-2bit-DQ-mtp", + "--model", "/Users/davidtai/models/DeepSeek-V4-Flash-0731-oQ2e-mtp", "--model-id", "deepseek-v4-0731-candidate", "--reasoning", "on", "--reasoning-effort", "low", diff --git a/services/deepseek-v4-0731/launch_candidate.sh b/services/deepseek-v4-0731/launch_candidate.sh index 137ef854..67df2665 100755 --- a/services/deepseek-v4-0731/launch_candidate.sh +++ b/services/deepseek-v4-0731/launch_candidate.sh @@ -5,12 +5,15 @@ umask 077 SERVICE_ROOT=/Users/davidtai/projects/OpenSourceWTF/.worktrees/dsv4-0731-service/services/deepseek-v4-0731 WORKTREE=/Users/davidtai/projects/OpenSourceWTF/.worktrees/dsv4-0731-service -MODEL=/Users/davidtai/models/DeepSeek-V4-Flash-2bit-DQ-mtp +MODEL=/Users/davidtai/models/DeepSeek-V4-Flash-0731-oQ2e-mtp PYTHON=/Users/davidtai/projects/OpenSourceWTF/.worktrees/dsv4-0731-service/.venv/bin/python PYTHON_TARGET=/Users/davidtai/.local/share/uv/python/cpython-3.12-macos-aarch64-none/bin/python3.12 ENTRY="$SERVICE_ROOT/candidate_entry.py" ENCODING="$SERVICE_ROOT/encoding" REVIEWED_REF=refs/tags/mtplx-dsv4-0731-reviewed +ARTIFACT_VALIDATOR_COMMIT=bbf02944aab3e17be754ba3c88d6aad3c488d10d +ARTIFACT_VALIDATOR_PATH=scripts/deepseek_v4_0731_artifact_check.py +ARTIFACT_VALIDATOR_BLOB_SHA256=672e3bafa8381c5264960d065730d9894b12f832eeb358922e0dd703042ac67b PORT=8081 die() { printf '%s\n' "deepseek-v4-0731 candidate: $1" >&2; exit 64; } @@ -31,10 +34,11 @@ fi [ -x "$PYTHON_TARGET" ] && [ ! -L "$PYTHON_TARGET" ] || die "trusted python target is missing or unsafe" [ "$(sha256 "$PYTHON_TARGET")" = 96793b100c947cdc81a38e8fb8c9c1889abccda9840ce1bef58d372bf3f2c263 ] || die "trusted python hash changed" [ -f "$ENTRY" ] && [ ! -L "$ENTRY" ] || die "candidate entrypoint is missing or unsafe" -[ "$(sha256 "$ENTRY")" = c7e2e79e45b3f2e8afd8453d5976e5234caea724813eae0e555c5f956bf725aa ] || die "candidate entrypoint hash changed" +[ "$(sha256 "$ENTRY")" = 029a177a932b2736fc02b90ebf8e240b576a28993ea0940ddba66a2719c26d1b ] || die "candidate entrypoint hash changed" [ -d "$MODEL" ] && [ ! -L "$MODEL" ] || die "pinned model path is missing or unsafe" -[ "$(sha256 "$MODEL/config.json")" = c8ff87fd5ee5c9587d0c937e9bfd3193e1a1621141aa367848a9610b3291fa6f ] || die "model configuration hash changed" -[ "$(sha256 "$MODEL/model.safetensors.index.json")" = c84d2b369f5d5023d0f2d183fc36a935a3981751414996243b65f069983e43d8 ] || die "model index hash changed" +[ "$(sha256 "$MODEL/config.json")" = 6d0297a4329d55dccf3cd48fd168efea8044996245195d518a9e8aaa14906d3e ] || die "model configuration hash changed" +[ "$(sha256 "$MODEL/model.safetensors.index.json")" = 9edcd0db7e6b8f0b8e02978d73c30083b2aa64c2e3a8fd77d3b776a4efb4bc91 ] || die "model index hash changed" +[ "$(sha256 "$MODEL/tokenizer_config.json")" = 6ac8c8dc065ed118161d02dd532749ae3f52c243deac27872134fae2f50d8547 ] || die "model tokenizer configuration hash changed" [ "$(sha256 "$ENCODING/SHA256SUMS")" = 6758dfda8a39afdd00d907606c42c1a268289c463351b9628ac07f4f916d7d0a ] || die "official encoding manifest hash changed" (cd "$ENCODING" && /usr/bin/shasum -a 256 -c SHA256SUMS >/dev/null) || die "official encoding/vector asset hash changed" @@ -43,6 +47,16 @@ current_commit=$(/usr/bin/git -C "$WORKTREE" rev-parse --verify HEAD) || die "wo [ "$current_commit" = "$reviewed_commit" ] || die "worktree is not the exact reviewed commit" [ -z "$(/usr/bin/git -C "$WORKTREE" status --porcelain=v1 --untracked-files=all)" ] || die "reviewed worktree is not clean" +# Execute the exact reviewed validator blob. It pins tokenizer.json plus every +# one of the 20 model shards and rejects unknown, missing, or changing files. +validator_blob_sha=$( + /usr/bin/git -C "$WORKTREE" cat-file blob "$ARTIFACT_VALIDATOR_COMMIT:$ARTIFACT_VALIDATOR_PATH" | + /usr/bin/shasum -a 256 | /usr/bin/awk '{print $1}' +) || die "reviewed artifact validator is unavailable" +[ "$validator_blob_sha" = "$ARTIFACT_VALIDATOR_BLOB_SHA256" ] || die "reviewed artifact validator hash changed" +/usr/bin/git -C "$WORKTREE" cat-file blob "$ARTIFACT_VALIDATOR_COMMIT:$ARTIFACT_VALIDATOR_PATH" | + "$PYTHON" - "$MODEL" >/dev/null || die "reviewed artifact validation failed" + exec /usr/bin/env -i \ HOME=/Users/davidtai \ LC_ALL=C \ diff --git a/services/deepseek-v4-0731/promote_cutover.py b/services/deepseek-v4-0731/promote_cutover.py index ef7155c2..ad3df61d 100755 --- a/services/deepseek-v4-0731/promote_cutover.py +++ b/services/deepseek-v4-0731/promote_cutover.py @@ -43,8 +43,8 @@ REVIEWED_REF = "refs/tags/mtplx-dsv4-0731-reviewed" CANDIDATE_PLIST_SHA256 = "93eac0d4eaac491c7f2f1d3a293ba38a3144ade59ee3afdf52b35cc9ec9bb101" ENCODING_ASSET_SET_SHA256 = "6758dfda8a39afdd00d907606c42c1a268289c463351b9628ac07f4f916d7d0a" -MODEL_CONFIG_SHA256 = "c8ff87fd5ee5c9587d0c937e9bfd3193e1a1621141aa367848a9610b3291fa6f" -MODEL_INDEX_SHA256 = "c84d2b369f5d5023d0f2d183fc36a935a3981751414996243b65f069983e43d8" +MODEL_CONFIG_SHA256 = "6d0297a4329d55dccf3cd48fd168efea8044996245195d518a9e8aaa14906d3e" +MODEL_INDEX_SHA256 = "9edcd0db7e6b8f0b8e02978d73c30083b2aa64c2e3a8fd77d3b776a4efb4bc91" class PromotionError(RuntimeError): @@ -96,7 +96,11 @@ def _smoke_stop(model_id: str) -> None: payload = json.loads(response.read().decode("utf-8")) choice = payload["choices"][0] content = choice["message"]["content"] - if choice.get("finish_reason") != "stop" or not isinstance(content, str) or "READY" not in content: + if ( + choice.get("finish_reason") != "stop" + or not isinstance(content, str) + or content.strip() != "READY" + ): raise ValueError("required READY/stop evidence absent") except (KeyError, OSError, TypeError, ValueError, urllib.error.URLError) as error: raise PromotionError("service smoke did not return READY with finish_reason=stop") from error @@ -119,26 +123,45 @@ def _launchctl_pid(label: str) -> int: return int(match.group(1)) -def attest_live(*, label: str, plist: Path) -> dict[str, Any]: - """Capture exact live identity without sending a generation prompt.""" +def attest_process_identity(*, label: str, plist: Path) -> dict[str, Any]: + """Bind one launchd label, plist, and 8080 listener before any HTTP probe.""" launch_pid = _launchctl_pid(label) listener_pid = _listener_pid(LIVE_PORT) if launch_pid != listener_pid: raise PromotionError("launchd PID and 8080 listener PID differ") + return { + "label": label, + "pid": launch_pid, + "listener_port": LIVE_PORT, + "plist_sha256": _sha256(plist), + } + + +def attest_live(*, label: str, plist: Path) -> dict[str, Any]: + """Capture exact live identity without sending a generation prompt.""" + process = attest_process_identity(label=label, plist=plist) models = _http_json(f"http://127.0.0.1:{LIVE_PORT}/v1/models") model_ids = [item.get("id") for item in models.get("data", []) if isinstance(item, dict)] if not model_ids or not all(isinstance(model_id, str) for model_id in model_ids): raise PromotionError("live /v1/models is not a valid service identity") return { "schema": "mtplx.live-identity.v1", - "label": label, - "pid": launch_pid, - "listener_port": LIVE_PORT, - "plist_sha256": _sha256(plist), + **process, "model_ids": model_ids, } +def _wait_for_process_identity(*, label: str, plist: Path) -> dict[str, Any]: + """Wait for launchd and the listener to converge without probing HTTP.""" + deadline = time.monotonic() + 600 + while time.monotonic() < deadline: + try: + return attest_process_identity(label=label, plist=plist) + except PromotionError: + time.sleep(0.5) + raise PromotionError("promoted launchd process did not acquire the 8080 listener") + + def _read_json_bytes(path: Path) -> tuple[dict[str, Any], bytes]: if not path.is_file() or path.is_symlink(): raise PromotionError("receipt is missing or unsafe") @@ -357,12 +380,16 @@ def promote(args: argparse.Namespace) -> None: try: _bootout(current["label"]) _bootstrap(target) + promoted = _wait_for_process_identity(label=args.production_label, plist=target) + if promoted["plist_sha256"] != promotion_target["plist_sha256"]: + raise PromotionError("promoted service plist identity changed during cutover") _verify_live_ready(candidate["candidate_smoke"]["candidate_model_ids"]) except BaseException: try: _bootout(args.production_label) finally: _bootstrap(prior_plist) + _wait_for_process_identity(label=current["label"], plist=prior_plist) _verify_live_ready(current["model_ids"]) raise diff --git a/services/deepseek-v4-0731/tests/test_service_surface.py b/services/deepseek-v4-0731/tests/test_service_surface.py index 2b74d21c..d9ccdd94 100644 --- a/services/deepseek-v4-0731/tests/test_service_surface.py +++ b/services/deepseek-v4-0731/tests/test_service_surface.py @@ -39,6 +39,29 @@ def test_candidate_config_pins_all_installation_identities() -> None: for key in ("model_config_sha256", "model_index_sha256", "trusted_python_sha256"): assert len(config[key]) == 64 assert config["reviewed_ref"] == "refs/tags/mtplx-dsv4-0731-reviewed" + assert config["artifact_validator_commit"] == "bbf02944aab3e17be754ba3c88d6aad3c488d10d" + assert config["artifact_validator_path"] == "scripts/deepseek_v4_0731_artifact_check.py" + assert config["artifact_validator_blob_sha256"] == ( + "672e3bafa8381c5264960d065730d9894b12f832eeb358922e0dd703042ac67b" + ) + validator = subprocess.check_output( + [ + "git", + "-C", + str(ROOT.parents[1]), + "cat-file", + "blob", + f'{config["artifact_validator_commit"]}:{config["artifact_validator_path"]}', + ] + ) + assert hashlib.sha256(validator).hexdigest() == config["artifact_validator_blob_sha256"] + assert config["model_path"] == "/Users/davidtai/models/DeepSeek-V4-Flash-0731-oQ2e-mtp" + assert config["model_config_sha256"] == ( + "6d0297a4329d55dccf3cd48fd168efea8044996245195d518a9e8aaa14906d3e" + ) + assert config["model_index_sha256"] == ( + "9edcd0db7e6b8f0b8e02978d73c30083b2aa64c2e3a8fd77d3b776a4efb4bc91" + ) assert len(config["encoding_assets"]) == 9 for relative, expected in config["encoding_assets"].items(): assert hashlib.sha256((ROOT / "encoding" / relative).read_bytes()).hexdigest() == expected @@ -85,8 +108,24 @@ def test_cutover_requires_receipts_lock_identity_and_explicit_promotion() -> Non "finally:", "_bootstrap(prior_plist)", "candidate_model_ids", + "attest_process_identity", ): assert required in source + bootstrap = source.index("_bootstrap(target)") + new_identity = source.index("_wait_for_process_identity(", bootstrap) + readiness = source.index("_verify_live_ready(", new_identity) + assert bootstrap < new_identity < readiness + assert 'content.strip() != "READY"' in source + assert "return attest_process_identity(label=label, plist=plist)" in source + + +def test_launcher_invokes_exact_reviewed_artifact_validator() -> None: + source = (ROOT / "launch_candidate.sh").read_text(encoding="utf-8") + assert "bbf02944aab3e17be754ba3c88d6aad3c488d10d" in source + assert "672e3bafa8381c5264960d065730d9894b12f832eeb358922e0dd703042ac67b" in source + assert "scripts/deepseek_v4_0731_artifact_check.py" in source + assert "git -C \"$WORKTREE\" cat-file blob" in source + assert '"$PYTHON" - "$MODEL"' in source def test_server_construction_installs_verified_0731_encoder() -> None: @@ -151,6 +190,13 @@ def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: assert extraction.tool_calls[0]["function"]["name"] == "get_weather" assert extraction.tool_calls[0]["id"].startswith("call_") + no_tool = server.omlx_extract_tool_calls_with_thinking("", "READY", tokenizer, []) + assert no_tool.parser_source == "deepseek_v4_0731_official" + assert no_tool.status == "no_tool" + assert no_tool.cleaned_text == "READY" + assert no_tool.tool_calls is None + assert server._parse_generated_tool_calls_or_content("READY", tools=[]) == (None, None) + translator = server._ToolAwareContentStreamTranslator( tools=[], argument_chunk_chars=16, tokenizer=tokenizer ) @@ -162,6 +208,35 @@ def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: assert translator.tool_calls[0]["function"]["name"] == "get_weather" assert any("tool_calls" in delta for delta in deltas) + preamble = "I will check." + nonstream = server.omlx_extract_tool_calls_with_thinking( + "", preamble + regular, tokenizer, [] + ) + translator = server._ToolAwareContentStreamTranslator( + tools=[], argument_chunk_chars=16, tokenizer=tokenizer + ) + # The ordinary preamble arrives before there is any evidence of DSML. + deltas = translator.feed("content", preamble) + deltas.extend(translator.feed("content", regular[:midpoint])) + deltas.extend(translator.feed("content", regular[midpoint:])) + deltas.extend(translator.finish()) + streamed_content = "".join(delta.get("content", "") for delta in deltas) + assert streamed_content == nonstream.cleaned_text == preamble + assert translator.tool_calls == nonstream.tool_calls + assert any("tool_calls" in delta for delta in deltas) + assert "<｜DSML｜" not in json.dumps(deltas, ensure_ascii=False) + + split_translator = server._ToolAwareContentStreamTranslator( + tools=[], argument_chunk_chars=16, tokenizer=tokenizer + ) + split_deltas: list[dict[str, object]] = [] + for character in preamble + regular: + split_deltas.extend(split_translator.feed("content", character)) + split_deltas.extend(split_translator.finish()) + assert "".join(str(delta.get("content", "")) for delta in split_deltas) == preamble + assert split_translator.tool_calls == nonstream.tool_calls + assert "<｜DSML｜" not in json.dumps(split_deltas, ensure_ascii=False) + @pytest.mark.parametrize("forbidden_key, forbidden_value", [ ("stdout", "must never be retained"), From f9c5aa851404c71cb5b0158918cada5a80353767 Mon Sep 17 00:00:00 2001 From: davidtai Date: Mon, 3 Aug 2026 16:46:51 -0500 Subject: [PATCH 19/24] service: harden DeepSeek 0731 stream and promotion gates --- pyproject.toml | 4 + services/deepseek-v4-0731/README.md | 9 ++ services/deepseek-v4-0731/candidate.json | 2 +- services/deepseek-v4-0731/candidate_entry.py | 55 +++++++ services/deepseek-v4-0731/launch_candidate.sh | 2 +- services/deepseek-v4-0731/promote_cutover.py | 38 ++++- .../tests/test_service_surface.py | 136 +++++++++++++++++- 7 files changed, 237 insertions(+), 9 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 5c6905fd..78262680 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -104,3 +104,7 @@ mtplx = ["templates/**/*.jinja"] testpaths = ["tests"] pythonpath = ["."] addopts = "-q" + +[tool.ruff.lint.per-file-ignores] +# Byte-exact vendored upstream source; candidate wrappers remain fully linted. +"services/deepseek-v4-0731/encoding/encoding_dsv4.py" = ["ALL"] diff --git a/services/deepseek-v4-0731/README.md b/services/deepseek-v4-0731/README.md index 5b6136b0..518b312f 100644 --- a/services/deepseek-v4-0731/README.md +++ b/services/deepseek-v4-0731/README.md @@ -46,3 +46,12 @@ probe is allowed. Receipts have an exact allowlist and recursively reject local paths, request content, tool schemas, secrets, argv/env, and captured process output. + +The reviewed, dedicated signer list lives at +`~/.config/mtplx/deepseek-v4-0731-allowed-signers` (owner-only mode `0600`), +with its SHA-256 pinned in `promote_cutover.py`. Sign a receipt without +printing key material using `/usr/bin/ssh-keygen -Y sign -f +~/.config/mtplx/deepseek-v4-0731-signing -n mtplx-deepseek-v4-0731 +candidate-receipt.json`; this dedicated key is not a reused login key. +Promotion rejects a missing, changed, wrongly owned, or group/world-writable +signer list; it does not discover trust at runtime. diff --git a/services/deepseek-v4-0731/candidate.json b/services/deepseek-v4-0731/candidate.json index 26a79213..2fdb8a3d 100644 --- a/services/deepseek-v4-0731/candidate.json +++ b/services/deepseek-v4-0731/candidate.json @@ -27,7 +27,7 @@ "trusted_python": "/Users/davidtai/projects/OpenSourceWTF/.worktrees/dsv4-0731-service/.venv/bin/python", "trusted_python_target": "/Users/davidtai/.local/share/uv/python/cpython-3.12-macos-aarch64-none/bin/python3.12", "trusted_python_sha256": "96793b100c947cdc81a38e8fb8c9c1889abccda9840ce1bef58d372bf3f2c263", - "candidate_entry_sha256": "029a177a932b2736fc02b90ebf8e240b576a28993ea0940ddba66a2719c26d1b", + "candidate_entry_sha256": "cc16dff9d2ebcd01af55cfc9ba36ca01ee832e3db264e5af311552292d127c97", "candidate_plist_sha256": "93eac0d4eaac491c7f2f1d3a293ba38a3144ade59ee3afdf52b35cc9ec9bb101", "served_model_id": "deepseek-v4-0731-candidate" } diff --git a/services/deepseek-v4-0731/candidate_entry.py b/services/deepseek-v4-0731/candidate_entry.py index 88bbab2b..b4a7e841 100755 --- a/services/deepseek-v4-0731/candidate_entry.py +++ b/services/deepseek-v4-0731/candidate_entry.py @@ -354,6 +354,59 @@ def resolve_deferred_content(self, *, has_tool_calls: bool): server._ToolAwareContentStreamTranslator = DSV40731StreamTranslator +def _install_no_tools_stream_sanitizer(server: ModuleType) -> None: + """Apply the official DSML sanitizer to no-tools SSE content too.""" + stock_factory = server._stream_splitter_for_state + + class DSV40731NoToolsStreamSplitter: + def __init__(self, stock: Any) -> None: + self._stock = stock + self._translator = server._ToolAwareContentStreamTranslator( + tools=[], argument_chunk_chars=1 + ) + + def start(self): + return self._translate(self._stock.start()) + + def feed(self, text: str): + return self._translate(self._stock.feed(text)) + + def finish(self, **kwargs): + chunks = self._translate(self._stock.finish(**kwargs)) + try: + chunks.extend(self._content_pairs(self._translator.finish())) + except CandidateConstructionError: + # A malformed completion cannot release the DSML suffix that + # the translator withheld while it waited for official parsing. + return chunks + return chunks + + def _translate(self, chunks: list[tuple[str, str]]): + translated: list[tuple[str, str]] = [] + for field, text in chunks: + if field == "content": + translated.extend(self._content_pairs(self._translator.feed(field, text))) + else: + translated.append((field, text)) + return translated + + @staticmethod + def _content_pairs(deltas: list[dict[str, Any]]): + return [ + ("content", text) + for delta in deltas + if isinstance((text := delta.get("content")), str) and text + ] + + def stream_splitter_for_state(*args, **kwargs): + stock = stock_factory(*args, **kwargs) + if kwargs.get("suppress_orphan_tool_markup"): + return DSV40731NoToolsStreamSplitter(stock) + return stock + + server._stream_splitter_for_state = stream_splitter_for_state + + def _install_reasoning_policy(server: ModuleType) -> None: def normalize(value: Any, *, default: str = "low") -> str: effort = str(value or default).strip().lower() @@ -410,6 +463,8 @@ def install_candidate_surface(server: ModuleType) -> dict[str, str]: _install_actual_tool_extractor(server, encoding) if hasattr(server, "_ToolAwareContentStreamTranslator"): _install_stream_translator(server, encoding) + if hasattr(server, "_stream_splitter_for_state"): + _install_no_tools_stream_sanitizer(server) _install_reasoning_policy(server) _install_construction_identity(server) server._DSV4_0731_ENCODER_INSTALLED = True diff --git a/services/deepseek-v4-0731/launch_candidate.sh b/services/deepseek-v4-0731/launch_candidate.sh index 67df2665..e6e14bc3 100755 --- a/services/deepseek-v4-0731/launch_candidate.sh +++ b/services/deepseek-v4-0731/launch_candidate.sh @@ -34,7 +34,7 @@ fi [ -x "$PYTHON_TARGET" ] && [ ! -L "$PYTHON_TARGET" ] || die "trusted python target is missing or unsafe" [ "$(sha256 "$PYTHON_TARGET")" = 96793b100c947cdc81a38e8fb8c9c1889abccda9840ce1bef58d372bf3f2c263 ] || die "trusted python hash changed" [ -f "$ENTRY" ] && [ ! -L "$ENTRY" ] || die "candidate entrypoint is missing or unsafe" -[ "$(sha256 "$ENTRY")" = 029a177a932b2736fc02b90ebf8e240b576a28993ea0940ddba66a2719c26d1b ] || die "candidate entrypoint hash changed" +[ "$(sha256 "$ENTRY")" = cc16dff9d2ebcd01af55cfc9ba36ca01ee832e3db264e5af311552292d127c97 ] || die "candidate entrypoint hash changed" [ -d "$MODEL" ] && [ ! -L "$MODEL" ] || die "pinned model path is missing or unsafe" [ "$(sha256 "$MODEL/config.json")" = 6d0297a4329d55dccf3cd48fd168efea8044996245195d518a9e8aaa14906d3e ] || die "model configuration hash changed" [ "$(sha256 "$MODEL/model.safetensors.index.json")" = 9edcd0db7e6b8f0b8e02978d73c30083b2aa64c2e3a8fd77d3b776a4efb4bc91 ] || die "model index hash changed" diff --git a/services/deepseek-v4-0731/promote_cutover.py b/services/deepseek-v4-0731/promote_cutover.py index ad3df61d..be1e93bf 100755 --- a/services/deepseek-v4-0731/promote_cutover.py +++ b/services/deepseek-v4-0731/promote_cutover.py @@ -16,10 +16,12 @@ import os import plistlib import re +import stat import subprocess import sys import time import urllib.error +import urllib.parse import urllib.request from contextlib import contextmanager from pathlib import Path @@ -36,9 +38,16 @@ r"[\"']?(?:prompt|messages|tools|secret|authorization)[\"']?\s*[:=])", re.I, ) +PATH_LIKE_VALUE = re.compile( + r"(?:^~[\\/]|(?:^|[\\/])\.\.(?:[\\/]|$)|\b[A-Za-z]:[\\/]|" + r"[A-Za-z][A-Za-z0-9+.-]*://|(?:^|\s)/)" +) ALLOWED_SIGNERS = Path("/Users/davidtai/.config/mtplx/deepseek-v4-0731-allowed-signers") +# Digest of the reviewed dedicated public signer list. +ALLOWED_SIGNERS_SHA256 = "003f258613fe308134ef184e52988a082a3376655d6b44f526017d7d71c7f843" SIGNING_IDENTITY = "mtplx-deepseek-v4-0731-candidate" SIGNING_NAMESPACE = "mtplx-deepseek-v4-0731" +ALLOWED_CANDIDATE_MODEL_IDS = frozenset({"deepseek-v4-0731-candidate"}) CANDIDATE_WORKTREE = Path("/Users/davidtai/projects/OpenSourceWTF/.worktrees/dsv4-0731-service") REVIEWED_REF = "refs/tags/mtplx-dsv4-0731-reviewed" CANDIDATE_PLIST_SHA256 = "93eac0d4eaac491c7f2f1d3a293ba38a3144ade59ee3afdf52b35cc9ec9bb101" @@ -180,8 +189,7 @@ def _read_json(path: Path) -> dict[str, Any]: def _verify_candidate_signature(receipt_bytes: bytes, signature: Path) -> None: - if not ALLOWED_SIGNERS.is_file() or ALLOWED_SIGNERS.is_symlink(): - raise PromotionError("pinned candidate allowed-signers file is missing or unsafe") + _assert_allowed_signers_trusted() if not signature.is_file() or signature.is_symlink(): raise PromotionError("detached candidate signature is missing or unsafe") result = subprocess.run( @@ -206,13 +214,30 @@ def _verify_candidate_signature(receipt_bytes: bytes, signature: Path) -> None: raise PromotionError("candidate receipt signature verification failed") +def _assert_allowed_signers_trusted() -> None: + """Trust exactly the reviewed signer list, owned by this operator.""" + if not ALLOWED_SIGNERS.is_file() or ALLOWED_SIGNERS.is_symlink(): + raise PromotionError("pinned candidate allowed-signers file is missing or unsafe") + try: + metadata = ALLOWED_SIGNERS.stat() + except OSError as error: + raise PromotionError("pinned candidate allowed-signers metadata is unavailable") from error + if metadata.st_uid != os.getuid(): + raise PromotionError("pinned candidate allowed-signers owner is unsafe") + if stat.S_IMODE(metadata.st_mode) & (stat.S_IWGRP | stat.S_IWOTH): + raise PromotionError("pinned candidate allowed-signers permissions are unsafe") + if hashlib.sha256(ALLOWED_SIGNERS.read_bytes()).hexdigest() != ALLOWED_SIGNERS_SHA256: + raise PromotionError("pinned candidate allowed-signers digest changed") + + def _contains_sensitive(value: Any) -> bool: if isinstance(value, dict): return any(SENSITIVE_KEY.search(str(key)) or _contains_sensitive(item) for key, item in value.items()) if isinstance(value, list): return any(_contains_sensitive(item) for item in value) if isinstance(value, str): - if any(marker in value for marker in ("/Users/", "/private/", "/tmp/")): + decoded = urllib.parse.unquote(value) + if PATH_LIKE_VALUE.search(decoded): return True if SENSITIVE_VALUE.search(value): return True @@ -271,9 +296,12 @@ def assert_candidate_receipt(payload: dict[str, Any]) -> None: if ( not isinstance(candidate_model_ids, list) or not candidate_model_ids - or not all(isinstance(model_id, str) and model_id for model_id in candidate_model_ids) + or not all( + isinstance(model_id, str) and model_id in ALLOWED_CANDIDATE_MODEL_IDS + for model_id in candidate_model_ids + ) ): - raise PromotionError("candidate smoke lacks nonempty candidate_model_ids") + raise PromotionError("candidate smoke has a disallowed model ID") def assert_live_identity(expected: dict[str, Any], current: dict[str, Any]) -> None: diff --git a/services/deepseek-v4-0731/tests/test_service_surface.py b/services/deepseek-v4-0731/tests/test_service_surface.py index d9ccdd94..c86bf56b 100644 --- a/services/deepseek-v4-0731/tests/test_service_surface.py +++ b/services/deepseek-v4-0731/tests/test_service_surface.py @@ -5,8 +5,10 @@ import hashlib import json import os +import stat import subprocess import sys +import tempfile from pathlib import Path from types import SimpleNamespace @@ -238,6 +240,136 @@ def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: assert "<｜DSML｜" not in json.dumps(split_deltas, ensure_ascii=False) +def test_no_tools_api_stream_split_chunks_never_release_dsml() -> None: + """The candidate's stream boundary sanitizes no-tools API content too.""" + from candidate_entry import install_candidate_surface + from mtplx.server import openai as openai_server + + class PassthroughSplitter: + def start(self): + return [] + + def feed(self, text: str): + return [("content", text)] + + def finish(self, **_kwargs): + return [] + + server = SimpleNamespace( + _encode_messages=lambda *_args, **_kwargs: [], + _parse_generated_tool_calls_or_content=lambda *_args, **_kwargs: (None, None), + omlx_extract_tool_calls_with_thinking=lambda *_args, **_kwargs: None, + _ToolAwareContentStreamTranslator=openai_server._ToolAwareContentStreamTranslator, + _stream_tool_call_deltas=openai_server._stream_tool_call_deltas, + _stream_splitter_for_state=lambda *_args, **_kwargs: PassthroughSplitter(), + ) + install_candidate_surface(server) + splitter = server._stream_splitter_for_state( + SimpleNamespace(), thinking_enabled=False, suppress_orphan_tool_markup=True + ) + wire_chunks: list[tuple[str, str]] = [] + raw = "Preamble.\n\n<｜DSML｜tool_calls>\n<｜DSML｜invoke name=\"x\">\n\n" + for character in raw: + wire_chunks.extend(splitter.feed(character)) + wire_chunks.extend(splitter.finish()) + assert "".join(text for field, text in wire_chunks if field == "content") == "Preamble." + assert "<｜DSML｜" not in json.dumps(wire_chunks, ensure_ascii=False) + + +def test_allowed_signers_is_digest_pinned_owned_and_not_writable() -> None: + import promote_cutover + + assert promote_cutover.ALLOWED_SIGNERS_SHA256 == ( + "003f258613fe308134ef184e52988a082a3376655d6b44f526017d7d71c7f843" + ) + with tempfile.TemporaryDirectory() as directory: + allowed = Path(directory) / "allowed-signers" + allowed.write_text("mtplx-deepseek-v4-0731-candidate ssh-ed25519 AAAA\n", encoding="utf-8") + allowed.chmod(0o600) + original_path = promote_cutover.ALLOWED_SIGNERS + original_digest = promote_cutover.ALLOWED_SIGNERS_SHA256 + try: + promote_cutover.ALLOWED_SIGNERS = allowed + promote_cutover.ALLOWED_SIGNERS_SHA256 = hashlib.sha256(allowed.read_bytes()).hexdigest() + promote_cutover._assert_allowed_signers_trusted() + + allowed.write_text("mutated\n", encoding="utf-8") + with pytest.raises(promote_cutover.PromotionError, match="digest"): + promote_cutover._assert_allowed_signers_trusted() + + allowed.write_text("mtplx-deepseek-v4-0731-candidate ssh-ed25519 AAAA\n", encoding="utf-8") + allowed.chmod(stat.S_IRUSR | stat.S_IWUSR | stat.S_IWGRP) + with pytest.raises(promote_cutover.PromotionError, match="permissions"): + promote_cutover._assert_allowed_signers_trusted() + + allowed.chmod(0o600) + original_getuid = promote_cutover.os.getuid + try: + promote_cutover.os.getuid = lambda: original_getuid() + 1 + with pytest.raises(promote_cutover.PromotionError, match="owner"): + promote_cutover._assert_allowed_signers_trusted() + finally: + promote_cutover.os.getuid = original_getuid + finally: + promote_cutover.ALLOWED_SIGNERS = original_path + promote_cutover.ALLOWED_SIGNERS_SHA256 = original_digest + + +@pytest.mark.parametrize( + "model_id", + [ + "/Users/davidtai/models/private", + "~/models/private", + r"C:\\Users\\davidtai\\models\\private", + "file:///Users/davidtai/models/private", + "deepseek-v4-0731/../../private", + "..%2Fprivate", + ], +) +def test_candidate_receipt_rejects_path_like_model_ids(model_id: str) -> None: + from promote_cutover import PromotionError, assert_candidate_receipt + + receipt = _passing_candidate_receipt() + receipt["candidate_smoke"]["candidate_model_ids"] = [model_id] + with pytest.raises(PromotionError, match="sensitive|model ID"): + assert_candidate_receipt(receipt) + + +@pytest.mark.parametrize("path_value", ["~/private", r"C:\\private", "file:///private", "a/../private"]) +def test_candidate_receipt_recursively_rejects_path_like_values(path_value: str) -> None: + from promote_cutover import PromotionError, assert_candidate_receipt + + receipt = _passing_candidate_receipt() + receipt["candidate_preflight"]["promotion_target"]["label"] = path_value + with pytest.raises(PromotionError, match="sensitive"): + assert_candidate_receipt(receipt) + + +def _passing_candidate_receipt() -> dict[str, object]: + return { + "schema": "mtplx.dsv4-0731-candidate.v1", + "candidate_preflight": { + "ok": True, + "label": "com.tea.deepseek-v4-0731.candidate", + "port": 8081, + "plist_sha256": "a" * 64, + "encoding_source_revision": "7872f01b1d1fe23eabc4c98b48bffcef5a386062", + "encoding_asset_set_sha256": "b" * 64, + "reviewed_commit": "c" * 40, + "model_config_sha256": "d" * 64, + "model_index_sha256": "e" * 64, + "promotion_target": {"label": "com.tea.deepseek-v4-0731.production", "plist_sha256": "b" * 64}, + }, + "candidate_smoke": { + "ok": True, + "models_ok": True, + "ready": True, + "finish_reason": "stop", + "candidate_model_ids": ["deepseek-v4-0731-candidate"], + }, + } + + @pytest.mark.parametrize("forbidden_key, forbidden_value", [ ("stdout", "must never be retained"), ("prompt", "must never be retained"), @@ -295,7 +427,7 @@ def test_scrubbed_passing_candidate_receipt_is_accepted() -> None: "models_ok": True, "ready": True, "finish_reason": "stop", - "candidate_model_ids": ["deepseek-v4-0731"], + "candidate_model_ids": ["deepseek-v4-0731-candidate"], }, } ) @@ -327,7 +459,7 @@ def test_candidate_receipt_rejects_unknown_nested_fields() -> None: "models_ok": True, "ready": True, "finish_reason": "stop", - "candidate_model_ids": ["deepseek-v4-0731"], + "candidate_model_ids": ["deepseek-v4-0731-candidate"], }, } with pytest.raises(PromotionError, match="strict receipt schema"): From 21c23874bac73fd4a7eb88c265e6f055474f0bc3 Mon Sep 17 00:00:00 2001 From: davidtai Date: Mon, 3 Aug 2026 16:55:07 -0500 Subject: [PATCH 20/24] service: preserve no-tools stream splitter metrics --- services/deepseek-v4-0731/candidate.json | 2 +- services/deepseek-v4-0731/candidate_entry.py | 4 ++++ services/deepseek-v4-0731/launch_candidate.sh | 2 +- .../tests/test_service_surface.py | 24 ++++++++++--------- 4 files changed, 19 insertions(+), 13 deletions(-) diff --git a/services/deepseek-v4-0731/candidate.json b/services/deepseek-v4-0731/candidate.json index 2fdb8a3d..4d3977d9 100644 --- a/services/deepseek-v4-0731/candidate.json +++ b/services/deepseek-v4-0731/candidate.json @@ -27,7 +27,7 @@ "trusted_python": "/Users/davidtai/projects/OpenSourceWTF/.worktrees/dsv4-0731-service/.venv/bin/python", "trusted_python_target": "/Users/davidtai/.local/share/uv/python/cpython-3.12-macos-aarch64-none/bin/python3.12", "trusted_python_sha256": "96793b100c947cdc81a38e8fb8c9c1889abccda9840ce1bef58d372bf3f2c263", - "candidate_entry_sha256": "cc16dff9d2ebcd01af55cfc9ba36ca01ee832e3db264e5af311552292d127c97", + "candidate_entry_sha256": "de4be5c6e248590f2ac9283692e7d58f39f7b2263a1ce5a037e356e2f5d2fa32", "candidate_plist_sha256": "93eac0d4eaac491c7f2f1d3a293ba38a3144ade59ee3afdf52b35cc9ec9bb101", "served_model_id": "deepseek-v4-0731-candidate" } diff --git a/services/deepseek-v4-0731/candidate_entry.py b/services/deepseek-v4-0731/candidate_entry.py index b4a7e841..52c5de0c 100755 --- a/services/deepseek-v4-0731/candidate_entry.py +++ b/services/deepseek-v4-0731/candidate_entry.py @@ -365,6 +365,10 @@ def __init__(self, stock: Any) -> None: tools=[], argument_chunk_chars=1 ) + @property + def reentry_count(self) -> int: + return int(self._stock.reentry_count) + def start(self): return self._translate(self._stock.start()) diff --git a/services/deepseek-v4-0731/launch_candidate.sh b/services/deepseek-v4-0731/launch_candidate.sh index e6e14bc3..5737b9b8 100755 --- a/services/deepseek-v4-0731/launch_candidate.sh +++ b/services/deepseek-v4-0731/launch_candidate.sh @@ -34,7 +34,7 @@ fi [ -x "$PYTHON_TARGET" ] && [ ! -L "$PYTHON_TARGET" ] || die "trusted python target is missing or unsafe" [ "$(sha256 "$PYTHON_TARGET")" = 96793b100c947cdc81a38e8fb8c9c1889abccda9840ce1bef58d372bf3f2c263 ] || die "trusted python hash changed" [ -f "$ENTRY" ] && [ ! -L "$ENTRY" ] || die "candidate entrypoint is missing or unsafe" -[ "$(sha256 "$ENTRY")" = cc16dff9d2ebcd01af55cfc9ba36ca01ee832e3db264e5af311552292d127c97 ] || die "candidate entrypoint hash changed" +[ "$(sha256 "$ENTRY")" = de4be5c6e248590f2ac9283692e7d58f39f7b2263a1ce5a037e356e2f5d2fa32 ] || die "candidate entrypoint hash changed" [ -d "$MODEL" ] && [ ! -L "$MODEL" ] || die "pinned model path is missing or unsafe" [ "$(sha256 "$MODEL/config.json")" = 6d0297a4329d55dccf3cd48fd168efea8044996245195d518a9e8aaa14906d3e ] || die "model configuration hash changed" [ "$(sha256 "$MODEL/model.safetensors.index.json")" = 9edcd0db7e6b8f0b8e02978d73c30083b2aa64c2e3a8fd77d3b776a4efb4bc91 ] || die "model index hash changed" diff --git a/services/deepseek-v4-0731/tests/test_service_surface.py b/services/deepseek-v4-0731/tests/test_service_surface.py index c86bf56b..198653b8 100644 --- a/services/deepseek-v4-0731/tests/test_service_surface.py +++ b/services/deepseek-v4-0731/tests/test_service_surface.py @@ -245,23 +245,16 @@ def test_no_tools_api_stream_split_chunks_never_release_dsml() -> None: from candidate_entry import install_candidate_surface from mtplx.server import openai as openai_server - class PassthroughSplitter: - def start(self): - return [] - - def feed(self, text: str): - return [("content", text)] - - def finish(self, **_kwargs): - return [] - server = SimpleNamespace( _encode_messages=lambda *_args, **_kwargs: [], _parse_generated_tool_calls_or_content=lambda *_args, **_kwargs: (None, None), omlx_extract_tool_calls_with_thinking=lambda *_args, **_kwargs: None, _ToolAwareContentStreamTranslator=openai_server._ToolAwareContentStreamTranslator, _stream_tool_call_deltas=openai_server._stream_tool_call_deltas, - _stream_splitter_for_state=lambda *_args, **_kwargs: PassthroughSplitter(), + _stream_splitter_for_state=lambda *_args, **kwargs: openai_server._ThinkingContentStreamSplitter( + thinking_enabled=kwargs["thinking_enabled"], + suppress_orphan_tool_markup=kwargs.get("suppress_orphan_tool_markup", False), + ), ) install_candidate_surface(server) splitter = server._stream_splitter_for_state( @@ -275,6 +268,15 @@ def finish(self, **_kwargs): assert "".join(text for field, text in wire_chunks if field == "content") == "Preamble." assert "<｜DSML｜" not in json.dumps(wire_chunks, ensure_ascii=False) + # The real endpoint performs both reads unconditionally during streaming + # finalization, including when the request declared no tools. + generated = {"stats": {}} + state = SimpleNamespace(last_metrics=[{}]) + generated["stats"]["reasoning_reentries"] = splitter.reentry_count + state.last_metrics[-1]["reasoning_reentries"] = splitter.reentry_count + assert generated["stats"]["reasoning_reentries"] == 0 + assert state.last_metrics[-1]["reasoning_reentries"] == 0 + def test_allowed_signers_is_digest_pinned_owned_and_not_writable() -> None: import promote_cutover From 2e9b7cf164ef8578a053044b0bd1c95e3d154259 Mon Sep 17 00:00:00 2001 From: davidtai Date: Mon, 3 Aug 2026 17:13:44 -0500 Subject: [PATCH 21/24] service: bind promotion and sanitize no-tools responses --- services/deepseek-v4-0731/README.md | 8 +- services/deepseek-v4-0731/candidate.json | 2 +- services/deepseek-v4-0731/candidate_entry.py | 22 ++ services/deepseek-v4-0731/launch_candidate.sh | 2 +- services/deepseek-v4-0731/promote_cutover.py | 307 +++++++++++++++--- .../tests/test_service_surface.py | 191 ++++++++++- 6 files changed, 477 insertions(+), 55 deletions(-) diff --git a/services/deepseek-v4-0731/README.md b/services/deepseek-v4-0731/README.md index 518b312f..6bb8e699 100644 --- a/services/deepseek-v4-0731/README.md +++ b/services/deepseek-v4-0731/README.md @@ -42,7 +42,13 @@ verify rollback. The same lock remains held through restoration and the real `/v1/models` plus exact `content.strip() == "READY"`/`finish_reason=stop` smoke. After the new plist is bootstrapped, its hash, launchd PID, and ownership of the 8080 listener are reattested under the lock before any HTTP identity or readiness -probe is allowed. +probe is allowed. The prior live identity is pinned to `com.tea.qwen` serving +`mtplx-qwen36-27b-optimized-quality`; the target is pinned to +`com.tea.deepseek-v4-0731.production`. Promotion parses `Label` and +`ProgramArguments` from one descriptor-read plist snapshot and requires the +loaded launchd job path, program, and arguments to match. The source inode and +bytes are rechecked before bootout, and rollback uses the held exact snapshot +rather than rereading a possibly replaced path. Receipts have an exact allowlist and recursively reject local paths, request content, tool schemas, secrets, argv/env, and captured process output. diff --git a/services/deepseek-v4-0731/candidate.json b/services/deepseek-v4-0731/candidate.json index 4d3977d9..778fd836 100644 --- a/services/deepseek-v4-0731/candidate.json +++ b/services/deepseek-v4-0731/candidate.json @@ -27,7 +27,7 @@ "trusted_python": "/Users/davidtai/projects/OpenSourceWTF/.worktrees/dsv4-0731-service/.venv/bin/python", "trusted_python_target": "/Users/davidtai/.local/share/uv/python/cpython-3.12-macos-aarch64-none/bin/python3.12", "trusted_python_sha256": "96793b100c947cdc81a38e8fb8c9c1889abccda9840ce1bef58d372bf3f2c263", - "candidate_entry_sha256": "de4be5c6e248590f2ac9283692e7d58f39f7b2263a1ce5a037e356e2f5d2fa32", + "candidate_entry_sha256": "35b268195eba1af59028f96dd5e6b474d76dcc42844e610743c48a55771d2268", "candidate_plist_sha256": "93eac0d4eaac491c7f2f1d3a293ba38a3144ade59ee3afdf52b35cc9ec9bb101", "served_model_id": "deepseek-v4-0731-candidate" } diff --git a/services/deepseek-v4-0731/candidate_entry.py b/services/deepseek-v4-0731/candidate_entry.py index 52c5de0c..6a07dadc 100755 --- a/services/deepseek-v4-0731/candidate_entry.py +++ b/services/deepseek-v4-0731/candidate_entry.py @@ -239,6 +239,26 @@ def extract( server.omlx_extract_tool_calls_with_thinking = extract +def _install_no_tools_nonstream_sanitizer(server: ModuleType, encoding: ModuleType) -> None: + """Route the candidate's no-tools JSON response through official parsing.""" + dsml_marker = f"<{encoding.dsml_token}" + + def strip_orphan_tool_markup(text: str) -> tuple[str, int]: + try: + extraction = server.omlx_extract_tool_calls_with_thinking( + "", text, None, [] + ) + except CandidateConstructionError: + marker = text.find(dsml_marker) + return (text[:marker].rstrip(), 1) if marker >= 0 else ("", 0) + cleaned = extraction.cleaned_text.strip() + if dsml_marker in cleaned: + raise CandidateConstructionError("official nonstream parser retained DSML markup") + return cleaned, int(dsml_marker in text) + + server._strip_orphan_tool_markup = strip_orphan_tool_markup + + def _install_stream_translator(server: ModuleType, encoding: ModuleType) -> None: """Keep scanning after visible prose while holding split DSML prefixes.""" @@ -465,6 +485,8 @@ def install_candidate_surface(server: ModuleType) -> dict[str, str]: server._parse_generated_tool_calls_or_content = _install_completion_parser(server, encoding) if hasattr(server, "omlx_extract_tool_calls_with_thinking"): _install_actual_tool_extractor(server, encoding) + if hasattr(server, "_strip_orphan_tool_markup"): + _install_no_tools_nonstream_sanitizer(server, encoding) if hasattr(server, "_ToolAwareContentStreamTranslator"): _install_stream_translator(server, encoding) if hasattr(server, "_stream_splitter_for_state"): diff --git a/services/deepseek-v4-0731/launch_candidate.sh b/services/deepseek-v4-0731/launch_candidate.sh index 5737b9b8..4ace40a7 100755 --- a/services/deepseek-v4-0731/launch_candidate.sh +++ b/services/deepseek-v4-0731/launch_candidate.sh @@ -34,7 +34,7 @@ fi [ -x "$PYTHON_TARGET" ] && [ ! -L "$PYTHON_TARGET" ] || die "trusted python target is missing or unsafe" [ "$(sha256 "$PYTHON_TARGET")" = 96793b100c947cdc81a38e8fb8c9c1889abccda9840ce1bef58d372bf3f2c263 ] || die "trusted python hash changed" [ -f "$ENTRY" ] && [ ! -L "$ENTRY" ] || die "candidate entrypoint is missing or unsafe" -[ "$(sha256 "$ENTRY")" = de4be5c6e248590f2ac9283692e7d58f39f7b2263a1ce5a037e356e2f5d2fa32 ] || die "candidate entrypoint hash changed" +[ "$(sha256 "$ENTRY")" = 35b268195eba1af59028f96dd5e6b474d76dcc42844e610743c48a55771d2268 ] || die "candidate entrypoint hash changed" [ -d "$MODEL" ] && [ ! -L "$MODEL" ] || die "pinned model path is missing or unsafe" [ "$(sha256 "$MODEL/config.json")" = 6d0297a4329d55dccf3cd48fd168efea8044996245195d518a9e8aaa14906d3e ] || die "model configuration hash changed" [ "$(sha256 "$MODEL/model.safetensors.index.json")" = 9edcd0db7e6b8f0b8e02978d73c30083b2aa64c2e3a8fd77d3b776a4efb4bc91 ] || die "model index hash changed" diff --git a/services/deepseek-v4-0731/promote_cutover.py b/services/deepseek-v4-0731/promote_cutover.py index be1e93bf..a70be9d4 100755 --- a/services/deepseek-v4-0731/promote_cutover.py +++ b/services/deepseek-v4-0731/promote_cutover.py @@ -19,17 +19,21 @@ import stat import subprocess import sys +import tempfile import time import urllib.error import urllib.parse import urllib.request from contextlib import contextmanager +from dataclasses import dataclass from pathlib import Path from typing import Any, Iterator LOCK_PATH = Path("/tmp/mtplx-gpu-exclusive.lock") CANDIDATE_LABEL = "com.tea.deepseek-v4-0731.candidate" +PRIOR_LIVE_LABEL = "com.tea.qwen" +PRODUCTION_LABEL = "com.tea.deepseek-v4-0731.production" CANDIDATE_PORT = 8081 LIVE_PORT = 8080 SENSITIVE_KEY = re.compile(r"(?:prompt|message|tool|secret|token|authorization|argv|env|stdout|stderr)", re.I) @@ -48,6 +52,8 @@ SIGNING_IDENTITY = "mtplx-deepseek-v4-0731-candidate" SIGNING_NAMESPACE = "mtplx-deepseek-v4-0731" ALLOWED_CANDIDATE_MODEL_IDS = frozenset({"deepseek-v4-0731-candidate"}) +ALLOWED_PRIOR_MODEL_IDS = ("mtplx-qwen36-27b-optimized-quality",) +ALLOWED_LAUNCHD_LABELS = frozenset({PRIOR_LIVE_LABEL, PRODUCTION_LABEL}) CANDIDATE_WORKTREE = Path("/Users/davidtai/projects/OpenSourceWTF/.worktrees/dsv4-0731-service") REVIEWED_REF = "refs/tags/mtplx-dsv4-0731-reviewed" CANDIDATE_PLIST_SHA256 = "93eac0d4eaac491c7f2f1d3a293ba38a3144ade59ee3afdf52b35cc9ec9bb101" @@ -60,10 +66,104 @@ class PromotionError(RuntimeError): pass -def _sha256(path: Path) -> str: - if not path.is_file() or path.is_symlink(): - raise PromotionError("attested plist is missing or unsafe") - return hashlib.sha256(path.read_bytes()).hexdigest() +@dataclass +class PlistSnapshot: + source_path: Path + path: Path + raw: bytes + sha256: str + label: str + program_arguments: tuple[str, ...] + source_device: int + source_inode: int + + def assert_source_unchanged(self) -> None: + raw, metadata = _read_regular_file(self.source_path, "source plist") + if ( + metadata.st_dev != self.source_device + or metadata.st_ino != self.source_inode + or hashlib.sha256(raw).hexdigest() != self.sha256 + ): + raise PromotionError("source plist changed since snapshot") + + def assert_snapshot_intact(self) -> None: + raw, _metadata = _read_regular_file(self.path, "rollback plist snapshot") + if raw != self.raw or hashlib.sha256(raw).hexdigest() != self.sha256: + raise PromotionError("rollback plist snapshot changed") + + +def _read_regular_file(path: Path, context: str) -> tuple[bytes, os.stat_result]: + flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) + try: + fd = os.open(path, flags) + except OSError as error: + raise PromotionError(f"{context} is missing or unsafe") from error + try: + metadata = os.fstat(fd) + if not stat.S_ISREG(metadata.st_mode): + raise PromotionError(f"{context} is not a regular file") + chunks: list[bytes] = [] + while chunk := os.read(fd, 1024 * 1024): + chunks.append(chunk) + return b"".join(chunks), metadata + finally: + os.close(fd) + + +def _parse_plist_identity(raw: bytes) -> tuple[str, tuple[str, ...]]: + try: + payload = plistlib.loads(raw) + except plistlib.InvalidFileException as error: + raise PromotionError("attested plist is not valid") from error + if not isinstance(payload, dict): + raise PromotionError("attested plist root is not a dictionary") + label = payload.get("Label") + arguments = payload.get("ProgramArguments") + if label not in ALLOWED_LAUNCHD_LABELS: + raise PromotionError("attested plist Label is not allowlisted") + if ( + not isinstance(arguments, list) + or not arguments + or not all(isinstance(argument, str) and argument for argument in arguments) + ): + raise PromotionError("attested plist ProgramArguments are invalid") + return str(label), tuple(arguments) + + +@contextmanager +def plist_snapshot(source: Path) -> Iterator[PlistSnapshot]: + """Hold exact descriptor-read plist bytes for identity and rollback.""" + source = Path(os.path.abspath(source)) + raw, metadata = _read_regular_file(source, "source plist") + label, program_arguments = _parse_plist_identity(raw) + fd, name = tempfile.mkstemp(prefix="mtplx-dsv4-0731-plist-", suffix=".snapshot") + snapshot_path = Path(name) + try: + os.fchmod(fd, 0o400) + view = memoryview(raw) + while view: + written = os.write(fd, view) + view = view[written:] + os.fsync(fd) + finally: + os.close(fd) + snapshot = PlistSnapshot( + source_path=source, + path=snapshot_path, + raw=raw, + sha256=hashlib.sha256(raw).hexdigest(), + label=label, + program_arguments=program_arguments, + source_device=metadata.st_dev, + source_inode=metadata.st_ino, + ) + try: + yield snapshot + finally: + try: + snapshot_path.unlink() + except FileNotFoundError: + pass def _command(*argv: str) -> str: @@ -123,18 +223,60 @@ def _listener_pid(port: int) -> int: return pids.pop() -def _launchctl_pid(label: str) -> int: +def _launchctl_job(label: str) -> dict[str, Any]: + if label not in ALLOWED_LAUNCHD_LABELS: + raise PromotionError("launchd label is not allowlisted") domain = f"gui/{os.getuid()}/{label}" output = _command("/bin/launchctl", "print", domain) - match = re.search(r"\bpid = (\d+)", output) - if not match: + + def scalar(name: str) -> str: + match = re.search(rf"^\s*{re.escape(name)} = (.+?)\s*$", output, re.MULTILINE) + if not match: + raise PromotionError(f"launchd job has no {name}") + return match.group(1) + + arguments_match = re.search( + r"^\s*arguments = \{\s*$(.*?)^\s*\}\s*$", + output, + re.MULTILINE | re.DOTALL, + ) + if not arguments_match: + raise PromotionError("launchd job has no arguments") + arguments = tuple( + line.strip() + for line in arguments_match.group(1).splitlines() + if line.strip() + ) + pid_text = scalar("pid") + if not pid_text.isdigit(): raise PromotionError("launchd service has no single running PID") - return int(match.group(1)) + return { + "pid": int(pid_text), + "path": Path(scalar("path")), + "program": scalar("program"), + "arguments": arguments, + } -def attest_process_identity(*, label: str, plist: Path) -> dict[str, Any]: +def _attest_process_snapshot( + *, + label: str, + snapshot: PlistSnapshot, + loaded_path: Path, +) -> dict[str, Any]: """Bind one launchd label, plist, and 8080 listener before any HTTP probe.""" - launch_pid = _launchctl_pid(label) + if label != snapshot.label: + raise PromotionError("supplied label differs from the plist Label") + snapshot.assert_snapshot_intact() + job = _launchctl_job(label) + if job["path"] != loaded_path: + raise PromotionError("launchd job path differs from the supplied plist") + if ( + job["program"] != snapshot.program_arguments[0] + or job["arguments"] != snapshot.program_arguments + ): + raise PromotionError("launchd ProgramArguments differ from the supplied plist") + launch_pid = int(job["pid"]) listener_pid = _listener_pid(LIVE_PORT) if launch_pid != listener_pid: raise PromotionError("launchd PID and 8080 listener PID differ") @@ -142,13 +284,36 @@ def attest_process_identity(*, label: str, plist: Path) -> dict[str, Any]: "label": label, "pid": launch_pid, "listener_port": LIVE_PORT, - "plist_sha256": _sha256(plist), + "plist_sha256": snapshot.sha256, } -def attest_live(*, label: str, plist: Path) -> dict[str, Any]: +def attest_process_identity(*, label: str, plist: Path) -> dict[str, Any]: + with plist_snapshot(plist) as snapshot: + identity = _attest_process_snapshot( + label=label, + snapshot=snapshot, + loaded_path=snapshot.source_path, + ) + snapshot.assert_source_unchanged() + return identity + + +def attest_live( + *, + label: str, + plist: Path | PlistSnapshot, + loaded_path: Path | None = None, +) -> dict[str, Any]: """Capture exact live identity without sending a generation prompt.""" - process = attest_process_identity(label=label, plist=plist) + if isinstance(plist, PlistSnapshot): + process = _attest_process_snapshot( + label=label, + snapshot=plist, + loaded_path=loaded_path or plist.source_path, + ) + else: + process = attest_process_identity(label=label, plist=plist) models = _http_json(f"http://127.0.0.1:{LIVE_PORT}/v1/models") model_ids = [item.get("id") for item in models.get("data", []) if isinstance(item, dict)] if not model_ids or not all(isinstance(model_id, str) for model_id in model_ids): @@ -160,11 +325,22 @@ def attest_live(*, label: str, plist: Path) -> dict[str, Any]: } -def _wait_for_process_identity(*, label: str, plist: Path) -> dict[str, Any]: +def _wait_for_process_identity( + *, + label: str, + plist: Path | PlistSnapshot, + loaded_path: Path | None = None, +) -> dict[str, Any]: """Wait for launchd and the listener to converge without probing HTTP.""" deadline = time.monotonic() + 600 while time.monotonic() < deadline: try: + if isinstance(plist, PlistSnapshot): + return _attest_process_snapshot( + label=label, + snapshot=plist, + loaded_path=loaded_path or plist.source_path, + ) return attest_process_identity(label=label, plist=plist) except PromotionError: time.sleep(0.5) @@ -287,8 +463,8 @@ def assert_candidate_receipt(payload: dict[str, Any]) -> None: if not isinstance(preflight.get("reviewed_commit"), str) or not re.fullmatch(r"[0-9a-f]{40}", preflight["reviewed_commit"]): raise PromotionError("candidate preflight has invalid reviewed_commit") target = _require_exact_keys(preflight["promotion_target"], {"label", "plist_sha256"}, "promotion target") - if not isinstance(target.get("label"), str): - raise PromotionError("candidate preflight lacks a separately reviewed promotion target") + if target.get("label") != PRODUCTION_LABEL: + raise PromotionError("candidate preflight has a disallowed production label") digest = target.get("plist_sha256") if not isinstance(digest, str) or not re.fullmatch(r"[0-9a-f]{64}", digest): raise PromotionError("candidate preflight lacks a valid promotion plist digest") @@ -311,6 +487,24 @@ def assert_live_identity(expected: dict[str, Any], current: dict[str, Any]) -> N "live attestation", ) fields = ("schema", "label", "pid", "listener_port", "plist_sha256", "model_ids") + if expected.get("schema") != "mtplx.live-identity.v1": + raise PromotionError("live attestation schema is not allowlisted") + if expected.get("label") != PRIOR_LIVE_LABEL: + raise PromotionError("live attestation label is not allowlisted") + if expected.get("model_ids") != list(ALLOWED_PRIOR_MODEL_IDS): + raise PromotionError("live attestation model IDs are not allowlisted") + if ( + not isinstance(expected.get("pid"), int) + or isinstance(expected.get("pid"), bool) + or expected["pid"] <= 0 + ): + raise PromotionError("live attestation has invalid pid") + if expected.get("listener_port") != LIVE_PORT: + raise PromotionError("live attestation has invalid listener port") + if not isinstance(expected.get("plist_sha256"), str) or not re.fullmatch( + r"[0-9a-f]{64}", expected["plist_sha256"] + ): + raise PromotionError("live attestation has invalid plist digest") if any(expected.get(field) != current.get(field) for field in fields): raise PromotionError("live service identity changed since its attestation") @@ -355,6 +549,8 @@ def _verify_live_ready(expected_model_ids: list[str]) -> None: def promote(args: argparse.Namespace) -> None: if args.promote is not True: raise PromotionError("refusing promotion without --promote") + if args.production_label != PRODUCTION_LABEL: + raise PromotionError("production label is not allowlisted") candidate, candidate_bytes = _read_json_bytes(args.candidate_receipt) _verify_candidate_signature(candidate_bytes, args.candidate_signature) expected_live = _read_json(args.live_attestation) @@ -386,40 +582,59 @@ def promote(args: argparse.Namespace) -> None: raise PromotionError("an absolute separately reviewed production plist is required") if not target.is_file() or target.is_symlink(): raise PromotionError("production plist is missing or unsafe") - promotion_target = preflight["promotion_target"] - if promotion_target["label"] != args.production_label or promotion_target["plist_sha256"] != _sha256(target): - raise PromotionError("production plist identity does not match the passing candidate preflight") - try: - target_label = plistlib.loads(target.read_bytes()).get("Label") - except (plistlib.InvalidFileException, OSError) as error: - raise PromotionError("production plist is not valid") from error - if target_label != args.production_label or args.production_label == str(expected_live.get("label")): - raise PromotionError("production label is unsafe or does not match its plist") - prior_plist = args.live_plist if not prior_plist.is_absolute(): raise PromotionError("live attestation does not name an absolute prior plist") - with exclusive_gpu_lock(): - current = attest_live(label=str(expected_live.get("label", "")), plist=prior_plist) - assert_live_identity(expected_live, current) - # No service is stopped until every receipt and identity check above has - # passed under the lock. Any post-cutover exception restores the exact - # attested plist before releasing that same lock. - try: - _bootout(current["label"]) - _bootstrap(target) - promoted = _wait_for_process_identity(label=args.production_label, plist=target) - if promoted["plist_sha256"] != promotion_target["plist_sha256"]: - raise PromotionError("promoted service plist identity changed during cutover") - _verify_live_ready(candidate["candidate_smoke"]["candidate_model_ids"]) - except BaseException: + with plist_snapshot(target) as target_snapshot, plist_snapshot(prior_plist) as prior_snapshot: + promotion_target = preflight["promotion_target"] + if ( + promotion_target["label"] != args.production_label + or promotion_target["plist_sha256"] != target_snapshot.sha256 + ): + raise PromotionError("production plist identity does not match the passing candidate preflight") + if ( + target_snapshot.label != args.production_label + or args.production_label == str(expected_live.get("label")) + ): + raise PromotionError("production label is unsafe or does not match its plist") + + with exclusive_gpu_lock(): + current = attest_live( + label=str(expected_live.get("label", "")), + plist=prior_snapshot, + loaded_path=prior_snapshot.source_path, + ) + assert_live_identity(expected_live, current) + prior_snapshot.assert_source_unchanged() + target_snapshot.assert_source_unchanged() + # No service is stopped until every receipt and identity check above + # has passed under the lock. Any post-cutover exception restores + # the exact descriptor-read prior snapshot under that same lock. try: - _bootout(args.production_label) - finally: - _bootstrap(prior_plist) - _wait_for_process_identity(label=current["label"], plist=prior_plist) - _verify_live_ready(current["model_ids"]) - raise + _bootout(current["label"]) + target_snapshot.assert_source_unchanged() + _bootstrap(target_snapshot.source_path) + promoted = _wait_for_process_identity( + label=args.production_label, + plist=target_snapshot, + loaded_path=target_snapshot.source_path, + ) + if promoted["plist_sha256"] != promotion_target["plist_sha256"]: + raise PromotionError("promoted service plist identity changed during cutover") + _verify_live_ready(candidate["candidate_smoke"]["candidate_model_ids"]) + except BaseException: + try: + _bootout(args.production_label) + finally: + prior_snapshot.assert_snapshot_intact() + _bootstrap(prior_snapshot.path) + _wait_for_process_identity( + label=current["label"], + plist=prior_snapshot, + loaded_path=prior_snapshot.path, + ) + _verify_live_ready(current["model_ids"]) + raise def main(argv: list[str] | None = None) -> int: diff --git a/services/deepseek-v4-0731/tests/test_service_surface.py b/services/deepseek-v4-0731/tests/test_service_surface.py index 198653b8..fd99b9bc 100644 --- a/services/deepseek-v4-0731/tests/test_service_surface.py +++ b/services/deepseek-v4-0731/tests/test_service_surface.py @@ -5,6 +5,7 @@ import hashlib import json import os +import plistlib import stat import subprocess import sys @@ -105,20 +106,76 @@ def test_cutover_requires_receipts_lock_identity_and_explicit_promotion() -> Non "assert_candidate_receipt", "assert_live_identity", "/v1/models", - "finish_reason", - "SENSITIVE_KEY", - "finally:", - "_bootstrap(prior_plist)", + "finish_reason", + "SENSITIVE_KEY", + "finally:", + "_bootstrap(prior_snapshot.path)", "candidate_model_ids", "attest_process_identity", ): assert required in source - bootstrap = source.index("_bootstrap(target)") + bootstrap = source.index("_bootstrap(target_snapshot.source_path)") new_identity = source.index("_wait_for_process_identity(", bootstrap) readiness = source.index("_verify_live_ready(", new_identity) assert bootstrap < new_identity < readiness assert 'content.strip() != "READY"' in source - assert "return attest_process_identity(label=label, plist=plist)" in source + assert "_bootstrap(prior_snapshot.path)" in source + unchanged = source.index("prior_snapshot.assert_source_unchanged()") + bootout = source.index('_bootout(current["label"])') + assert unchanged < bootout + + +def test_process_attestation_rejects_plist_not_loaded_by_launchd(monkeypatch, tmp_path) -> None: + import promote_cutover + + plist = tmp_path / "prior.plist" + plist.write_bytes( + plistlib.dumps( + { + "Label": "com.tea.qwen", + "ProgramArguments": ["/bin/false", "--reviewed"], + } + ) + ) + + def fake_command(*argv: str) -> str: + if argv[0] == "/bin/launchctl": + return f"""gui/501/com.tea.qwen = {{ +\tpath = {plist} +\tprogram = /bin/true +\targuments = {{ +\t\t/bin/true +\t\t--unrelated +\t}} +\tpid = 4242 +}}""" + return "p4242\n" + + monkeypatch.setattr(promote_cutover, "_command", fake_command) + with pytest.raises(promote_cutover.PromotionError, match="ProgramArguments"): + promote_cutover.attest_process_identity(label="com.tea.qwen", plist=plist) + + +def test_prior_plist_snapshot_detects_replacement_and_preserves_rollback_bytes(tmp_path) -> None: + from promote_cutover import PromotionError, plist_snapshot + + prior = tmp_path / "prior.plist" + original = plistlib.dumps( + {"Label": "com.tea.qwen", "ProgramArguments": ["/bin/false"]} + ) + prior.write_bytes(original) + with plist_snapshot(prior) as snapshot: + replacement = tmp_path / "attacker.plist" + replacement.write_bytes( + plistlib.dumps( + {"Label": "com.tea.qwen", "ProgramArguments": ["/bin/true"]} + ) + ) + os.replace(replacement, prior) + with pytest.raises(PromotionError, match="changed since snapshot"): + snapshot.assert_source_unchanged() + snapshot.assert_snapshot_intact() + assert snapshot.path.read_bytes() == original def test_launcher_invokes_exact_reviewed_artifact_validator() -> None: @@ -278,6 +335,79 @@ def test_no_tools_api_stream_split_chunks_never_release_dsml() -> None: assert state.last_metrics[-1]["reasoning_reentries"] == 0 +@pytest.mark.parametrize("malformed", [False, True]) +def test_no_tools_nonstream_endpoint_uses_official_dsml_sanitizer(monkeypatch, malformed: bool) -> None: + from fastapi.testclient import TestClient + from candidate_entry import install_candidate_surface + from mtplx.server import openai as openai_server + + tests_root = ROOT.parents[1] / "tests" + monkeypatch.syspath_prepend(str(tests_root)) + from test_server_openai import _fake_generation, _fake_state # noqa: PLC0415 + + class Tokenizer: + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: + assert add_special_tokens is False + return [ord(character) for character in text] + + def decode(self, tokens, **_kwargs) -> str: + return "".join(chr(int(token)) for token in tokens) + + state = _fake_state() + state.args.stats_footer = False + state.runtime.tokenizer = Tokenizer() + vector = (ROOT / "encoding/tests/test_output_1.txt").read_text(encoding="utf-8") + marker = "<｜Assistant｜>" + start = vector.find(marker) + len(marker) + end = vector.find("<｜User｜>", start) + _thinking, dsml = vector[start:end].split("", 1) + raw = "Preamble.\n\n" + ("<｜DSML｜tool_calls><｜DSML｜invoke" if malformed else dsml) + + replaced = ( + "_encode_messages", + "_parse_generated_tool_calls_or_content", + "omlx_extract_tool_calls_with_thinking", + "_ToolAwareContentStreamTranslator", + "_stream_splitter_for_state", + "_strip_orphan_tool_markup", + "_normalize_reasoning_effort", + "_reasoning_effort_for_state", + "_apply_chat_template_profile", + "_template_hash", + "_template_supports_scoped_reasoning", + "_DSV4_0731_ENCODER_INSTALLED", + ) + missing = object() + originals = {name: getattr(openai_server, name, missing) for name in replaced} + try: + install_candidate_surface(openai_server) + monkeypatch.setattr( + openai_server, + "_run_generation", + lambda *_args, **_kwargs: _fake_generation(raw), + ) + response = TestClient(openai_server.create_app(state)).post( + "/v1/chat/completions", + headers={"x-mtplx-cache-mode": "bypass"}, + json={ + "messages": [{"role": "user", "content": "Do not use tools."}], + "enable_thinking": False, + "max_tokens": 64, + }, + ) + finally: + for name, original in originals.items(): + if original is missing: + delattr(openai_server, name) + else: + setattr(openai_server, name, original) + + assert response.status_code == 200 + content = response.json()["choices"][0]["message"]["content"] + assert content == "Preamble." + assert "<｜DSML｜" not in response.text + + def test_allowed_signers_is_digest_pinned_owned_and_not_writable() -> None: import promote_cutover @@ -347,6 +477,55 @@ def test_candidate_receipt_recursively_rejects_path_like_values(path_value: str) assert_candidate_receipt(receipt) +@pytest.mark.parametrize( + "label", + [ + "com.tea.prod(/etc/x)", + "../com.tea.prod", + "file://com.tea.prod", + r"C:\\com.tea.prod", + "~/com.tea.prod", + "com.tea.other", + ], +) +def test_candidate_receipt_rejects_nonallowlisted_production_label(label: str) -> None: + from promote_cutover import PromotionError, assert_candidate_receipt + + receipt = _passing_candidate_receipt() + receipt["candidate_preflight"]["promotion_target"]["label"] = label + with pytest.raises(PromotionError, match="sensitive|production label"): + assert_candidate_receipt(receipt) + + +def test_live_receipt_rejects_nonallowlisted_label_and_model_id() -> None: + from promote_cutover import PromotionError, assert_live_identity + + live = { + "schema": "mtplx.live-identity.v1", + "label": "com.tea.qwen", + "pid": 42, + "listener_port": 8080, + "plist_sha256": "a" * 64, + "model_ids": ["mtplx-qwen36-27b-optimized-quality"], + } + for field, value in ( + ("label", "com.tea.qwen(/etc/x)"), + ("model_ids", ["../private-model"]), + ): + altered = {**live, field: value} + with pytest.raises(PromotionError, match="allowlist"): + assert_live_identity(altered, altered) + + for field, value in ( + ("pid", "42"), + ("listener_port", 9999), + ("plist_sha256", "com.tea.prod(/etc/x)"), + ): + altered = {**live, field: value} + with pytest.raises(PromotionError, match="invalid"): + assert_live_identity(altered, altered) + + def _passing_candidate_receipt() -> dict[str, object]: return { "schema": "mtplx.dsv4-0731-candidate.v1", From bf2118cf313b03330a016af9da57d1477734702d Mon Sep 17 00:00:00 2001 From: davidtai Date: Mon, 3 Aug 2026 17:25:15 -0500 Subject: [PATCH 22/24] service: make DeepSeek cutover snapshots durable --- services/deepseek-v4-0731/README.md | 14 +- services/deepseek-v4-0731/promote_cutover.py | 170 +++++++++++++++--- .../tests/test_service_surface.py | 115 +++++++++++- 3 files changed, 260 insertions(+), 39 deletions(-) diff --git a/services/deepseek-v4-0731/README.md b/services/deepseek-v4-0731/README.md index 6bb8e699..4211cf74 100644 --- a/services/deepseek-v4-0731/README.md +++ b/services/deepseek-v4-0731/README.md @@ -45,10 +45,16 @@ After the new plist is bootstrapped, its hash, launchd PID, and ownership of the probe is allowed. The prior live identity is pinned to `com.tea.qwen` serving `mtplx-qwen36-27b-optimized-quality`; the target is pinned to `com.tea.deepseek-v4-0731.production`. Promotion parses `Label` and -`ProgramArguments` from one descriptor-read plist snapshot and requires the -loaded launchd job path, program, and arguments to match. The source inode and -bytes are rechecked before bootout, and rollback uses the held exact snapshot -rather than rereading a possibly replaced path. +`ProgramArguments` from one descriptor read, then writes those exact bytes to a +content-addressed, owner-only (`0400` file in a `0700` directory), same-filesystem +snapshot and fsyncs it before any service action. Both promotion and rollback +bootstrap that durable snapshot path rather than the mutable source path, and +post-bootstrap attestation requires launchd's loaded path, current snapshot +bytes and metadata, program, and arguments to still match. A snapshot remains +on disk for as long as launchd references it, including after all Python context +managers exit. The unreferenced side is removed only after a successful +`launchctl print` proves that its label is absent or loaded from another path; +an ambiguous probe preserves the snapshot and fails closed. Receipts have an exact allowlist and recursively reject local paths, request content, tool schemas, secrets, argv/env, and captured process output. diff --git a/services/deepseek-v4-0731/promote_cutover.py b/services/deepseek-v4-0731/promote_cutover.py index a70be9d4..5daec184 100755 --- a/services/deepseek-v4-0731/promote_cutover.py +++ b/services/deepseek-v4-0731/promote_cutover.py @@ -54,6 +54,7 @@ ALLOWED_CANDIDATE_MODEL_IDS = frozenset({"deepseek-v4-0731-candidate"}) ALLOWED_PRIOR_MODEL_IDS = ("mtplx-qwen36-27b-optimized-quality",) ALLOWED_LAUNCHD_LABELS = frozenset({PRIOR_LIVE_LABEL, PRODUCTION_LABEL}) +SNAPSHOT_DIR_NAME = ".mtplx-dsv4-0731-snapshots" CANDIDATE_WORKTREE = Path("/Users/davidtai/projects/OpenSourceWTF/.worktrees/dsv4-0731-service") REVIEWED_REF = "refs/tags/mtplx-dsv4-0731-reviewed" CANDIDATE_PLIST_SHA256 = "93eac0d4eaac491c7f2f1d3a293ba38a3144ade59ee3afdf52b35cc9ec9bb101" @@ -87,10 +88,26 @@ def assert_source_unchanged(self) -> None: raise PromotionError("source plist changed since snapshot") def assert_snapshot_intact(self) -> None: - raw, _metadata = _read_regular_file(self.path, "rollback plist snapshot") - if raw != self.raw or hashlib.sha256(raw).hexdigest() != self.sha256: + _assert_snapshot_directory(self.path.parent, source_device=self.source_device) + raw, metadata = _read_regular_file(self.path, "durable plist snapshot") + if ( + raw != self.raw + or hashlib.sha256(raw).hexdigest() != self.sha256 + or metadata.st_dev != self.source_device + or metadata.st_uid != os.getuid() + or stat.S_IMODE(metadata.st_mode) != 0o400 + ): raise PromotionError("rollback plist snapshot changed") + def cleanup_if_unloaded(self) -> None: + """Remove this snapshot only after launchd no longer references it.""" + job = _launchctl_job_if_loaded(self.label) + if job is not None and job["path"] == self.path: + raise PromotionError("durable plist snapshot is still loaded") + self.assert_snapshot_intact() + self.path.unlink() + _fsync_directory(self.path.parent) + def _read_regular_file(path: Path, context: str) -> tuple[bytes, os.stat_result]: flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) @@ -130,40 +147,116 @@ def _parse_plist_identity(raw: bytes) -> tuple[str, tuple[str, ...]]: return str(label), tuple(arguments) -@contextmanager -def plist_snapshot(source: Path) -> Iterator[PlistSnapshot]: - """Hold exact descriptor-read plist bytes for identity and rollback.""" - source = Path(os.path.abspath(source)) - raw, metadata = _read_regular_file(source, "source plist") - label, program_arguments = _parse_plist_identity(raw) - fd, name = tempfile.mkstemp(prefix="mtplx-dsv4-0731-plist-", suffix=".snapshot") - snapshot_path = Path(name) +def _fsync_directory(path: Path) -> None: + flags = os.O_RDONLY | getattr(os, "O_DIRECTORY", 0) | getattr(os, "O_NOFOLLOW", 0) + try: + fd = os.open(path, flags) + except OSError as error: + raise PromotionError("snapshot directory is missing or unsafe") from error + try: + os.fsync(fd) + finally: + os.close(fd) + + +def _assert_snapshot_directory(path: Path, *, source_device: int) -> None: + try: + metadata = path.lstat() + except OSError as error: + raise PromotionError("snapshot directory is missing or unsafe") from error + if ( + not stat.S_ISDIR(metadata.st_mode) + or path.is_symlink() + or metadata.st_dev != source_device + or metadata.st_uid != os.getuid() + or stat.S_IMODE(metadata.st_mode) != 0o700 + ): + raise PromotionError("snapshot directory metadata is unsafe") + + +def _materialize_durable_snapshot( + *, source: Path, raw: bytes, metadata: os.stat_result, label: str, digest: str +) -> Path: + if source.parent.name == SNAPSHOT_DIR_NAME: + snapshot_dir = source.parent + else: + snapshot_dir = source.parent / SNAPSHOT_DIR_NAME + directory_created = False + try: + snapshot_dir.mkdir(mode=0o700) + directory_created = True + except FileExistsError: + pass + if directory_created: + _fsync_directory(source.parent) + _assert_snapshot_directory(snapshot_dir, source_device=metadata.st_dev) + snapshot_path = snapshot_dir / f"{label}-{digest}.plist" + if snapshot_path.exists() or snapshot_path.is_symlink(): + existing, existing_metadata = _read_regular_file( + snapshot_path, "durable plist snapshot" + ) + if ( + existing != raw + or existing_metadata.st_dev != metadata.st_dev + or existing_metadata.st_uid != os.getuid() + or stat.S_IMODE(existing_metadata.st_mode) != 0o400 + ): + raise PromotionError("existing durable plist snapshot is unsafe") + return snapshot_path + + fd, temporary_name = tempfile.mkstemp( + prefix=".mtplx-dsv4-0731-write-", suffix=".tmp", dir=snapshot_dir + ) + temporary_path = Path(temporary_name) try: os.fchmod(fd, 0o400) view = memoryview(raw) while view: written = os.write(fd, view) + if written <= 0: + raise PromotionError("durable plist snapshot write did not progress") view = view[written:] os.fsync(fd) - finally: os.close(fd) + fd = -1 + os.replace(temporary_path, snapshot_path) + _fsync_directory(snapshot_dir) + finally: + if fd >= 0: + os.close(fd) + try: + temporary_path.unlink() + except FileNotFoundError: + pass + return snapshot_path + + +@contextmanager +def plist_snapshot(source: Path) -> Iterator[PlistSnapshot]: + """Materialize exact plist bytes durably beside the reviewed source.""" + source = Path(os.path.abspath(source)) + raw, metadata = _read_regular_file(source, "source plist") + label, program_arguments = _parse_plist_identity(raw) + digest = hashlib.sha256(raw).hexdigest() + snapshot_path = _materialize_durable_snapshot( + source=source, + raw=raw, + metadata=metadata, + label=label, + digest=digest, + ) snapshot = PlistSnapshot( source_path=source, path=snapshot_path, raw=raw, - sha256=hashlib.sha256(raw).hexdigest(), + sha256=digest, label=label, program_arguments=program_arguments, source_device=metadata.st_dev, source_inode=metadata.st_ino, ) - try: - yield snapshot - finally: - try: - snapshot_path.unlink() - except FileNotFoundError: - pass + snapshot.assert_snapshot_intact() + yield snapshot def _command(*argv: str) -> str: @@ -223,12 +316,7 @@ def _listener_pid(port: int) -> int: return pids.pop() -def _launchctl_job(label: str) -> dict[str, Any]: - if label not in ALLOWED_LAUNCHD_LABELS: - raise PromotionError("launchd label is not allowlisted") - domain = f"gui/{os.getuid()}/{label}" - output = _command("/bin/launchctl", "print", domain) - +def _parse_launchctl_job(output: str) -> dict[str, Any]: def scalar(name: str) -> str: match = re.search(rf"^\s*{re.escape(name)} = (.+?)\s*$", output, re.MULTILINE) if not match: @@ -258,6 +346,32 @@ def scalar(name: str) -> str: } +def _launchctl_job(label: str) -> dict[str, Any]: + if label not in ALLOWED_LAUNCHD_LABELS: + raise PromotionError("launchd label is not allowlisted") + domain = f"gui/{os.getuid()}/{label}" + return _parse_launchctl_job(_command("/bin/launchctl", "print", domain)) + + +def _launchctl_job_if_loaded(label: str) -> dict[str, Any] | None: + """Return a loaded job, distinguishing absence from an unsafe probe failure.""" + if label not in ALLOWED_LAUNCHD_LABELS: + raise PromotionError("launchd label is not allowlisted") + domain = f"gui/{os.getuid()}/{label}" + result = subprocess.run( + ["/bin/launchctl", "print", domain], + check=False, + capture_output=True, + text=True, + ) + if result.returncode: + diagnostic = f"{result.stdout}\n{result.stderr}" + if "Could not find service" in diagnostic: + return None + raise PromotionError("could not safely determine whether launchd still references snapshot") + return _parse_launchctl_job(result.stdout) + + def _attest_process_snapshot( *, label: str, @@ -613,15 +727,16 @@ def promote(args: argparse.Namespace) -> None: try: _bootout(current["label"]) target_snapshot.assert_source_unchanged() - _bootstrap(target_snapshot.source_path) + _bootstrap(target_snapshot.path) promoted = _wait_for_process_identity( label=args.production_label, plist=target_snapshot, - loaded_path=target_snapshot.source_path, + loaded_path=target_snapshot.path, ) if promoted["plist_sha256"] != promotion_target["plist_sha256"]: raise PromotionError("promoted service plist identity changed during cutover") _verify_live_ready(candidate["candidate_smoke"]["candidate_model_ids"]) + prior_snapshot.cleanup_if_unloaded() except BaseException: try: _bootout(args.production_label) @@ -634,6 +749,7 @@ def promote(args: argparse.Namespace) -> None: loaded_path=prior_snapshot.path, ) _verify_live_ready(current["model_ids"]) + target_snapshot.cleanup_if_unloaded() raise diff --git a/services/deepseek-v4-0731/tests/test_service_surface.py b/services/deepseek-v4-0731/tests/test_service_surface.py index fd99b9bc..b902d379 100644 --- a/services/deepseek-v4-0731/tests/test_service_surface.py +++ b/services/deepseek-v4-0731/tests/test_service_surface.py @@ -106,15 +106,15 @@ def test_cutover_requires_receipts_lock_identity_and_explicit_promotion() -> Non "assert_candidate_receipt", "assert_live_identity", "/v1/models", - "finish_reason", - "SENSITIVE_KEY", - "finally:", - "_bootstrap(prior_snapshot.path)", + "finish_reason", + "SENSITIVE_KEY", + "finally:", + "_bootstrap(prior_snapshot.path)", "candidate_model_ids", "attest_process_identity", ): assert required in source - bootstrap = source.index("_bootstrap(target_snapshot.source_path)") + bootstrap = source.index("_bootstrap(target_snapshot.path)") new_identity = source.index("_wait_for_process_identity(", bootstrap) readiness = source.index("_verify_live_ready(", new_identity) assert bootstrap < new_identity < readiness @@ -156,19 +156,43 @@ def fake_command(*argv: str) -> str: promote_cutover.attest_process_identity(label="com.tea.qwen", plist=plist) -def test_prior_plist_snapshot_detects_replacement_and_preserves_rollback_bytes(tmp_path) -> None: +def test_prior_plist_snapshot_is_durable_reignorable_and_safely_cleaned( + monkeypatch, tmp_path +) -> None: + import promote_cutover from promote_cutover import PromotionError, plist_snapshot prior = tmp_path / "prior.plist" original = plistlib.dumps( - {"Label": "com.tea.qwen", "ProgramArguments": ["/bin/false"]} + { + "Label": "com.tea.qwen", + "ProgramArguments": ["/bin/false"], + "EnvironmentVariables": {"SAFE": "reviewed"}, + "KeepAlive": True, + "StandardOutPath": "/tmp/reviewed.out", + "StandardErrorPath": "/tmp/reviewed.err", + } ) prior.write_bytes(original) with plist_snapshot(prior) as snapshot: + snapshot_path = snapshot.path + assert snapshot_path.stat().st_dev == prior.stat().st_dev + assert stat.S_IMODE(snapshot_path.stat().st_mode) == 0o400 + assert snapshot_path.stat().st_uid == os.getuid() + assert not snapshot_path.is_symlink() + assert stat.S_IMODE(snapshot_path.parent.stat().st_mode) == 0o700 + snapshot.assert_source_unchanged() replacement = tmp_path / "attacker.plist" replacement.write_bytes( plistlib.dumps( - {"Label": "com.tea.qwen", "ProgramArguments": ["/bin/true"]} + { + "Label": "com.tea.qwen", + "ProgramArguments": ["/bin/false"], + "EnvironmentVariables": {"SAFE": "attacker"}, + "KeepAlive": False, + "StandardOutPath": "/tmp/attacker.out", + "StandardErrorPath": "/tmp/attacker.err", + } ) ) os.replace(replacement, prior) @@ -177,6 +201,81 @@ def test_prior_plist_snapshot_detects_replacement_and_preserves_rollback_bytes(t snapshot.assert_snapshot_intact() assert snapshot.path.read_bytes() == original + assert snapshot_path.is_file() + + def fake_command(*argv: str) -> str: + if argv[0] == "/bin/launchctl": + return f"""gui/501/com.tea.qwen = {{ +\tpath = {snapshot_path} +\tprogram = /bin/false +\targuments = {{ +\t\t/bin/false +\t}} +\tpid = 4242 +}}""" + return "p4242\n" + + monkeypatch.setattr(promote_cutover, "_command", fake_command) + repeated = promote_cutover.attest_process_identity( + label="com.tea.qwen", plist=snapshot_path + ) + assert repeated["plist_sha256"] == hashlib.sha256(original).hexdigest() + monkeypatch.setattr( + promote_cutover, + "_launchctl_job_if_loaded", + lambda _label: {"path": snapshot_path}, + ) + with pytest.raises(PromotionError, match="still loaded"): + snapshot.cleanup_if_unloaded() + monkeypatch.setattr( + promote_cutover, "_launchctl_job_if_loaded", lambda _label: None + ) + snapshot.cleanup_if_unloaded() + assert not snapshot_path.exists() + + +def test_snapshot_cleanup_distinguishes_absent_job_from_probe_failure( + monkeypatch, tmp_path +) -> None: + import promote_cutover + + plist = tmp_path / "prior.plist" + plist.write_bytes( + plistlib.dumps( + { + "Label": "com.tea.qwen", + "ProgramArguments": ["/bin/false"], + } + ) + ) + with promote_cutover.plist_snapshot(plist) as snapshot: + snapshot_path = snapshot.path + + monkeypatch.setattr( + promote_cutover.subprocess, + "run", + lambda *_args, **_kwargs: SimpleNamespace( + returncode=70, + stdout="", + stderr="launchd transport failure", + ), + ) + with pytest.raises(promote_cutover.PromotionError, match="safely determine"): + snapshot.cleanup_if_unloaded() + assert snapshot_path.is_file() + + monkeypatch.setattr( + promote_cutover.subprocess, + "run", + lambda *_args, **_kwargs: SimpleNamespace( + returncode=113, + stdout="", + stderr='Could not find service "com.tea.qwen" in domain for user gui: 501', + ), + ) + snapshot.cleanup_if_unloaded() + assert not snapshot_path.exists() + def test_launcher_invokes_exact_reviewed_artifact_validator() -> None: source = (ROOT / "launch_candidate.sh").read_text(encoding="utf-8") From 6ea59bd2e686d5c7439fda04feed2d064fd151b3 Mon Sep 17 00:00:00 2001 From: davidtai Date: Mon, 3 Aug 2026 17:35:23 -0500 Subject: [PATCH 23/24] service: commit cutover before snapshot cleanup --- services/deepseek-v4-0731/README.md | 7 +- services/deepseek-v4-0731/promote_cutover.py | 37 ++++- .../tests/test_service_surface.py | 151 +++++++++++++++++- 3 files changed, 187 insertions(+), 8 deletions(-) diff --git a/services/deepseek-v4-0731/README.md b/services/deepseek-v4-0731/README.md index 4211cf74..764aec1f 100644 --- a/services/deepseek-v4-0731/README.md +++ b/services/deepseek-v4-0731/README.md @@ -54,7 +54,12 @@ bytes and metadata, program, and arguments to still match. A snapshot remains on disk for as long as launchd references it, including after all Python context managers exit. The unreferenced side is removed only after a successful `launchctl print` proves that its label is absent or loaded from another path; -an ambiguous probe preserves the snapshot and fails closed. +an ambiguous probe preserves the snapshot and fails closed. Successful backend +readiness is the cutover commit point: prior-snapshot cleanup happens afterward, +and a cleanup failure reports a warning without stopping the verified production +service or entering rollback. Backend PID attestation is bound specifically to +the listener on `127.0.0.1:8080`; a same-port gateway on another interface is +ignored, while wildcard listeners and multiple loopback owners are rejected. Receipts have an exact allowlist and recursively reject local paths, request content, tool schemas, secrets, argv/env, and captured process output. diff --git a/services/deepseek-v4-0731/promote_cutover.py b/services/deepseek-v4-0731/promote_cutover.py index 5daec184..88f5effa 100755 --- a/services/deepseek-v4-0731/promote_cutover.py +++ b/services/deepseek-v4-0731/promote_cutover.py @@ -310,10 +310,28 @@ def _smoke_stop(model_id: str) -> None: def _listener_pid(port: int) -> int: output = _command("/usr/sbin/lsof", "-nP", f"-iTCP:{port}", "-sTCP:LISTEN", "-Fpn") - pids = {int(line[1:]) for line in output.splitlines() if line.startswith("p") and line[1:].isdigit()} - if len(pids) != 1: - raise PromotionError("listener identity is absent or ambiguous") - return pids.pop() + current_pid: int | None = None + loopback_pids: set[int] = set() + wildcard_pids: set[int] = set() + loopback_endpoint = f"127.0.0.1:{port}" + wildcard_endpoints = { + f"*:{port}", + f"0.0.0.0:{port}", + f"[::]:{port}", + f":::{port}", + } + for line in output.splitlines(): + if line.startswith("p"): + current_pid = int(line[1:]) if line[1:].isdigit() else None + elif line.startswith("n") and current_pid is not None: + endpoint = line[1:] + if endpoint == loopback_endpoint: + loopback_pids.add(current_pid) + elif endpoint in wildcard_endpoints: + wildcard_pids.add(current_pid) + if wildcard_pids or len(loopback_pids) != 1: + raise PromotionError("exact loopback listener identity is absent or ambiguous") + return loopback_pids.pop() def _parse_launchctl_job(output: str) -> dict[str, Any]: @@ -736,7 +754,6 @@ def promote(args: argparse.Namespace) -> None: if promoted["plist_sha256"] != promotion_target["plist_sha256"]: raise PromotionError("promoted service plist identity changed during cutover") _verify_live_ready(candidate["candidate_smoke"]["candidate_model_ids"]) - prior_snapshot.cleanup_if_unloaded() except BaseException: try: _bootout(args.production_label) @@ -751,6 +768,16 @@ def promote(args: argparse.Namespace) -> None: _verify_live_ready(current["model_ids"]) target_snapshot.cleanup_if_unloaded() raise + # Readiness is the cutover commit point. Snapshot reclamation is + # post-commit housekeeping and must never re-enter rollback after + # it has removed the only prior rollback path. + try: + prior_snapshot.cleanup_if_unloaded() + except Exception as error: + print( + f"promotion committed; prior snapshot cleanup was incomplete: {error}", + file=sys.stderr, + ) def main(argv: list[str] | None = None) -> int: diff --git a/services/deepseek-v4-0731/tests/test_service_surface.py b/services/deepseek-v4-0731/tests/test_service_surface.py index b902d379..dfabc6cf 100644 --- a/services/deepseek-v4-0731/tests/test_service_surface.py +++ b/services/deepseek-v4-0731/tests/test_service_surface.py @@ -10,6 +10,7 @@ import subprocess import sys import tempfile +from contextlib import nullcontext from pathlib import Path from types import SimpleNamespace @@ -149,13 +150,45 @@ def fake_command(*argv: str) -> str: \t}} \tpid = 4242 }}""" - return "p4242\n" + return "p4242\nn127.0.0.1:8080\n" monkeypatch.setattr(promote_cutover, "_command", fake_command) with pytest.raises(promote_cutover.PromotionError, match="ProgramArguments"): promote_cutover.attest_process_identity(label="com.tea.qwen", plist=plist) +def test_listener_identity_selects_exact_loopback_backend(monkeypatch) -> None: + import promote_cutover + + monkeypatch.setattr( + promote_cutover, + "_command", + lambda *_argv: ( + "p3098\nn10.8.0.2:8080\n" + "p14242\nn127.0.0.1:8080\n" + ), + ) + + assert promote_cutover._listener_pid(8080) == 14242 + + +@pytest.mark.parametrize( + "listeners", + [ + "p3098\nn*:8080\np14242\nn127.0.0.1:8080\n", + "p14242\nn127.0.0.1:8080\np14243\nn127.0.0.1:8080\n", + ], +) +def test_listener_identity_rejects_wildcard_or_duplicate_loopback_owners( + monkeypatch, listeners: str +) -> None: + import promote_cutover + + monkeypatch.setattr(promote_cutover, "_command", lambda *_argv: listeners) + with pytest.raises(promote_cutover.PromotionError, match="loopback listener"): + promote_cutover._listener_pid(8080) + + def test_prior_plist_snapshot_is_durable_reignorable_and_safely_cleaned( monkeypatch, tmp_path ) -> None: @@ -213,7 +246,7 @@ def fake_command(*argv: str) -> str: \t}} \tpid = 4242 }}""" - return "p4242\n" + return "p4242\nn127.0.0.1:8080\n" monkeypatch.setattr(promote_cutover, "_command", fake_command) repeated = promote_cutover.attest_process_identity( @@ -234,6 +267,120 @@ def fake_command(*argv: str) -> str: assert not snapshot_path.exists() +def test_post_commit_snapshot_cleanup_failure_does_not_roll_back_production( + monkeypatch, tmp_path, capsys +) -> None: + import promote_cutover + + prior_plist = tmp_path / "prior.plist" + prior_plist.write_bytes( + plistlib.dumps( + { + "Label": promote_cutover.PRIOR_LIVE_LABEL, + "ProgramArguments": ["/bin/false", "--prior"], + } + ) + ) + target_plist = tmp_path / "production.plist" + target_plist.write_bytes( + plistlib.dumps( + { + "Label": promote_cutover.PRODUCTION_LABEL, + "ProgramArguments": ["/bin/false", "--production"], + } + ) + ) + reviewed_commit = "c" * 40 + candidate = _passing_candidate_receipt() + preflight = candidate["candidate_preflight"] + preflight.update( + { + "plist_sha256": promote_cutover.CANDIDATE_PLIST_SHA256, + "encoding_asset_set_sha256": promote_cutover.ENCODING_ASSET_SET_SHA256, + "reviewed_commit": reviewed_commit, + "model_config_sha256": promote_cutover.MODEL_CONFIG_SHA256, + "model_index_sha256": promote_cutover.MODEL_INDEX_SHA256, + "promotion_target": { + "label": promote_cutover.PRODUCTION_LABEL, + "plist_sha256": hashlib.sha256(target_plist.read_bytes()).hexdigest(), + }, + } + ) + candidate_receipt = tmp_path / "candidate.json" + candidate_receipt.write_text(json.dumps(candidate), encoding="utf-8") + live = { + "schema": "mtplx.live-identity.v1", + "label": promote_cutover.PRIOR_LIVE_LABEL, + "pid": 14242, + "listener_port": 8080, + "plist_sha256": hashlib.sha256(prior_plist.read_bytes()).hexdigest(), + "model_ids": list(promote_cutover.ALLOWED_PRIOR_MODEL_IDS), + } + live_attestation = tmp_path / "live.json" + live_attestation.write_text(json.dumps(live), encoding="utf-8") + events: list[tuple[str, object]] = [] + + monkeypatch.setattr(promote_cutover, "_verify_candidate_signature", lambda *_args: None) + monkeypatch.setattr(promote_cutover, "_command", lambda *_args: reviewed_commit) + monkeypatch.setattr(promote_cutover, "exclusive_gpu_lock", nullcontext) + monkeypatch.setattr(promote_cutover, "attest_live", lambda **_kwargs: dict(live)) + monkeypatch.setattr( + promote_cutover, + "_bootout", + lambda label: events.append(("bootout", label)), + ) + monkeypatch.setattr( + promote_cutover, + "_bootstrap", + lambda path: events.append(("bootstrap", Path(path))), + ) + monkeypatch.setattr( + promote_cutover, + "_wait_for_process_identity", + lambda **kwargs: { + "plist_sha256": kwargs["plist"].sha256, + }, + ) + monkeypatch.setattr( + promote_cutover, + "_verify_live_ready", + lambda model_ids: events.append(("ready", tuple(model_ids))), + ) + + def fail_after_unlink(snapshot) -> None: + events.append(("cleanup", snapshot.label)) + if snapshot.label == promote_cutover.PRIOR_LIVE_LABEL: + snapshot.path.unlink() + raise OSError("injected failure immediately after unlink") + + monkeypatch.setattr(promote_cutover.PlistSnapshot, "cleanup_if_unloaded", fail_after_unlink) + args = SimpleNamespace( + promote=True, + candidate_receipt=candidate_receipt, + candidate_signature=tmp_path / "candidate.json.sig", + live_attestation=live_attestation, + live_plist=prior_plist, + production_plist=target_plist, + production_label=promote_cutover.PRODUCTION_LABEL, + ) + + promote_cutover.promote(args) + + assert ("ready", tuple(candidate["candidate_smoke"]["candidate_model_ids"])) in events + assert ("cleanup", promote_cutover.PRIOR_LIVE_LABEL) in events + assert ("bootout", promote_cutover.PRODUCTION_LABEL) not in events + assert any( + event == "bootstrap" + and path.name.startswith(promote_cutover.PRODUCTION_LABEL) + for event, path in events + ) + assert not any( + event == "bootstrap" and path.name.startswith(promote_cutover.PRIOR_LIVE_LABEL) + for event, path in events + ) + assert "promotion committed; prior snapshot cleanup was incomplete" in capsys.readouterr().err + + def test_snapshot_cleanup_distinguishes_absent_job_from_probe_failure( monkeypatch, tmp_path ) -> None: From b871c368514439682f6b2b0dce6a8756ffbbe072 Mon Sep 17 00:00:00 2001 From: davidtai Date: Thu, 13 Aug 2026 00:31:41 -0500 Subject: [PATCH 24/24] feat: deploy optimized DeepSeek V4 0731 K3 --- mtplx/cli.py | 10 +++ mtplx/commands/public.py | 45 +++++++--- mtplx/deepseek_v4_dspark_generation.py | 20 ++++- mtplx/runtime.py | 18 +++- mtplx/server/openai.py | 36 ++++++-- services/deepseek-v4-0731/README.md | 45 +++++++++- .../com.tea.deepseek-v4.plist | 22 +++++ .../deepseek-v4-0731/launch_production.sh | 65 ++++++++++++++ services/deepseek-v4-0731/production_entry.py | 84 +++++++++++++++++++ .../tests/test_service_surface.py | 24 ++++++ tests/test_deepseek_v4_dspark_generation.py | 25 +++++- tests/test_public_cli.py | 55 ++++++++++++ tests/test_runtime_deepseek_v4_dspark.py | 7 +- tests/test_server_openai.py | 78 +++++++++++++++++ 14 files changed, 506 insertions(+), 28 deletions(-) create mode 100644 services/deepseek-v4-0731/com.tea.deepseek-v4.plist create mode 100755 services/deepseek-v4-0731/launch_production.sh create mode 100755 services/deepseek-v4-0731/production_entry.py diff --git a/mtplx/cli.py b/mtplx/cli.py index e75d1466..aca3e54d 100644 --- a/mtplx/cli.py +++ b/mtplx/cli.py @@ -661,6 +661,16 @@ def _add_mtp_toggle_args(parser: argparse.ArgumentParser) -> None: "greedy AR. Requires explicit --depth 2 and MTP." ), ) + parser.add_argument( + "--deepseek-v4-0731-optimized", + action="store_true", + help=( + "Select the construction-bound DeepSeek-V4-Flash-0731 optimized " + "DSpark stack for an explicit depth from 1 through 3. K2 uses the " + "measured physical-M3 target route; K1 and K3 use the native target " + "shape route. This lane is not token-exact against serial greedy AR." + ), + ) SCHEDULER_MODE_CHOICES = ( diff --git a/mtplx/commands/public.py b/mtplx/commands/public.py index 3876d890..4abedf63 100644 --- a/mtplx/commands/public.py +++ b/mtplx/commands/public.py @@ -657,24 +657,38 @@ def _validate_public_depth(args: Any, *, printer=print) -> int | None: return None -def _deepseek_v4_0731_k2_entrypoint_error(args: Any) -> dict[str, str] | None: - """Reject an explicit K2 selection before any public model construction.""" +def _deepseek_v4_0731_entrypoint_error(args: Any) -> dict[str, str] | None: + """Reject an invalid optimized selection before public model construction.""" - if not bool(getattr(args, "deepseek_v4_0731_k2", False)): + legacy_k2 = bool(getattr(args, "deepseek_v4_0731_k2", False)) + optimized = bool(getattr(args, "deepseek_v4_0731_optimized", False)) + if not (legacy_k2 or optimized): return None + if legacy_k2 and optimized: + return { + "error": "DeepSeek-V4-0731 optimized selection accepts one entrypoint flag", + "detail": "Choose either the legacy K2 flag or the depth-selectable flag.", + } cli_flags = getattr(args, "_cli_flags", set()) or set() - if "depth" not in cli_flags or int(getattr(args, "depth", 3)) != 2: + depth = int(getattr(args, "depth", 3)) + if legacy_k2 and ("depth" not in cli_flags or depth != 2): return { "error": "DeepSeek-V4-0731 K2 requires explicit --depth 2", "detail": "The selected construction owns exactly two future drafts.", } + if optimized and ("depth" not in cli_flags or depth not in {1, 2, 3}): + return { + "error": "DeepSeek-V4-0731 optimized requires explicit --depth 1, 2, or 3", + "detail": "The selected construction binds one fixed depth before loading.", + } if ( _generation_mode_from_args(args) != GENERATION_MODE_MTP or getattr(args, "load_mtp", True) is False or bool(getattr(args, "no_mtp", False)) ): + label = "K2" if legacy_k2 else "optimized" return { - "error": "DeepSeek-V4-0731 K2 requires MTP generation", + "error": f"DeepSeek-V4-0731 {label} requires MTP generation", "detail": "Remove target-only AR or --no-load-mtp/--no-mtp options.", } return None @@ -684,8 +698,12 @@ def _runtime_load_kwargs(args: Any) -> dict[str, Any]: kwargs: dict[str, Any] = { "mtp": getattr(args, "load_mtp", True) is not False, } - if bool(getattr(args, "deepseek_v4_0731_k2", False)): + legacy_k2 = bool(getattr(args, "deepseek_v4_0731_k2", False)) + optimized = bool(getattr(args, "deepseek_v4_0731_optimized", False)) + if legacy_k2 or optimized: kwargs["deepseek_v4_0731_k2"] = True + if optimized: + kwargs["deepseek_v4_0731_depth"] = int(args.depth) return kwargs @@ -8352,7 +8370,7 @@ def cmd_serve_public(args: Any) -> int: if depth_error is not None: return depth_error generation_mode = _generation_mode_from_args(args) - k2_error = _deepseek_v4_0731_k2_entrypoint_error(args) + k2_error = _deepseek_v4_0731_entrypoint_error(args) if k2_error is not None: _print_command_error( k2_error, @@ -8650,7 +8668,11 @@ def cmd_serve_public(args: Any) -> int: if draft_sampler_override is not None: draft_sampler = draft_sampler_override cli_flags = getattr(args, "_cli_flags", set()) or set() - if dspark_request_defaults and ("depth" not in cli_flags or int(args.depth) != 2): + if ( + dspark_request_defaults + and not bool(getattr(args, "deepseek_v4_0731_optimized", False)) + and ("depth" not in cli_flags or int(args.depth) != 2) + ): _print_command_error( { "error": "DeepSeek-V4-0731 DSpark requires explicit --depth 2", @@ -8743,6 +8765,8 @@ def cmd_serve_public(args: Any) -> int: ] if bool(getattr(args, "deepseek_v4_0731_k2", False)): cmd.append("--deepseek-v4-0731-k2") + if bool(getattr(args, "deepseek_v4_0731_optimized", False)): + cmd.append("--deepseek-v4-0731-optimized") for attr, flag in ( ("max_active_requests", "--max-active-requests"), ("decode_batch_max", "--decode-batch-max"), @@ -9360,7 +9384,7 @@ def _generate_one_shot_public( profile = get_profile(getattr(args, "profile", None) or DEFAULT_PROFILE_NAME) apply_profile_env(profile.name) generation_mode = _generation_mode_from_args(args) - k2_error = _deepseek_v4_0731_k2_entrypoint_error(args) + k2_error = _deepseek_v4_0731_entrypoint_error(args) if k2_error is not None: return 2, k2_error, [] draft_lm_head = ( @@ -11749,6 +11773,7 @@ def _with_server_policy_args(target: Any, source: Any) -> Any: ("api_key_file", None), ("api_key_source", "none"), ("deepseek_v4_0731_k2", False), + ("deepseek_v4_0731_optimized", False), ("default_presence_penalty", 0.0), ("default_frequency_penalty", 0.0), ("paged_kv_quantization", "off"), @@ -12408,7 +12433,7 @@ def _quickstart_run_terminal_chat_body( profile = get_profile(getattr(args, "profile", None) or DEFAULT_PROFILE_NAME) apply_profile_env(profile.name) generation_mode = _generation_mode_from_args(args) - k2_error = _deepseek_v4_0731_k2_entrypoint_error(args) + k2_error = _deepseek_v4_0731_entrypoint_error(args) if k2_error is not None: _print_command_error( k2_error, diff --git a/mtplx/deepseek_v4_dspark_generation.py b/mtplx/deepseek_v4_dspark_generation.py index 1fd0aed2..71cb569a 100644 --- a/mtplx/deepseek_v4_dspark_generation.py +++ b/mtplx/deepseek_v4_dspark_generation.py @@ -26,7 +26,12 @@ class DeepseekV4DSparkBackend: minimum_proposal_target_position: int = 2 @classmethod - def bind(cls, model: Any) -> "DeepseekV4DSparkBackend": + def bind( + cls, + model: Any, + *, + supported_depths: tuple[int, ...] = (2,), + ) -> "DeepseekV4DSparkBackend": """Validate ownership once and bind direct hot-path callables.""" dspark = getattr(model, "_dspark", None) inner = getattr(model, "model", None) @@ -39,13 +44,24 @@ def bind(cls, model: Any) -> "DeepseekV4DSparkBackend": or not callable(model) ): raise ValueError("DeepSeek-V4 DSpark backend cannot bind model ownership") - if len(tuple(getattr(dspark, "stages", ()))) != 3: + stage_count = len(tuple(getattr(dspark, "stages", ()))) + if stage_count != 3: raise ValueError("DeepSeek-V4 DSpark backend requires three owned stages") + selected_depths = tuple(int(depth) for depth in supported_depths) + if ( + not selected_depths + or len(set(selected_depths)) != len(selected_depths) + or any(depth < 1 or depth > stage_count for depth in selected_depths) + ): + raise ValueError( + "DeepSeek-V4 DSpark construction depths must be unique values from 1 to 3" + ) return cls( dspark=dspark, embed_tokens=embed_tokens, lm_head=lm_head, target_forward=model, + supported_depths=selected_depths, ) def make_cache(self, rt: Any) -> Any: diff --git a/mtplx/runtime.py b/mtplx/runtime.py index f8fe3403..4c8356f6 100644 --- a/mtplx/runtime.py +++ b/mtplx/runtime.py @@ -651,6 +651,7 @@ def load( proj_quant: str | None = None, proj_requant: str | None = None, deepseek_v4_0731_k2: bool = False, + deepseek_v4_0731_depth: int | None = None, ) -> MTPLXRuntime: """Load an MLX model and optionally inject native MTP support. @@ -662,6 +663,11 @@ def load( """ if deepseek_v4_0731_k2 and not mtp: raise ValueError("DeepSeek-V4-0731 K2 construction requires mtp=True") + selected_0731_depth = ( + 2 if deepseek_v4_0731_depth is None else int(deepseek_v4_0731_depth) + ) + if deepseek_v4_0731_k2 and selected_0731_depth not in {1, 2, 3}: + raise ValueError("DeepSeek-V4-0731 optimized depth must be 1, 2, or 3") path = Path(model_path) k2_config = None if deepseek_v4_0731_k2: @@ -1091,7 +1097,10 @@ def load( k2_prepared = (target_prepared, ffn_prepared) target_prepared.publish() ffn_prepared.publish() - block_speculative_backend = DeepseekV4DSparkBackend.bind(model) + block_speculative_backend = DeepseekV4DSparkBackend.bind( + model, + supported_depths=(selected_0731_depth,), + ) except Exception as failure: rollback_failures = _rollback_deepseek_v4_0731_k2( k2_prepared, @@ -1104,7 +1113,12 @@ def load( ) from failure raise deepseek_v4_0731_k2_receipt = { - "target_protocol": "primary_plus_two_drafts_physical_m3", + "target_protocol": ( + "primary_plus_two_drafts_physical_m3" + if selected_0731_depth == 2 + else f"primary_plus_{selected_0731_depth}_drafts_native_m{selected_0731_depth + 1}" + ), + "selected_depth": selected_0731_depth, "exact_vs_serial_greedy": False, "target": target_prepared.receipt, "dspark_ffn": ffn_prepared.receipt, diff --git a/mtplx/server/openai.py b/mtplx/server/openai.py index fed6cf42..f1e33b3a 100644 --- a/mtplx/server/openai.py +++ b/mtplx/server/openai.py @@ -1763,20 +1763,33 @@ def _validate_mtp_batch_settings(args: argparse.Namespace) -> None: _require_mlx_lm_arrays_cache_fix() -def _validate_deepseek_v4_0731_k2_entrypoint(args: argparse.Namespace) -> None: - if not bool(getattr(args, "deepseek_v4_0731_k2", False)): +def _validate_deepseek_v4_0731_entrypoint(args: argparse.Namespace) -> None: + legacy_k2 = bool(getattr(args, "deepseek_v4_0731_k2", False)) + optimized = bool(getattr(args, "deepseek_v4_0731_optimized", False)) + if not (legacy_k2 or optimized): return + if legacy_k2 and optimized: + raise ValueError( + "DeepSeek-V4-0731 optimized selection accepts only one entrypoint flag" + ) cli_flags = getattr(args, "_cli_flags", set()) or set() - if "depth" not in cli_flags or int(getattr(args, "depth", 3)) != 2: + depth = int(getattr(args, "depth", 3)) + if legacy_k2 and ("depth" not in cli_flags or depth != 2): raise ValueError( "DeepSeek-V4-0731 K2 requires explicit --depth 2 before model load" ) + if optimized and ("depth" not in cli_flags or depth not in {1, 2, 3}): + raise ValueError( + "DeepSeek-V4-0731 optimized requires explicit --depth 1, 2, or 3 " + "before model load" + ) if ( getattr(args, "load_mtp", True) is False or str(getattr(args, "generation_mode", "mtp")) != "mtp" or bool(getattr(args, "stock_ar", False)) ): - raise ValueError("DeepSeek-V4-0731 K2 requires MTP generation") + label = "K2" if legacy_k2 else "optimized" + raise ValueError(f"DeepSeek-V4-0731 {label} requires MTP generation") def _global_stateless_session_cache_bypass( @@ -1837,7 +1850,7 @@ def archive_cold_tier() -> dict[str, Any]: class ServerState: def __init__(self, args: argparse.Namespace) -> None: _validate_mtp_batch_settings(args) - _validate_deepseek_v4_0731_k2_entrypoint(args) + _validate_deepseek_v4_0731_entrypoint(args) self.args = args try: args.paged_kv_quantization = normalize_paged_kv_quantization( @@ -1960,8 +1973,11 @@ def __init__(self, args: argparse.Namespace) -> None: _startup_line(" Model load in progress (this may take a minute).") load_heartbeat = _startup_heartbeat("Model still loading") construction_options = {} - if bool(getattr(args, "deepseek_v4_0731_k2", False)): + if bool(getattr(args, "deepseek_v4_0731_k2", False)) or bool( + getattr(args, "deepseek_v4_0731_optimized", False) + ): construction_options["deepseek_v4_0731_k2"] = True + construction_options["deepseek_v4_0731_depth"] = int(args.depth) try: self.runtime = self.model_scheduler.submit_foreground( load, @@ -28864,6 +28880,14 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: "DSpark K2 stack. Requires explicit --depth 2 and MTP." ), ) + parser.add_argument( + "--deepseek-v4-0731-optimized", + action="store_true", + help=( + "Select the construction-bound DeepSeek-V4-Flash-0731 optimized " + "DSpark stack at an explicit depth from 1 through 3." + ), + ) parser.add_argument( "--max-response-tokens", "--max-tokens", diff --git a/services/deepseek-v4-0731/README.md b/services/deepseek-v4-0731/README.md index 764aec1f..d533eefe 100644 --- a/services/deepseek-v4-0731/README.md +++ b/services/deepseek-v4-0731/README.md @@ -1,9 +1,48 @@ -# DeepSeek-V4-Flash-0731 isolated candidate service +# DeepSeek-V4-Flash-0731 service -This directory owns a separate candidate only. It does not change the loaded -production service, and its launchd identity is +This directory owns both the reviewed isolated candidate and the pinned local +production profile. The candidate remains separate under `com.tea.deepseek-v4-0731.candidate` on `127.0.0.1:8081`. +## Production profile + +`production_entry.py`, `launch_production.sh`, and +`com.tea.deepseek-v4.plist` define the local production service on +`127.0.0.1:8080`. The launcher acquires the exclusive GPU lock before model +construction and pins the exact model hashes, DeepSeek-V4 topology, MLX 0.32.0, +official 0731 encoder, 262,144-token context, greedy sampling, K3, and the +construction-bound optimized DSpark route. Cline should use model ID +`mtplx-deepseek-v4-flash-0731-2.4bit-k3`. + +The following warm production probes were measured under that lock on an Apple +M5 Max with 128 GB unified memory. MLX active memory stayed flat at 86.73 GiB +across the probes; the process-wide 139.71 GiB MLX peak includes model loading +and first compilation and is not concurrent resident memory. + +| probe | prompt tokens | output tokens | prefill tok/s | decode tok/s | TTFT s | accepted / drafted | active GiB | +|---|---:|---:|---:|---:|---:|---:|---:| +| exact-text smoke | 17 | 20 | 20.85 | 44.10 | 0.92 | 14 / 15 | 86.73 | +| coding response | 40 | 109 | 44.59 | 38.63 | 1.00 | 75 / 101 | 86.73 | +| native tool call | 295 | 55 | 154.79 | 43.16 | 2.01 | 39 / 45 | 86.73 | +| streamed tool call | 295 | 55 | 166.22 | 43.48 | 1.88 | 39 / 45 | 86.73 | + +The exact-text smoke deterministically repeated the requested marker twice, so +it is a liveness/performance probe rather than an instruction-following pass. +The coding response was coherent, and both nonstream and stream probes emitted +a valid OpenAI tool call with incrementally valid streamed arguments. A replay +of the prior 540 KB Cline transcript remained compute-bound for more than 1,323 +seconds and exceeded the 900-second client timeout. System free-memory pressure +oscillated between 14% and 44% with zero throttled pages instead of growing +monotonically. That establishes bounded chunk cleanup, but not acceptable +long-context prefill latency; the replay is not listed as a successful TPS row. + +The historical K0-K3 diagnostic and the promoted physical-M3 K2 bracket remain +in `docs/perf/receipts/deepseek-v4-0731-dspark.md`. Those runs used different +target geometry and must not be compared directly with this native-M4 K3 +production profile. + +## Candidate profile + The `encoding/` directory vendors the exact official Python encoder and four input/output vectors from `deepseek-ai/DeepSeek-V4-Flash-0731@7872f01b1d1fe23eabc4c98b48bffcef5a386062`. diff --git a/services/deepseek-v4-0731/com.tea.deepseek-v4.plist b/services/deepseek-v4-0731/com.tea.deepseek-v4.plist new file mode 100644 index 00000000..a753e0f2 --- /dev/null +++ b/services/deepseek-v4-0731/com.tea.deepseek-v4.plist @@ -0,0 +1,22 @@ + + + + + Label + com.tea.deepseek-v4 + ProgramArguments + + /Users/davidtai/projects/OpenSourceWTF/.worktrees/dsv4-0731-prod/services/deepseek-v4-0731/launch_production.sh + + RunAtLoad + + KeepAlive + + ThrottleInterval + 15 + StandardOutPath + /Users/davidtai/Library/Logs/com.tea.deepseek-v4.log + StandardErrorPath + /Users/davidtai/Library/Logs/com.tea.deepseek-v4.log + + diff --git a/services/deepseek-v4-0731/launch_production.sh b/services/deepseek-v4-0731/launch_production.sh new file mode 100755 index 00000000..1cb42218 --- /dev/null +++ b/services/deepseek-v4-0731/launch_production.sh @@ -0,0 +1,65 @@ +#!/bin/zsh +set -euo pipefail + +ROOT=/Users/davidtai/projects/OpenSourceWTF/.worktrees/dsv4-0731-prod +SERVICE_ROOT=$ROOT/services/deepseek-v4-0731 +MODEL=/Users/davidtai/models/DeepSeek-V4-Flash-0731-2.4bit-mixed +PYTHON=/Users/davidtai/projects/OpenSourceWTF/mtplx-hy3-ssd/.venv/bin/python +ENTRY=$SERVICE_ROOT/production_entry.py +LOCK=/tmp/mtplx-gpu-exclusive.lock + +die() { print -u2 -- "[deepseek-v4-0731] preflight failed: $*"; exit 1; } +sha256() { /usr/bin/shasum -a 256 "$1" | /usr/bin/awk '{print $1}'; } + +[[ -x "$PYTHON" ]] || die "MLX 0.32 interpreter is unavailable" +[[ -f "$ENTRY" ]] || die "production entrypoint is unavailable" +[[ -d "$MODEL" && ! -L "$MODEL" ]] || die "pinned model is unavailable" +[[ "$(sha256 "$MODEL/config.json")" == 44735712733fcf8f299bdf1faa1d87fac88f1917efe1d3876d6d4c582f79a68f ]] || die "model config changed" +[[ "$(sha256 "$MODEL/model.safetensors.index.json")" == f1332b2b209769c2db335954c2651652a8048e7d7dbf60296c2f2c0198715861 ]] || die "model index changed" +[[ "$(sha256 "$MODEL/tokenizer_config.json")" == 6ac8c8dc065ed118161d02dd532749ae3f52c243deac27872134fae2f50d8547 ]] || die "tokenizer config changed" + +"$PYTHON" - "$MODEL" <<'PY' || die "mlx==0.32.0 or model invariant differs" +import importlib.metadata +import json +import sys +from pathlib import Path + +model = Path(sys.argv[1]) +if importlib.metadata.version("mlx") != "0.32.0": + raise SystemExit("MLX must be exactly 0.32.0") +config = json.loads((model / "config.json").read_text(encoding="utf-8")) +expected = { + "model_type": "deepseek_v4", + "num_hidden_layers": 43, + "dspark_block_size": 5, + "dspark_target_layer_ids": [40, 41, 42], + "dspark_markov_rank": 256, +} +for key, value in expected.items(): + if config.get(key) != value: + raise SystemExit(f"unexpected {key}") +if int(config.get("max_position_embeddings", 0)) < 262144: + raise SystemExit("model context is below 262144") +PY + +export HF_HUB_OFFLINE=1 +export PYTHONNOUSERSITE=1 +export PYTHONPATH="$ROOT" +export MTPLX_MEMORY_LIMIT_BYTES=111669149696 +export MTPLX_WIRED_LIMIT_BYTES=111669149696 + +exec "$PYTHON" - "$LOCK" "$PYTHON" "$ENTRY" <<'PY' +import fcntl +import os +import sys + +lock_path, python, entry = sys.argv[1:] +lock_fd = os.open(lock_path, os.O_CREAT | os.O_RDWR, 0o600) +try: + fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB) +except BlockingIOError as error: + raise SystemExit("GPU lock is already held") from error +os.set_inheritable(lock_fd, True) +os.environ["MTPLX_GPU_LOCK_FD"] = str(lock_fd) +os.execv(python, [python, entry]) +PY diff --git a/services/deepseek-v4-0731/production_entry.py b/services/deepseek-v4-0731/production_entry.py new file mode 100755 index 00000000..3f2ebed8 --- /dev/null +++ b/services/deepseek-v4-0731/production_entry.py @@ -0,0 +1,84 @@ +#!/usr/bin/env python3 +"""Production entrypoint for the optimized DeepSeek-V4-Flash-0731 service.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +from candidate_entry import CandidateConstructionError, install_candidate_surface + + +MODEL = Path("/Users/davidtai/models/DeepSeek-V4-Flash-0731-2.4bit-mixed") +MODEL_ID = "mtplx-deepseek-v4-flash-0731-2.4bit-k3" + + +def serve_argv() -> list[str]: + return [ + "serve", + "--host", + "127.0.0.1", + "--port", + "8080", + "--model", + str(MODEL), + "--model-id", + MODEL_ID, + "--backend-id", + "deepseek_mtp", + "--context-window", + "262144", + "--temperature", + "0", + "--top-p", + "1", + "--top-k", + "0", + "--generation-mode", + "mtp", + "--load-mtp", + "--depth", + "3", + "--deepseek-v4-0731-optimized", + "--verify-strategy", + "batched", + "--verify-core", + "stock", + "--warmup-tokens", + "0", + "--max-active-requests", + "1", + "--session-cache-mode", + "off", + "--ssd-session-cache", + "off", + "--prefill-chunk-tokens", + "512", + "--mlx-cache-limit", + "536870912", + "--reasoning", + "on", + "--reasoning-effort", + "low", + "--reasoning-parser", + "qwen3", + "--tool-prompt-mode", + "native", + "--chat-template-profile", + "tokenizer", + "--no-stats-footer", + ] + + +def main() -> int: + if sys.argv[1:]: + raise CandidateConstructionError("production entrypoint accepts no arguments") + from mtplx.server import openai as server + + install_candidate_surface(server) + server.main(serve_argv()[1:]) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/services/deepseek-v4-0731/tests/test_service_surface.py b/services/deepseek-v4-0731/tests/test_service_surface.py index dfabc6cf..05091f66 100644 --- a/services/deepseek-v4-0731/tests/test_service_surface.py +++ b/services/deepseek-v4-0731/tests/test_service_surface.py @@ -80,6 +80,30 @@ def test_candidate_config_pins_all_installation_identities() -> None: ] +def test_production_profile_pins_coherent_0731_k3_mlx032_and_256k() -> None: + from production_entry import MODEL, MODEL_ID, serve_argv + + assert MODEL == Path( + "/Users/davidtai/models/DeepSeek-V4-Flash-0731-2.4bit-mixed" + ) + assert MODEL_ID == "mtplx-deepseek-v4-flash-0731-2.4bit-k3" + argv = serve_argv() + assert argv[:1] == ["serve"] + assert argv[argv.index("--port") + 1] == "8080" + assert argv[argv.index("--backend-id") + 1] == "deepseek_mtp" + assert argv[argv.index("--context-window") + 1] == "262144" + assert argv[argv.index("--depth") + 1] == "3" + assert "--deepseek-v4-0731-optimized" in argv + assert argv[argv.index("--session-cache-mode") + 1] == "off" + assert argv[argv.index("--ssd-session-cache") + 1] == "off" + + launcher = (ROOT / "launch_production.sh").read_text(encoding="utf-8") + assert "/tmp/mtplx-gpu-exclusive.lock" in launcher + assert "fcntl.LOCK_EX | fcntl.LOCK_NB" in launcher + assert "mlx==0.32.0" in launcher + assert "DeepSeek-V4-Flash-0731-2.4bit-mixed" in launcher + + def test_command_override_is_rejected_except_for_nonstarting_fixture() -> None: launcher = ROOT / "launch_candidate.sh" fixture_env = {"PATH": os.environ["PATH"], "MTPLX_DSV4_0731_TEST_FIXTURE": "1"} diff --git a/tests/test_deepseek_v4_dspark_generation.py b/tests/test_deepseek_v4_dspark_generation.py index 89d073c9..e86227ca 100644 --- a/tests/test_deepseek_v4_dspark_generation.py +++ b/tests/test_deepseek_v4_dspark_generation.py @@ -136,7 +136,7 @@ def __call__(self, *args, **kwargs): class _DSparkRuntime(MTPLXRuntime): - def __init__(self): + def __init__(self, *, supported_depths=(2,)): model = _DSparkModel() super().__init__( model=model, @@ -148,7 +148,9 @@ def __init__(self): self.deepseek_v4_dspark_enabled = True model._dspark.owner = self model.owner = self - self.block_speculative_backend = DeepseekV4DSparkBackend.bind(model) + self.block_speculative_backend = DeepseekV4DSparkBackend.bind( + model, supported_depths=supported_depths + ) self.target_cache = _TargetCache() self.prefill_hidden = None self.proposal_cache = None @@ -380,6 +382,25 @@ def test_dspark_depth_two_verifies_primary_and_two_drafts_in_one_m3_call(): assert out.stats.repair_time_s == 0.0 +def test_dspark_construction_selected_depth_three_verifies_one_physical_m4(): + rt = _DSparkRuntime(supported_depths=(3,)) + + out = generate_mtpk( + rt, + [9, 10], + max_tokens=4, + sampler=SamplerConfig(temperature=0.0), + speculative_depth=3, + stop_token_ids=set(), + ) + + assert out.tokens == [11, 12, 13, 14] + assert rt.target_forward_widths == [1, 1, 4, 1] + assert rt.proposal_widths == [4] + assert out.stats.accepted_drafts == 2 + assert out.stats.accepted_by_depth == [1, 1, 0] + + def test_dspark_proposal_restore_precedes_the_full_accepted_prefix_commit(): """The proposal's poisoned rings never reach the accepted-prefix commit.""" rt = _DSparkRuntime() diff --git a/tests/test_public_cli.py b/tests/test_public_cli.py index 1abc11f5..b5d89bd9 100644 --- a/tests/test_public_cli.py +++ b/tests/test_public_cli.py @@ -1145,6 +1145,61 @@ def test_serve_forwards_explicit_0731_k2_construction_option( assert "--depth 2" in command +def test_serve_forwards_explicit_0731_optimized_construction_option( + monkeypatch, tmp_path, capsys +): + model_dir = tmp_path / "DeepSeek-V4-0731" + model_dir.mkdir() + (model_dir / "config.json").write_text( + json.dumps({"model_type": "deepseek_v4", "dspark_block_size": 5}) + ) + + payload = _serve_dry_run_payload_for_model( + monkeypatch, + capsys, + model_dir, + extra_args=("--deepseek-v4-0731-optimized", "--depth", "3"), + ) + + command = payload["server_command"] + assert "--deepseek-v4-0731-optimized" in command + assert "--depth 3" in command + + +@pytest.mark.parametrize( + "extra_args", + [ + ("--deepseek-v4-0731-optimized",), + ("--deepseek-v4-0731-optimized", "--depth", "3", "--no-load-mtp"), + ( + "--deepseek-v4-0731-optimized", + "--depth", + "3", + "--generation-mode", + "ar", + ), + ], +) +def test_serve_rejects_invalid_0731_optimized_entrypoint_selection( + monkeypatch, tmp_path, capsys, extra_args +): + model_dir = tmp_path / "DeepSeek-V4-0731" + model_dir.mkdir() + (model_dir / "config.json").write_text( + json.dumps({"model_type": "deepseek_v4", "dspark_block_size": 5}) + ) + + payload = _serve_dry_run_payload_for_model( + monkeypatch, + capsys, + model_dir, + extra_args=extra_args, + expected_code=2, + ) + + assert "DeepSeek-V4-0731 optimized" in payload["error"] + + @pytest.mark.parametrize( "extra_args", [ diff --git a/tests/test_runtime_deepseek_v4_dspark.py b/tests/test_runtime_deepseek_v4_dspark.py index a957de3a..370ade3f 100644 --- a/tests/test_runtime_deepseek_v4_dspark.py +++ b/tests/test_runtime_deepseek_v4_dspark.py @@ -420,6 +420,7 @@ def test_k2_option_publishes_one_construction_transaction(monkeypatch, tmp_path) ] assert loaded.deepseek_v4_0731_k2_receipt == { "target_protocol": "primary_plus_two_drafts_physical_m3", + "selected_depth": 2, "exact_vs_serial_greedy": False, "target": {"candidate": "target"}, "dspark_ffn": {"candidate": "ffn"}, @@ -436,7 +437,7 @@ def test_k2_option_restores_both_staged_stacks_when_backend_binding_fails( monkeypatch.setattr( DeepseekV4DSparkBackend, "bind", - lambda _model: (_ for _ in ()).throw(RuntimeError("bind failed")), + lambda _model, **_kwargs: (_ for _ in ()).throw(RuntimeError("bind failed")), ) with pytest.raises(RuntimeError, match="bind failed"): @@ -520,7 +521,7 @@ def test_k2_rollback_attempts_target_and_o_lora_after_ffn_restore_failure( monkeypatch.setattr( DeepseekV4DSparkBackend, "bind", - lambda _model: (_ for _ in ()).throw(RuntimeError("bind failed")), + lambda _model, **_kwargs: (_ for _ in ()).throw(RuntimeError("bind failed")), ) with pytest.raises(ExceptionGroup) as caught: @@ -581,7 +582,7 @@ def install_gather(_model, **kwargs): monkeypatch.setattr( DeepseekV4DSparkBackend, "bind", - lambda _model: (_ for _ in ()).throw(RuntimeError("bind failed")), + lambda _model, **_kwargs: (_ for _ in ()).throw(RuntimeError("bind failed")), ) with pytest.raises(ExceptionGroup) as caught: diff --git a/tests/test_server_openai.py b/tests/test_server_openai.py index 49ab6ee3..a2a1897d 100644 --- a/tests/test_server_openai.py +++ b/tests/test_server_openai.py @@ -693,6 +693,22 @@ def test_server_parser_accepts_explicit_0731_k2_construction_option(): assert "depth" in selected._cli_flags +def test_server_parser_accepts_optimized_0731_k3_construction_option(): + selected = parse_args( + [ + "--warmup-tokens", + "0", + "--deepseek-v4-0731-optimized", + "--depth", + "3", + ] + ) + + assert selected.deepseek_v4_0731_optimized is True + assert "deepseek-v4-0731-optimized" in selected._cli_flags + openai._validate_deepseek_v4_0731_entrypoint(selected) + + @pytest.mark.parametrize( "argv", [ @@ -714,6 +730,31 @@ def test_server_rejects_invalid_0731_k2_selection_before_load(monkeypatch, argv) openai.ServerState(args) +@pytest.mark.parametrize( + "argv", + [ + ["--deepseek-v4-0731-optimized"], + ["--deepseek-v4-0731-optimized", "--depth", "4"], + ["--deepseek-v4-0731-optimized", "--depth", "3", "--no-load-mtp"], + ["--deepseek-v4-0731-optimized", "--depth", "3", "--generation-mode", "ar"], + ], +) +def test_server_rejects_invalid_0731_optimized_selection_before_load( + monkeypatch, argv +): + monkeypatch.setattr( + openai, + "load", + lambda *_args, **_kwargs: pytest.fail( + "invalid optimized 0731 selection reached load" + ), + ) + args = parse_args(["--warmup-tokens", "0", *argv]) + + with pytest.raises(ValueError, match="DeepSeek-V4-0731 optimized"): + openai.ServerState(args) + + def test_server_parser_resolves_api_key_file_before_env(monkeypatch, tmp_path): api_key_file = tmp_path / "api-key" api_key_file.write_text("file-secret\n", encoding="utf-8") @@ -11886,6 +11927,43 @@ def stop_after_load(model, mtp, contract, **kwargs): openai.ServerState(args) assert captured["deepseek_v4_0731_k2"] is True + assert captured["deepseek_v4_0731_depth"] == 2 + + +def test_server_state_maps_0731_optimized_k3_to_runtime_installer(monkeypatch): + captured = {} + monkeypatch.setattr(openai, "apply_profile_env", lambda _profile, **_kwargs: None) + monkeypatch.setattr(openai, "profile_env_status", lambda _profile, **_kwargs: {}) + monkeypatch.setattr(openai, "_fast_path_env_status", lambda: {}) + monkeypatch.setattr(openai, "_mlx_runtime_status", lambda: {"ok": True}) + monkeypatch.setattr( + openai, + "_configure_mlx_cache_limit", + lambda _args: {"configured": False}, + ) + + def stop_after_load(model, mtp, contract, **kwargs): + captured.update(kwargs) + raise RuntimeError("stop after load") + + monkeypatch.setattr(openai, "load", stop_after_load) + args = parse_args( + [ + "--model", + "models/DeepSeek-V4-Flash-0731", + "--warmup-tokens", + "0", + "--deepseek-v4-0731-optimized", + "--depth", + "3", + ] + ) + + with pytest.raises(RuntimeError, match="stop after load"): + openai.ServerState(args) + + assert captured["deepseek_v4_0731_k2"] is True + assert captured["deepseek_v4_0731_depth"] == 3 def test_normalize_stop_sequences_accepts_string_list_and_caps_at_four():