Skip to content

Add Nemotron 3.5 Lightning DFlash serving (stacked on #248) - #250

Draft
davidtai wants to merge 2 commits into
youssofal:mainfrom
davidtai:feat/nemotron35-dflash
Draft

Add Nemotron 3.5 Lightning DFlash serving (stacked on #248)#250
davidtai wants to merge 2 commits into
youssofal:mainfrom
davidtai:feat/nemotron35-dflash

Conversation

@davidtai

Copy link
Copy Markdown
Contributor

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

uv run --frozen --with pytest pytest -q \
  tests/test_modelopt_nvfp4.py \
  tests/test_dflash_nemotron.py \
  tests/test_artifacts.py::test_dflash_pair_bundle_inspects_as_native_runtime \
  tests/test_artifacts.py::test_gemma4_pair_bundle_inspects_as_assistant_runtime \
  tests/test_compressed_tensors.py \
  tests/test_server_openai.py::test_gemma4_request_uses_draft_block_before_depth

uv run mtplx inspect <dflash-pair-bundle> --require-mtp --json
uv run mtplx serve --model <dflash-pair-bundle> --mtp

Local real-model verification used:

  • target nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4@e0b753dc24903ad4d62f5696077da22020eca89a
  • drafter nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4-DFlash@bbea7f89d6221e2242a2f00f958afce59d32fab7
  • 18 focused tests passed
  • DFlash matched independent greedy target AR for 32/32 generated tokens
  • standard mtplx serve returned 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 8

Notable decisions

  • Preserve NVIDIA's released block size 8 as a construction-time invariant; generic native-MTP depth defaults do not resize the installed DFlash lane per request.
  • Validate topology, capture layers, and quantization layout during construction and conversion, not inside the measured per-layer path.
  • Requantize ModelOpt NVFP4/mixed-FP8 weights to standard MLX affine 4-bit while retaining pinned source revisions in the converted configs.
  • Use snapshot/rollback plus committed-prefix repair for Nemotron-H's hybrid recurrent/attention cache after partial acceptance.
  • Keep the initial serving contract greedy-only. No non-greedy or performance-speedup claim is made.
  • Skip retokenized SessionBank postcommit for DFlash until this external backend exposes a safe final-state snapshot contract.

davidtai and others added 2 commits August 10, 2026 17:09
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
@youssofal

Copy link
Copy Markdown
Owner

Stack noted. Review lands after #248 does; the dc68a6c split makes the Nemotron leg easy to isolate.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants