Train a DFlash block drafter

EAGLE-3 drafts one token per forward pass of its draft layer, so a draft of 3 tokens costs 3 passes. DFlash drafts a whole block in one pass. It is a small block-diffusion model: starting from the last accepted token and the target’s hidden states, it predicts every position in a block of 16 in parallel. The target then verifies the block in one pass, as it does any draft.

When the target accepts long runs, a block drafter pays for one draft pass where EAGLE-3 pays for several. That favors code, structured output, and low-concurrency serving. When the target rejects early in the block, the extra positions are wasted work.

This tutorial fine-tunes Red Hat’s public DFlash head for Qwen3-8B, on the same dataset as Train an EAGLE-3 draft head.

The DFlash recipe

Each drafter family is a recipe: the head to start from, the target layers it reads, the draft length to serve with, and its training flags.

@dataclass(frozen=True)
class Recipe:
    method: str  # speculators' --speculator-type, and vLLM's speculative method
    init_from: str  # Hugging Face head to fine-tune; "" trains from scratch
    target_layers: tuple[int, ...]  # Target layers the drafter reads
    num_speculative_tokens: int  # Draft tokens per verification step when serving
    lr: float
    epochs: int
    train_args: tuple[str, ...] = field(default_factory=tuple)  # Method-specific flags
    vllm_memory: float = 0.45  # Share of a shared GPU for the hidden-state server
    seq_len: int = 4096  # Training sequence length, in tokens

DFLASH = Recipe(
    method="dflash",
    init_from="RedHatAI/Qwen3-8B-speculator.dflash",
    target_layers=(2, 10, 18, 26, 34),
    # A block of 16 positions: the bonus token plus 15 drafted tokens.
    num_speculative_tokens=15,
    lr=1e-4,
    epochs=2,
    train_args=(
        "--block-size", "16",
        # Anchors are the positions a block is drafted from; fewer anchors per sequence
        # trade throughput for memory on a 48 GB GPU.
        "--max-anchors", "512",
        # D-PACE weights each block position by how often it can be accepted.
        "--loss-fn", "ce",
        "--per-position-loss-weight", "dpace",
        "--gradient-checkpointing",
    ),
    # Block drafters predict every block position over the vocabulary, so on a shared 48 GB
    # GPU they train on shorter sequences to keep the logits in memory.
    seq_len=2048,
)

• Five target layers, not three. DFlash reads layers 2, 10, 18, 26, and 34 of Qwen3-8B. The hidden-state server captures exactly the recipe’s layers, and the trainer reads the same list, so the two can never disagree.
• 15 draft tokens. A block of 16 positions starts with the token the target already produced, so vLLM serves it with num_speculative_tokens=15.
• D-PACE loss weighting. Later positions in a block are accepted only when every earlier one was, so they matter less. D-PACE weights each position’s loss by how often it can be accepted.
• Sized for one 48 GB GPU. The hidden-state server takes 45% of the GPU, so the trainer gets the rest. The public head predicts over Qwen3’s full 151,936-token vocabulary, and the logits for one 4,096-token sequence alone take 2.3 GB in fp32. The recipe trains on 2,048-token sequences, uses 512 anchors per sequence (positions a training block is drafted from), and checkpoints activations. On a 4-GPU node, raise all three.

Train

train_drafter is the same task as for EAGLE-3. The method argument selects the recipe:

@driver_env.task
async def main_dflash(
    num_chat: int = 4000,
    num_agent: int = 4000,
    num_shards: int = 4,
    max_steps: int | None = None,
) -> Dir:
    """Fine-tune the public DFlash head on the shared dataset."""
    # Same inputs as the EAGLE-3 pipeline, so every data shard is a cache hit.
    dataset = await build_dataset(num_chat, num_agent, num_shards)
    return await train_drafter(data=dataset.files, method="dflash", max_steps=max_steps)

build_dataset has the same inputs as in the EAGLE-3 pipeline, so if you ran that first, every data shard is a cache hit and this pipeline starts at training.

python dflash.py --smoke   # 20 training steps
python dflash.py           # the full run

The training report charts draft accuracy at each block position. Expect accuracy to fall steadily across the block, and even position 0 to trail EAGLE-3’s: in the full run, 78% against EAGLE-3’s 95%. DFlash drafts every position of the block in one pass, without seeing its own earlier guesses. It wins on speed by proposing 15 tokens per step, not by drafting each one better.

Benchmark draft models compares the fine-tuned drafter with the public DFlash head and with every other family.

Serve the DFlash drafter

serve.py builds the app for every family from its recipe:

python serve.py --method dflash

Serve and retrain draft models covers the app, the agent that sends it traffic, and nightly retraining. Every step there takes --method dflash.