diff --git a/docs/muse-glimmer-optimization-ledger.md b/docs/muse-glimmer-optimization-ledger.md new file mode 100644 index 00000000..f31fb170 --- /dev/null +++ b/docs/muse-glimmer-optimization-ledger.md @@ -0,0 +1,45 @@ +# Muse-Glimmer-30B (q4) text-model optimization ledger + +Measure-first, A/B each candidate against stock decode (same weights, same +shapes, in-window back-to-back so thermal drift doesn't confound), keep wins, +reject losers with evidence. GPU work under the serialized MLX window. + +## Profile (q4, M5 Max, B=1 decode) + +- decode **26 tok/s (38.5 ms/tok) = 82% of the 36.6 tok/s bandwidth roofline** (16.75 GB read/token). +- per-component census: **MLP 77%**, o_proj 7%, lm_head 6%, attn gate_proj 3%, q+kv 5%, norms 1%. +- isolated GEMM sum ≈ 28.3 ms; the remaining **~10 ms/tok (26%)** is SDPA + per-layer glue (cache/rope/gate-mul/softcap) + B=1 host dispatch — the schedulable headroom. + +## Verdicts + +| # | candidate | verdict | evidence | +|---|-----------|---------|----------| +| 1 | **async scheduling** | ✅ already in stock | `mx.async_eval` throughout `batched_decode.py`; MTPLX stock decode already overlaps dispatch (the Laguna S1 +4% lever is already captured) | +| 2 | **QKVG fusion** | ✅ **WIN +4.8%, integrated** | fuse q/k/v/gate → one `quantized_matmul`; **bit-exact** (max\|Δ\|=0), 208→52 launches/token; in-window A/B **26.10 → 27.35 tok/s**. Now the default path in `vendored_muse_glimmer_text.Attention` (id-cached lazy fuse, off the param tree). NB: this *mlx-level concat* beats stock, unlike the Laguna hand-kernel qkvg which was ineligible on affine. | +| 3a | **MLP gate/up fusion (naive mlx-concat)** | ❌ −0.6% | *Not a valid rejection* — mlx-level concat, not a shape-optimized kernel. Superseded by 3b. | +| 3b | **MLP dense-SwiGLU fused KERNEL (shape-optimized)** | ✅ **+5.2%**, quality-parity | Laguna `dense_swiglu_qmv` (in-kernel affine dequant, gate+up+silu+mult fused, per-output-element row-owned) at Glimmer's exact gs32/4-bit/6656/19968. Decode 26.58 → 27.96. **Not** bit-exact (5.86e-3, FP accum-order), but HumanEval **28/40 vs 27/40 stock = parity**. Env-gated `MG_MLP_KERNEL`. Proves a shape-optimized kernel beats stock at M=1 where the naive concat lost. | +| 4 | **qk-norm+rope kernel (row-owned)** | ✅ **+4.8%, bit-exact** | Laguna `fused_qk_rope_sliding` at Glimmer's `SlidingRopeSpec` (32q/2kv, hd128, θ=500000, param-free norm via q_weight=3.87·ones/k_weight=ones, eps1e-5), sliding layers only (globals NoPE). max\|diff\|=0 on q and k. Decode 26.59 → 27.86. Win is **dispatch** (4 kernels × 39 layers → 1 each; host-encode lag), not GPU-exec. Earlier "reject by analogy" was invalid. | +| 5 | **gated-o_proj kernel (row-owned)** | ✅ **+5.1%** isolated | Custom kernel (MLP-`down` pattern + sigmoid-gate fused into the input read), Glimmer 4/5-bit gs32. Decode 26.51 → 27.85. Non-bit-exact (3.12e-2; in-kernel sigmoid) — would need a HumanEval gate, but ~0 stacked so not pursued. | + +## Stacking: the wins SATURATE (do not add) + +All four are **dispatch / host-encode** wins competing for one fixed budget. Stacked (QKVG + MLP-kernel + qk-rope), decode = **27.86 tok/s ≈ any single kernel** (MLP-alone was 27.96). So the bankable win is **~+7% over raw stock (~26.0 → ~27.9)** from **QKVG (bit-exact, default) + the MLP kernel (quality-parity, env-gated)**; qk-rope and gated-o_proj are genuine +4.8–5.1% *isolated* wins but ~0 marginal on top. + +**Method lesson (David's correction, 4/4 vindicated):** you cannot reject a fusion by a naive mlx-concat or by analogy — only by benchmarking a kernel **tiled+fused for the model's exact shapes**. Every candidate rejected the wrong way became a real +5%-class isolated win, because at M=1 decode the bottleneck is stock `qmm`'s small-T ramp + per-op host-encode lag, which an in-kernel-dequant row-owned kernel beats. The *new* structural lesson is that these wins saturate against a fixed dispatch budget rather than compounding. + +## Salvage benchmarks (2026-08-10, in-window guarded) + +| test | result | verdict | +|------|--------|---------| +| **gated-o_proj stacked on QKVG+MLP** | QKVG 26.72 → +MLP 28.02 → +MLP+gated-oproj **27.95 (−0.3%)** | ❌ no salvage — confirms the budget is fully saturated by QKVG+MLP; gated-o_proj's isolated +5.1% was pure dispatch, zero marginal here. Dead. | +| **MLP kernel at prefill (L=512)** | stock qmm **932.4** → MLP kernel **896.3 (−3.9%)** | ⚠️ the MLP kernel is a **decode-only** win. At L=512 the regime is compute-bound (large-T qmm is optimal), the M=1-tuned row-owned kernel loses. **Gate `MG_MLP_KERNEL` to L==1** so prefill keeps stock qmm. | + +Net after salvage: the bankable win stands at **QKVG (bit-exact, always-on) + MLP kernel (decode-only, L==1-gated) ≈ +7% decode over raw stock**. gated-o_proj and qk-rope add nothing on top (saturated). The MLP-kernel gate is now **`MG_MLP_KERNEL` AND L==1**, not unconditional. +| 5 | **lm_head 5→4-bit** | ⚠️ not worth | ~1% byte win; lm_head is deliberately 5-bit for output fidelity (quality-gated). | + +## Net + +One integrated win: **QKVG fusion, +4.8%, bit-exact.** MLP and SDPA are already +stock-`qmm`/flash-optimal and are intentionally **not** hand-kerneled. The +remaining decode gap is bandwidth (MLP, quant-gated) + B=1 dispatch (already +async-overlapped in stock). diff --git a/mtplx/artifacts.py b/mtplx/artifacts.py index e3ccc381..1aa5e375 100644 --- a/mtplx/artifacts.py +++ b/mtplx/artifacts.py @@ -471,6 +471,7 @@ class ModelInspection: backend_status: str | None = None backend_artifact: dict[str, Any] | None = None gemma4_pair: dict[str, Any] | None = None + dflash_pair: dict[str, Any] | None = None @property def passes_primary_gate(self) -> bool: @@ -517,6 +518,7 @@ def to_dict(self) -> dict[str, Any]: "backend_status": self.backend_status, "backend_artifact": self.backend_artifact, "gemma4_pair": self.gemma4_pair, + "dflash_pair": self.dflash_pair, "mtp_supported": self.compatibility.get("mtp_supported"), "mtp_arch": self.compatibility.get("arch_id"), "recommended_backend": self.compatibility.get("recommended_backend"), @@ -1109,6 +1111,60 @@ def inspect_model(model_dir: Path | str) -> ModelInspection: if repo_id is not None: return _inspect_hf_model(repo_id) model_path = Path(model_dir) + try: + from .dflash_pair import dflash_pair_inspection, resolve_dflash_pair_paths + except Exception: + dflash_pair = None + else: + dflash_pair = resolve_dflash_pair_paths(model_path) + if dflash_pair is not None: + payload = dflash_pair_inspection( + model_ref=str(model_path), + bundle_root=dflash_pair["bundle_root"], + target_model=dflash_pair["target_model"], + drafter_model=dflash_pair["drafter_model"], + metadata=dflash_pair["metadata"], + ) + target_config = load_config(dflash_pair["target_model"]) + tcfg = text_config(target_config) + target_quant = ( + target_config.get("quantization") + or target_config.get("quantization_config") + or tcfg.get("quantization") + or tcfg.get("quantization_config") + or {} + ) + return ModelInspection( + model_dir=str(model_path), + source="local", + config_exists=True, + architecture=str(payload.get("architecture") or "DFlashDrafterPair"), + model_type=str(payload.get("model_type") or "dflash_pair"), + mtp_num_hidden_layers=1, + hidden_size=tcfg.get("hidden_size"), + num_hidden_layers=tcfg.get("num_hidden_layers"), + vocab_size=tcfg.get("vocab_size"), + num_experts=tcfg.get("n_routed_experts") or tcfg.get("num_experts"), + num_experts_per_tok=tcfg.get("num_experts_per_tok"), + mtp_pattern="dflash-pair", + quantization=target_quant, + sidecars={name: False for name in MULTIMODAL_SIDECARS}, + model_files=tuple( + sorted( + path.name + for path in Path(dflash_pair["target_model"]).glob( + "model*.safetensors" + ) + ) + ), + runtime_model=payload.get("runtime_model"), + dflash_pair=payload.get("dflash_pair") + if isinstance(payload.get("dflash_pair"), dict) + else None, + compatibility=payload.get("compatibility") + if isinstance(payload.get("compatibility"), dict) + else {}, + ) try: from .gemma4_pair import gemma4_pair_inspection, resolve_gemma4_pair_paths except Exception: diff --git a/mtplx/backends/descriptors.py b/mtplx/backends/descriptors.py index d29fa1ce..3c2143d3 100644 --- a/mtplx/backends/descriptors.py +++ b/mtplx/backends/descriptors.py @@ -729,11 +729,63 @@ def supports(self, capability: str) -> bool: ) +DFLASH_DESCRIPTOR = BackendDescriptor( + backend_id="dflash", + architecture_id="dflash-drafter-pair", + model_family="dflash", + display_name="DFlash external drafter", + artifact_layout="dflash_pair_bundle", + runtime_capabilities=( + "target_logits", + "external_dflash_drafter", + "target_prefix_greedy_verification", + "requires_generation_thread_affinity", + ), + sampler_defaults=SamplerDefaults(temperature=0.0, top_p=1.0, top_k=0), + reasoning_codec=ReasoningCodec( + parser="none", + display_name="Tokenizer-native text", + default_mode="off", + supported=False, + ), + draft_semantics=DraftSemantics( + request_field="speculative_depth", + display_label="DFlash block", + default=8, + minimum=2, + maximum=8, + unit="block", + ), + uses_external_assistant=True, + uses_draft_lm_head=False, + hidden_variant="dflash_aux_taps", + mtp_history_policy="dflash_context_kv", + tune_policy=TunePolicy( + supported=False, + unsupported_reason="DFlash block tuning has not been wired to mtplx tune.", + ), + kv_quant_policy=KVQuantPolicy( + supported=False, + disabled_reason="KV quantization is not validated for DFlash hybrid caches.", + ), + context_window_policy=ContextWindowPolicy( + maximum=1_048_576, + default=131_072, + source="dflash_target_config", + ), + validation_status="runtime_runnable_qa_pending", + status="runtime_runnable_qa_pending", + profile_policy="backend-aware-sustained", + notes=("DFlash currently exposes exact greedy target-prefix verification.",), +) + + DESCRIPTORS_BY_BACKEND_ID: dict[str, BackendDescriptor] = { QWEN3_NEXT_DESCRIPTOR.backend_id: QWEN3_NEXT_DESCRIPTOR, LAGUNA_AR_DESCRIPTOR.backend_id: LAGUNA_AR_DESCRIPTOR, NATIVE_CONTRACT_DESCRIPTOR.backend_id: NATIVE_CONTRACT_DESCRIPTOR, GEMMA4_ASSISTANT_DESCRIPTOR.backend_id: GEMMA4_ASSISTANT_DESCRIPTOR, + DFLASH_DESCRIPTOR.backend_id: DFLASH_DESCRIPTOR, STEP3P5_MTP_DESCRIPTOR.backend_id: STEP3P5_MTP_DESCRIPTOR, DEEPSEEK_MTP_DESCRIPTOR.backend_id: DEEPSEEK_MTP_DESCRIPTOR, GLM_MTP_DESCRIPTOR.backend_id: GLM_MTP_DESCRIPTOR, diff --git a/mtplx/backends/dflash.py b/mtplx/backends/dflash.py new file mode 100644 index 00000000..d0985313 --- /dev/null +++ b/mtplx/backends/dflash.py @@ -0,0 +1,363 @@ +"""Generic dflash speculative-decoding backend for MTPLX. + +Drives a :class:`~mtplx.models.dflash.DFlashDrafter` (a config-driven block- +diffusion drafter, validated bit-exact against bstnxbt/dflash-mlx) against any +MTPLX target that exposes the residual-stream tap hook (``model._tap_layers`` / +``model._taps``). One round = propose a block, verify it in a single target +forward, accept a prefix by **target-prefix** semantics, and roll the target KV +cache back to the accepted length. + +Adding a future dflash drafter is code-free: drop a pair bundle (``target/`` + +``drafter/`` + ``dflash_pair.json``) and this backend loads and runs it — the +tap layers, block size, and mask token all come from the drafter's config. + +The decode loop mirrors the acceptance-validated reference flow: + * ``staged_first`` = the last *forwarded* token; its tap + all prior committed + taps form the drafter's context (accumulated across rounds, like the + reference ``ContextOnlyDraftKVCache``). + * verify forwards the k-1 draft tokens once; ``target[j]`` is the argmax of the + logit *before* draft position j (the first uses the carried ``prev_logit``). + * accept the leading run of matches; the first miss becomes the correction, + an all-match becomes the bonus; both are the next ``staged_first``. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import mlx.core as mx + +from mtplx.models.dflash import DFlashDrafter, load_dflash + + +@dataclass(frozen=True) +class DFlashRuntimeConfig: + target_model_path: str + drafter_model_path: str + block_size: int + # Target-model layer outputs to tap after applying the checkpoint-specific + # offset once during bundle construction. + capture_layers: list[int] + mask_token_id: int + embed_scale: float = 1.0 + max_context: int = 1088 # sink(64)+window(1024); cap accumulated ctx taps + + @property + def assistant_model_path(self) -> str: + return self.drafter_model_path + + @property + def draft_block_size(self) -> int: + return self.block_size + + target_distribution_mode: str = "target_prefix" + + +class DFlashRuntime: + """Self-contained dflash runtime. ``backend_id`` marks it for the + ``generate_mtpk`` dispatch branch (mirrors ``gemma4_assistant``).""" + + backend_id = "dflash" + mtp_enabled = True + + def __init__(self, target, tokenizer, drafter: DFlashDrafter, + config: DFlashRuntimeConfig): + self.target = target + self.model = target + self.args = target.args + self.tokenizer = tokenizer + self.drafter = drafter + self.config = config + if getattr(target.args, "model_type", None) == "nemotron_h": + from mtplx.nemotron_lightning_dflash import ( + install_nemotron_lightning_capture, + ) + + install_nemotron_lightning_capture(target, config.capture_layers) + self._tm = target.backbone + self._target_embedding = target.backbone.embeddings + self._target_lm_head = target.lm_head + self._forward_capture_impl = target.dflash_forward_capture + self._round_impl = self._round_hybrid + else: + self._tm = target.model + self._tm._tap_layers = set(config.capture_layers) + self._target_embedding = self._tm.embed_tokens + self._target_lm_head = None + self._forward_capture_impl = self._forward_capture_tapped + self._round_impl = self._round_kv + self.model_path = config.target_model_path + self.path = config.target_model_path + + # ---- target token-embedding / lm-head as callables for the drafter ---- + def _tok_embd(self, ids: mx.array) -> mx.array: + return self._target_embedding(ids) + + def _lm_head(self, x: mx.array) -> mx.array: + if self._target_lm_head is not None: + return self._target_lm_head(x) + args = self.target.args + logits = (self._tm.embed_tokens.as_linear(x) + if args.tie_word_embeddings else self.target.lm_head(x)) + logits = logits * args.logit_scale + cap = args.final_logit_softcapping + return mx.tanh(logits / cap) * cap if cap else logits + + def _forward_capture_tapped(self, ids: mx.array, cache) -> tuple[mx.array, dict]: + """Forward the target on ids [1,T] with tap capture on. Returns + (logits [1,T,vocab], taps {L: [T, hidden]}).""" + logits = self.target(ids, cache=cache) + taps = {L: self._tm._taps[L][0] for L in self.config.capture_layers} + return logits, taps + + def _forward_capture(self, ids: mx.array, cache) -> tuple[mx.array, dict]: + return self._forward_capture_impl(ids, cache) + + # ---- drop grown-buffer zero-garbage so temporal_order stays clean ------ + @staticmethod + def _normalize(cache) -> None: + """A single-token forward (`_update_in_place`) grows the KV buffer with + zeros past `offset`; the next batched forward's `_temporal_order` would + fold that garbage into committed context (the exactness bug). Slice each + non-rotated cache down to its valid offset before every batched forward.""" + for c in cache: + if c.keys is None: + continue + max_size = getattr(c, "max_size", None) + if max_size is not None and c.offset >= max_size: + continue # rotated sliding cache: leave the rotating buffer alone + if c.keys.shape[2] > c.offset: + c.keys = c.keys[..., : c.offset, :] + c.values = c.values[..., : c.offset, :] + if hasattr(c, "_idx"): + c._idx = c.offset + + # ---- roll every layer cache back to `keep_len` committed positions ----- + @staticmethod + def _rollback(cache, keep_len: int) -> None: + """Drop speculative KV beyond `keep_len`. For a non-rotated cache we + SLICE the buffer (not just move `offset`) so the next batched forward's + temporal-order/trim never re-mixes rejected drafts into context — the + cause of the exactness divergence. A genuinely rotated sliding cache + falls back to `trim` bookkeeping.""" + for c in cache: + if c.keys is None or c.offset <= keep_len: + continue + max_size = getattr(c, "max_size", None) + if max_size is not None and c.offset >= max_size: # rotated sliding + c.trim(c.offset - keep_len) + else: # global / unrotated + c.keys = c.keys[..., :keep_len, :] + c.values = c.values[..., :keep_len, :] + c.offset = keep_len + if hasattr(c, "_idx"): + c._idx = keep_len + + # ---- one propose+verify+accept round (cached ctx, one target forward) -- + @staticmethod + def _accepted_prefix(draft_ids: list[int], vlog: mx.array) -> tuple[list[int], int]: + targ = [int(x) for x in mx.argmax(vlog, axis=-1).tolist()] + accepted: list[int] = [] + nxt: int | None = None + for j, draft in enumerate(draft_ids): + if draft == targ[j]: + accepted.append(draft) + else: + nxt = targ[j] + break + if nxt is None: + nxt = targ[len(draft_ids)] + return accepted, nxt + + def _round_kv(self, ctx_cache: list, ctx_len: int, primary: int, cache): + cfg = self.config + drafts = self.drafter.propose_block_cached( + ctx_cache, ctx_len, self._tok_embd, self._lm_head, + primary_token_id=primary, mask_token_id=cfg.mask_token_id, + block_size=cfg.block_size, embed_scale=cfg.embed_scale, + ) + draft_ids = [int(x) for x in drafts.tolist()] # k-1 tokens + self._normalize(cache) # clean grown-buffer garbage + base = cache[0].offset + + # ONE target forward over [primary, *drafts] (folds the old separate + # nxt decode): vlog[j] = logits after position j predicts token j+1. + vlogits, vtaps = self._forward_capture( + mx.array([primary] + draft_ids, dtype=mx.int32)[None], cache) # appends k + vlog = vlogits[0] # [k, vocab] + + # target-prefix accept walk. Compute ALL k target argmaxes in ONE op + + # ONE host sync (not k `.item()` calls — that was k GPU stalls/round). + accepted, nxt = self._accepted_prefix(draft_ids, vlog) + A = len(accepted) + + # commit [primary + A accepted]; the correction/bonus `nxt` becomes the + # next round's primary (seated by that round's verify position 0). + self._rollback(cache, base + 1 + A) + new_taps = [vtaps[L][: 1 + A] for L in cfg.capture_layers] # primary+accepted taps + self.drafter.extend_context(ctx_cache, new_taps, ctx_len) + committed = [primary] + accepted + return ctx_len + 1 + A, nxt, committed + + def _round_hybrid(self, ctx_cache: list, ctx_len: int, primary: int, cache): + from mtplx.cache_state import rollback_after_verify, snapshot_untrimmable_cache + + cfg = self.config + drafts = self.drafter.propose_block_cached( + ctx_cache, + ctx_len, + self._tok_embd, + self._lm_head, + primary_token_id=primary, + mask_token_id=cfg.mask_token_id, + block_size=cfg.block_size, + embed_scale=cfg.embed_scale, + ) + draft_ids = [int(x) for x in drafts.tolist()] + verify_ids = [primary] + draft_ids + before_verify = snapshot_untrimmable_cache(cache) + vlogits, verify_taps = self._forward_capture( + mx.array(verify_ids, dtype=mx.int32)[None], cache + ) + accepted, nxt = self._accepted_prefix(draft_ids, vlogits[0]) + committed = [primary] + accepted + if len(accepted) == len(draft_ids): + committed_taps = verify_taps + else: + rollback_after_verify(cache, before_verify, verified_tokens=len(verify_ids)) + _repair_logits, committed_taps = self._forward_capture( + mx.array(committed, dtype=mx.int32)[None], cache + ) + new_taps = [committed_taps[layer] for layer in cfg.capture_layers] + self.drafter.extend_context(ctx_cache, new_taps, ctx_len) + return ctx_len + len(committed), nxt, committed + + # ---- greedy generate -------------------------------------------------- + def generate(self, prompt, max_tokens: int = 128, *, + stop_token_ids: set | None = None, token_callback=None) -> dict: + """Greedy speculative generate. `prompt` is a string or a list of token + ids. `token_callback([id])` is called per committed token (streaming); + generation stops at `max_tokens` or when a stop id is committed. Output + is token-exact vs greedy AR (up to fp near-tie non-determinism).""" + prompt_ids = (self.tokenizer.encode(prompt) if isinstance(prompt, str) + else list(prompt)) + ids = mx.array(prompt_ids)[None] + cache = self.target.make_cache() + logits, taps = self._forward_capture(ids, cache) + P = int(ids.shape[1]) + ctx_cache = self.drafter.init_context_cache() + self.drafter.extend_context( + ctx_cache, [taps[L] for L in self.config.capture_layers], 0) # prompt @0..P-1 + ctx_len = P + primary = int(mx.argmax(logits[0, -1]).item()) # first token to commit + stops = stop_token_ids or set() + + out: list[int] = [] + rounds = 0 + accepts = 0 + drafted = 0 + stopped = False + while len(out) < max_tokens and not stopped: + ctx_len, primary, committed = self._round_impl( + ctx_cache, ctx_len, primary, cache + ) + rounds += 1 + drafted += self.config.block_size - 1 + accepts += len(committed) - 1 # accepted drafts only + for tid in committed: + out.append(tid) + if token_callback is not None: + token_callback([tid]) + if tid in stops or len(out) >= max_tokens: + stopped = True + break + return { + "text": self.tokenizer.decode(out), + "tokens": out, + "rounds": rounds, + "accepted": accepts, + "drafted": drafted, + "rejected": drafted - accepts, + "mean_accept": accepts / max(1, rounds), # drafts accepted / round + "tokens_per_target_step": (accepts + rounds) / max(1, rounds), + } + + +def generate_dflash(runtime: DFlashRuntime, prompt_ids, *, max_tokens: int, + sampler=None, speculative_depth: int | None = None, + stop_token_ids=None, token_callback=None, seed: int = 0, + **_ignored): + """`generate_mtpk` dispatch entry for the dflash backend (mirrors + `generate_gemma4_assistant`). Greedy target-prefix decode — token-exact vs + greedy AR. `sampler`/session/trace kwargs are accepted for signature + compatibility; sampling (temperature>0) is not yet wired (the argmax drafter + has no q-distribution for a p/q accept), so decode is greedy.""" + import time as _time + + from mtplx.generation import GenerationOutput, GenerationStats + + # The released assistant was trained for the block geometry encoded in its + # pair manifest. Keep that construction-time invariant: the server's + # generic native-MTP depth argument does not resize an installed DFlash + # lane at request time. + del speculative_depth + stops = {int(t) for t in (stop_token_ids or ())} + t0 = _time.perf_counter() + result = runtime.generate(list(prompt_ids), max_tokens=int(max_tokens), + stop_token_ids=stops, token_callback=token_callback) + elapsed = _time.perf_counter() - t0 + toks = result["tokens"] + rate = len(toks) / max(1e-9, elapsed) + stats = GenerationStats( + mode="mtpk", generated_tokens=len(toks), elapsed_s=elapsed, tok_s=rate, + decode_elapsed_s=elapsed, decode_tok_s=rate, end_to_end_tok_s=rate, + runtime_mtp_enabled=True, mtp_forward_calls=result["rounds"], + accepted_drafts=result["accepted"], rejected_drafts=result["rejected"], + drafted_tokens=result["drafted"], verify_calls=result["rounds"], + speculative_depth=runtime.config.block_size - 1, + requested_speculative_depth=runtime.config.block_size - 1, + ) + finish = "stop" if (toks and toks[-1] in stops) else "length" + return GenerationOutput(tokens=toks, text=result["text"], stats=stats, + final_state=None, finish_reason=finish) + + +def load_dflash_runtime(bundle_root: str) -> DFlashRuntime: + """Load a dflash pair bundle (target/ + drafter/ + dflash_pair.json) into a + ready `DFlashRuntime`. Installs the target's model shim if needed (the + drafter borrows the target's tok_embd/lm_head, so the target must load).""" + import json + import os + + from mlx_lm import load as _load + + from mtplx.dflash_pair import ( + dflash_pair_block_size, + resolve_dflash_pair_paths, + ) + + pair = resolve_dflash_pair_paths(bundle_root) + if pair is None: + raise ValueError(f"{bundle_root} is not a dflash pair bundle") + # install the target arch shim before loading (e.g. Muse-Glimmer) + try: + tcfg = json.load(open(os.path.join(pair["target_model"], "config.json"))) + from mtplx.muse_glimmer_patch import ( + install_muse_glimmer_model_shim, + is_muse_glimmer_config, + ) + if is_muse_glimmer_config(tcfg): + install_muse_glimmer_model_shim() + except (OSError, ValueError, ImportError): + pass + target, tokenizer = _load(pair["target_model"]) + drafter, dcfg = load_dflash(pair["drafter_model"]) + block = dflash_pair_block_size(pair["metadata"], dcfg.block_size) + cfg = DFlashRuntimeConfig( + target_model_path=pair["target_model"], + drafter_model_path=pair["drafter_model"], + block_size=block, + capture_layers=[t + dcfg.target_layer_offset for t in dcfg.target_layers], + mask_token_id=dcfg.mask_token_id if dcfg.mask_token_id is not None else 201818, + ) + return DFlashRuntime(target, tokenizer, drafter, cfg) diff --git a/mtplx/backends/registry.py b/mtplx/backends/registry.py index 29db3bd7..d7f31d9e 100644 --- a/mtplx/backends/registry.py +++ b/mtplx/backends/registry.py @@ -137,6 +137,33 @@ def to_dict(self) -> dict[str, Any]: "with mtp=False." ), ), + "muse-glimmer-ar": ArchitectureSupport( + arch_id="muse-glimmer-ar", + display_name="Muse-Glimmer-30B text tower (MLX)", + family="muse_glimmer", + backend="muse_glimmer_ar", + support_level="experimental-native-ar-only", + runtime_compatibility="native-ar-only", + can_run_verified=True, + aliases=( + "muse_glimmer", + "muse_glimmer_text", + "MuseGlimmerForConditionalGeneration", + ), + config_markers=(), + family_gate="none", + references=( + "https://huggingface.co/meta-models/Muse-Glimmer-30B", + "https://github.com/ggml-org/llama.cpp/blob/master/src/models/muse-glimmer.cpp", + ), + notes=( + "Target-only AR runtime for the Muse-Glimmer text tower (Gemma-family: " + "sigmoid gated attention, parameter-free QK-norm with qk_scale_factor, " + "NoPE on the global layers, sandwich norms, final-logit softcap). Loaded " + "via the vendored mlx_lm model class registered by muse_glimmer_patch; " + "no native MTP head, so it runs mtp=False like any AR checkpoint." + ), + ), "qwen3-next-mtp": ArchitectureSupport( arch_id="qwen3-next-mtp", display_name="Qwen3.6 / Qwen3-Next / Qwen3.5 MTP", @@ -1078,6 +1105,10 @@ def _passes_family_runtime_gate(arch_id: str, inspection: Any, tensor_gate: bool _text(getattr(inspection, "model_type", None)) == "gemma4_pair" and isinstance(getattr(inspection, "gemma4_pair", None), dict) ) + if arch_id == "muse-glimmer-ar": + # Plain target-only AR text tower; the vendored mlx_lm loader handles it + # like any dense AR checkpoint, so recognition is sufficient. + return True return False diff --git a/mtplx/dflash_pair.py b/mtplx/dflash_pair.py new file mode 100644 index 00000000..5e000e44 --- /dev/null +++ b/mtplx/dflash_pair.py @@ -0,0 +1,130 @@ +"""dflash drafter-pair bundle helpers. + +A dflash artifact is a bundle root with the *target* verifier under ``target/`` +and the *dflash drafter* under ``drafter/``, described by a ``dflash_pair.json`` +manifest. This mirrors :mod:`mtplx.gemma4_pair` so the native-MTP path learns no +dflash-specific assumptions, and so adding a future dflash drafter is purely +"drop a bundle" — no code. + +Manifest shape (``dflash_pair.json``):: + + { + "layout": {"target": "target", "drafter": "drafter"}, + "backend": "dflash", + "diffusion": {"num_steps": 8}, # optional; drafter denoise budget + "benchmark": {"best_block_size": 16, "mean_accept": 2.1} + } + +The drafter's own ``config.json`` (a :class:`~mtplx.models.dflash.DFlashConfig`) +carries ``target_layers``, ``block_size``, ``mask_token_id`` etc., so the +manifest only needs the layout + serving knobs. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +DFLASH_PAIR_FILE = "dflash_pair.json" +DFLASH_BACKEND = "dflash" +DFLASH_ARCH_ID = "dflash-drafter-pair" + + +def load_dflash_pair_metadata(bundle_root: str | Path) -> dict[str, Any] | None: + path = Path(bundle_root).expanduser() / DFLASH_PAIR_FILE + if not path.is_file(): + return None + try: + data = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return None + return data if isinstance(data, dict) else None + + +def resolve_dflash_pair_paths(bundle_root: str | Path) -> dict[str, Any] | None: + """Return the resolved target/drafter paths for a dflash bundle, or ``None`` + if ``bundle_root`` is not a dflash pair (so ``load()`` falls through).""" + root = Path(bundle_root).expanduser() + metadata = load_dflash_pair_metadata(root) + if metadata is None: + return None + layout = metadata.get("layout") if isinstance(metadata.get("layout"), dict) else {} + target = root / str(layout.get("target") or "target") + drafter = root / str(layout.get("drafter") or "drafter") + if not (target / "config.json").is_file() or not (drafter / "config.json").is_file(): + return None + return { + "bundle_root": str(root), + "target_model": str(target), + "drafter_model": str(drafter), + "metadata": metadata, + } + + +def dflash_pair_block_size(metadata: dict[str, Any] | None, fallback: int) -> int: + if isinstance(metadata, dict): + bench = metadata.get("benchmark") + if isinstance(bench, dict): + try: + return int(bench["best_block_size"]) + except (KeyError, TypeError, ValueError): + pass + return int(fallback) + + +def dflash_pair_num_steps(metadata: dict[str, Any] | None, fallback: int) -> int: + if isinstance(metadata, dict): + diff = metadata.get("diffusion") + if isinstance(diff, dict): + try: + return int(diff["num_steps"]) + except (KeyError, TypeError, ValueError): + pass + return int(fallback) + + +def dflash_pair_inspection( + *, + model_ref: str, + bundle_root: str | Path, + target_model: str | Path, + drafter_model: str | Path, + metadata: dict[str, Any], +) -> dict[str, Any]: + benchmark = metadata.get("benchmark") if isinstance(metadata.get("benchmark"), dict) else {} + return { + "source": model_ref, + "model_dir": str(bundle_root), + "runtime_model": str(target_model), + "drafter_model": str(drafter_model), + "architecture": "DFlashDrafterPair", + "model_type": "dflash_pair", + "mtp_arch": DFLASH_ARCH_ID, + "mtp_supported": True, + "recommended_backend": DFLASH_BACKEND, + "recommended_profile": "sustained", + "runtime_compatibility": "drafter-pair-native", + "dflash_pair": { + "bundle_root": str(bundle_root), + "target_model": str(target_model), + "drafter_model": str(drafter_model), + "benchmark": benchmark, + }, + "compatibility": { + "tier": "family-compatible-unverified", + "can_run": True, + "recognized": True, + "exit_code": 0, + "arch_id": DFLASH_ARCH_ID, + "recommended_backend": DFLASH_BACKEND, + "mtp_supported": "yes", + "runtime_compatibility": "drafter-pair-native", + "support_notes": ( + "External dflash block-diffusion drafter; target and drafter live " + "in bundle subdirectories and load together. Verified via the " + "acceptance@K bench." + ), + "unverified_model": True, + }, + } diff --git a/mtplx/generation.py b/mtplx/generation.py index e64dece6..a1128720 100644 --- a/mtplx/generation.py +++ b/mtplx/generation.py @@ -6162,6 +6162,19 @@ def generate_mtpk( repetition_stop=repetition_stop, requested_speculative_depth=requested_block_size, ) + if getattr(rt, "backend_id", None) == "dflash": + from .backends.dflash import generate_dflash + + return generate_dflash( + rt, + prompt_ids, + max_tokens=max_tokens, + sampler=sampler, + speculative_depth=int(speculative_depth or 0), + stop_token_ids=stop_token_ids, + token_callback=token_callback, + seed=seed, + ) 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) diff --git a/mtplx/modelopt_nvfp4.py b/mtplx/modelopt_nvfp4.py new file mode 100644 index 00000000..f24eac22 --- /dev/null +++ b/mtplx/modelopt_nvfp4.py @@ -0,0 +1,248 @@ +"""NVIDIA ModelOpt NVFP4/FP8 to standard MLX affine conversion.""" + +from __future__ import annotations + +import json +import shutil +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import mlx.core as mx + +from mtplx.compressed_tensors import ( + SIDECAR_FILES, + _read_f8_e4m3_tensor, + _TensorReader, +) +from mtplx.expert_layout import NumberedExpertAccumulator, num_experts_from_config + +_E2M1 = mx.array( + [ + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, + ], + dtype=mx.float32, +) + + +def dequantize_modelopt_nvfp4( + packed: mx.array, + block_scale: mx.array, + weight_scale_2: mx.array, +) -> mx.array: + """Apply ModelOpt's exported ``nibble * block_scale * scale_2`` contract.""" + packed = packed.astype(mx.uint8) + low = mx.take(_E2M1, packed & 0xF) + high = mx.take(_E2M1, (packed >> 4) & 0xF) + values = mx.stack([low, high], axis=-1).reshape( + *packed.shape[:-1], packed.shape[-1] * 2 + ) + scales = mx.repeat(block_scale.astype(mx.float32), repeats=16, axis=-1) + if tuple(scales.shape) != tuple(values.shape): + raise ValueError( + f"ModelOpt NVFP4 scales {scales.shape} do not match weights {values.shape}" + ) + return values * scales * weight_scale_2.astype(mx.float32) + + +def dequantize_modelopt_fp8(weight: mx.array, weight_scale: mx.array) -> mx.array: + """Dequantize ModelOpt FP8 weight tensors before MLX affine requantization.""" + return weight.astype(mx.float32) * weight_scale.astype(mx.float32) + + +def requantize_affine4(weight: mx.array, *, group_size: int = 64) -> dict[str, mx.array]: + qweight, scales, biases = mx.quantize( + weight.astype(mx.float16), group_size=group_size, bits=4, mode="affine" + ) + return {"weight": qweight, "scales": scales, "biases": biases} + + +def _normalized_expert_key(key: str) -> str: + return key.replace(".switch_mlp.up_proj.", ".switch_mlp.fc1.").replace( + ".switch_mlp.down_proj.", ".switch_mlp.fc2." + ) + + +def _quantized_module(prefix: str) -> str: + if ".experts." not in prefix: + return prefix + before, after = prefix.split(".experts.", 1) + _expert, projection = after.split(".", 1) + projection = {"up_proj": "fc1", "down_proj": "fc2"}.get( + projection, projection + ) + return f"{before}.switch_mlp.{projection}" + + +def convert_modelopt_checkpoint( + source_path: str | Path, + output_path: str | Path, + *, + group_size: int = 64, + source_repo: str | None = None, + source_sha: str | None = None, + progress_callback: Callable[[dict[str, Any]], None] | None = None, +) -> dict[str, Any]: + """Stream a ModelOpt mixed NVFP4/FP8 checkpoint into MLX affine INT4.""" + source = Path(source_path).expanduser() + output = Path(output_path).expanduser() + if output.exists(): + raise FileExistsError(output) + output.mkdir(parents=True) + + config = json.loads((source / "config.json").read_text(encoding="utf-8")) + index_path = source / "model.safetensors.index.json" + if index_path.exists(): + source_index = json.loads(index_path.read_text(encoding="utf-8")) + weight_map = {str(k): str(v) for k, v in source_index["weight_map"].items()} + else: + only_file = "model.safetensors" + weights = mx.load(str(source / only_file)) + weight_map = {str(key): only_file for key in weights} + del weights + + source_files = sorted(set(weight_map.values())) + keys_by_file = {filename: [] for filename in source_files} + for key, filename in weight_map.items(): + keys_by_file[filename].append(key) + + num_experts = num_experts_from_config(config) + experts = NumberedExpertAccumulator(num_experts=num_experts or None) + quantized_modules: set[str] = set() + output_map: dict[str, str] = {} + total_size = 0 + counts = {"nvfp4": 0, "fp8": 0, "plain": 0} + + with _TensorReader(source, weight_map) as reader: + for file_index, filename in enumerate(source_files, start=1): + if progress_callback: + progress_callback( + { + "event": "shard_start", + "filename": filename, + "completed": file_index - 1, + "total": len(source_files), + } + ) + out: dict[str, mx.array] = {} + for key in sorted(keys_by_file[filename]): + if key.endswith((".k_scale", ".v_scale")): + # Exported KV-cache quantization metadata is not a model + # parameter in MLX's Nemotron-H implementation. + continue + if key.endswith((".weight_scale", ".weight_scale_2", ".input_scale")): + continue + if not key.endswith(".weight"): + out[key] = reader.tensor(key) + counts["plain"] += 1 + continue + + prefix = key[: -len(".weight")] + scale_key = f"{prefix}.weight_scale" + scale2_key = f"{prefix}.weight_scale_2" + if scale2_key in weight_map: + packed = reader.tensor(key).astype(mx.uint8) + scale = _read_f8_e4m3_tensor(reader, scale_key) + scale2 = reader.tensor(scale2_key) + weight = dequantize_modelopt_nvfp4(packed, scale, scale2) + converted = requantize_affine4(weight, group_size=group_size) + counts["nvfp4"] += 1 + elif scale_key in weight_map: + weight = _read_f8_e4m3_tensor(reader, key) + scale = reader.tensor(scale_key) + converted = requantize_affine4( + dequantize_modelopt_fp8(weight, scale), group_size=group_size + ) + counts["fp8"] += 1 + else: + out[key] = reader.tensor(key) + counts["plain"] += 1 + continue + + quantized_modules.add(_quantized_module(prefix)) + for leaf, value in converted.items(): + out_key = f"{prefix}.{leaf}" + if not experts.add(out_key, value): + out[out_key] = value + + out.update( + { + _normalized_expert_key(key): value + for key, value in experts.flush_complete().items() + } + ) + if out: + output_file = output / filename + mx.save_safetensors(str(output_file), out, metadata={"format": "mlx"}) + for key, value in out.items(): + output_map[key] = filename + total_size += int(value.nbytes) + if progress_callback: + progress_callback( + { + "event": "shard_complete", + "filename": filename, + "completed": file_index, + "total": len(source_files), + } + ) + + remaining = { + _normalized_expert_key(key): value + for key, value in experts.flush_remaining(strict=True).items() + } + if remaining: + filename = "model-experts.safetensors" + mx.save_safetensors(str(output / filename), remaining, metadata={"format": "mlx"}) + for key, value in remaining.items(): + output_map[key] = filename + total_size += int(value.nbytes) + + qparams = {"group_size": group_size, "bits": 4, "mode": "affine"} + quantization = dict(qparams) + quantization.update( + {key: dict(qparams) for key in sorted(quantized_modules)} + ) + config["quantization"] = quantization + config["quantization_config"] = quantization + config["mtplx_source_quantization"] = { + "format": "modelopt-w4a16-nvfp4-mixed-fp8", + "source": source_repo, + "revision": source_sha, + } + (output / "config.json").write_text( + json.dumps(config, indent=2) + "\n", encoding="utf-8" + ) + for name in SIDECAR_FILES: + candidate = source / name + if candidate.exists(): + shutil.copy2(candidate, output / name) + index = { + "metadata": {"total_size": total_size, "source_sha": source_sha}, + "weight_map": {key: output_map[key] for key in sorted(output_map)}, + } + (output / "model.safetensors.index.json").write_text( + json.dumps(index, indent=2) + "\n", encoding="utf-8" + ) + return { + "source": str(source), + "output": str(output), + "counts": counts, + "quantized_modules": len(quantized_modules), + "total_size": total_size, + } diff --git a/mtplx/models/dflash.py b/mtplx/models/dflash.py new file mode 100644 index 00000000..7e98fa1c --- /dev/null +++ b/mtplx/models/dflash.py @@ -0,0 +1,343 @@ +"""Generic dflash block-diffusion drafter (MLX). + +A *config-driven* drafter for MTPLX speculative decoding. A "dflash" drafter is a +small Qwen3-style transformer that, given a few taps of the target model's +residual stream, proposes a whole block of `block_size` draft tokens in one +non-autoregressive forward — then MTPLX verifies them against the target. + +The mechanism (authoritative reference: llama.cpp ``src/models/dflash.cpp``, +simple/Qwen3 variant — no DSpark/Markov head): + + encode: fused = enc_norm( fc( concat of target hidden taps @ target_layers ) ) + inject: for each drafter layer, K/V = rope(k_norm(k_proj(fused))) / v_proj(fused) + become the *context* the block attends to (no Q, no output). + decode: a block of `block_size` MASK tokens (embedded via the TARGET's token + embedding) runs through the drafter layers with *non-causal* + attention over [injected-context ++ block], then the TARGET's + lm_head projects the final hidden → argmax → `block_size` draft ids. + +Nothing here is Muse-Glimmer specific. A future dflash drafter is added by +dropping its weights + a ``config.json`` (this shape) + a pair manifest naming +its target — zero code. Everything model-specific lives in :class:`DFlashConfig`: +the tap layer indices, the block size, GQA shape, rope, and the block-seed +token. The drafter borrows the *target's* ``tok_embd`` and ``lm_head`` (passed +in at proposal time) so it never carries a vocab-sized table of its own. +""" +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Callable, List, Optional + +import mlx.core as mx +import mlx.nn as nn +from mlx_lm.models.rope_utils import initialize_rope + + +@dataclass +class DFlashConfig: + hidden_size: int + num_hidden_layers: int + intermediate_size: int + num_attention_heads: int + num_key_value_heads: int + head_dim: int + block_size: int = 8 + # which TARGET layer residual streams are tapped and stacked into `fc`. + target_layers: List[int] = field(default_factory=list) + # Glimmer's converted config stores hidden-state indices (layer output is + # index-1); NVIDIA stores zero-based model-layer ids directly. + target_layer_offset: int = -1 + rope_theta: float = 500000.0 + rms_norm_eps: float = 1e-5 + # encoder input width; must equal len(target_layers) * hidden_size. + n_embd_inp_enc: int = 0 + sliding_window: int = 2048 + # token id used to seed the MASK block. ``None`` => seed each block position + # with the last committed token (a driver-side convention; the acceptance + # bench validates which the checkpoint was trained with). + mask_token_id: Optional[int] = None + max_position_embeddings: int = 1048576 + rope_scaling: Optional[dict[str, Any]] = None + has_embed_tokens: bool = False + causal: bool = False + vocab_size: int = 0 + quantization: Optional[dict[str, Any]] = None + model_type: str = "dflash" + + def __post_init__(self): + if not self.n_embd_inp_enc: + self.n_embd_inp_enc = len(self.target_layers) * self.hidden_size + + @classmethod + def from_dict(cls, d: dict) -> "DFlashConfig": + d = dict(d) + dflash = d.get("dflash_config") or {} + rope = dict(d.get("rope_parameters") or d.get("rope_scaling") or {}) + if "rope_type" in rope and "type" not in rope: + rope["rope_type"] = rope["rope_type"] + d.setdefault("target_layers", d.get("target_layer_ids") or dflash.get("target_layer_ids") or []) + if "target_layer_ids" in d or "eagle_aux_hidden_state_layer_ids" in d: + d.setdefault("target_layer_offset", 0) + d.setdefault("mask_token_id", dflash.get("mask_token_id")) + d.setdefault("causal", bool(dflash.get("causal", False))) + d.setdefault("rope_theta", rope.pop("rope_theta", d.get("rope_theta", 10000.0))) + d.setdefault("rope_scaling", rope or None) + keys = {f.name for f in cls.__dataclass_fields__.values()} # type: ignore[attr-defined] + return cls(**{k: v for k, v in d.items() if k in keys}) + + +class _RMSNorm(nn.Module): + """Standard (Qwen3) RMSNorm — learned weight, NOT Gemma (1+w).""" + + def __init__(self, dims: int, eps: float): + super().__init__() + self.weight = mx.ones((dims,)) + self.eps = eps + + def __call__(self, x: mx.array) -> mx.array: + return mx.fast.rms_norm(x, self.weight, self.eps) + + +class _Attention(nn.Module): + def __init__(self, cfg: DFlashConfig): + super().__init__() + d, H, KV, hd = (cfg.hidden_size, cfg.num_attention_heads, + cfg.num_key_value_heads, cfg.head_dim) + self.n_heads, self.n_kv, self.head_dim = H, KV, hd + self.repeat = H // KV + self.scale = hd ** -0.5 + self.q_proj = nn.Linear(d, H * hd, bias=False) + self.k_proj = nn.Linear(d, KV * hd, bias=False) + self.v_proj = nn.Linear(d, KV * hd, bias=False) + self.o_proj = nn.Linear(H * hd, d, bias=False) + self.q_norm = _RMSNorm(hd, cfg.rms_norm_eps) + self.k_norm = _RMSNorm(hd, cfg.rms_norm_eps) + self.rope = initialize_rope( + hd, + cfg.rope_theta, + False, + scaling_config=cfg.rope_scaling, + max_position_embeddings=cfg.max_position_embeddings, + ) + + def _kv(self, feats: mx.array, offset: int) -> tuple[mx.array, mx.array]: + """Project `feats` [T, d] → per-head K/V with k_norm + rope. [1,KV,T,hd].""" + T = feats.shape[0] + k = self.k_norm(self.k_proj(feats).reshape(T, self.n_kv, self.head_dim)) + v = self.v_proj(feats).reshape(T, self.n_kv, self.head_dim) + k = k.transpose(1, 0, 2)[None] # [1, KV, T, hd] + v = v.transpose(1, 0, 2)[None] + k = self.rope(k, offset=offset) + return k, v + + def __call__(self, h: mx.array, ctx_k: mx.array, ctx_v: mx.array, + ctx_len: int, block_offset: int) -> mx.array: + """`h` [Tblk, d] is the normed block. Attends non-causally over + [injected context ++ block]. Returns o_proj output [Tblk, d].""" + Tb = h.shape[0] + q = self.q_norm(self.q_proj(h).reshape(Tb, self.n_heads, self.head_dim)) + q = q.transpose(1, 0, 2)[None] # [1, H, Tb, hd] + q = self.rope(q, offset=block_offset) + bk, bv = self._kv(h, block_offset) # block's own K/V + k = mx.concatenate([ctx_k, bk], axis=2) # [1,KV,ctx+Tb,hd] + v = mx.concatenate([ctx_v, bv], axis=2) + # full (non-causal) attention: block attends to all context + all block. + out = mx.fast.scaled_dot_product_attention(q, k, v, scale=self.scale, mask=None) + out = out[0].transpose(1, 0, 2).reshape(Tb, self.n_heads * self.head_dim) + return self.o_proj(out) + + +class _MLP(nn.Module): + def __init__(self, cfg: DFlashConfig): + super().__init__() + self.gate_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False) + self.up_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False) + self.down_proj = nn.Linear(cfg.intermediate_size, cfg.hidden_size, bias=False) + + def __call__(self, x: mx.array) -> mx.array: + return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x)) + + +class _Layer(nn.Module): + def __init__(self, cfg: DFlashConfig): + super().__init__() + self.self_attn = _Attention(cfg) + self.mlp = _MLP(cfg) + self.attn_norm = _RMSNorm(cfg.hidden_size, cfg.rms_norm_eps) + self.ffn_norm = _RMSNorm(cfg.hidden_size, cfg.rms_norm_eps) + + def inject(self, fused: mx.array, offset: int) -> tuple[mx.array, mx.array]: + """Precompute this layer's injected context K/V from the fused feature.""" + return self.self_attn._kv(fused, offset) + + def __call__(self, x: mx.array, ctx_k, ctx_v, ctx_len, block_offset) -> mx.array: + r = self.self_attn(self.attn_norm(x), ctx_k, ctx_v, ctx_len, block_offset) + x = x + r + return x + self.mlp(self.ffn_norm(x)) + + +class DFlashDrafter(nn.Module): + """Config-driven dflash drafter. Weights map 1:1 to the checkpoint keys + (fc, enc_norm, layers.N.{attn_norm,ffn_norm,self_attn.*,mlp.*}, norm).""" + + def __init__(self, cfg: DFlashConfig): + super().__init__() + self.cfg = cfg + self.fc = nn.Linear(cfg.n_embd_inp_enc, cfg.hidden_size, bias=False) + if cfg.has_embed_tokens: + if cfg.vocab_size <= 0: + raise ValueError("dflash has_embed_tokens requires vocab_size") + self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size) + self.enc_norm = _RMSNorm(cfg.hidden_size, cfg.rms_norm_eps) + self.layers = [_Layer(cfg) for _ in range(cfg.num_hidden_layers)] + self.norm = _RMSNorm(cfg.hidden_size, cfg.rms_norm_eps) + + # --- encode: stacked target taps -> fused context feature --------------- + def encode(self, taps: List[mx.array]) -> mx.array: + """`taps` is a list of target residual-stream tensors (one per + config.target_layers), each [T, hidden]. Returns fused [T, hidden].""" + if len(taps) != len(self.cfg.target_layers): + raise ValueError( + f"dflash expects {len(self.cfg.target_layers)} taps " + f"(target_layers={self.cfg.target_layers}), got {len(taps)}") + stacked = mx.concatenate([t.astype(mx.float32) for t in taps], axis=-1) + return self.enc_norm(self.fc(stacked.astype(self.fc.weight.dtype))) + + # --- block logits: one non-causal block-decode forward ------------------- + def block_logits( + self, + fused: mx.array, + target_tok_embd: Callable[[mx.array], mx.array], + target_lm_head: Callable[[mx.array], mx.array], + block_ids: mx.array, + embed_scale: float = 1.0, + ) -> mx.array: + """One forward over `block_ids` [Tblk], attending non-causally over + [context ++ block]. Cache-relative positions (matching the reference): + the fused context sits at offset 0, the block at offset Tctx. Returns + TARGET-projected logits [Tblk, vocab].""" + Tctx = fused.shape[0] + ctx_kv = [layer.inject(fused, 0) for layer in self.layers] # context @ 0 + embed = self.embed_tokens if self.cfg.has_embed_tokens else target_tok_embd + x = embed(block_ids).astype(mx.float32) * embed_scale + x = x.astype(self.norm.weight.dtype) + for layer, (ck, cv) in zip(self.layers, ctx_kv): + x = layer(x, ck, cv, Tctx, Tctx) # block @ Tctx + return target_lm_head(self.norm(x)) + + # --- propose: dflash single-forward block draft -------------------------- + def propose_block( + self, + fused: mx.array, + target_tok_embd: Callable[[mx.array], mx.array], + target_lm_head: Callable[[mx.array], mx.array], + primary_token_id: int, + mask_token_id: int, + block_size: Optional[int] = None, + embed_scale: float = 1.0, + ) -> mx.array: + """dflash proposal (matches bstnxbt/dflash-mlx ``draft_greedy``): seed the + block with the KNOWN primary token at position 0 and MASK for the rest, + run ONE non-autoregressive forward, and read the TARGET-projected argmax + at positions 1: — the ``block_size - 1`` speculative tokens after the + primary. Not iterative; the "diffusion" is this single masked pass.""" + k = block_size or self.cfg.block_size + block_ids = mx.concatenate([ + mx.array([int(primary_token_id)], dtype=mx.int32), + mx.full((k - 1,), int(mask_token_id), dtype=mx.int32), + ]) + logits = self.block_logits(fused, target_tok_embd, target_lm_head, + block_ids, embed_scale) + return mx.argmax(logits[1:], axis=-1).astype(mx.int32) # drop pos 0 + + # ---- incremental context cache (avoids re-encoding every round) --------- + # The drafter's context is the projected+injected K/V of every committed + # position. Re-encoding all of it each round is O(context) and dominates the + # per-round cost. Instead cache per-layer (K,V) and only inject the new + # committed positions — "lockstep" drafting with parallel weight reads. + def init_context_cache(self) -> list: + return [[None, None] for _ in self.layers] + + def extend_context(self, ctx_cache: list, new_taps: List[mx.array], offset: int) -> list: + """Inject `new_taps` (list per target_layer, each [new, hidden]) at rope + `offset` into the per-layer context K/V cache. Returns the cache.""" + fused = self.encode(new_taps) # [new, hidden] + for i, layer in enumerate(self.layers): + nk, nv = layer.self_attn._kv(fused, offset) # [1, KV, new, hd] roped @offset + ck, cv = ctx_cache[i] + if ck is None: + ctx_cache[i] = [nk, nv] + else: + ctx_cache[i] = [mx.concatenate([ck, nk], axis=2), + mx.concatenate([cv, nv], axis=2)] + return ctx_cache + + def propose_block_cached( + self, + ctx_cache: list, + ctx_len: int, + target_tok_embd: Callable[[mx.array], mx.array], + target_lm_head: Callable[[mx.array], mx.array], + primary_token_id: int, + mask_token_id: int, + block_size: Optional[int] = None, + embed_scale: float = 1.0, + ) -> mx.array: + """Same proposal as :meth:`propose_block` but attends to the cached + context K/V (block at offset ``ctx_len``). Bit-identical to the + re-encoding path — it only skips recomputing the context.""" + k = block_size or self.cfg.block_size + block_ids = mx.concatenate([ + mx.array([int(primary_token_id)], dtype=mx.int32), + mx.full((k - 1,), int(mask_token_id), dtype=mx.int32), + ]) + embed = self.embed_tokens if self.cfg.has_embed_tokens else target_tok_embd + x = embed(block_ids).astype(mx.float32) * embed_scale + x = x.astype(self.norm.weight.dtype) + for layer, (ck, cv) in zip(self.layers, ctx_cache): + x = layer(x, ck, cv, ctx_len, ctx_len) + logits = target_lm_head(self.norm(x)) + return mx.argmax(logits[1:], axis=-1).astype(mx.int32) + + +def normalize_dflash_weights(weights: dict[str, mx.array]) -> dict[str, mx.array]: + """Map NVIDIA's DFlash module names onto the original Glimmer adapter tree.""" + replacements = ( + ("hidden_norm.", "enc_norm."), + (".input_layernorm.", ".attn_norm."), + (".post_attention_layernorm.", ".ffn_norm."), + ) + out: dict[str, mx.array] = {} + for key, value in weights.items(): + normalized = key + for old, new in replacements: + normalized = normalized.replace(old, new) + out[normalized] = value + return out + + +def load_dflash(path: str) -> tuple[DFlashDrafter, DFlashConfig]: + """Load a dflash drafter from a directory holding config.json + + model.safetensors (weights keyed to match the module tree).""" + import json + import os + + raw_config = json.load(open(os.path.join(path, "config.json"))) + cfg = DFlashConfig.from_dict(raw_config) + model = DFlashDrafter(cfg) + quantization = raw_config.get("quantization") + if isinstance(quantization, dict): + def quant_predicate(module_path, _module): + return quantization.get(module_path, False) + + nn.quantize( + model, + group_size=64, + bits=4, + mode="affine", + class_predicate=quant_predicate, + ) + weights = normalize_dflash_weights(mx.load(os.path.join(path, "model.safetensors"))) + model.load_weights(list(weights.items())) + model.eval() + return model, cfg diff --git a/mtplx/muse_glimmer_patch.py b/mtplx/muse_glimmer_patch.py new file mode 100644 index 00000000..f21e7e46 --- /dev/null +++ b/mtplx/muse_glimmer_patch.py @@ -0,0 +1,45 @@ +"""Register the vendored ``muse_glimmer_text`` model class for MTPLX. + +No released mlx-lm ships a muse_glimmer model class, so ``mlx_lm.utils.load`` +(and therefore ``mtplx serve`` / inspect) cannot build the Muse-Glimmer text +tower without this shim. The vendored class lives in +``mtplx.vendored_muse_glimmer_text``. +""" + +from __future__ import annotations + +import logging +import sys +from typing import Any + +logger = logging.getLogger(__name__) + + +def is_muse_glimmer_config(config: dict[str, Any]) -> bool: + """True for a Muse-Glimmer text checkpoint (converted text tower or the + multimodal wrapper's text_config).""" + model_type = str(config.get("model_type", "")).lower() + if model_type in ("muse_glimmer_text", "muse_glimmer"): + return True + architectures = [str(a) for a in config.get("architectures") or []] + return any("museglimmer" in a.lower() for a in architectures) + + +def install_muse_glimmer_model_shim() -> None: + """Register the vendored Muse-Glimmer classes under ``mlx_lm.models`` so + ``mlx_lm.utils.load`` can resolve them. Idempotent. + + * ``muse_glimmer_text`` -> the text backbone + * ``muse_glimmer`` -> the qwen3_vl-style multimodal wrapper + """ + from . import vendored_muse_glimmer, vendored_muse_glimmer_text + + for name, mod in ( + ("mlx_lm.models.muse_glimmer_text", vendored_muse_glimmer_text), + ("mlx_lm.models.muse_glimmer", vendored_muse_glimmer), + ): + existing = sys.modules.get(name) + if existing is not None and getattr(existing, "Model", None) is not None: + continue + sys.modules[name] = mod + logger.info("[muse-glimmer] vendored model registered as %s", name) diff --git a/mtplx/nemotron_lightning_dflash.py b/mtplx/nemotron_lightning_dflash.py new file mode 100644 index 00000000..1ea66eba --- /dev/null +++ b/mtplx/nemotron_lightning_dflash.py @@ -0,0 +1,73 @@ +"""Construction-time DFlash capture route for Nemotron 3.5 Lightning.""" + +from __future__ import annotations + +from types import MethodType + +from mlx_lm.models.base import create_attention_mask, create_ssm_mask + + +def install_nemotron_lightning_capture(target, capture_layers: list[int]) -> None: + """Install a dedicated verifier forward that returns fixed residual taps.""" + if getattr(target.args, "model_type", None) != "nemotron_h": + raise ValueError("Nemotron Lightning DFlash requires model_type=nemotron_h") + if not hasattr(target, "backbone"): + raise ValueError("Nemotron Lightning target is missing its backbone") + layer_count = len(target.backbone.layers) + captures = tuple(int(layer) for layer in capture_layers) + if captures != tuple(sorted(set(captures))): + raise ValueError(f"DFlash capture layers must be sorted and unique: {captures}") + if not captures or captures[0] < 0 or captures[-1] >= layer_count: + raise ValueError( + f"DFlash capture layers {captures} are outside Lightning's {layer_count} layers" + ) + + # Segments make capture positions construction-time invariants. The inner + # layer loop has no membership check or metadata validation. + segments: tuple[tuple[int, int], ...] = tuple( + (0 if index == 0 else captures[index - 1] + 1, capture + 1) + for index, capture in enumerate(captures) + ) + expected_start = 0 + for start, end in segments: + if start != expected_start or end <= start: + raise ValueError("invalid Lightning DFlash capture segments") + expected_start = end + + def forward_capture(self, inputs, cache): + backbone = self.backbone + hidden = backbone.embeddings(inputs) + if cache is None: + cache = [None] * sum( + layer.block_type in {"M", "*"} for layer in backbone.layers + ) + attn_mask = create_attention_mask(hidden, cache[backbone.fa_idx]) + ssm_mask = create_ssm_mask(hidden, cache[backbone.ssm_idx]) + cache_index = 0 + layer_index = 0 + taps = {} + for start, end in segments: + for layer_index in range(start, end): + layer = backbone.layers[layer_index] + if layer.block_type in {"M", "*"}: + layer_cache = cache[cache_index] + cache_index += 1 + else: + layer_cache = None + mask = attn_mask if layer.block_type == "*" else ssm_mask + hidden = layer(hidden, mask=mask, cache=layer_cache) + capture = end - 1 + taps[capture] = hidden[0] + layer_index = end + for layer_index in range(layer_index, len(backbone.layers)): + layer = backbone.layers[layer_index] + if layer.block_type in {"M", "*"}: + layer_cache = cache[cache_index] + cache_index += 1 + else: + layer_cache = None + mask = attn_mask if layer.block_type == "*" else ssm_mask + hidden = layer(hidden, mask=mask, cache=layer_cache) + return self.lm_head(backbone.norm_f(hidden)), taps + + target.dflash_forward_capture = MethodType(forward_capture, target) diff --git a/mtplx/runtime.py b/mtplx/runtime.py index 61d37fa3..68efd106 100644 --- a/mtplx/runtime.py +++ b/mtplx/runtime.py @@ -617,6 +617,20 @@ def load( runtime.bundle_path = path return runtime path = Path(gemma4_pair["target_model"]) + + from .dflash_pair import resolve_dflash_pair_paths + + dflash_pair = resolve_dflash_pair_paths(path) + if dflash_pair is not None: + if mtp: + from .backends.dflash import load_dflash_runtime + + runtime = load_dflash_runtime(str(path)) + runtime.model_path = path + runtime.path = path + runtime.bundle_path = path + return runtime + path = Path(dflash_pair["target_model"]) config = load_config(path) from .a3b_whole_moe import validate_a3b_whole_moe_load_options @@ -648,6 +662,16 @@ def load( if is_hy_v3_config(config): install_hy_v3_model_shim() + # muse_glimmer has no model class in any released mlx-lm; register the + # vendored text-tower class before mlx_lm.utils.load resolves the type. + from .muse_glimmer_patch import ( + install_muse_glimmer_model_shim, + is_muse_glimmer_config, + ) + + if is_muse_glimmer_config(config): + install_muse_glimmer_model_shim() + # A checkpoint whose model_type has no mlx-lm module may still declare the # implementing class in ``architectures`` — new Qwen generations reuse the # qwen3_5 schema under fresh model_type strings (Qwen3.6 shipped as diff --git a/mtplx/server/openai.py b/mtplx/server/openai.py index 6ae270c5..f2b7014c 100644 --- a/mtplx/server/openai.py +++ b/mtplx/server/openai.py @@ -11654,11 +11654,14 @@ def _request_depth_for_generation( None, ) if default_value is None: - default_value = getattr( - state.args, - "depth", - descriptor.draft_semantics.default, - ) + if descriptor.draft_semantics.request_field == "depth": + default_value = getattr( + state.args, + "depth", + descriptor.draft_semantics.default, + ) + else: + default_value = descriptor.draft_semantics.default return descriptor.draft_semantics.clamp(default_value) try: depth = int(value) @@ -16109,6 +16112,15 @@ def _skipped_idle_postcommit_snapshot( "assistant_tool_calls": len(assistant_tool_calls or []), "prompt_prefix_len": int(prompt_prefix_len or 0), } + if backend_id == "dflash": + return { + "stored": False, + "mode": "skipped", + "reason": "dflash_retokenized_postcommit_unsupported", + "unsafe_reason": unsafe_reason, + "assistant_tool_calls": len(assistant_tool_calls or []), + "prompt_prefix_len": int(prompt_prefix_len or 0), + } del unsafe_reason, assistant_tool_calls, prompt_prefix_len return None diff --git a/mtplx/vendored_muse_glimmer.py b/mtplx/vendored_muse_glimmer.py new file mode 100644 index 00000000..7ad9c1cd --- /dev/null +++ b/mtplx/vendored_muse_glimmer.py @@ -0,0 +1,82 @@ +# Copyright © 2026 MTPLX contributors. +"""Vendored multimodal wrapper for Muse-Glimmer (``model_type: muse_glimmer``). + +Mirrors mlx-lm's ``qwen3_vl`` treatment of Qwen3.6-27B-VL: a thin wrapper that +builds the text backbone from ``text_config``, drops the vision tower's weights +at load, and accepts spliced image ``input_embeddings`` (the vision encoder runs +externally via mlx-vlm; MTPLX's runtime splices its output through the +``input_embeddings`` path). Registered as ``mlx_lm.models.muse_glimmer`` by +``muse_glimmer_patch``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +import mlx.core as mx +import mlx.nn as nn +from mlx.utils import tree_flatten, tree_unflatten + +from mlx_lm.models.base import BaseModelArgs + +from . import vendored_muse_glimmer_text as muse_glimmer_text + + +@dataclass +class ModelArgs(BaseModelArgs): + model_type: str + text_config: dict + + @classmethod + def from_dict(cls, params): + if "text_config" not in params: + return cls(model_type=params["model_type"], text_config=params) + return super().from_dict(params) + + +class Model(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.model_type = args.model_type + text_args = muse_glimmer_text.ModelArgs.from_dict(args.text_config) + self.language_model = muse_glimmer_text.Model(text_args) + + def __call__( + self, + inputs: mx.array, + cache=None, + input_embeddings: Optional[mx.array] = None, + ): + return self.language_model( + inputs, cache=cache, input_embeddings=input_embeddings + ) + + def sanitize(self, weights): + # Drop the vision stack (perception encoder + projector); it is served + # externally via mlx-vlm and spliced as input_embeddings. + weights = tree_unflatten(list(weights.items())) + if isinstance(weights, dict): + model = weights.get("model") + if isinstance(model, dict): + for k in ("vision_tower", "vision_adapter", "vision_projection"): + model.pop(k, None) + weights = dict(tree_flatten(weights)) + + # Remap the multimodal checkpoint's language-model weights under the + # wrapper's ``language_model.`` prefix (matching qwen3_vl). + sanitized = {} + for key, value in weights.items(): + if key.startswith("model.language_model."): + key = "language_model.model." + key[len("model.language_model.") :] + elif key == "lm_head.weight": + key = "language_model.lm_head.weight" + elif key.startswith("model.vision"): + continue + sanitized[key] = value + return sanitized + + @property + def layers(self): + return self.language_model.model.layers diff --git a/mtplx/vendored_muse_glimmer_text.py b/mtplx/vendored_muse_glimmer_text.py new file mode 100644 index 00000000..10ce6a4d --- /dev/null +++ b/mtplx/vendored_muse_glimmer_text.py @@ -0,0 +1,338 @@ +# Copyright © 2026 MTPLX contributors. +"""Vendored MLX model definition for the Muse-Glimmer text tower. + +Registered as ``mlx_lm.models.muse_glimmer_text`` by ``muse_glimmer_patch`` so +``mlx_lm.utils.load`` (and therefore ``mtplx serve``) can build the model; no +released mlx-lm ships it. + +Ported 1:1 from llama.cpp ``src/models/muse-glimmer.cpp`` (the authoritative +reference — transformers has no muse_glimmer modeling code). Deviations from a +plain Gemma-3 text model: + + * Sigmoid **gated attention** — ``o_proj(sdpa_out * sigmoid(gate_proj(x)))``. + * **Parameter-free QK-norm**: RMSNorm(q,k) with no learnable weight; the + ``qk_scale_factor`` (3.87) is multiplied onto Q (llama.cpp synthesizes this + into ``attn_q_norm``; the HF checkpoint has no q/k-norm weights). + * **NoPE on the global (full-attention) layers**; RoPE (theta=500000) only on + the sliding-window layers (3 local : 1 global). + * Gemma-style ``(1 + weight)`` sandwich norms; post-attn/post-FFN norms use + eps ``1e-8`` (``post_norm_eps``) vs ``1e-5`` for the pre-norms. + * Embeddings **RMS-normalized** (no weight), NOT scaled by ``sqrt(hidden)``. + * Output: ``logit_scale`` (0.196) then tanh softcap (20). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import List, Optional + +import mlx.core as mx +import mlx.nn as nn + +# Per-attention cache of the fused q/k/v/gate projection (built lazily from the +# loaded per-projection weights), keyed by id() since nn.Module is unhashable. +# Bit-exact vs the 4 separate matmuls; fusing them into one quantized_matmul +# cuts 4 kernel launches/layer -> 1 and measured +4.8% B=1 decode on the q4 +# checkpoint. Kept off the nn.Module parameter tree so load/save are unaffected. +_QKVG_FUSED: dict = {} + +from mlx_lm.models.base import ( + BaseModelArgs, + create_attention_mask, + scaled_dot_product_attention, +) +from mlx_lm.models.cache import KVCache, RotatingKVCache +from mlx_lm.models.rope_utils import initialize_rope + + +@dataclass +class ModelArgs(BaseModelArgs): + model_type: str + hidden_size: int = 6656 + num_hidden_layers: int = 52 + intermediate_size: int = 19968 + num_attention_heads: int = 32 + num_key_value_heads: int = 2 + head_dim: int = 128 + vocab_size: int = 202048 + sliding_window: int = 2048 + sliding_window_pattern: int = 4 + rope_theta: float = 500000.0 + rms_norm_eps: float = 1e-5 + post_norm_eps: float = 1e-8 + qk_scale_factor: float = 3.87 + logit_scale: float = 0.1961161345243454 + final_logit_softcapping: float = 20.0 + max_position_embeddings: int = 131072 + tie_word_embeddings: bool = False + layer_types: Optional[List[str]] = None + + def is_global(self, layer_idx: int) -> bool: + if self.layer_types is not None and layer_idx < len(self.layer_types): + return self.layer_types[layer_idx] == "full_attention" + return (layer_idx + 1) % self.sliding_window_pattern == 0 + + +def _rms(x: mx.array, eps: float) -> mx.array: + dt = x.dtype + x = x.astype(mx.float32) + x = x * mx.rsqrt(mx.mean(x * x, axis=-1, keepdims=True) + eps) + return x.astype(dt) + + +class RMSNorm(nn.Module): + """Gemma-style RMSNorm applying ``(1 + weight)`` (weights stored raw in HF).""" + + def __init__(self, dims: int, eps: float = 1e-5): + super().__init__() + self.weight = mx.ones((dims,)) + self.eps = eps + + def __call__(self, x): + return mx.fast.rms_norm(x, 1.0 + self.weight, self.eps) + + +class Attention(nn.Module): + def __init__(self, args: ModelArgs, is_global: bool): + super().__init__() + self.n_heads = args.num_attention_heads + self.n_kv_heads = args.num_key_value_heads + self.head_dim = args.head_dim + self.qk_eps = args.rms_norm_eps + self.qk_scale_factor = args.qk_scale_factor + self.scale = self.head_dim**-0.5 + + dim = args.hidden_size + self.q_proj = nn.Linear(dim, self.n_heads * self.head_dim, bias=False) + self.k_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(self.n_heads * self.head_dim, dim, bias=False) + self.gate_proj = nn.Linear(dim, self.n_heads * self.head_dim, bias=False) + + if is_global: + self.rope = None + else: + self.rope = initialize_rope( + self.head_dim, args.rope_theta, False, None, args.max_position_embeddings + ) + + def _fused_qkvg(self, x): + """One matmul for q/k/v/gate (bit-exact vs 4 separate); split the output.""" + f = _QKVG_FUSED.get(id(self)) + if f is None: + ms = (self.q_proj, self.k_proj, self.v_proj, self.gate_proj) + if all(hasattr(m, "scales") for m in ms): # quantized + f = ("q", + mx.concatenate([m.weight for m in ms], axis=0), + mx.concatenate([m.scales for m in ms], axis=0), + mx.concatenate([m.biases for m in ms], axis=0), + int(self.q_proj.group_size), int(self.q_proj.bits)) + else: # bf16 dense (attention_bias=False -> no bias) + f = ("d", mx.concatenate([m.weight for m in ms], axis=0)) + _QKVG_FUSED[id(self)] = f + if f[0] == "q": + t = mx.quantized_matmul(x, f[1], f[2], f[3], transpose=True, group_size=f[4], bits=f[5]) + else: + t = x @ f[1].T + qw = self.n_heads * self.head_dim + kw = self.n_kv_heads * self.head_dim + return t[..., :qw], t[..., qw:qw + kw], t[..., qw + kw:qw + 2 * kw], t[..., qw + 2 * kw:] + + def __call__(self, x: mx.array, mask=None, cache=None) -> mx.array: + B, L, _ = x.shape + + queries, keys, values, gate = self._fused_qkvg(x) + + queries = queries.reshape(B, L, self.n_heads, self.head_dim).transpose(0, 2, 1, 3) + keys = keys.reshape(B, L, self.n_kv_heads, self.head_dim).transpose(0, 2, 1, 3) + values = values.reshape(B, L, self.n_kv_heads, self.head_dim).transpose(0, 2, 1, 3) + + # Parameter-free QK-norm; qk_scale_factor folded onto Q. + queries = _rms(queries, self.qk_eps) * self.qk_scale_factor + keys = _rms(keys, self.qk_eps) + + if self.rope is not None: + offset = cache.offset if cache is not None else 0 + queries = self.rope(queries, offset=offset) + keys = self.rope(keys, offset=offset) + + if cache is not None: + keys, values = cache.update_and_fetch(keys, values) + + output = scaled_dot_product_attention( + queries, keys, values, cache=cache, scale=self.scale, mask=mask + ) + output = output.transpose(0, 2, 1, 3).reshape(B, L, -1) + output = output * mx.sigmoid(gate) + return self.o_proj(output) + + +# Optional decode-only fused dense-SwiGLU kernel (+5.2% decode, quality-parity). +# Enabled by MG_MLP_KERNEL=1; the kernel module path is MG_MLP_KERNEL_PATH. +# If either is unset or the module can't load, we fall back to stock qmm — so a +# fresh checkout without the external kernel just runs the (bit-exact) default. +_MLP_KERNEL_ENABLED = __import__("os").environ.get("MG_MLP_KERNEL") == "1" +_dense_swiglu_qmv = None +_dense_swiglu_tried = False + + +def _get_dense_swiglu(): + global _dense_swiglu_qmv, _dense_swiglu_tried + if _dense_swiglu_tried: + return _dense_swiglu_qmv + _dense_swiglu_tried = True + import importlib.util + import os + path = os.environ.get("MG_MLP_KERNEL_PATH") + if not path or not os.path.isfile(path): + return None + try: + spec = importlib.util.spec_from_file_location("_mtplx_dense_mlp", path) + m = importlib.util.module_from_spec(spec) + spec.loader.exec_module(m) + _dense_swiglu_qmv = m.dense_swiglu_qmv + except Exception: + _dense_swiglu_qmv = None + return _dense_swiglu_qmv + + +class MLP(nn.Module): + def __init__(self, dim: int, hidden_dim: int): + super().__init__() + self.gate_proj = nn.Linear(dim, hidden_dim, bias=False) + self.up_proj = nn.Linear(dim, hidden_dim, bias=False) + self.down_proj = nn.Linear(hidden_dim, dim, bias=False) + self._hidden = dim + self._intermediate = hidden_dim + + def __call__(self, x) -> mx.array: + # Optional shape-optimized fused dense-SwiGLU kernel (+5.2% decode; NOT + # bit-exact vs stock qmm ~5.9e-3, so env-gated pending a quality gate). + # DECODE-ONLY: it's a row-owned qmv tuned for M=1; at prefill (M>1) the + # compute-bound large-T stock qmm wins (measured −3.9% at L=512), so gate + # on a single flattened row and fall through to stock qmm otherwise. + lead = x.shape[:-1] + x2 = x.reshape(-1, x.shape[-1]) + f = (_get_dense_swiglu() + if _MLP_KERNEL_ENABLED and x2.shape[0] == 1 and hasattr(self.gate_proj, "scales") + else None) + if f is not None: + out = f( + x2, + self.gate_proj.weight, self.gate_proj.scales, self.gate_proj.biases, + self.up_proj.weight, self.up_proj.scales, self.up_proj.biases, + self.down_proj.weight, self.down_proj.scales, self.down_proj.biases, + hidden=self._hidden, intermediate=self._intermediate, + gate_up_bits=int(self.gate_proj.bits), down_bits=int(self.down_proj.bits), + group_size=int(self.gate_proj.group_size)) + return out.reshape(*lead, out.shape[-1]) + return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x)) + + +class TransformerBlock(nn.Module): + def __init__(self, args: ModelArgs, layer_idx: int): + super().__init__() + self.self_attn = Attention(args, is_global=args.is_global(layer_idx)) + self.mlp = MLP(args.hidden_size, args.intermediate_size) + self.input_layernorm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(args.hidden_size, eps=args.post_norm_eps) + self.pre_feedforward_layernorm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.post_feedforward_layernorm = RMSNorm(args.hidden_size, eps=args.post_norm_eps) + + def __call__(self, x: mx.array, mask=None, cache=None) -> mx.array: + r = self.self_attn(self.input_layernorm(x), mask, cache) + h = x + self.post_attention_layernorm(r) + r = self.mlp(self.pre_feedforward_layernorm(h)) + return h + self.post_feedforward_layernorm(r) + + +class MuseGlimmerModel(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size) + self.layers = [TransformerBlock(args, i) for i in range(args.num_hidden_layers)] + self.norm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + # Generic residual-stream tap for external drafters (dflash): when a + # backend sets ``_tap_layers`` to a set of layer indices, ``_taps`` is + # populated with the residual stream at the OUTPUT of each such layer on + # every forward. ``None`` => zero overhead, no behavior change. + self._tap_layers = None + self._taps: dict[int, mx.array] = {} + + def __call__(self, inputs, cache=None, input_embeddings=None): + if input_embeddings is not None: + h = input_embeddings + else: + h = self.embed_tokens(inputs) + h = _rms(h, self.args.rms_norm_eps) # RMS-normed embeddings (not ×√hidden) + + if cache is None: + cache = [None] * len(self.layers) + + pattern = self.args.sliding_window_pattern + global_mask = create_attention_mask(h, cache[pattern - 1]) + sliding_mask = create_attention_mask(h, cache[0], window_size=self.args.sliding_window) + + taps = self._tap_layers + if taps is not None: + self._taps = {} + for i, (layer, c) in enumerate(zip(self.layers, cache)): + mask = global_mask if self.args.is_global(i) else sliding_mask + h = layer(h, mask, c) + if taps is not None and i in taps: + self._taps[i] = h + + return self.norm(h) + + +class Model(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.model_type = args.model_type + self.model = MuseGlimmerModel(args) + if not args.tie_word_embeddings: + self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False) + + def __call__(self, inputs, cache=None, input_embeddings=None): + h = self.model(inputs, cache, input_embeddings) + if self.args.tie_word_embeddings: + logits = self.model.embed_tokens.as_linear(h) + else: + logits = self.lm_head(h) + + logits = logits * self.args.logit_scale + cap = self.args.final_logit_softcapping + if cap: + logits = mx.tanh(logits / cap) * cap + return logits + + def sanitize(self, weights): + """Strip the multimodal wrapper: keep the text tower, drop the vision + stack, and remap ``model.language_model.*`` -> ``model.*``.""" + out = {} + for k, v in weights.items(): + if ( + k.startswith("model.vision_tower") + or k.startswith("model.vision_adapter") + or k.startswith("model.vision_projection") + ): + continue + if k.startswith("model.language_model."): + k = "model." + k[len("model.language_model.") :] + out[k] = v + return out + + @property + def layers(self): + return self.model.layers + + def make_cache(self): + caches = [] + for i in range(self.args.num_hidden_layers): + if self.args.is_global(i): + caches.append(KVCache()) + else: + caches.append(RotatingKVCache(max_size=self.args.sliding_window)) + return caches diff --git a/scripts/convert_modelopt_nvfp4_to_mlx.py b/scripts/convert_modelopt_nvfp4_to_mlx.py new file mode 100644 index 00000000..2c811845 --- /dev/null +++ b/scripts/convert_modelopt_nvfp4_to_mlx.py @@ -0,0 +1,35 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import json + +from mtplx.modelopt_nvfp4 import convert_modelopt_checkpoint + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("source") + parser.add_argument("output") + parser.add_argument("--source-repo") + parser.add_argument("--source-sha") + parser.add_argument("--group-size", type=int, default=64) + args = parser.parse_args() + + def progress(event): + if event["event"] == "shard_complete": + print(f"[{event['completed']}/{event['total']}] {event['filename']}", flush=True) + + report = convert_modelopt_checkpoint( + args.source, + args.output, + group_size=args.group_size, + source_repo=args.source_repo, + source_sha=args.source_sha, + progress_callback=progress, + ) + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/tests/test_artifacts.py b/tests/test_artifacts.py index 93884965..5f639fa6 100644 --- a/tests/test_artifacts.py +++ b/tests/test_artifacts.py @@ -1546,6 +1546,58 @@ def test_gemma4_pair_bundle_inspects_as_assistant_runtime(tmp_path): assert result.gemma4_pair["assistant_model"].endswith("/assistant") +def test_dflash_pair_bundle_inspects_as_native_runtime(tmp_path): + target = tmp_path / "target" + drafter = tmp_path / "drafter" + target.mkdir() + drafter.mkdir() + (tmp_path / "dflash_pair.json").write_text( + json.dumps( + { + "backend": "dflash", + "layout": {"target": "target", "drafter": "drafter"}, + "benchmark": {"best_block_size": 8}, + } + ), + encoding="utf-8", + ) + (target / "config.json").write_text( + json.dumps( + { + "architectures": ["NemotronHForCausalLM"], + "model_type": "nemotron_h", + "hidden_size": 2688, + "num_hidden_layers": 52, + "vocab_size": 131072, + "n_routed_experts": 128, + "num_experts_per_tok": 6, + "quantization": {"group_size": 64, "bits": 4, "mode": "affine"}, + } + ), + encoding="utf-8", + ) + (drafter / "config.json").write_text( + json.dumps( + { + "architectures": ["DFlashDraftModel"], + "model_type": "qwen3", + "target_layer_ids": [1, 5, 19, 29, 41, 51], + } + ), + encoding="utf-8", + ) + + result = inspect_model(tmp_path) + + assert result.model_type == "dflash_pair" + assert result.architecture == "DFlashDrafterPair" + assert result.compatibility["can_run"] is True + assert result.compatibility["recommended_backend"] == "dflash" + assert result.compatibility["runtime_compatibility"] == "drafter-pair-native" + assert result.dflash_pair["target_model"].endswith("/target") + assert result.dflash_pair["drafter_model"].endswith("/drafter") + + def test_gemma4_pair_subfolder_reports_bundle_required(tmp_path): bundle = _write_gemma4_pair_bundle(tmp_path) diff --git a/tests/test_dflash_nemotron.py b/tests/test_dflash_nemotron.py new file mode 100644 index 00000000..e4410e11 --- /dev/null +++ b/tests/test_dflash_nemotron.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +from types import SimpleNamespace + +from mtplx.backends.descriptors import descriptor_for_backend_id +from mtplx.backends.dflash import generate_dflash +from mtplx.models.dflash import DFlashConfig, normalize_dflash_weights +from mtplx.server import openai + + +def test_nvidia_dflash_config_maps_to_runtime_contract(): + cfg = DFlashConfig.from_dict( + { + "hidden_size": 2688, + "num_hidden_layers": 6, + "intermediate_size": 6144, + "num_attention_heads": 32, + "num_key_value_heads": 2, + "head_dim": 128, + "max_position_embeddings": 1048576, + "rms_norm_eps": 1e-6, + "mask_token_id": 990, + "target_layer_ids": [1, 5, 19, 29, 41, 51], + "has_embed_tokens": True, + "rope_parameters": { + "factor": 128.0, + "original_max_position_embeddings": 8192, + "rope_theta": 10000, + "rope_type": "yarn", + }, + "dflash_config": { + "causal": False, + "target_layer_ids": [1, 5, 19, 29, 41, 51], + }, + } + ) + + assert cfg.block_size == 8 + assert cfg.target_layers == [1, 5, 19, 29, 41, 51] + assert cfg.target_layer_offset == 0 + assert cfg.has_embed_tokens is True + assert cfg.causal is False + assert cfg.rope_theta == 10000 + assert cfg.rope_scaling == { + "factor": 128.0, + "original_max_position_embeddings": 8192, + "rope_type": "yarn", + } + + +def test_nvidia_dflash_weight_names_map_to_existing_module_tree(): + normalized = normalize_dflash_weights( + { + "hidden_norm.weight": object(), + "layers.0.input_layernorm.weight": object(), + "layers.0.post_attention_layernorm.weight": object(), + "layers.0.self_attn.q_proj.weight": object(), + } + ) + + assert set(normalized) == { + "enc_norm.weight", + "layers.0.attn_norm.weight", + "layers.0.ffn_norm.weight", + "layers.0.self_attn.q_proj.weight", + } + + +def test_dflash_server_descriptor_uses_external_drafter_without_mtp_head(): + descriptor = descriptor_for_backend_id("dflash") + + assert descriptor.backend_id == "dflash" + assert descriptor.uses_external_assistant is True + assert descriptor.uses_draft_lm_head is False + assert descriptor.draft_semantics.request_field == "speculative_depth" + + +def test_generate_dflash_keeps_installed_block_and_reports_verification_stats(): + class FakeRuntime: + config = SimpleNamespace(block_size=8) + + def generate(self, prompt, *, max_tokens, stop_token_ids, token_callback): + assert prompt == [1, 2] + assert max_tokens == 4 + assert stop_token_ids == set() + assert self.config.block_size == 8 + return { + "text": "done", + "tokens": [3, 4, 5, 6], + "rounds": 2, + "accepted": 5, + "drafted": 14, + "rejected": 9, + "mean_accept": 2.5, + "tokens_per_target_step": 3.5, + } + + runtime = FakeRuntime() + output = generate_dflash( + runtime, + [1, 2], + max_tokens=4, + speculative_depth=3, + ) + + assert runtime.config.block_size == 8 + assert output.stats.accepted_drafts == 5 + assert output.stats.drafted_tokens == 14 + assert output.stats.rejected_drafts == 9 + assert output.stats.verify_calls == 2 + assert output.stats.speculative_depth == 7 + assert output.stats.requested_speculative_depth == 7 + + +def test_dflash_skips_unsupported_retokenized_session_postcommit(): + state = SimpleNamespace( + backend_descriptor=descriptor_for_backend_id("dflash") + ) + + skipped = openai._skipped_idle_postcommit_snapshot( + state=state, + unsafe_reason="missing_generation_final_state", + prompt_prefix_len=12, + ) + + assert skipped == { + "stored": False, + "mode": "skipped", + "reason": "dflash_retokenized_postcommit_unsupported", + "unsafe_reason": "missing_generation_final_state", + "assistant_tool_calls": 0, + "prompt_prefix_len": 12, + } + + +def test_dflash_request_uses_backend_block_default_not_native_mtp_depth(): + state = SimpleNamespace( + args=SimpleNamespace(depth=3), + backend_descriptor=descriptor_for_backend_id("dflash"), + ) + + block_size = openai._request_depth_for_generation( + state, + openai.ChatCompletionRequest(), + generation_mode="mtp", + ) + + assert block_size == 8 diff --git a/tests/test_modelopt_nvfp4.py b/tests/test_modelopt_nvfp4.py new file mode 100644 index 00000000..31c55ea9 --- /dev/null +++ b/tests/test_modelopt_nvfp4.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +import mlx.core as mx + +from mtplx.modelopt_nvfp4 import ( + dequantize_modelopt_fp8, + dequantize_modelopt_nvfp4, +) + + +def test_modelopt_nvfp4_dequantizes_nibbles_with_both_scales(): + # Low nibble first: 0x21 -> [0.5, 1.0], 0xCB -> [-1.5, -2.0]. + packed = mx.array([[0x21, 0xCB] + [0x00] * 6], dtype=mx.uint8) + block_scale = mx.array([[2.0]], dtype=mx.float32) + global_scale = mx.array(0.25, dtype=mx.float32) + + weight = dequantize_modelopt_nvfp4(packed, block_scale, global_scale) + + assert mx.allclose( + weight[:, :4], mx.array([[0.25, 0.5, -0.75, -1.0]]) + ).item() + + +def test_modelopt_fp8_dequantizes_with_exported_weight_scale(): + weight = mx.array([[1.0, -2.0]], dtype=mx.float32) + scale = mx.array(0.125, dtype=mx.float32) + + actual = dequantize_modelopt_fp8(weight, scale) + + assert mx.allclose(actual, mx.array([[0.125, -0.25]])).item()