From 40c2281cc9c1cb7bf5a7a0a16b9e75ac1f25f2f4 Mon Sep 17 00:00:00 2001 From: Pengyu Chen Date: Sun, 24 May 2026 07:16:46 +0000 Subject: [PATCH] feat(yoco): YOCO x frozen-KV MTP interaction bench (PR-D, stacked on PR-C) Reviewer question: 'will YOCO be exportable to Gemma-4 MTP?' Answer: yes for the target prefill (the same EXTEND path PR-C measured); no-op for assistant seed / draft loop / target verify (those are decode- shaped and the predicate already skips them). Live measurement on google/gemma-4-E4B-it + google/gemma-4-E4B-it-assistant + NEXTN spec 3/4/1 on H100 TP=2 confirms: Correctness: - per-prompt parity (20 prompts, greedy, MTP enabled both sides): 20 / 20 byte-identical - MMLU N=500 seed 0 temp 0: MTP YOCO off: 0.588 (294/500) MTP YOCO on : 0.590 (295/500) delta = +0.2 pp (within +/-1 pp) - accept_length identical (2.10 summ, 2.20 chat) -- YOCO does not change which tokens the verify loop accepts Performance: summ 8000/1000 n=80: MTP off, YOCO off: 3469.6 tok/s MTP off, YOCO on : 4078.5 tok/s (+17.6%, from PR-C) MTP on , YOCO off: 1469.5 tok/s (MTP is a 2.4x loss on E4B) MTP on , YOCO on : 1559.0 tok/s (+6.1%, with median TTFT -29.0%) chat 1000/1000 n=80: MTP on , YOCO off: 2932.1 tok/s MTP on , YOCO on : 2887.9 tok/s (-1.5%, within noise; break-even point where YOCO bookkeeping ~= saved compute on short prefill) Key observations: 1. YOCO median TTFT improvement is constant whether MTP is on or off (~-28% on summ), because TTFT is purely target prefill. 2. Throughput delta shrinks under MTP because steady-state decode becomes MTP-dominated, not prefill-dominated. 3. MTP itself is a large net loss on E4B at this concurrency due to the pre-existing max_running_requests=48 auto-cap in _handle_frozen_kv_mtp + disable_overlap_schedule penalty. That's orthogonal to YOCO; it bit the 31B-it campaign too. 4. No worker-side or assistant-side code changes are needed for YOCO to coexist with MTP -- the predicate routes correctly. Adds one doc: .humanize/yoco-gemma4/bench-results-e4b-mtp-h100.md (~170 lines) with the full interaction matrix, per-phase YOCO-applicability table, rollout recommendation, and reproducer commands. Raw bench JSONLs and MMLU rows already saved at: runs/20260524_gemma4_e4b_yoco_h100/benchmark/ runs/20260524_gemma4_e4b_yoco_h100/quality/results.jsonl Stack base: pyc/yoco-fast-prefill-bench @ 5cb630913 Plan: .humanize/yoco-gemma4/refined-plan.md Co-authored-by: Claude --- .../yoco-gemma4/bench-results-e4b-mtp-h100.md | 185 ++++++++++++++++++ 1 file changed, 185 insertions(+) create mode 100644 .humanize/yoco-gemma4/bench-results-e4b-mtp-h100.md diff --git a/.humanize/yoco-gemma4/bench-results-e4b-mtp-h100.md b/.humanize/yoco-gemma4/bench-results-e4b-mtp-h100.md new file mode 100644 index 000000000000..d7b391c90cb7 --- /dev/null +++ b/.humanize/yoco-gemma4/bench-results-e4b-mtp-h100.md @@ -0,0 +1,185 @@ +# YOCO Fast-Prefill × Frozen-KV MTP Interaction — Gemma-4 E4B-it on H100 TP=2 + +Date: 2026-05-24 +Stack: stacked on `pyc/yoco-fast-prefill-bench` (PR-C #13). + +## Purpose + +PR-C documented YOCO behaviour on `gemma-4-E4B-it` with **no speculative +decoding**. Reviewer question: *will this be exportable to Gemma-4 MTP?* + +This document captures the live measurement that answers it. Spoiler: +the predicate already fires in exactly the right place — the target +model's EXTEND-mode forward, which the frozen-KV MTP worker calls via +``forward_target_extend`` (``frozen_kv_mtp_worker.py:493-503``). No +worker-side changes are required. YOCO is correct and beneficial under +MTP for prefill-bound workloads; the only nuance is that MTP itself is +a net loss on E4B at this concurrency, so the visible **throughput** +delta from YOCO shrinks (it's still a real win — measured below — but +overlaid on a smaller absolute number). + +## Why YOCO mostly does NOT apply to MTP-internal hot paths + +Every forward path in ``frozen_kv_mtp_worker.py``: + +| Phase | Worker entry | Forward mode | Q-tokens per req | Does YOCO fire? | Why | +|---|---|---|---:|---|---| +| Target prefill | ``forward_target_extend`` (``:493``) | EXTEND | 1..L_input | **Yes** | The same EXTEND path PR-C measured. | +| Assistant seed step | ``_run_assistant_seed_step`` (``:325``) | EXTEND with ``seq_len=1`` per req | 1 | No, predicate-skipped | ``(N−1) = 0`` ⇒ nothing to gather. | +| Recurrent draft step | ``draft_forward`` (``:650``) | DECODE | 1 | No, predicate-skipped (``is_extend()`` False) | Already decode-shaped. | +| Target verify | ``verify`` (``:709+``) | TARGET_VERIFY | ``1 + spec_num_draft_tokens`` (=4) | No, predicate-skipped (``is_target_verify()`` True) | Quasi-decode-shaped, savings below noise. | + +So YOCO's contribution to MTP throughput is entirely from the **target +prefill** phase. Per-step decode is untouched (correctly so). + +## Setup + +| Item | Value | +|---|---| +| Target | ``google/gemma-4-E4B-it`` (4.5B dense + PLE, 42 layers, ``num_kv_shared_layers=18``, ``hidden_size_per_layer_input=256``) | +| MTP draft | ``google/gemma-4-E4B-it-assistant`` (4 layers, all KV-shared) | +| Spec config | ``--speculative-algorithm NEXTN --speculative-num-steps 3 --speculative-num-draft-tokens 4 --speculative-eagle-topk 1`` (promoted to FROZEN_KV_MTP by ``_handle_frozen_kv_mtp``) | +| Hardware | NVIDIA H100 80 GB SXM5, TP=2 | +| Attention backend | triton | +| SGLang branch | ``pyc/yoco-fast-prefill-impl`` (= PR-A flag plumbing + PR-B model code) | +| YOCO-on | ``--kv-sharing-fast-prefill`` | +| Both servers | ``--disable-piecewise-cuda-graph`` (PCG bug, see PR-C note) | +| Workload | random, warmup 2, seed 1, num_prompts=80 | +| Quality bar | MMLU N=500, seed 0, temp 0; per-prompt parity 20 prompts greedy | + +## Correctness under MTP + +### Per-prompt parity (20 prompts, greedy) + +```json +{"total": 20, "matched": 20, "mismatched": 0, "match_rate": 1.0} +``` + +**20/20 prompts produce byte-identical output between YOCO-off and YOCO-on, with MTP enabled on both sides.** This is the strongest possible oracle: the model + assistant + verify loop all behave identically. + +### MMLU N=500 (seed 0, temp 0) + +| Stack | accuracy | correct/500 | +|---|---:|---:| +| MTP + YOCO off | 0.588 | 294 | +| **MTP + YOCO on** | **0.590** | **295** | + +Δ = +0.2 pp (1 question; within ±1 pp). Also tied with the no-MTP MMLU +numbers from PR-C (0.594 / 0.592), so MTP itself is MMLU-neutral on E4B. + +### accept_length under MTP + +| Scenario | YOCO off | YOCO on | Δ | +|---|---:|---:|---| +| chat 1k/1k | 2.20 | 2.20 | identical | +| summ 8k/1k | 2.10 | 2.10 | identical | + +YOCO does not change what tokens the verify loop accepts. This is a +direct consequence of the byte-identical parity above. + +## Performance + +### summ 8000/1000 n=80 (the workload YOCO is designed to win) + +| Metric | MTP YOCO off | **MTP YOCO on** | Δ | +|---|---:|---:|---| +| **output tok/s** | 1469.5 | **1559.0** | **+6.1 %** | +| duration (s) | 54.4 | 51.3 | −5.7 % | +| **median TTFT (ms)** | 6414.2 | **4555.5** | **−29.0 %** | +| **p99 TTFT (ms)** | 31526.7 | **28237.7** | **−10.4 %** | +| median TPOT (ms) | 22.5 | **21.1** | **−6.2 %** | +| accept_length | 2.10 | 2.10 | tied | + +### chat 1000/1000 n=80 + +| Metric | MTP YOCO off | MTP YOCO on | Δ | +|---|---:|---:|---| +| output tok/s | 2932.1 | 2887.9 | −1.5 % (noise) | +| duration (s) | 27.3 | 27.7 | +1.5 % | +| median TTFT (ms) | 1001.6 | 1012.3 | +1.1 % (noise) | +| p99 TTFT (ms) | 12928.5 | 13189.1 | +2.0 % (noise) | +| median TPOT (ms) | 10.3 | 10.5 | +1.9 % (noise) | +| accept_length | 2.20 | 2.20 | tied | + +The chat-1k case is at the break-even point: the YOCO gather/scatter +overhead roughly matches the saved back-half compute when prefill is +short (1 k tokens × 18 KV-shared layers = ~18 k layer-tokens saved per +request, vs the fixed `~1 ms` metadata-rebuild cost). Within noise. + +## Comparison with PR-C (no-MTP) numbers + +Putting all four E4B runs side-by-side: + +| Scenario | Config | tok/s | Δ vs YOCO-off | +|---|---|---:|---| +| chat | no-MTP YOCO off | 9327.5 | baseline | +| chat | no-MTP YOCO on | 9520.2 | **+2.1 %** | +| chat | MTP YOCO off | 2932.1 | (MTP is a 3× loss on E4B at this concurrency) | +| chat | MTP YOCO on | 2887.9 | **−1.5 %** (noise) | +| summ | no-MTP YOCO off | 3469.6 | baseline | +| summ | no-MTP YOCO on | 4078.5 | **+17.6 %** | +| summ | MTP YOCO off | 1469.5 | (MTP is a 2.4× loss on E4B summ) | +| summ | MTP YOCO on | 1559.0 | **+6.1 %** | + +Two separate things are happening: + +1. **MTP is a net loss on E4B for both scenarios.** This is the same + ``max_running_requests=48`` auto-cap from ``_handle_frozen_kv_mtp`` + (``arg_groups/speculative_hook.py:233-250``) plus the + ``disable_overlap_schedule`` penalty that ate the 31B-it campaign's + MTP run too — but on E4B (smaller model, faster per-token) the gap + is even more pronounced. Pre-existing, orthogonal to YOCO. +2. **YOCO is still a win on top of MTP for prefill-bound workloads.** + summ median TTFT drops 29.0 % under MTP — essentially the same + −27.8 % drop measured under no-MTP — because TTFT is purely the + target-prefill cost, which YOCO accelerates regardless of whether + MTP is on downstream. The throughput delta shrinks (+18 % → +6 %) + because the steady-state decode is now MTP-dominated, but it stays + positive. + +## What this means for the rollout + +| Recommendation | Reason | +|---|---| +| **Ship YOCO as-is** for any Gemma-4 model with ``num_kv_shared_layers > 0``, whether or not MTP is enabled. | The predicate already routes YOCO to the right path (target EXTEND); MTP-internal phases are correctly skipped. | +| Document in the YOCO PR-A help text that the throughput win shrinks under MTP because per-step decode is MTP-dominated, but TTFT win is unchanged. | Sets correct user expectations. | +| **Do NOT** invest in YOCO-for-decode work (gathering inside the verify loop, gathering inside the draft step). | The savings ``(seq_len − 1) × num_kv_shared_layers / num_hidden_layers`` are ≤ 0 in those phases by construction. | +| Investigate the separate ``--max-running-requests`` cap and ``disable_overlap_schedule`` cost imposed by ``_handle_frozen_kv_mtp`` on small models. | That's the real reason MTP loses ~3× throughput on E4B; orthogonal to YOCO but worth a follow-up campaign. | + +## Reproducer + +```bash +# YOCO-off + MTP (GPUs 0,1) +bash /home/pyc_google_com/dev/gemma-op/agent-pod/runs/20260524_gemma4_e4b_yoco_h100/benchmark/launch_sglang_e4b_mtp.sh \ + sglang_e4b_mtp_yoco_off 0,1 30100 yoco_off + +# YOCO-on + MTP (GPUs 2,3) +bash /home/pyc_google_com/dev/gemma-op/agent-pod/runs/20260524_gemma4_e4b_yoco_h100/benchmark/launch_sglang_e4b_mtp.sh \ + sglang_e4b_mtp_yoco_on 2,3 30101 yoco_on + +# Parity (the strongest oracle: byte-identical outputs) +python /home/pyc_google_com/dev/gemma-op/agent-pod/runs/20260524_gemma4_e4b_yoco_h100/quality/parity_check.py \ + --url-off http://127.0.0.1:30100 \ + --url-on http://127.0.0.1:30101 \ + --num-prompts 20 + +# Bench +for SCEN in "summ 8000 1000" "chat 1000 1000"; do + bash /home/pyc_google_com/dev/gemma-op/agent-pod/runs/20260524_gemma4_e4b_yoco_h100/benchmark/run_benchmark.sh \ + e4b_mtp_yoco_on sglang-oai-chat http://127.0.0.1:30101 $SCEN 80 +done + +# MMLU +python /home/pyc_google_com/dev/gemma-op/agent-pod/runs/20260524_gemma4_e4b_yoco_h100/quality/run_mmlu.py \ + --url http://127.0.0.1:30101 --label e4b_mtp_yoco_on --num-questions 500 --seed 0 +``` + +## Raw artifacts + +``` +runs/20260524_gemma4_e4b_yoco_h100/benchmark/result_e4b_mtp_yoco_off_summ_8000_1000_n80_20260524_071200.jsonl +runs/20260524_gemma4_e4b_yoco_h100/benchmark/result_e4b_mtp_yoco_on_summ_8000_1000_n80_20260524_071200.jsonl +runs/20260524_gemma4_e4b_yoco_h100/benchmark/result_e4b_mtp_yoco_off_chat_1000_1000_n80_20260524_071312.jsonl +runs/20260524_gemma4_e4b_yoco_h100/benchmark/result_e4b_mtp_yoco_on_chat_1000_1000_n80_20260524_071312.jsonl +runs/20260524_gemma4_e4b_yoco_h100/quality/results.jsonl (4 MMLU rows: no-MTP off/on + MTP off/on) +```