Add Nemotron 3.5 Lightning DFlash serving (stacked on #248) - #250
Draft
davidtai wants to merge 2 commits into
Draft
Add Nemotron 3.5 Lightning DFlash serving (stacked on #248)#250davidtai wants to merge 2 commits into
davidtai wants to merge 2 commits into
Conversation
Muse-Glimmer-30B (Gemma-family multimodal; text tower served) as an MLX arch, plus a generic dflash block-diffusion speculative backend that drafts for it. Muse-Glimmer text tower (vendored_muse_glimmer_text.py, arch from llama.cpp src/models/muse-glimmer.cpp): sigmoid gated attention, parameter-free QK-norm with qk_scale_factor folded onto Q, NoPE on global layers, Gemma (1+w) sandwich norms, RMS-normed embeddings, logit softcap. Registered via a sys.modules shim (muse_glimmer_patch.py) + ARCHITECTURE_CATALOG/family-gate entries (registry.py) so `mtplx serve` accepts it. qwen3_vl-style multimodal wrapper drops the vision stack (vendored_muse_glimmer.py). Decode optimization: bit-exact QKVG fusion (default) + optional decode-only fused dense-SwiGLU kernel (env-gated, graceful fallback). Ledger in docs/. dflash backend (models/dflash.py + backends/dflash.py): config-driven DFlashDrafter (block-diffusion drafter, forward bit-exact vs the bstnxbt/ dflash-mlx reference) + DFlashRuntime doing propose -> single-forward verify -> target-prefix accept -> cache-normalize rollback -> incremental context cache. Decode is token-exact vs greedy AR (up to fp near-tie non-determinism). Wired through the serve path: runtime.load() resolves a dflash pair bundle (dflash_pair.py) to a DFlashRuntime, and generate_mtpk() dispatches to it. Adding a future dflash drafter is a drop-in bundle (target/ + drafter/ + dflash_pair.json), no code. Draft: measured ~2.5x on predictable code, ~1.0x (break-even) on diverse code (acceptance-bound); best block size tracks acceptance. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01AzsPk4X2dkhrNRUDuCE1K1
Owner
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stack dependency
This PR is stacked on #248, which supplies the Glimmer DFlash backend used as the integration base. Until #248 merges, GitHub will show both commits in this PR. The Nemotron-specific change is
dc68a6c0.What changed
Adds native MTPLX serving support for NVIDIA Nemotron 3.5 Lightning 30B-A3B with NVIDIA's released DFlash drafter. The change adds a streaming NVIDIA ModelOpt NVFP4/mixed-FP8 to MLX affine-4 conversion path, maps the released DFlash config and weights onto the existing Glimmer adapter, installs a fixed construction-time capture route for Nemotron-H's hybrid Mamba/attention trunk, and integrates DFlash pair inspection, backend metadata, server defaults, telemetry, and tests.
Why
NVIDIA publishes the Lightning target and DFlash assistant in ModelOpt formats and with a hybrid Nemotron-H target topology. The existing Glimmer adapter establishes the block-diffusion algorithm, but the Lightning checkpoint needs its own quantization conversion, direct layer-output taps, YaRN/config mapping, and hybrid-cache rollback/repair to preserve exact greedy target-prefix verification.
How to verify
Local real-model verification used:
nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4@e0b753dc24903ad4d62f5696077da22020eca89anvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4-DFlash@bbea7f89d6221e2242a2f00f958afce59d32fab7mtplx servereturned HTTP 200 from health, models, and chat endpoints; the chat request recorded 112 drafted tokens, 16 accepted drafts, 16 verify calls, and requested block size 8Notable decisions