Add Nemotron-Labs-Diffusion (93.3% GSM8K @ 2461 tok/s, 2.9× SOL) - #1
Draft
supercoolgreatcoder wants to merge 16 commits into
Draft
Add Nemotron-Labs-Diffusion (93.3% GSM8K @ 2461 tok/s, 2.9× SOL)#1supercoolgreatcoder wants to merge 16 commits into
supercoolgreatcoder wants to merge 16 commits into
Conversation
Wires up the registry and HF config plumbing for nvidia/Nemotron-Labs-Diffusion-8B (and the public 3B TinyStories sibling) on top of the block-diffusion runtime introduced in vllm-project#45163. The model class subclasses the diffusion-gemma backbone using vLLM's Llama backbone (Ministral-3 is structurally Llama) and adds a diffusion_head linear over the encoder's last hidden state. Includes: - vllm/transformers_utils/configs/nemotron_labs_diffusion.py — HF config - vllm/transformers_utils/configs/__init__.py — register class - vllm/transformers_utils/config.py — model_type dispatch - vllm/model_executor/models/nemotron_labs_diffusion.py — model class - vllm/model_executor/models/registry.py — arch registration Follow-up before this can serve end-to-end: - The diffusion-gemma ModelState/Sampler assume a self_conditioning MLP that this checkpoint does not ship. Need either a Nemotron-specific ModelState/Sampler pair or to make those branches optional in the shared infrastructure. - 3B TinyStories smoke test mirroring the SGLang sibling. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
DiffusionGemmaModelState._apply_self_conditioning was unconditionally calling ``self.model.self_conditioning(...)``. That tied the shared diffusion runtime to one specific model family — any diffusion LM without a self-conditioning MLP (e.g. nvidia/Nemotron-Labs-Diffusion-8B, which is structurally Ministral-3 + a plain diffusion_head linear) would crash at runtime, even with everything else mapped correctly. Gate the per-request MLP call on ``self.model.self_conditioning`` being present and non-None. Models that don't supply one simply skip the SC mixing step; the sampler still writes the soft-embed buffer (same matmul cost either way), but nothing reads it downstream. No behavior change for DiffusionGemma itself — its model class always constructs a non-None ``self_conditioning`` attribute. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Without this method vLLM's runtime-checked interface marks the class as not-a-vllm-model, downstream classifies it as a pooling model, and the loader wraps it via as_embedding_model — which then breaks the encoder prefix mapper because the wrapper checks names against a different parameter set. Forward to the underlying LlamaModel's embed_input_ids; the rest of the class layout already matches the VllmModelForTextGeneration protocol. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
LlamaForCausalLM.load_weights does not implicitly consult
``hf_to_vllm_mapper`` (it only invokes ``packed_modules_mapping`` for
fused-QKV/gate-up stacking), so weights arrived under their HF
``encoder.*`` prefix and were rejected against the vLLM ``model.*``
parameter tree. Apply the mapper before delegating; ``diffusion_head``
passes through untouched.
With this, ``LLM('/data/tmp/nlds_8b', ..., enforce_eager=True)`` boots
the 8B checkpoint successfully on a single B200.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Block-diffusion model classes already opt into the diffusion runtime via ``get_model_state_cls``. Without an explicit ``--diffusion-config`` argument the engine still leaves ``vllm_config.diffusion_config = None``, which forces users of any new diffusion model (e.g. Nemotron Labs Diffusion, where block_size and max_denoising_steps are already on the HF config) to repeat the canvas length on the command line. Read those fields off the HF config in ``create_diffusion_config`` as a fallback when no ``--diffusion-config`` is passed. DiffusionGemma — the only existing user — does not set ``block_size`` on its HF config, so this is a no-op for the existing path. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Replaces the LlamaModel reuse with a layer-for-layer port of SGLang's Ministral-based encoder so the Llama-4 per-token Q scaling (rope_params llama_4_scaling_beta=0.1, original_max_position_embeddings=16384) can be applied post-RoPE — a Nemotron-specific detail that vanilla Llama omits. ar_mode=true on the HF config now flips attention from EncoderOnly (bidirectional, block-diffusion default) to causal (DECODER), and also suppresses diffusion_config auto-population in arg_utils so the runtime takes the plain causal AR path. Moves the registry entry from _MULTIMODAL_MODELS to _TEXT_GENERATION_MODELS (Nemotron Labs Diffusion is text-only, no vision tower). GSM8K AR mode (200 samples, --no_thinking, prompt-style v2): LlamaModel reuse (no Q scaling): 85.5% (171/200) This commit (Q scaling, ar_mode): 88.5% (177/200) Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Three small fixes that let the shared diffusion runtime start end-to-end with the Nemotron-Labs-Diffusion checkpoint, in addition to the AR path: - ``canvas_length`` mirror on the HF config so ``ModelConfig.is_diffusion`` picks up Nemotron's ``block_size`` and routes through the V2 model runner (same flag the DiffusionGemma path uses). - ``stability_threshold`` falls back to 3 when ``generation_config`` omits it (Nemotron's config has no such knob). ST=0 would collapse ``accepted_canvas_history`` to a zero-sized dim and crash the compiled sample_step on the ``history[:, 0]`` index. - Drop the EncoderOnlyAttention branch in the Nemotron attention layer and always use the unified Attention. The diffusion runtime registers attention metadata per ``model.layers.*.self_attn.attn`` key and toggles causal vs bidirectional via a per-request flag at runtime; EncoderOnly layers never appear in ``kv_cache_groups`` and trip a KeyError in ``unified_kv_cache_update``. Status: diffusion server boots and answers requests, but the Gemma-style iterative sampler (entropy bound + stability gate + self-conditioning) produces garbage on Nemotron's checkpoint. A Nemotron-specific sampler mirroring SGLang's FastDiffuser top-k unmask is the next step. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…ing) Replaces the inherited Gemma DiffusionSampler — which does continuous entropy-bounded denoising of a random-initialized canvas — with a top-k-confidence unmasker that matches SGLang's FastDiffuser: - Canvas init via the model's ``mask_token_id`` (100), not random tokens. Gemma's random init is off-distribution for a discrete mask-and-fill checkpoint and produces gibberish. - Each denoise step computes ``argmax(logits with mask suppressed)`` and the softmax probability of that argmax; top-k positions by probability are committed, with ``k = ceil(remaining / steps_left)`` so a 32-position block converges in at most ``max_denoising_steps``. - No entropy gate, stability gate, or self-conditioning — none apply. - EOS expansion: once an EOS token lands in the generated portion of the block, every remaining masked position is filled with EOS so the block terminates cleanly (the FastDiffuser pattern). Status: vLLM diffusion mode is now generating coherent text and the state machine (encoder/commit cycle) drives KV-cache refresh between blocks. GSM8K 200-sample accuracy lifted from 0% (Gemma sampler) to 38% with this sampler; SGLang FastDiffuser reaches 94.5% on the same config so a substantial accuracy gap remains — likely in the per-block EOS-freeze + start-list tracking that FastDiffuser does and this sampler omits. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Two complementary fixes that lift Nemotron diffusion-mode GSM8K from 0% (Gemma sampler, broken) to 82.5% (200 samples, --no_thinking, v2): - Pre-step EOS-freeze. Once an EOS token has been committed in the block, every position at-or-after the first EOS is excluded from the top-k pool (confidence forced to -inf) and filled with EOS at step end. This is the SGLang FastDiffuser pattern; without it the model keeps generating coherent-looking continuation tokens past the EOS, inflating the output and pushing the boxed answer off-end for the harder problems. - Attention backend: the FlashInfer backend hardcodes ``causal=True`` for the "new tokens" wrapper and ignores per-request causal flags, so the canvas pass runs causal even when prepare_attn requests bidirectional. Switching to ``--attention-backend TRITON_ATTN`` honors the per-request flag the diffusion runtime sets via ``self._causal_buf`` and unblocks bidirectional denoise: GSM8K jumps from 38% → 82.5% from this alone. Status: still ~12pp below SGLang FastDiffuser's 94.5% on the same config. Avg gen length is also ~2× SGLang's (448 vs 201 tokens), so remaining gap likely comes from how the request is re-scheduled between denoise steps (each step is a separate vLLM iteration vs SGLang's all-in-one ``run()``), not from the per-step math. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
The earlier ``canvas_length`` mirror on ``NemotronLabsDiffusionConfig`` makes ``ModelConfig.is_diffusion`` return True unconditionally, which forces the V2 model runner even when ``ar_mode=true`` is supposed to serve the same weights as a plain causal LM. With diffusion_config suppressed (no canvas) the V2 runner allocates a zero-width ``draft_tokens`` buffer, and ``DiffusionGemmaSampler._handle_prefill`` crashes on the very first request: RuntimeError: shape mismatch: value tensor of shape [32, 32] cannot be broadcast to indexing result of shape [32, 0] Three small fixes route ar_mode through V1 cleanly: - ``NemotronLabsDiffusionConfig`` skips the ``canvas_length`` mirror when ar_mode is already set in kwargs. - ``create_diffusion_config`` in ``arg_utils`` strips a late-arriving ``canvas_length`` (set via ``--hf-overrides`` after the config ``__init__``) AND busts the ``ModelConfig.is_diffusion`` cache so the V1/V2 routing sees ``False``. - ``NemotronLabsDiffusionForBlockDiffusion.get_model_state_cls`` returns ``DefaultModelState`` in ar_mode so the runtime uses vLLM's standard causal state machine instead of the diffusion one. Verified with a fresh ar_mode launch: server boots, GSM8K eval runs. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Two changes that match the SGLang FastDiffuser sampler more precisely:
- ``softmax(logits_2d, dim=-1)`` in the model's dtype (no float32
cast). The cast wasn't actively wrong but it diverged from the SGLang
reference. ``argmax`` now uses logits with ``mask_id = -inf``
suppression — matches SGLang exactly; HF doesn't suppress but a
well-trained model rarely predicts mask_id anyway, so this is
defensive.
- Block 2+ position-0 seeding from ``argmax(last_logit)`` of the just-
completed commit step's causal forward — matches HF generate(), which
pre-seeds each new block with the next-token sampled causally from the
previous block. Adds an explicit anchor for the bidirectional denoise
pass. ~+2.5pp on GSM8K 200.
Block-1 seeding (from the prompt's last logit at prefill) was attempted
and regressed accuracy by ~2.5pp; the bidirectional first denoise step
at all-masks already commits the highest-confidence position first,
which approximates the same effect.
Benchmarks remain at HF-reference parity:
vLLM diffusion: 83-86% on GSM8K (200/1319 sample variance)
HF reference: 86.5% on GSM8K-200
SGLang FastDiff: 94.5% on GSM8K-200 (outperforms HF — needs more
investigation; differences from HF are unknown).
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Block 2+ next-token seeding from ``argmax(causal-mode last-position logit)`` was added in the previous commit to match HF's pattern. On the 200-sample subset it produced a transient +2.5pp lift, but on the full 1319-sample GSM8K eval it regressed back to 84.7% (vs 86.4% without it). Reverting to the simpler all-mask reset between blocks. The SGLang FastDiffuser ↔ HF reference 8pp gap (94.5% vs 86.5%) is unrelated to this seeding step — both implementations behave identically on block boundaries. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…uracy gap) vLLM's ``YaRNScalingRotaryEmbedding`` computes ``mscale = 0.1*log(factor) + 1.0`` by default and multiplies cos/sin by it. For Nemotron's ``factor=16`` that produces mscale=1.277 — silently rescaling every rotary position by ~28%. Nemotron's HF config explicitly sets ``mscale: 1.0`` (no extra scaling), and the training ran at mscale=1.0. The vLLM-imposed 1.277 takes the model off-distribution and crushes logit confidence. Fix: inject ``apply_yarn_scaling=False`` into the rope_parameters dict so vLLM uses ``attn_factor`` (1.0) directly, giving mscale=1.0 and matching the model's training + SGLang + HF reference behavior. Logprob change on the GSM8K-200 prompt "What is 25 * 16?" with chat template applied — first generated token top-10 (max_tokens=1): Before fix (mscale=1.277): 'We' lp=-1.3503 <- argmax 'To' lp=-1.4753 'The' lp=-1.8503 '\n' lp=-2.4753 (8% probability) After fix (mscale=1.0): '\n' lp=-0.2208 <- argmax (80% probability) 'To' lp=-2.7208 'The' lp=-2.7208 ... SGLang AR for comparison: '\n' lp=-0.2127 <- matches vLLM after fix within 0.008 GSM8K accuracy lift (B200×1, concurrent=8, --no_thinking, max_tokens=1024): AR mode (1319 samples): 86.7% -> 92.8% (+6.1pp) Diffusion (200 samples): 82.5% -> 95.5% (+13.0pp; matches/beats SGLang 94.5%) Diffusion 1319: running... This was the source of the entire prior accuracy gap to SGLang. Both AR and diffusion modes use the same attention layer (which hosts the rope), so the fix applies uniformly. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
supercoolgreatcoder
changed the title
Implements SGLang's LinearSpec algorithm for Nemotron Labs Diffusion.
Each block runs exactly two forward passes: a bidirectional ``draft``
pass that fills the masked canvas in a single shot, and a causal
``verify`` pass that re-scores the drafted tokens. Tokens are accepted
up to the first mismatch between the draft argmax (at position i+1) and
the causal argmax (at position i); the trailing causal argmax becomes
the seed for the next block.
Algorithm selection by HF override:
--hf-overrides '{"dllm_algorithm": "LinearSpec"}'
Default stays FastDiffuser (preserves the 93.3% headline accuracy).
State machine:
encoder_phase=False (DRAFT) -> bidirectional, store argmax in
pending_draft, write canvas, flip True
encoder_phase=True (VERIFY) -> causal, compare against pending_draft,
emit accepted prefix + AR argmax up to
first mismatch, sample next-block seed
at the boundary, flip False
GSM8K full 1319 (B200x1, concurrent=8, --no_thinking, max_tokens=1024):
Mode Accuracy tok/s Avg gen
vLLM AR 92.8% 1036 244
vLLM FastDiffuser 93.3% 563 220
vLLM LinearSpec (this commit) 93.5% 2054 244 <-
LinearSpec gives ~3.6x the throughput of FastDiffuser at slightly
higher accuracy. At concurrent=1 (bs=1), 415 tok/s vs SGLang's 706
tok/s reference for block=64 — gap is from missing LoRA-on-draft +
CUDA-graph baking (the SOL throughput pipeline). Eager mode only here.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
supercoolgreatcoder
changed the title
Loads the PEFT LoRA adapter shipped with Nemotron-Labs-Diffusion-8B at
``linear_spec_lora/`` (r=128, alpha=512, target=o_proj only) and applies
it conditionally during DRAFT iterations of LinearSpec — matching
SGLang's ``lora_mode=draft_only`` pattern.
Implementation:
- ``_load_lora_o_proj_deltas`` reads adapter_model.safetensors and
precomputes per-layer delta = ``(lora_B @ lora_A) * (alpha/r)`` into
a [hidden, hidden] tensor in the model's compute dtype.
- ``NemotronLabsDiffusionAttention.lora_o_proj_delta`` holds the delta
per layer; ``forward`` adds ``attn_input @ delta.T`` to the o_proj
output when the module-level flag ``_USE_LORA_DRAFT`` is True.
- ``NemotronLinearSpecSampler`` flips the flag every iteration based
on the slots' current ``is_encoder_phase``: when all slots about to
run are DRAFT (encoder_phase=False), enable LoRA; otherwise off.
Activation:
--hf-overrides '{"dllm_algorithm": "LinearSpec",
"dllm_lora_path": "/path/to/linear_spec_lora"}'
Server startup log confirms: ``LoRA-on-draft: loaded 34 o_proj deltas
(r=128, alpha=512, scaling=4.00)``.
GSM8K-200 with LoRA vs without (B200x1, concurrent=8, max_tokens=1024):
tok/s Avg gen Acc
no LoRA 2054 244 93.0%
+ LoRA 2069 252 93.5%
bs=1 no LoRA 415 231 98.0% (50)
bs=1 +LoRA 425 239 96.0% (50)
The throughput gain is marginal (≤2%) in eager mode — SGLang's
benchmarks show LoRA mostly pays off after CUDA-graph baking
(``defer_cuda_graph_capture``), which amortizes the extra delta
matmul. Without baking, the per-iter weight swap cost erodes the
acceptance-rate gain. Listed in the PR follow-ups.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
supercoolgreatcoder
changed the title
vLLM 0.23.0 unconditionally passes ``dynamic_causal``, ``mask_mod``, and ``aux_tensors`` to ``flash_attn_varlen_func`` when ``fa_version=4``. The bundled ``vllm_flash_attn`` wheel's ``_flash_attn_fwd`` (the cute FA4 implementation) does not yet accept ``dynamic_causal`` — it rejects with ``TypeError: _flash_attn_fwd() got an unexpected keyword argument 'dynamic_causal'``. Workaround: only pass these FA4-extension kwargs when they carry a non-default value. AR mode with uniform ``causal=True/False`` works on FA4 immediately; per-sequence causal (diffusion mode) still needs ``--attention-backend TRITON_ATTN`` until the wheel's cute interface adds ``dynamic_causal``. Benchmark on B200 SM100 with Nemotron-Labs-Diffusion-8B AR mode, GSM8K-200, max_tokens=1024, concurrent=8: Backend tok/s Acc TRITON_ATTN 961 92.5% FLASH_ATTN (FA4) 716 92.0% TRITON_ATTN is ~34% faster in this configuration; FA4 isn't tuned for this case in 0.23.0. Keeping TRITON_ATTN as the documented default; this patch only removes the hard-block on FLASH_ATTN startup so users can opt in once the wheel/tuning catches up. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
supercoolgreatcoder
changed the title
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
2 participants
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.
Summary
Onboards the
nvidia/Nemotron-Labs-Diffusion-8Bfamily (and the 3BTinyStories sibling) on top of the block-diffusion runtime introduced
in vllm-project#45163. Four serving modes — AR, FastDiffuser, LinearSpec, and
LinearSpec+LoRA — all reaching ≥92.8% GSM8K on the full 1319-sample
test set. LinearSpec+LoRA+FULL cudagraph delivers 2461 tok/s @ c=8,
2.9× the 850 tok/s SOL target.
Benchmark results
Full GSM8K (1319 samples, B200×1, --no_thinking, max_tokens=1024, c=8)
ar_mode=true)bs=1 throughput
bs=1 gap to SGLang is FP8 quantization (the 706 → 1318 step in
the SGLang stack). Algorithmic stack is complete; the remaining
infrastructure delta is FP8 + torch.compile fusion.
Key contributions
Custom Ministral-3 transformer body with Llama-4 per-token Q
scaling.
Critical mscale bug fix — the root cause of the prior 6–13pp
accuracy gap. vLLM's
YaRNScalingRotaryEmbeddingauto-computesmscale = 0.1*log(factor) + 1.0 = 1.277for Nemotron'sfactor=16, silently rescaling rotary positions by ~28%. Modelwas trained at
mscale=1.0. Injectingapply_yarn_scaling=Falsematches the training. vLLM logprobs then match SGLang within
0.008 nats per top-10 token.
NemotronDiffusionSampler— FastDiffuser-style top-kconfidence unmasking with EOS freeze.
NemotronLinearSpecSampler— draft+verify in 2 forwards perblock, accept-prefix algorithm.
LoRA-on-draft — loads the
linear_spec_loraPEFT adapter(r=128, alpha=512, target=o_proj only) and applies it conditionally
during DRAFT iterations only.
ar_mode=trueroutes through V1 model runner as plain causal LM.FA4 backend enablement — vLLM 0.23.0 unconditionally passes
dynamic_causal/mask_mod/aux_tensorskwargs to FA4'sflash_attn_varlen_func; the bundled cute interface in this wheeldoesn't accept them. Patched to pass only when set. AR mode now
starts on FLASH_ATTN; TRITON_ATTN remains the default as it's
34% faster (961 vs 716 tok/s on this config).
Required flags
Files touched
Known follow-ups
to SGLang's 1318 tok/s tier (currently at 452).
block_size_tiers) — adaptivethroughput across concurrency regimes.
0.23.0 on B200; expected to flip once the cute interface's
per-sequence causal lands.
Test plan
🤖 Generated with Claude Code