NVIDIA-NeMo / NVIDIA-NeMo/RL

Speculative decoding in rollout: interest in Domino / DFlash / DSpark draft methods, and online-trained MTP?

Open
#3,301 0 comments 0 reactions 0 assignees View on GitHub
community-request Feature waiting-on-maintainers
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

## Context

The rollout path already supports vLLM-native speculative decoding via `policy.generation.vllm_kwargs.speculative_config`, with first-class code, recipes and tests for:

- **EAGLE3**, both a fixed offline draft and an online-trained draft (`policy.draft.enabled: true`, draft weights refit into vLLM each step), per `docs/guides/eagle3-speculative-decoding.md`.
- **DeepSeek MTP** (`method: deepseek_mtp` / `mtp`), with V3 recipes and tests.

I am the primary contributor of speculative decoding in NeMo Automodel (EAGLE3 / MTP / Domino / DFlash / DSpark drafters), and I would like to help extend the rollout side here. Before I open PRs I want to check scope and interest.

## Questions

**1. Online MTP-head training.** MTP is supported for inference (frozen draft refit), but the online drafter-training loop looks EAGLE3-specific (`EagleModel` wrapping the ModelOpt `EagleModule`). Is training an MTP head alongside the policy something you'd want, and would you accept a PR generalizing the `policy.draft` path to a second draft architecture?

**2. Domino / DFlash / DSpark.** These are Automodel-native draft methods and are currently absent here. Because the rollout spec-decode is a pass-through of vLLM's `speculative_config`, supporting them has a hard prerequisite: vLLM has no upstream inference path for these methods today, so an engine-side bridge (custom `method` + drafter loading) must land first. Is adding these in scope for this repo, or would you rather keep rollout limited to methods vLLM supports natively?

## Offer

I am happy to implement whichever of these you consider in scope, starting with a design note so we agree on the config surface (`speculative_config.method` values, and any `policy.draft` generalization) before code. Let me know the preferred direction.

Contributor guide

Open the contributing guide

Research direction

Start with docs/guides/eagle3-speculative-decoding.md and the existing EAGLE3 and DeepSeek MTP recipes and tests. Read the rollout policy.generation.vllm_kwargs.speculative_config and policy.draft paths, including EagleModel and the ModelOpt EagleModule, then determine whether online MTP training or new draft methods are in scope. Done means the maintainers agree on the design and config surface before implementation.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.