How I Trained a DFlash Drafter for Speculative Decoding

If you prefer video format

How I Trained a DFlash Drafter for Speculative Decoding

Running a capable local LLM is often easy. Making it responsive enough for interactive use, coding, or agent workflows is much harder.

I wanted to improve the decode throughput of a Qwen3.8 27B target model running in my local setup. Rather than changing the target model, I trained a small DFlash draft model and used speculative decoding. This post focuses on the practical workflow: how to prepare training data, choose an online or offline pipeline, train and inspect a drafter, export it, serve it, and benchmark it honestly.

This is not a universal copy-paste recipe. Exact package names, flags, configurations, and supported architectures move quickly. Treat the commands as a concrete template, pin the versions that work in your environment, and verify the current documentation for your chosen framework.

What we are building

A speculative decoding system uses two models:

• A target model. This is the model whose output behavior you want to preserve.
• A draft model, or drafter. This smaller model proposes several next tokens quickly.

At generation time, the drafter proposes a block of tokens. The target model verifies that proposal. It accepts the matching prefix and corrects the first disagreement. Because the target can verify multiple proposed tokens together, speculative decoding can improve output throughput without changing the target model's output distribution.

The important point is that a drafter is not a general-purpose replacement for the target model. It is a specialized accelerator coupled to a particular target model and serving implementation.

Why DFlash is interesting

Many speculative approaches draft autoregressively: token one, then token two, then token three. A DFlash drafter uses block diffusion to predict a block of future tokens in one parallel forward pass.

A simplified view looks like this:

Prompt
  |
  +--> Target model produces context features
  |
  +--> DFlash drafter proposes a token block in parallel
  |
  +--> Target model verifies the proposed block
  |
  +--> Accept matching prefix and continue

DFlash uses the target model's internal context features to help the small draft model make better proposals. In particular, its design fuses hidden states from the target and injects that context into the drafter's attention keys and values. This lets the drafter spend its capacity on predicting the next block rather than reconstructing the entire context from scratch.

This is why two metrics must always be considered together:

• Acceptance. How many consecutive draft tokens the target accepts.
• Draft speed. How quickly the drafter can create proposed blocks.

A slow but accurate drafter may not help. A fast drafter with poor acceptance may add overhead. The metric that matters at the end is measured end-to-end output throughput under a realistic workload.

Prerequisites and planning

Before starting a run, confirm these items.

1. A compatible target model

Your target must be supported by the training code and the serving runtime. The draft model is architecture-specific in practice: tokenizer vocabulary, hidden sizes, RoPE behavior, attention layout, and target hidden-state interfaces all matter.

Do not assume that a drafter trained for one Qwen version can be paired with another. Pair the exact target checkpoint named by the draft model or training recipe.

2. A suitable GPU environment

Drafter training can be memory-bandwidth and storage intensive. You will commonly have all of these on disk at once:

• The target model
• Source prompt data
• Regenerated or distilled responses
• Training checkpoints
• Exported Hugging Face artifacts
• Logs and temporary caches

Provision more storage than your first estimate. For large models, 1 TB is a safer starting point than trying to expand a rented instance halfway through a run.

Use a persistent remote session for every long operation:

tmux new -s dflash-train

# Detach with Ctrl-b then d
# Reattach later
tmux attach -t dflash-train

Back up progress outside the rented machine at regular intervals:

rsync -avP /workspace/dflash-runs/ user@your-local-host:/data/dflash-backups/

3. A reproducible software stack

Do not blindly install the latest version of every library. The intersection of PyTorch, CUDA, Transformers, FlashAttention, vLLM, SGLang, and a speculative-decoding training framework can be fragile.

Start with an isolated environment and record the working versions immediately.

python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip

python --version
pip freeze | tee requirements.lock.txt
nvidia-smi

If the project provides a Docker image, development container, or pinned lockfile, prefer it. A known-good container can save days of dependency debugging.

The training pipeline

At a high level, the workflow has five stages:

1. Collect prompts
2. Generate target-model responses
3. Train the drafter
4. Export the checkpoint
5. Serve and benchmark

The training data should resemble the work you actually ask the target model to do. A coding assistant benefits from coding prompts. An agent server benefits from tool calls, structured output, planning traces, and the system prompts it sees in production.

A generic instruction dataset can produce a usable drafter, but task-aligned prompts are more likely to improve acceptance on your own workload.

Step 1: Prepare prompt data

A simple starting format is JSONL: one JSON object per line.

{"messages":[{"role":"user","content":"Write a Python function that retries an HTTP request with exponential backoff."}]}
{"messages":[{"role":"user","content":"Explain the difference between Kubernetes Deployment and StatefulSet."}]}

For chat models, preserve the structure that the target model expects. Do not flatten messages into ad hoc strings unless you also control and validate the exact chat template.

Here is a small Python script that samples a larger JSONL dataset reproducibly:

import json
import random
from pathlib import Path

source = Path("prompts-all.jsonl")
destination = Path("prompts-sample-30k.jsonl")
sample_size = 30_000
seed = 42

lines = source.read_text().splitlines()
random.Random(seed).shuffle(lines)
selected = lines[:sample_size]

with destination.open("w") as f:
    for line in selected:
        json.loads(line)  # Fail early if the input is malformed.
        f.write(line + "\n")

print(f"Wrote {len(selected)} examples to {destination}")

Before spending GPU money, inspect examples manually. Look for malformed conversations, duplicate prompts, giant contexts, private data, template artifacts, and prompts that do not represent your intended workload.

head -n 3 prompts-sample-30k.jsonl | jq .
wc -l prompts-sample-30k.jsonl

Use the target model to create labels

The drafter should learn the behavior of the target model, not the behavior of whatever dataset originally supplied an answer. Send each prompt to the target model and save the target's generated response. This is the distillation stage.

First launch an OpenAI-compatible target server. This SGLang command is illustrative; use the version and flags validated for your target model.

python -m sglang.launch_server \
  --model-path Qwen/Qwen3.8-27B \
  --host 0.0.0.0 \
  --port 30000 \
  --tp-size 1 \
  --dtype bfloat16 \
  --mem-fraction-static 0.80

Then run the framework's data-regeneration script in a second terminal. The command below shows the inputs you generally need: source prompts, output path, target endpoint, concurrency, generation settings, and resumability.

python scripts/regenerate_train_data.py \
  --input prompts-sample-30k.jsonl \
  --output distilled-qwen3.8-27b.jsonl \
  --base-url http://127.0.0.1:30000/v1 \
  --model Qwen/Qwen3.8-27B \
  --concurrency 64 \
  --temperature 0.7 \
  --max-tokens 2048 \
  --resume

The most important practical feature here is --resume or its equivalent. Distillation can run for many hours. Spot instances can disappear, servers can crash, and local networks can drop. A resumable output turns a disaster into an interruption.

If your target has a supported built-in speculative mode, enabling it during data generation may reduce the time needed to build the distilled dataset. Validate that doing so preserves the target behavior you want to imitate.

Step 2: Choose offline or online training

There are two broad ways to feed target-model information into a drafter.

Offline training

Offline training separates the work into two phases:

Distilled data
  -> Run target model and save hidden states
  -> Train drafter from saved hidden states

Typical workflow:

python scripts/prepare_hidden_states.py \
  --model Qwen/Qwen3.8-27B \
  --input distilled-qwen3.8-27b.jsonl \
  --output-dir hidden-states-qwen3.8-27b

torchrun --nproc_per_node=8 train_dflash.py \
  --config configs/qwen3.8-27b-dflash.yaml \
  --data-path distilled-qwen3.8-27b.jsonl \
  --hidden-state-path hidden-states-qwen3.8-27b

Advantages:

• Easier mental model: prepare features, then train.
• Target and drafter do not have to share live runtime memory during the training step.
• Easier to retry only one stage.

Disadvantages:

• Hidden states can be very large.
• You pay for a separate extraction pass.
• Data preparation can become the bottleneck.

Online training

Online training keeps the target and drafter live together. The training process obtains target features as it trains, often through a shared KV-cache or shared-memory mechanism.

Distilled data
  -> Live target model
  -> Shared context features or KV cache
  -> Live drafter training

A generic layout looks like this:

# Terminal 1: start the target model with the shared-cache integration.
python -m sglang.launch_server \
  --model-path Qwen/Qwen3.8-27B \
  --port 30000 \
  --enable-shared-kv-cache

# Terminal 2: run training.
torchrun --nproc_per_node=8 train_dflash.py \
  --config configs/qwen3.8-27b-dflash-online.yaml \
  --data-path distilled-qwen3.8-27b.jsonl \
  --target-base-url http://127.0.0.1:30000/v1

Advantages:

• Avoids a separate hidden-state extraction workflow.
• Can use a large GPU more efficiently.
• Faster iteration when the environment is already working.

Disadvantages:

• More moving parts.
• Greater runtime memory demand.
• More sensitive to version compatibility between the training process, serving runtime, and shared cache.

My practical rule is simple: choose offline when you want the least complicated path or have limited runtime memory. Choose online when you have enough GPU headroom and are willing to spend time validating the shared-cache setup.

Step 3: Configure DFlash training

A DFlash configuration usually controls the drafter architecture, block size, target feature layers, optimizer settings, dataset handling, and checkpoint cadence.

Here is an illustrative YAML configuration. Field names vary by framework, so use it as a map of the decisions you need to make.

model:
  target_model: Qwen/Qwen3.8-27B
  drafter_type: dflash
  num_hidden_layers: 5
  hidden_size: 2048
  num_attention_heads: 16
  num_key_value_heads: 4

training:
  output_dir: runs/qwen3.8-27b-dflash
  max_steps: 20000
  per_device_train_batch_size: 1
  gradient_accumulation_steps: 8
  learning_rate: 0.0001
  warmup_ratio: 0.03
  bf16: true
  gradient_checkpointing: true
  save_steps: 1000
  logging_steps: 20

sequence:
  max_length: 4096
  block_size: 16
  num_anchors: 256

loss:
  self_logit_distillation: true
  loss_decay_factor: 7.0
  answer_only_loss: true

data:
  train_file: distilled-qwen3.8-27b.jsonl
  chat_template: templates/qwen3-train.jinja

The parameters that deserve attention

Block size

Block size is the number of token positions in a draft block. Larger blocks create more opportunity for acceleration, but later tokens are generally harder to predict and may be rejected.

Start with the block size recommended by the implementation and the corresponding reference checkpoint. Do not change it casually at serving time. The drafter is trained around a particular block structure.

Drafter depth and size

A larger drafter may improve acceptance, but it consumes more memory and can reduce draft speed. DFlash is attractive because it can predict a block in parallel, but that does not eliminate the cost of a deeper network.

Start with an upstream architecture configuration that is known to export and serve correctly. Only tune drafter capacity after you have a benchmark harness.

Number of steps

A small experimental run is useful for validating the pipeline, but it is not a production training schedule. In my experiment, I used a short staged run: an initial phase, then a lower learning rate phase. That was enough to prove the full path worked, but not necessarily enough to maximize generalization.

Save checkpoints often enough that a failed rental instance does not waste a day of work:

training:
  save_steps: 1000
  save_total_limit: 3

Assistant-only loss and chat templates

If you only want to train on assistant responses, ensure that the data pipeline can identify those tokens accurately. Some frameworks require chat templates with explicit generation markers in order to produce an assistant token mask.

Do not assume a model's default template has the required markers. Test tokenization before beginning a large run.

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3.8-27B")
messages = [
    {"role": "user", "content": "Explain a binary search."},
    {"role": "assistant", "content": "Binary search repeatedly halves a sorted search space."},
]

encoded = tokenizer.apply_chat_template(
    messages,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
    return_assistant_tokens_mask=True,
)

print(encoded.keys())
print(encoded.get("assistant_masks"))

If no assistant mask is created, fix the template before training. Otherwise, you may waste substantial compute learning to predict user messages, system prompts, or formatting tokens that do not improve generation behavior.

Step 4: Launch and monitor training

A typical multi-GPU launch uses torchrun.

torchrun \
  --standalone \
  --nproc_per_node=8 \
  train_dflash.py \
  --config configs/qwen3.8-27b-dflash-online.yaml

For distributed jobs, log the exact command, config, git commit, container tag, and hardware details alongside every run.

git rev-parse HEAD | tee runs/qwen3.8-27b-dflash/git-commit.txt
pip freeze | tee runs/qwen3.8-27b-dflash/requirements.txt
nvidia-smi -q | tee runs/qwen3.8-27b-dflash/nvidia-smi.txt
cp configs/qwen3.8-27b-dflash-online.yaml runs/qwen3.8-27b-dflash/

What to monitor

At minimum, watch these signals:

• Training loss should trend downward without exploding or becoming NaN.
• Draft token accuracy should generally rise.
• GPU utilization should be high enough to justify the rental.
• Memory usage should leave a small safety margin.
• Checkpoints should be written successfully.
• Data loading should not starve the GPU.

Accuracy in a training log is useful, but it is not the final serving metric. A per-token training accuracy number does not directly tell you how many consecutive tokens will be accepted in a live speculative loop.

Also watch for overfitting. An implausibly perfect score on a narrow dataset can mean the drafter memorized a distribution that does not match your real prompts.

A simple live log filter can help during debugging:

tail -f runs/qwen3.8-27b-dflash/trainer.log | \
  grep --line-buffered -E "step|loss|acc|accuracy|grad_norm|nan"

Step 5: Export to Hugging Face format

Training checkpoints are often framework-specific. Serving runtimes usually need a Hugging Face compatible export containing at least a config.json and one or more model.safetensors files.

A typical export command looks like this:

python scripts/export_hf_checkpoint.py \
  --model_path runs/qwen3.8-27b-dflash/checkpoint-20000 \
  --export_path exports/qwen3.8-27b-dflash

Inspect the exported directory before attempting to serve it:

find exports/qwen3.8-27b-dflash -maxdepth 2 -type f -printf "%f\n" | sort
ls -lh exports/qwen3.8-27b-dflash
jq . exports/qwen3.8-27b-dflash/config.json | head -n 60

Important configuration fields vary by implementation, but generally include architecture metadata, vocabulary size, block size, attention dimensions, and the target pairing assumptions.

Do a cheap load test before copying a large checkpoint to another machine:

from transformers import AutoModel

model = AutoModel.from_pretrained(
    "exports/qwen3.8-27b-dflash",
    trust_remote_code=True,
    device_map="cpu",
)
print(type(model).__name__)

Step 6: Serve the target and drafter

Serving integration changes quickly. Always begin with the current documentation for your runtime and DFlash implementation.

An example vLLM-style configuration commonly looks like this:

vllm serve Qwen/Qwen3.8-27B \
  --tensor-parallel-size 1 \
  --gpu-memory-utilization 0.90 \
  --speculative-config '{
    "method": "dflash",
    "model": "/models/qwen3.8-27b-dflash",
    "num_speculative_tokens": 15
  }'

An SGLang-style deployment can look like this:

python -m sglang.launch_server \
  --model-path Qwen/Qwen3.8-27B \
  --speculative-algorithm DFLASH \
  --speculative-draft-model-path /models/qwen3.8-27b-dflash \
  --tp-size 1 \
  --dtype bfloat16 \
  --trust-remote-code

The draft block is often represented as an anchor token plus masked positions internally. A runtime flag named num_speculative_tokens may count only proposed tokens, while the training block_size can include the anchor. Check the implementation's terminology and use the reference serving command where possible.

Test correctness before measuring speed:

curl http://127.0.0.1:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "Qwen/Qwen3.8-27B",
    "messages": [{"role": "user", "content": "Give me three concise Linux troubleshooting tips."}],
    "temperature": 0,
    "max_tokens": 128
  }'

Look for failed remote-code imports, incompatible model configs, tokenizer mismatches, cache-layout errors, and unsupported attention backends. Solve those before chasing benchmark numbers.

Benchmark the right way

Do not benchmark one prompt once and call the result a speedup. A useful benchmark compares baseline and speculative runs under the same conditions.

Keep these fixed:

• Target model and quantization
• Hardware and GPU count
• Runtime version and container image
• Prompt dataset
• Prompt length distribution
• Output length limit
• Sampling parameters
• Concurrency
• Warm-up procedure
• Measurement window

Record both decode throughput and request-level behavior. Depending on your use case, useful outputs include:

• Output tokens per second
• Time to first token
• End-to-end latency
• Accepted draft tokens per verification step
• Acceptance length
• GPU memory use
• Error rate

Here is a small benchmark client that sends prompts concurrently and estimates generated output tokens. It is not a replacement for a full benchmark tool, but it is useful for a controlled smoke test.

import asyncio
import time
from openai import AsyncOpenAI

client = AsyncOpenAI(base_url="http://127.0.0.1:8000/v1", api_key="not-needed")
prompts = [
    "Write a Python function to merge two sorted lists.",
    "Explain how a Kubernetes readiness probe differs from a liveness probe.",
    "Show a PostgreSQL query that finds duplicate email addresses.",
]

async def run_one(prompt):
    started = time.perf_counter()
    response = await client.chat.completions.create(
        model="Qwen/Qwen3.8-27B",
        messages=[{"role": "user", "content": prompt}],
        temperature=0,
        max_tokens=256,
    )
    elapsed = time.perf_counter() - started
    text = response.choices[0].message.content or ""
    approx_tokens = len(text.split()) * 1.3
    return elapsed, approx_tokens, text

async def main():
    results = await asyncio.gather(*(run_one(p) for p in prompts))
    total_seconds = sum(elapsed for elapsed, _, _ in results)
    total_tokens = sum(tokens for _, tokens, _ in results)
    print(f"Approximate aggregate throughput: {total_tokens / total_seconds:.2f} tokens/sec")

asyncio.run(main())

For serious measurements, use the tokenizer rather than an approximate token count, use enough prompts to smooth out variance, and run the exact same harness against baseline and speculative servers.

A result table should disclose enough context to be interpreted:

Target model: Qwen3.8 27B
Drafter: custom DFlash checkpoint
Hardware: one GPU, exact model recorded in run metadata
Prompt set: held-out mixed coding and agent prompts
Temperature: 0
Max output tokens: 512
Concurrency: 1

Mode                    Output tokens/sec    Notes
Baseline                14.0                 No speculative decoding
DFlash via SGLang       18.5                 Runtime version recorded
DFlash via vLLM         20.0                 Runtime version recorded

The numbers themselves matter less than the methodology. If the runtime versions or kernels differ, say so. If you compiled a runtime from a pull request, say so. If the speculative path only helps at a particular concurrency or output length, show it.

Common failure modes

The run dies after an SSH disconnect

Use tmux or a similar terminal multiplexer. Make data generation and checkpointing resumable. Keep logs on disk.

You run out of disk

Check disk space before and during the run:

df -h
du -sh /workspace/* | sort -h

Checkpoints, model caches, and hidden states can grow much faster than expected. Clean only files you fully understand. Never delete the latest known-good checkpoint until the export has loaded successfully.

Version mismatches

Symptoms include broken imports, serialization failures, CUDA kernel errors, incompatible model classes, and runtime errors while serving an exported checkpoint.

Fix this by pinning working versions, recording git commits, and rebuilding from a known-good environment. Avoid partial upgrades in the middle of an experiment.

High training accuracy but weak serving gain

This is the classic drafter failure mode. Possible causes include:

• The training prompts do not match serving prompts.
• The drafter is too slow relative to its acceptance benefit.
• The draft block is too long for the model's quality level.
• You measured a training proxy rather than live acceptance.
• The serving runtime is not using the expected optimized speculative path.
• The benchmark is dominated by prefill or short responses rather than decode.

Treat the entire pipeline as the system under optimization, not just the checkpoint.

Serving works but results differ across runtimes

vLLM and SGLang may have different implementations, kernels, versions, cache behavior, batching choices, and speculative support maturity. Benchmark each runtime separately, record exact versions, and avoid claiming a direct runtime comparison unless the setups are truly matched.

Practical next steps

Once the generic pipeline works, the most interesting improvement is data specialization.

A general-purpose drafter learns from a broad prompt distribution. But a local assistant may spend most of its time doing only a few things:

• Coding and debugging
• Tool calling
• Shell commands
• Structured JSON responses
• Infrastructure troubleshooting
• A fixed set of system prompts

You can log and sanitize your own requests and target outputs, deduplicate them, split them into train and held-out evaluation sets, and continue training a strong base drafter on that distribution.

Be careful with privacy. Strip secrets, tokens, personal data, customer data, and proprietary code before any dataset leaves your trusted environment. Keep a fixed held-out set so you can distinguish real generalization from memorization.

Final checklist

Before publishing or relying on a custom drafter, verify the following:

[ ] Target and draft checkpoints are an intended compatible pair
[ ] Training and serving dependency versions are recorded
[ ] Prompt data resembles the real serving workload
[ ] Distillation is resumable and outputs have been validated
[ ] Chat template and assistant-only loss behavior have been tested
[ ] Checkpoints are backed up off the rented instance
[ ] Hugging Face export loads successfully
[ ] Baseline and speculative servers use the same benchmark harness
[ ] Benchmark settings are documented
[ ] You measure end-to-end throughput, not acceptance alone
[ ] You have tested held-out prompts, not only training-like examples

Resources

• DFlash project and code
• DFlash project page
• DFlash paper
• NVIDIA Model Optimizer DFlash training documentation
• SGLang
• vLLM

Closing thoughts

Training a drafter is a useful systems exercise because it touches the whole LLM stack: data design, target-model behavior, GPU infrastructure, distributed training, checkpoint formats, serving runtimes, and benchmark methodology.

The best lesson is not that every model needs a custom drafter. It is that speculative decoding is only valuable when the complete system improves. Build a reliable pipeline first, measure the real workload, and then specialize the drafter where it makes sense.