[levanter] Optimize v6e prefill-heavy decode gap

[dlwh]

Description

Explain and reduce the v6e prefill-heavy gap for prefill_b8_i2048_o128_n1.

Historical comparison row: Levanter 1006.83 decode tok/s and 17116.08 total tok/s versus vLLM 1262.90 decode tok/s and 21469.31 total tok/s, ratio 0.797. That measured Levanter row had prefill chunks 4096,4096,4096,4096 and one decode iteration with 0.802s total, 0.623s device, 0.178s host, and 0.174s submit.

The corrected Levanter-only diagnostic completed at /dlwh/qwen3-v6e8-prefilldiag-drain-prefill-b8-i2048-o128-n1-20260606-1532 from PR #6185 commit 82bb6dbfb. It validated diagnostic prefill drain and produced 1022 decode iteration tokens for normal, no_lm_head, and lm_head_no_sampling rows.

Important interpretation correction: production decode_submit_seconds was previously measured from the start of the outer iteration before prefill drain, so the 0.174s/0.184s submit field included prefill admission work. PR #6185 commit dd993c03c moved the submit timer to immediately before _run_generation_loop(...), commit 6d74fc7ea added decode iteration/device throughput fields, and commit f55913d1 now separates prefill drain from the generation loop with prefill_drain_*, generation_*, and derived generation_tokens_per_second benchmark fields.

Corrected backend=both row: /dlwh/qwen3-v6e8-prefillcorr-20260607-0023 completed successfully from #6185 head 91f6ec06a. vLLM measured 1212.15 decode tok/s and 20606.54 total tok/s. Levanter measured 898.94 decode tok/s and 15282.00 total tok/s, ratio 0.742, target fail. The corrected Levanter fields available in that run show prefill chunks 4096,4096,4096,4096, prefill admission wall times 0.060,0.055,0.055,0.055s, decode iteration 0.852s total / 0.677s device / 0.174s host / 0.002s submit / 0.002s extract, 1022 decode iteration tokens, 1200.223 decode-iteration tok/s, and 1509.350 device tok/s. Because that row predates f55913d1, it does not yet include the new generation-only token/time split.

Post-f55913d1 backend=both row: /dlwh/qwen3-v6e8-prefillpostsplit-20260607-0419 completed successfully from #6185 head d63d1edf. vLLM measured 1264.03 decode tok/s and 21488.50 total tok/s. Levanter measured 957.72 decode tok/s and 16281.24 total tok/s, ratio 0.758, target fail. The new split shows hot prefill chunks 4096,4096,4096,4096 with prefill wall times 0.060,0.056,0.056,0.056s, decode iteration 0.800s total / 0.623s device / 0.177s host / 0.002s submit / 0.002s extract, 1022 decode iteration tokens, 0.171s prefill drain for 6 tokens, and 0.629s generation for 1016 tokens. Derived throughput: 1277.749 decode-iteration tok/s, 1639.896 decode-device tok/s, and 1614.816 generation tok/s.

Follow-up code changes:

• [levanter] Add multi-prefill admission for serving #6185 commit f7275e5a stops logging full InferenceRequest dataclasses at INFO and instead logs bounded request metadata. In the prefill-heavy measured row, the old path serialized eight 2048-token prompt lists before the first prefill admission; this was a concrete host-side measured-path overhead source. Validation passed with inference-server pytest (19 passed, 1 skipped) plus focused pre-commit/Pyrefly on the touched files.
• [levanter] Add multi-prefill admission for serving #6185 commit f55913d1 adds per-iteration prefill_drain_seconds_per_iteration, prefill_drain_tokens_per_iteration, generation_seconds_per_iteration, generation_host_seconds_per_iteration, generation_tokens_per_iteration, and generation_tokens_per_second to avoid conflating prefill-drain work with generation-loop throughput in future [levanter] Optimize v6e prefill-heavy decode gap #6229 rows. Validation passed with test_engine.py and test_qwen3_tpu_inference_parity_bench.py (63 passed) plus focused pre-commit/Pyrefly.

Current interpretation: the remaining failure is attributed to end-to-end prefill-heavy wall-clock overhead rather than a proven decode-device regression. The issue remains open as a performance target because the post-f55913d1 row still fails the end-to-end parity target at 0.758, even though generation-only throughput is above the vLLM decode row. The remaining question is whether to optimize hot prefill/admission/serving wall-clock for this weak prefill-heavy regime, or deprioritize it relative to the core RL rollout rows.

Definition of Done

• Rerun one corrected Levanter-only diagnostic on the fixed harness.
• Attribute the original row enough to avoid a false decode-kernel conclusion: the old submit timing was contaminated by prefill-drain work, and new benchmark fields separate prefill-drain, generation-loop, device, host, submit, and extraction timing.
• Collect a post-f55913d1 row for prefill_b8_i2048_o128_n1 if this regime remains a priority, so the result includes generation_* and prefill_drain_* fields after bounded request logging.
• Decide the next optimization or close/deprioritize this issue based on that post-f55913d1 row: prefill/host accounting, page/admission settings, or no further work if the row is acceptable for RL rollout priorities.
• Update the epic and handoff with terminal job results from the completed diagnostics/reruns.

Current Next Step

Do not launch another TPU run automatically. The next step is a decision, not more evidence collection: either optimize the remaining hot prefill/admission/serving wall-clock overhead for this regime, or close/deprioritize it if core RL rollout rows matter more than prefill-heavy parity. The post-split row should be interpreted using generation_tokens_per_second=1614.816, prefill_drain_seconds_per_iteration=0.171, and generation_host_seconds_per_iteration=0.006, not only the end-to-end decode_tokens_per_second=957.72.

Parent

• [Epic] Levanter TPU RL rollout backend parity #6227

[dlwh]

🤖 Initial progress update: PR #6185 now includes diagnostic harness fix 82bb6dbfb, which makes _generate_diagnostic() drain pending prefill admissions before diagnostic decode. Focused validation passed locally (lib/levanter/tests/inference/test_engine.py, 12 passed) plus the required pre-commit file slice.

A corrected single Levanter-only v6e diagnostic has been delegated from that head with prefix qwen3-v6e8-prefilldiag-drain-prefill-b8-i2048-o128-n1. The babysitter should update this issue and epic #6227 only on launch failure, concrete non-capacity failure/fix, or terminal benchmark results.

[claude]

🤖 Not actionable as a PR right now — blocked on attribution data.

Research

This issue is a performance attribution task, not a self-contained bug fix. The Definition of Done is explicit: "File or implement the next code change only after the attribution is clean." The corrected Levanter-only v6e diagnostic was already delegated (prefix `qwen3-v6e8-prefilldiag-drain-prefill-b8-i2048-o128-n1`) and the babysitter is instructed to update only on terminal results or concrete failure. No attribution data is yet in the issue.

The harness artifact that motivated the rerun is already fixed in PR #6185 commit `82bb6dbfb` (diagnostic prefill-drain). The next code change cannot be designed without seeing the corrected diagnostic split between device decode, host/submit overhead, prefill scheduling, and LM-head/sampling.

Relevant code

• `engine.py` - serving engine
• `jit_scheduler.py` - prefill/decode admission
• `tests/inference/test_engine.py` - diagnostic harness coverage
• PR [levanter] Add multi-prefill admission for serving #6185 commit `82bb6dbfb` - `_generate_diagnostic()` prefill-drain fix
• Parent epic [Epic] Levanter TPU RL rollout backend parity #6227

Why no PR

Opening a code-change PR now would mean guessing which of {device decode, host/submit, prefill scheduling, LM-head/sampling} dominates the `0.797` ratio for `prefill_b8_i2048_o128_n1`. The agent has no TPU access to run the benchmark and the delegated rerun hasn't reported back. Either outcome — pinning the gap on a specific component or finding a second harness artifact — drives a different fix.

Next step (for a human or the babysitter)

1. Wait for the delegated `qwen3-v6e8-prefilldiag-drain-prefill-b8-i2048-o128-n1` run to land terminal results in this issue.
2. With the corrected decode token grouping (`1022`, not `510,256,256`), compute the residual gap and attribute it.
3. Reopen this issue for an implementation PR scoped to the attributed component, or close it out if the gap collapses under the drain fix alone.

[dlwh]

🤖 Corrected v6e-8 Levanter-only prefill-heavy diagnostic completed.

Iris job: /dlwh/qwen3-v6e8-prefilldiag-drain-prefill-b8-i2048-o128-n1-20260606-1532
Commit: 82bb6dbfb
Shape: Qwen/Qwen3-8B, dense_qwen3_8b_v2, prefill_b8_i2048_o128_n1, backend levanter only, v6e-8
Terminal state: JOB_STATE_SUCCEEDED / exit 0

Key rows:

• levanter:auto: decode 906.62 tok/s, total 15412.56 tok/s, decode iter 0.813s total / 0.623s device / 0.190s host / 0.184s submit, 1022 decode iter toks
• levanter:auto:no_lm_head: decode 1244.03 tok/s, total 21148.47 tok/s, decode iter 0.751s total / 0.573s device / 0.178s host / 0.003s submit, 1022 decode iter toks
• levanter:auto:lm_head_no_sampling: decode 1154.22 tok/s, total 19621.78 tok/s, decode iter 0.815s total / 0.637s device / 0.178s host / 0.003s submit, 1022 decode iter toks

Drain fix validated for this diagnostic: all three rows show prefill chunks 4096,4096,4096,4096 and 1022 decode iter toks.

[dlwh]

🤖 Follow-up correction on the v6e prefill-heavy diagnostic interpretation.

The corrected diagnostic result remains valid, but code inspection found that production generate() was measuring decode_submit_seconds from the start of the outer iteration, before pending prefills were drained. That means the levanter:auto 0.184s submit field in /dlwh/qwen3-v6e8-prefilldiag-drain-prefill-b8-i2048-o128-n1-20260606-1532 includes prefill admission work and is not directly comparable to the diagnostic rows' 0.003s submit field.

Pushed PR #6185 commit dd993c03c to start the production submit timer immediately before _run_generation_loop(...), after prefill drain. Serving behavior is unchanged; this fixes benchmark attribution for future rows.

Updated PR #6214 handoff at b865dfdfd: the clean signal is now 1022 / 0.813s production decode iteration throughput, close to the original vLLM 1262.90 tok/s row. Next step is corrected benchmark timing attribution before assigning any residual gap to LM-head/sampling or kernels.

[dlwh]

🤖 Progress update: #6185 now includes 6d74fc7ea (Report pure decode iteration throughput). Future parity summaries and summary.json rows report decode_iteration_tokens_per_second and decode_device_tokens_per_second alongside the existing wall-clock decode_tokens_per_second, so v6e prefill-heavy rows can distinguish end-to-end admission/extraction overhead from the pure measured decode loop and device time.

Validation run locally on the #6185 branch:

• uv run --package marin-levanter --group test pytest lib/levanter/tests/test_qwen3_tpu_inference_parity_bench.py lib/levanter/tests/inference/test_engine.py -q -> 61 passed
• ./infra/pre-commit.py --files lib/levanter/scripts/bench/bench_qwen3_tpu_inference_parity.py lib/levanter/tests/test_qwen3_tpu_inference_parity_bench.py --fix -> OK
• commit hook -> OK, including Pyrefly

[dlwh]

🤖 Research update: #6185 head 91f6ec06a already has the corrected production submit timer and pure decode throughput fields; prefill_b8_i2048_o128_n1 is selectable with --matrix dense_qwen3_8b_v2 --case prefill_b8_i2048_o128_n1. The current #6229 comparison row predates those fields, so a kernel/LM-head PR would still be guessing.

Relevant code:

• engine.py#L1617-L1672 - submit timing after prefill drain
• bench_qwen3_tpu_inference_parity.py#L2083-L2089 - pure decode/device TPS
• bench_qwen3_tpu_inference_parity.py#L240-L254 - target dense case
• bench_qwen3_tpu_inference_parity.py#L2547-L2549 - single-case filter

Proposed next step: with top-level/user run authorization, collect one v6e-8 backend=both row from #6185 head using the same launch wrapper plus --matrix dense_qwen3_8b_v2 --case prefill_b8_i2048_o128_n1 --warmup-rounds 2 --measure-rounds 1 --max-pages 512 --tensor-parallel-size 8, then compare decode_tokens_per_second, decode_iteration_tokens_per_second, and decode_device_tokens_per_second. No code change from this pass; #6185 already contains the smallest harness support needed.

[dlwh]

🤖 Launched the corrected v6e-8 backend=both prefill-heavy rerun requested by this issue.

Iris job: /dlwh/qwen3-v6e8-prefillcorr-20260607-0023
PR/head: #6185 91f6ec06a
Shape: Qwen/Qwen3-8B, dense_qwen3_8b_v2, prefill_b8_i2048_o128_n1, backend both, v6e-8, TP=8, --max-pages 512, two warmups, one measured round.

Immediate Iris status: JOB_STATE_RUNNING, one running task, no failures or preemptions. A babysitter subagent and the top-level heartbeat are watching for terminal results or concrete non-capacity failures. No additional runs were launched.

[dlwh]

🤖 Corrected v6e-8 backend=both rerun completed successfully.

Iris job: /dlwh/qwen3-v6e8-prefillcorr-20260607-0023
PR/head: #6185 91f6ec06a
Shape: prefill_b8_i2048_o128_n1, backend both, v6e-8, TP=8, --max-pages 512, two warmups, one measured round.

Results:

• vLLM: 1212.15 decode tok/s, 20606.54 total tok/s
• Levanter: 898.94 decode tok/s, 15282.00 total tok/s
• ratio: 0.742 decode / 0.742 total, target fail
• Levanter prefill chunks: 4096,4096,4096,4096; measured prefill wall times 0.060,0.055,0.055,0.055s
• Levanter decode iteration: 0.852s total / 0.677s device / 0.174s host / 0.002s submit / 0.002s extract, 1022 tokens
• corrected throughput fields: 1200.223 decode-iteration tok/s, 1509.350 decode-device tok/s

Interpretation: the row still fails end-to-end, but the corrected fields do not point at decode-device throughput. The remaining bottleneck is prefill-heavy wall-clock overhead around prefill/host/measurement. Updated #6229 body, epic #6227, and handoff PR #6214 at 1fe859e30.