Serve and retrain draft models
A drafter trained on last month’s traffic drafts worse for this month’s. New tools, new prompt templates, and new response formats lower its acceptance length, and with it the speedup, while outputs stay correct. Nothing fails loudly; responses get slower.
The fix is a loop. The application’s own model calls are the training data, because every response the target served is, by definition, what the target says. This page serves a drafter of any family, runs an agent against it, and retrains each family’s drafter every night on the agent’s traces.
Serve the target with a drafter
drafter_app builds a VLLMAppEnvironment for any family. It passes the drafter as
draft_model_path, and the family’s speculative method and draft length from its recipe:
# The pipeline task whose latest output each family's app serves.
PIPELINES = {
"eagle3": "specdec-driver.main",
"dflash": "specdec-driver.main_dflash",
"dspark": "specdec-driver.main_dspark",
}
def drafter_app(method: str) -> VLLMAppEnvironment:
"""A vLLM app serving Qwen3-8B with the latest drafter of one family."""
return VLLMAppEnvironment(
name=f"qwen3-8b-{method}",
image=gpu_image,
model_hf_path=TARGET_MODEL,
model_id=TARGET_MODEL,
draft_model_path=flyte.app.RunOutput(type="directory", task_name=PIPELINES[method]),
speculative_config={
"method": method,
"num_speculative_tokens": RECIPES[method].num_speculative_tokens,
},
resources=ONE_GPU,
# Scale to zero after 30 idle minutes: an idle app holds no GPU.
scaling=flyte.app.Scaling(replicas=(0, 1), scaledown_after=1800),
requires_auth=False,
secrets=API_KEY,
extra_args=[*TOOL_ARGS, "--max-model-len", "8192", "--api-key", "$VLLM_API_KEY"],
)
APPS = {method: drafter_app(method) for method in PIPELINES}
llm_app = APPS["eagle3"]
-
One app per family.
qwen3-8b-eagle3,qwen3-8b-dflash, andqwen3-8b-dsparkserve the same target with different drafters, so you can route traffic to the family that wins your benchmark. -
flyte.app.RunOutputresolves to theDirreturned by the latest successful run of the family’s training pipeline, so each app serves whatever that pipeline last trained.--run <name>pins a specific run. -
Both models load from disk. With a draft model set, the plugin turns off weight streaming and downloads both models to the container, so give the app enough
disk. -
The endpoint is public with an API key.
requires_auth=Falseopens the URL, and vLLM’s--api-keychecks every request against thespecdec-api-keysecret. Create it first:flyte create secret specdec-api-key --project <project> --domain <domain>Enter any random string when prompted.
Deploy an app:
python serve.py # EAGLE-3
python serve.py --method dflash # or dsparkThe first start takes several minutes: vLLM loads both models and captures CUDA graphs for the draft and verify passes.
Record the agent’s traces
The application is a tool-calling agent that runs as a task and calls one of the apps.
depends_on deploys every app with the task environment. Each scales to zero, so an app
holds no GPU until traffic reaches it:
agent_env = flyte.TaskEnvironment(
name="specdec-agent",
image=cpu_image,
resources=flyte.Resources(cpu=2, memory="4Gi"),
secrets=API_KEY,
# Deploys every app with the agent. Each scales to zero until it gets traffic.
depends_on=list(APPS.values()),
)
Every model call is a flyte.trace. The trace records each call’s inputs and output in the
run, so you can inspect what the agent asked and what the model answered in the UI. The job
also returns its episodes as a file:
@flyte.trace
async def call_model(endpoint: str, messages: list[dict], tools: list[dict]) -> dict:
"""One model call. As a trace, its inputs and output are recorded and visible in the UI."""
import os
import httpx
async with httpx.AsyncClient() as client:
complete = http_complete(client, endpoint, api_key=os.environ["VLLM_API_KEY"])
return await complete(messages, tools)
@agent_env.task(report=True)
async def agent_job(
num_episodes: int = 100, method: str = "eagle3", seed: int = 0, concurrency: int = 16
) -> File:
"""Serve agent requests through one family's app and return the episodes as traces."""
endpoint = APPS[method].endpoint
tasks = load_agent_tasks("traffic", num_episodes, seed=seed)
semaphore = asyncio.Semaphore(concurrency)
latencies: list[float] = []
async def one(task: dict) -> dict | None:
async with semaphore:
start = time.perf_counter()
try:
trace = await run_agent(partial(call_model, endpoint), task)
except Exception as e:
print(f"episode failed: {e}")
return None
latencies.append(time.perf_counter() - start)
return trace
traces = [t for t in await asyncio.gather(*(one(t) for t in tasks)) if t]
path = "/tmp/traces.jsonl"
with open(path, "w") as f:
for trace in traces:
f.write(json.dumps(trace) + "\n")
await flyte.report.replace.aio(
data_reports.traffic_report(traces, len(tasks) - len(traces), latencies, title=f"Agent traffic through the {method} app")
)
await flyte.report.flush.aio()
return await File.from_local(path)
The job’s report shows what the agent did: episodes, failures, tool calls, latency per episode, the most called tools, and a rendered example. Send traffic through an app:
python serve.py --agent 500 # through the EAGLE-3 app
python serve.py --method dflash --agent 500 # through the DFlash appEach model call appears in the run as its own call_model trace, next to the report. This
run sent 250 episodes through the EAGLE-3 app in under a minute:
These are the traces the retraining loop consumes. Every response in them is the target’s, whichever drafter served it, so one stream of traffic trains every family. In your own application, any task or app that logs its prompts, tool schemas, and the model’s responses in OpenAI chat format can feed the same loop. Training data needs no hidden states: the trainer recomputes them from the text.
Traces contain whatever your users send. Filter or redact personal data before traces reach training, and keep trace files under the same access controls as your logs.
Retrain nightly, promote only on improvement
retrain runs every night once per family, from 03:00 UTC, an hour apart so their GPU jobs
do not start at once. It finds the agent runs from the past day with flyte.remote, and the
family’s production drafter from its newest successful retrain or training pipeline run:
async def collect_traces(since: datetime, until: datetime) -> list[dict]:
"""Read the trace files returned by every successful `agent_job` run in the window."""
traces = []
runs = remote.Run.listall.aio(
in_phase=("succeeded",),
task_name="specdec-agent.agent_job",
created_at=remote.TimeFilter(after=since, before=until),
limit=500,
)
async for run in runs:
trace_file = File.from_existing_remote(await _output_uri(run))
with open(await trace_file.download()) as f:
traces.extend(json.loads(line) for line in f)
return traces
async def latest_drafter(method: str) -> Dir:
"""The family's drafter in production: its newest retrain output, else its pipeline's."""
async for run in remote.Run.listall.aio(
in_phase=("succeeded",), task_name="specdec-flywheel.retrain",
sort_by=("created_at", "desc"), limit=50,
):
if dict(await run.inputs.aio()).get("method", "eagle3") == method:
return Dir.from_existing_remote(await _output_uri(run))
async for run in remote.Run.listall.aio(
in_phase=("succeeded",), task_name=PIPELINES[method],
sort_by=("created_at", "desc"), limit=1,
):
return Dir.from_existing_remote(await _output_uri(run))
raise RuntimeError(f"No trained {method} drafter yet. Run its training pipeline first.")
It then warm-starts from that drafter, trains one epoch on the new traces plus replayed chat data, and benchmarks the candidate against the current drafter on the newest 10% of traces, which training did not see:
@flywheel_env.task(
report=True,
# One nightly trigger per family, staggered so their GPU jobs do not all start at once.
triggers=tuple(
flyte.Trigger(
name=f"nightly-{method}-retrain",
automation=flyte.Cron(f"0 {3 + i} * * *"),
inputs={"window_end": flyte.TriggerTime, "method": method},
)
for i, method in enumerate(PIPELINES)
),
)
async def retrain(
window_end: datetime,
method: str = "eagle3",
window_hours: int = 24,
min_traces: int = 200,
min_gain: float = 0.02,
max_steps: int | None = None,
) -> Dir:
"""Fine-tune a family's production drafter on yesterday's traces; promote it if it wins."""
since = window_end - timedelta(hours=window_hours)
window = f"{method}: traces from {since:%Y-%m-%d %H:%M} to {window_end:%Y-%m-%d %H:%M} UTC"
current = await latest_drafter(method)
traces = await collect_traces(since, window_end)
async def finish(drafter: Dir, promoted: bool, reason: str, **details) -> Dir:
await flyte.report.replace.aio(
gate_report(window, len(traces), promoted=promoted, reason=reason, **details)
)
await flyte.report.flush.aio()
return drafter
if len(traces) < min_traces:
return await finish(
current, False, f"only {len(traces)} traces, below min_traces={min_traces}",
n_train=0, n_eval_prompts=0, acceptance=None,
)
cut = int(len(traces) * 0.9)
train_file = await _write_jsonl(traces[:cut], "/tmp/train-traces.jsonl")
# Replay general chat so the head keeps drafting well outside the agent workload.
# Same inputs as the pipeline's first shard, so this is a cache hit, not a GPU job.
shards = await prepare_prompts(num_chat=4000, num_agent=4000, num_shards=4)
replay = (await regenerate_chat(shards.chat[0])).data
candidate = await train_drafter(
data=[train_file, replay], method=method, warm_start=current, epochs=1, max_steps=max_steps
)
# Gate: benchmark both heads on the most recent, held-out traffic.
eval_prompts = _eval_prompts(traces[cut:])
eval_path = "/tmp/eval-prompts.json"
with open(eval_path, "w") as f:
json.dump({"agent-recent": eval_prompts}, f)
prompts = await File.from_local(eval_path)
with flyte.group("gate"):
results = await asyncio.gather(
*(
benchmark_server(drafter_config(n, method, d), prompts, [1])
for n, d in (("current", current), ("candidate", candidate))
)
)
al = {}
for f in results:
with open(await f.download()) as fh:
r = json.load(fh)
al[r["config"]] = r["workloads"]["agent-recent"]["1"]["acceptance_length"]
details = dict(n_train=cut, n_eval_prompts=len(eval_prompts), acceptance=al)
if al["candidate"] >= al["current"] * (1 + min_gain):
return await finish(candidate, True, f"acceptance length improved by at least {min_gain:.0%}", **details)
return await finish(current, False, f"gain below the {min_gain:.0%} threshold", **details)
Each cycle’s report records the decision: the family and trace window, how many traces trained the candidate and how many prompts tested it, both acceptance lengths, and whether the candidate was promoted.
In this cycle, the candidate trained on 225 traces and drafted no better than the current head on the 46 held-out prompts, so the gate kept the current head. The traffic came from the same tool-calling dataset the head was fine-tuned on, so there was nothing new to learn. The loop pays off when traffic drifts away from what the drafter last trained on.
Design choices that keep the loop safe:
- The gate measures acceptance length, not training loss. Loss and acceptance length can disagree. The gate serves both drafters in vLLM and measures what production would see.
- Evaluation uses the newest traffic. Holding out the most recent traces checks the candidate against where traffic is heading, not where it was.
- Replay prevents forgetting. Each cycle mixes in regenerated general chat. The replay shard has the same inputs as the pipeline’s, so it is a cache hit and costs no GPU time.
- A warm start keeps the drafter’s architecture. Starting from an existing drafter,
train_drafterdrops the flags that shape a new model, such as DSpark’s layer count, and trains the drafter as it is. - Too little traffic is a no-op. Below
min_traces, the task returns the current drafter unchanged. - A losing candidate is discarded. The task returns the current drafter, so the newest
retrainoutput for a family is always the drafter that should be in production.
Register the triggers, or run one cycle now:
python flywheel.py # deploys the environment and its nightly triggers
python flywheel.py --now --method dspark # one cycle immediately, for one familyTo roll an app forward, redeploy it pinned to the retrain run:
python serve.py --method dspark --run <retrain-run>.
A drafter is trained on one target’s hidden states. If you fine-tune, quantize, or upgrade the target model, retrain every drafter before serving them together.
Version drafters and traces as artifacts
On Union, publish each promoted drafter as a version of its family’s artifact, such as
qwen3-8b-dflash-drafter, with a model card. The card shows the family, its recipe, the run
that produced it, and its benchmark results: acceptance length, speedup, and agreement with
no speculation, per workload.
def model_card(method: str, summary: dict | None, source_run: str) -> str:
"""A model card: what the drafter is, how it was trained, and what it measured."""
recipe = RECIPES[method]
parts = [
ui.hero(
f"Qwen3-8B {FAMILY[method]} draft model",
f"A {FAMILY[method]} drafter for Qwen3-8B. Outputs are unchanged; only decoding speed differs.",
[FAMILY[method], "speculators", "vLLM", f"{recipe.num_speculative_tokens} speculative tokens"],
),
ui.section(
"Details",
ui.table(
["field", "value"],
[
["target model", TARGET_MODEL],
["initialized from", recipe.init_from or "random weights"],
["target layers", ", ".join(map(str, recipe.target_layers))],
["produced by run", source_run or "unknown"],
["serve with", '{"method": "%s", "num_speculative_tokens": %d}' % (method, recipe.num_speculative_tokens)],
],
numeric=set(),
),
),
]
if summary:
for workload, data in summary["workloads"].items():
rows = data["rows"]
parts.append(
ui.section(
f"Benchmark: {workload}",
ui.charts(
ui.bar_chart("Speedup at concurrency 1", {n: r["speedup"]["1"] for n, r in rows.items()}, "{:.2f}x", 1.0),
ui.bar_chart(
"Acceptance length",
{n: r["acceptance_length"] for n, r in rows.items() if n != "baseline"},
),
),
ui.table(
["config", "acceptance length", "speedup c=1", "matches baseline before divergence"],
[
[n, f"{r['acceptance_length']:.2f}", f"{r['speedup']['1']:.2f}x", f"{r['prefix_agreement']:.0%}"]
for n, r in rows.items()
],
),
)
)
parts.append(
ui.section(
"Limitations",
lede=(
"Trained for Qwen3-8B with thinking disabled. Retrain after any change to the target's "
"weights, quantization, or chat template. Speedups shrink as serving concurrency rises."
),
)
)
return ui.page(*parts)
@promote_env.task(produces_artifacts=True)
async def publish_drafter(
drafter: Dir, method: str = "eagle3", summary: File | None = None, source_run: str = ""
) -> Dir:
"""Publish a drafter as the next version of its family's artifact, with a model card."""
results = None
if summary is not None:
with open(await summary.download()) as f:
results = json.load(f)
card = await artifacts.Card.create_from.aio(
content=model_card(method, results, source_run), format="html", card_type="model"
)
return artifacts.new(
drafter,
artifacts.Metadata.create_model_metadata(
name=drafter_artifact(method),
description=f"{FAMILY[method]} draft model for {TARGET_MODEL}",
card=card,
framework="PyTorch",
architecture=FAMILY[method],
task="speculative-decoding",
serial_format="safetensors",
attrs={"target_model": TARGET_MODEL, "method": method, "source_run": source_run},
),
)
@promote_env.task(produces_artifacts=True)
async def publish_traces(traces: list[File]) -> File:
"""Merge agent trace files into one data artifact, with a data card."""
rows = []
for f in traces:
with open(await f.download()) as fh:
rows.extend(json.loads(line) for line in fh)
path = "/tmp/traces.jsonl"
with open(path, "w") as fh:
for r in rows:
fh.write(json.dumps(r) + "\n")
tools = Counter(
c["function"]["name"]
for r in rows
for m in r["messages"]
if m["role"] == "assistant"
for c in m.get("tool_calls", [])
)
turns = [sum(m["role"] == "assistant" for m in r["messages"]) for r in rows]
card_html = ui.page(
ui.hero(
"Agent traces for draft head training",
f"Tool-calling episodes served by {TARGET_MODEL}. Every assistant turn is target output.",
["OpenAI chat format", "tools included", "jsonl"],
),
ui.tiles({
"episodes": f"{len(rows):,}",
"source files": f"{len(traces):,}",
"tool calls": f"{sum(tools.values()):,}",
"distinct tools": f"{len(tools):,}",
}),
ui.charts(
ui.histogram("Assistant turns per episode", turns, bins=6),
ui.bar_chart("Most called tools", dict(tools.most_common(8)), "{:.0f}"),
),
ui.section(
"Schema",
ui.pretty_json({"messages": "[OpenAI chat messages, with tool_calls]", "tools": "[OpenAI tool schemas]"}),
),
ui.section("Example", ui.chat(rows[0]["messages"]) if rows else ""),
)
card = await artifacts.Card.create_from.aio(content=card_html, format="html", card_type="data")
return artifacts.new(
await File.from_local(path),
artifacts.Metadata(
name=TRACES_ARTIFACT,
description="Agent traces collected from the speculative decoding app",
card=card,
kind="data",
attrs={"episodes": str(len(rows)), "target_model": TARGET_MODEL},
),
)
ArtifactValue without a version resolves to
the latest version, and pinning a version rolls back:def artifact_app(method: str, version: str | None = None):
"""The family's app, serving its drafter artifact: the latest version, or a pinned one."""
app = APPS[method]
return app.clone_with(
name=f"{app.name}-promoted",
draft_model_path=flyte.app.ArtifactValue(
name=drafter_artifact(method), type="directory", version=version
),
)
def roll_back(method: str, version: str) -> None:
"""Serve an earlier drafter. The target model and everything else stay the same."""
pinned = artifact_app(method, version)
flyte.deploy(pinned)
App.get(name=pinned.name).activate()
python promote.py --method dflash --run <training-run> --benchmarks <compare-run>
python promote.py --traces <agent-job-run> ... # publish traces, with a data card
python promote.py --method dflash --serve # serve the latest drafter
python promote.py --method dflash --rollback <version>When a learned drafter is not the answer
On highly repetitive agent traffic, where outputs copy long spans from the prompt or earlier
turns, suffix decoding often accepts more tokens than any learned drafter, with no
training. The SuffixDecoding paper reports about 6.2 accepted tokens per step on an agentic
text-to-SQL workload, against about 3.2 to 3.6 for EAGLE-3. The benchmark
includes a suffix configuration so you can check this on your own traffic before investing
in a training loop.