MiMo-V2.5 + DFlash and a 4-bit NVFP4 KV cache in one vLLM instance — on the v0.24.0 release (2× DGX Spark)

Accelerated Computing DGX Spark / GB10 User Forum DGX Spark / GB10
1

Follow-up to the earlier MiMo-V2.5-DFlash thread. I got DFlash speculative decoding running together with a 4-bit NVFP4 KV cache in a single vLLM instance on a 2× DGX Spark (GB10) pair — and ported the whole thing to the v0.24.0 release. Full build + mods + recipe:

👉 GitHub - DoctorMasterNewb/vLLM-MiMo-V2.5-DFlash-NVFP4Kv · GitHub

Why this is awkward in stock vLLM

DFlash is a block-diffusion drafter — it fills a block of 8 masked tokens in one forward using non-causal attention. Two consequences:

  1. The draft’s KV registers into vLLM’s global paged allocator, which insists on one physical page size across layers. A bf16 draft page won’t unify with a 4-bit NVFP4 target page — the wall behind upstream #41559 (“DFlash fundamentally incompatible with all KV-cache quantization”).
  2. So --kv-cache-dtype nvfp4 + DFlash blows up.

The way around it: run the DFlash drafter as a vLLM custom_class proposer — a standalone module whose KV never enters the global allocator. The NVFP4 target keeps its full pool; the drafter runs decoupled beside it (attention forward reimplemented to match qwen3_dflash: partial rotary dim-64 @ θ=5e6, attention sinks, value-scale, mask-embedding sidecar, dense per-request context buffer). Decoupled → works with any target KV dtype.

What v0.24.0 changed (the useful part for everyone)

v0.24.0 upstreams most of the old MiMo mod stack: MiMoV2Omni model + config/arch registry, DiffKV backend (#41797), MXFP8 kernels, native DFlash / custom_class. What still needs a startup mod for the lukealonso/MiMo-V2.5-NVFP4 export:

mod why
fix-mimo-config transformers still doesn’t know model_type=mimo_v2
fix-mimo-qkv-mxfp8 fused qkv is `[Q_all
fix-mimo-…-merger vision merger mlp bias=True (checkpoint ships the bias → else KeyError)
fix-mimo-…-packed packed_modules_mapping on the Omni wrapper (else MXFP8 read as bf16 → garbage)
nvfp4-kv-diffkv NVFP4 KV for DiffKV; backend drops in, the triton kernel needed one line (v0.24.0’s compute_kv_seq_mask gained a seq_len arg)
dflash-custom-proposer the decoupled proposer — applies unchanged from the old build

Maintainer takeaway: on a version bump, most MiMo mods upstream or apply unchanged. The churn is checkpoint-specific quant loaders + triton arg drift.

Numbers (2× GB10, TP=2, 0.24.1.dev0)

Structured (temp 0): ~2.37 acc/block — JSON 3.33 (45 tok/s), code 2.81 (38), math 2.41 (32).

Deep context (llama-benchy, prose):

depth tg tok/s (peak) acc/block
short 21.7 (31) 1.13
100k 7.85 (12) 1.02
200k 5.03 (10.5) 1.17
  • NVFP4 KV pool ≈ 1.35M tokens (~2.8× fp8). Needle recall at 26k; coherent throughout.
  • Prose acceptance is lower — DFlash collapses on unpredictable prose, measure your own workload.
  • Honest tradeoff vs fp8: fp8-KV DiffKV builds decode deep context faster (~100k@12 / 200k@10) since fp8 attention is cheaper than 4-bit; NVFP4 buys the bigger pool at some decode cost. Both now coexist with DFlash via the decoupled proposer, so pick per use case.

Repo

GitHub - DoctorMasterNewb/vLLM-MiMo-V2.5-DFlash-NVFP4Kv · GitHub — Dockerfile (v0.24.0 GB10, sm_121a/CUDA 13.2), all seven mods, serve recipe, README with the full breakdown. Lineage based on eugr/spark-vllm-docker; fp8-KV DiffKV approach cross-referenced from this forum — thanks all. Experimental; validate before leaning on it.

3 Likes