[Tracking] Speculative decoding: implemented features and remaining gaps

[Tracking] Speculative decoding: implemented features and remaining gaps

Snapshot of main as of 2026-07-07, updated 2026-07-22. One place to see what the speculative-decoding subsystem (nemo_automodel/components/speculative/, the train_* recipes, examples/speculative/) already ships, which PR added each piece, what is in flight, and what is not implemented yet.

Update 2026-07-10: #2953, #2959, #2966, and #2976 merged since the original snapshot.
Update 2026-07-11: fp8 draft training and LoRA/PEFT drafts (#2963, merged 2026-07-08) and draft-side context parallelism (#2918) are now landed, so they move to "Implemented". Sequence-packing coverage for the remaining drafts is now in flight (#3002, #3003, #3004, #3005, #3023). The "Not implemented" list below is corrected accordingly.
Update 2026-07-14: sequence packing for DeepSeek MLA EAGLE-3, EAGLE-1/2, DFlash, DSpark Qwen3, Domino, and JetSpec (#3002, #3003, #3004, #3005, #3023), DSpark training-time acceptance metrics (#2957), and the variable-prefix DFlash plus EAGLE-3 LK objectives (#3032) are now merged.
Update 2026-07-14: a design RFC for the disaggregated streaming target backend is now open at #3062 (see the "Not implemented yet" item below).
Update 2026-07-22: Gemma4 EAGLE-3 target support and its example configs (#3071, #3073, #3077, #3079), unified DFlash / Domino / JetSpec validation metrics (#3072), periodic real accept-length eval during training (#3037), the on-policy regeneration loop (#3042), and the fp8 draft convergence smoke (#3081) are merged, so they move to "Implemented" and their entries are removed from the gap list. The first VLM speculative training PRs (#3166, #3173) and the streaming data-plane series (#3084, #3085, #3094) are now in flight. The remaining gaps are re-ranked below.

1. Implemented

Algorithms

Algorithm	Recipe	Key PRs
EAGLE-1 / EAGLE-2	train_eagle1.py / train_eagle2.py	#2263, #2264, feature-noise augmentation #2470, TP #2829, sequence packing #3003
EAGLE-3 (TTT unroll, soft-CE, optional LK objective)	train_eagle3.py	recipe #2224, FA2 draft backend #2254, sequence packing #2444, target TP #2827, target CP #2465, training-time simulated accept-length metric #2959, LK objectives #3032
EAGLE-3.1 (fc_norm / norm_output)	train_eagle3.py toggles	#2322, config plumbing fix #2897
P-EAGLE (parallel drafting, COD)	train_eagle3.py (parallel_drafting: true)	#2376, sequence partitioning #2409, activation checkpointing #2458
DFlash (block-parallel draft, optional variable-prefix objective)	train_dflash.py	#2406, target TP #2830, sequence packing #3004, variable-prefix D2SD objective #3032, validation metrics #3072; D-PACE remains in review in #2572
Domino (DFlash + GRU correction head)	train_domino.py	#2819, TP gather fix #2899, sequence packing #3023, validation metrics #3072
JetSpec (block-causal DFlash + forward-KL)	train_jetspec.py	#2867, sequence packing #3023, validation metrics #3072
DSpark (semi-autoregressive + Markov / confidence heads)	train_dspark.py	#2810, DeepSeek-V4 #2866, MiniMax-M3 VL #2877, GLM-5.2 #2885, registry dispatch #2909, offline cache #2924, distributed multi-node offline precompute #2953, FSDP2-sharded dense target #2976, sequence packing #3005, acceptance metrics #2957 and #3050, SGLang-compatible target_layer_ids #3083
MTP	model-level, not a speculative/ recipe	nemotron_v3, qwen3_5/3_6, deepseek_v4, minimax_m3_vl, step3p7 all carry MTP heads (e.g. #2510, #2486, #2727)

Draft / target coverage

Algorithm	Registered targets
EAGLE-1/2	Llama, Phi-3 (#2312), Qwen3 dense (#2313), Qwen3-MoE (#2317)
EAGLE-3 / 3.1 / P-EAGLE	the four above, plus gpt-oss (#2399), DeepSeek-V3 MLA draft (#2849, eager-only, no packing), and Gemma4 (#3071, with E2B / E4B / 31B / 26B-A4B example configs in #3073, #3077, #3079)
DFlash / Domino / JetSpec	Qwen3 dense + Qwen3-MoE only
DSpark	Qwen3 dense/MoE, DeepSeek-V4 (#2866), GLM-5.2 (#2885), Gemma4 (#2909 registry; FSDP2-sharded dense target for large models, validated on Gemma4-31B, #2976), MiniMax-M3 VL (#2877, text draft with multimodal data path)

Low-precision and adapter drafts

• fp8 draft training across every spec-decode recipe (#2963): the SFT fp8: block is accepted by EAGLE-1/2, EAGLE-3 / P-EAGLE, DFlash / Domino / JetSpec, and DSpark, reusing components/quantization/fp8.py. A shared apply_draft_fp8 helper swaps the draft's nn.Linear to torchao Float8Linear before the DDP / FSDP2 wrap; DSpark's FSDP2 path adds enable_fsdp_float8_all_gather; the frozen target is never converted. emulate: true allows numerics testing on pre-Hopper GPUs. Convergence on fp8-capable (SM89+) hardware is now covered by the smoke in #3081.
• LoRA / PEFT draft adaptation for EAGLE-3 (#2963): the SFT peft: block trains only lora_A / lora_B on a frozen base draft, with the checkpointer on the adapter-only PEFT path. recipe_args.draft_weights_path warm-starts the draft from a consolidated safetensors export. Unsupported combinations fail fast (peft: on DFlash / Domino / JetSpec / DSpark / EAGLE-1/2, or with parallel_drafting / fp8: in EAGLE-3, raises).

Data pipeline

• Target backends for EAGLE-3 training: colocated (default), remote target on separate GPUs over HTTP + NCCL (#2398, #2459, #2479), SGLang target backend (#2449), vLLM target backend (#2798).
• Offline precomputed hidden-state cache: EAGLE-3 (precompute_eagle3.py, #2373 plus follow-ups), opt-in top-k target_probs compression (#2847); DSpark (precompute_dspark.py, #2924; distributed multi-node sharded precompute for targets that OOM on one node, #2953).
• Dataset regeneration against the serving target: offline (regenerate.py, #2307; reasoning-mode control #2444; chat_template_kwargs fix #2921) and online (regen_loop.py, #3042), where a detached worker on a reserved GPU regenerates shards on a step cadence and all data-parallel ranks swap the dataloader in lockstep at the next epoch boundary. ShareGPT auto-conversion (#2375).
• Draft-vocab compression (d2t/t2d): token-map caching #2371, remap embedded in the checkpoint #2447.
• Sequence packing now covers dense and DeepSeek MLA EAGLE-3, EAGLE-1/2, DFlash, DSpark Qwen3, Domino, and JetSpec (#2444, #3002, #3003, #3004, #3005, #3023). Remaining packing gaps are DSpark non-Qwen3, multimodal and offline paths, plus composition with draft CP.

Parallelism

Algorithm	Draft	Target	Notes
EAGLE family	DDP; context parallelism (ring, zig-zag option) #2918	FSDP2, TP (#2827/#2829), CP (#2465)	TP and CP are colocated-backend only and mutually exclusive; CP excludes packing
DFlash family	DDP	FSDP2, TP (#2830)	no CP, no EP for MoE targets
DSpark	FSDP2 or DDP	FSDP2 + EP (DeepEP); optional FSDP2-sharded frozen dense target for large models that don't fit replicated (#2976)	multi-node EP validated (DeepSeek-V4 ep=32, GLM-5.2 ep=64); PP must be 1

Training-time metrics

• EAGLE-3: simulated accept length train/val tau_sim from per-TTT-step accuracies (#2959).
• DSpark: acceptance and confidence metrics (#2957), with unmeasured positional rates omitted rather than reported as zero (#3050).
• DFlash / Domino / JetSpec: token-weighted validation loss and accuracy reduced across ranks, block-weighted acceptance length, Domino final-head and base-head diagnostics, optional W&B logging (#3072).
• All families: periodic real acceptance length (train/tau_real, #3037). On a step cadence, rank 0 snapshots the draft, a detached worker on a reserved GPU converts it through serve_vllm, boots a vLLM server, and runs the bench_vllm workload, so the engine's own spec-decode counters are logged alongside the proxy metrics without the trainer importing vllm.

Serving and benchmarking

Engine	Servable drafts	Tooling
SGLang	EAGLE-1/2/3/3.1	serve_sglang.py (#2294), bench_sglang.py (#2372)
vLLM	EAGLE-3, P-EAGLE (#2841), DFlash + JetSpec (#2913)	serve_vllm.py, bench_vllm.py (#2912)
neither	Domino (GRU head has no engine runtime), DSpark (runtime not yet released upstream)

Multi-dataset acceptance-length benchmark sweep across standard eval sets, aggregated reporting (#2966).

Docs: subsystem README #2448, refresh covering DSpark/Domino/JetSpec/vLLM #2914, fp8 and LoRA draft sections #2963.

2. In flight

• #3166 multimodal EAGLE draft core (modality-aware draft over projected image embeddings, frozen VLM feature wrapper, hidden-state and soft-token supervision)
• #3173 ViSpec two-stage vision-language draft training for Qwen2.5-VL and Qwen3-VL targets, including multimodal target-logit regeneration (tracked by #3172, validation plan in #3167)
• #3084, #3085, #3094 the disaggregated streaming data plane from the #3062 RFC: data-plane contracts and local feature store, EAGLE-3 producer and loader, async producer with a shared-directory feature store
• #2572 DFlash D-PACE loss
• #2802 nemotron_v3 MTP validation acceptance metrics
• #2656 qwen3_5 dense MTP SDPA + packing fix

3. Not implemented yet

Roughly ranked by user impact.

• DFlash-family targets beyond Qwen3. DFlash / Domino / JetSpec register only Qwen3ForCausalLM and Qwen3MoeForCausalLM. Every other target family, including the newer Qwen generations and non-Qwen dense and MoE models, needs a draft class and a registry entry before these three algorithms are usable outside one model family. This is now the narrowest coverage in the subsystem.
• More EAGLE-3 target families. The registry covers 7 architectures after #3071. Mixtral, GLM, Llama4, Qwen3-Next, QwQ, and a small MLA smoke target such as DeepSeek-V2-Lite still have no draft support. VLM targets are being handled by #3166 and #3173, so far scoped to the Qwen-VL family.
• train-with-decode. #3042 regenerates data online against the frozen target, and #3037 periodically measures the real acceptance length of a draft snapshot in a live engine, but neither trains on draft-produced rollouts. The missing mode is a closed loop where the draft proposes, the target verifies, the draft weights are hot-synced into the engine on a cadence, and the training signal comes from the accept / reject outcome.
• Acceptance metrics for EAGLE-1/2. #2959, #2957, and #3072 gave every other family a simulated or block-weighted acceptance metric. train_eagle1.py and train_eagle2.py still log only loss and top-1 accuracy, so their runs cannot be compared against the other recipes. Optionally add best-checkpoint selection keyed on the acceptance metric.
• Remaining sequence-packing gaps. DSpark packing is gated to the Qwen3 draft (the recipe raises for any other DSpark draft class), the multimodal and offline-cache paths carry no packing metadata, and packing is mutually exclusive with context parallelism in every recipe's _validate_packing_gates.
• DSpark / Domino inference. Neither is servable on a stock engine; vLLM ships a dflash method (used by #2913) but no DSpark runtime is released yet, and the Domino GRU correction head has no engine runtime at all. This means neither algorithm can report a real end-to-end speedup today. Mostly an upstream-engine dependency, tracked here for completeness.
• A third serving backend and sub-fp8 drafts. serve_sglang and serve_vllm are the only export paths; there is no TensorRT-LLM target. Draft quantization stops at fp8, with no nvfp4 or int4 draft path.
• Draft-side long-context parallelism beyond #2918. The merged draft CP is a hand-written ring; a Ulysses + ring (USP-style) option would compose better with packing and non-Llama drafts.
• Online large-vocab memory optimization. Top-k target_probs compression (#2847) covers the offline cache; the online path still materializes full-vocab target distributions.
• Weights-only warm start across draft families. Generalize draft_weights_path beyond EAGLE-3 and reset optimizer, scheduler, RNG, and training progress when starting a new run from existing draft weights.
• Engine-side offline capture. precompute_eagle3.py and precompute_dspark.py run the target through the AutoModel forward. Capturing hidden states from a serving engine instead would raise precompute throughput and keep the cached features bit-consistent with what the same engine serves at inference time.
• A first-class MTP training path. MTP heads exist at the model level in several architectures, but there is no speculative/ recipe, no acceptance evaluation, and no shared tooling with the draft families. #2802 covers the validation metrics for one model.

Validation debt

These are verification gaps, not missing user-facing features.

• Add checkpoint-load and speculative output-equivalence smokes for each advertised serving engine. tests/unit_tests/speculative/test_serve_vllm.py and test_serve_sglang.py cover the config rewrite in isolation; nothing asserts that an exported draft loads in the engine and produces the expected accepted tokens.
• Add colocated, remote, and offline target-feature parity gates.

Recently closed (kept for history)

• Gemma4 EAGLE-3 target support delivered by #3071, with E2B, E4B, 31B, and 26B-A4B example configs in #3073, #3077, and #3079.
• Consistent validation metrics across the DFlash family delivered by #3072 (token-weighted distributed validation, block-weighted acceptance length, Domino base-head diagnostics, optional W&B logging).
• Periodic real acceptance-length eval during training delivered by #3037 (train/tau_real from a detached vLLM worker on a reserved GPU).
• On-policy data regeneration during training delivered by #3042 (step-cadence regeneration worker plus lockstep dataloader swap).
• fp8 draft convergence validation on SM89+ delivered by #3081.
• fp8 draft training and LoRA/PEFT drafts delivered by #2963 (merged 2026-07-08).
• Training-time acceptance metrics delivered for EAGLE-3 by #2959 and for DSpark by #2957, refined by #3050.
• Draft-side context parallelism delivered by #2918 (ring attention, zig-zag option).
• Sequence packing landed for the EAGLE-3 dense draft (#2444), DeepSeek MLA EAGLE-3 (#3002), EAGLE-1/2 (#3003), DFlash (#3004), DSpark Qwen3 (#3005), and Domino plus JetSpec (#3023).
• Variable-prefix DFlash and EAGLE-3 LK objectives delivered by #3032. D-PACE remains in review in #2572.

Happy to take any of these; comments and corrections welcome.

Choose one specific gap first, then read its named recipe entry point, relevant registry entries, and the examples/speculative/ area. This tracking issue does not name a single bounded change or test; done would require implementing the selected capability and updating this status list.