Train a DSpark drafter

DSpark builds on DFlash’s block drafting with two extra heads:

• A Markov head. DFlash predicts every position in a block independently, so position 5 cannot see what was drafted at position 4. The Markov head adds a low-rank dependency between neighboring tokens in the block, which raises accuracy at later positions, where DFlash is weakest.
• A confidence head. It predicts, for each position, whether the target will accept it. The serving engine can use that to stop verifying a block at the first position that is likely to be rejected.

This tutorial trains a DSpark drafter for Qwen3-8B from scratch, on the same dataset as Train an EAGLE-3 draft head.

The DSpark recipe

DSPARK = Recipe(
    method="dspark",
    init_from="",  # No permissively licensed DSpark head for Qwen3-8B: train from scratch.
    target_layers=(2, 10, 18, 26, 34),
    # A block of 8 positions: the bonus token plus 7 drafted tokens.
    num_speculative_tokens=7,
    lr=3e-4,
    epochs=5,  # From random weights: more passes over the data.
    train_args=(
        "--block-size", "8",
        # From scratch, every draft weight trains, with fp32 Adam state. Three decoder layers
        # and 256 anchors fit next to the hidden-state server on one 48 GB GPU.
        "--max-anchors", "256",
        "--num-layers", "3",
        "--draft-vocab-size", "32000",
        # The Markov head adds a low-rank dependency between neighboring block tokens.
        "--markov-rank", "256",
        "--markov-head-type", "vanilla",
        # The confidence head predicts acceptance, so serving can stop a draft early.
        "--enable-confidence-head",
        "--confidence-head-with-markov",
        "--confidence-head-alpha", "1.0",
        "--loss-fn", '{"ce": 0.1, "tv": 0.9}',
        "--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,
)

• Three draft layers and the DFlash target layers. With no head to load, the recipe sets the architecture: three decoder layers reading target layers 2, 10, 18, 26, and 34. Trained from scratch, every draft weight carries fp32 optimizer state, so three layers and 256 anchors per sequence fit next to the hidden-state server on one 48 GB GPU. On a 4-GPU node, use five layers and more anchors.
• A block of 8. Shorter blocks than DFlash’s 16 suit a drafter trained on little data, because accuracy at late positions is the last thing it learns. vLLM serves it with num_speculative_tokens=7.
• A reduced draft vocabulary. The draft head predicts over the 32,000 most frequent target tokens in the training data, which shrinks its output layer by about 5x. prepare-data computes the token frequencies.
• A mixed loss. Cross-entropy plus total variation distance, weighted 0.1 and 0.9. Total variation tracks acceptance more directly than cross-entropy does.
• A higher learning rate and more epochs. 3e-4 for 5 epochs, as for any model trained from random weights.

Train

@driver_env.task
async def main_dspark(
    num_chat: int = 4000,
    num_agent: int = 4000,
    num_shards: int = 4,
    max_steps: int | None = None,
) -> Dir:
    """Train a DSpark drafter from scratch 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="dspark", max_steps=max_steps)

The dataset is the same cached one as for EAGLE-3 and DFlash.

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

Training from scratch takes the longest of the three families: about 30,000 steps over five epochs, nearly seven hours on one L40S. The report shows the loss falling across all five epochs and which epoch’s checkpoint was kept:

Benchmark draft models compares the DSpark drafter with the public and fine-tuned EAGLE-3 and DFlash drafters.

Serve the DSpark drafter

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

python serve.py --method dspark

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

Choose a draft model family

• Start with EAGLE-3. It has the most mature tooling, a public head for most popular targets, and holds its speedup best as concurrency rises.
• Try DFlash for long accepted runs. Code, JSON, and low-concurrency serving reward block drafting. Fine-tune a public DFlash head where one exists.
• Train DSpark when you have the data. Its Markov and confidence heads improve on DFlash at later block positions, but a head trained from scratch needs far more data than a fine-tune to beat a public head.
• Measure on your own traffic. The benchmark in these tutorials runs every family on the same held-out prompts and GPUs. Rank them by acceptance length and speedup on your workload, not on a paper’s.