# M3.6 F4 — eliminate per-step `cuStreamSynchronize` in decode loop

> Trace-evidence-licensed candidate per
> [`docs/plans/M3.6-decode-throughput-no-retraction-tax.md`](M3.6-decode-throughput-no-retraction-tax.md)
> Phase 1 + Phase 2 wins entries. F1/F2/F3 (scheduler policy) are
> falsified; F4 is the remaining candidate before kernel-level work.
>
> **Evidence summary**: ARLE Phase 1 nsys trace shows
> `cuStreamSynchronize` consumes 48.6% of CPU API time (avg 65.7 ms,
> max 274 ms, 666 calls in 60 s). vLLM s14 control hits 3.30× faster
> ITL (20 ms vs 66 ms) and 18% faster per-row work. Phase 1 +
> Phase 2 bench evidence is sufficient to license F4 without an
> additional vLLM nsys trace.

## ITL gap = single-point + scheduling, both contribute

User question (2026-05-07): "ITL p50 这么慢的原因是什么 ... 是单点
问题 还是 调度问题"

Direct decomposition from Phase 1 trace:

| Per-decode-tick component | ARLE measured | vLLM (inferred from ITL) |
|---|---:|---:|
| `step_total` (CPU wall) | 82 ms | ~20 ms |
| Of which: GPU kernel work | ~76 ms (cuda_gpu_kern_sum / 551 ticks) | hidden — overlapped |
| Of which: CPU launch + setup | ~6 ms | included in 20 ms (overlapped) |
| Of which: `cuStreamSynchronize` | 65 ms | ~0 ms |
| Per-row-per-step (= ITL / batch) | 2.16 ms | 1.43 ms |

Both axes contribute to the 3.30× ITL gap (66 ms vs 20 ms):

- **Single-point (kernel) component**: per-row-per-step gap is
  2.16 ms vs 1.43 ms = **vLLM 51% faster per row**.
  `decode_attention_varlen_quantized_partial_kernel` consumes
  41.6% of GPU time at avg 4.55 ms / call vs vLLM's Triton
  attention which is provably faster at the same effective batch
  (vLLM ITL 20 ms with 14 rows = ~22 ms of GPU time per step,
  which on a 36-layer model is ~0.6 ms/layer — much faster than
  ARLE's per-layer attention).
- **Scheduling component**: ARLE step_total = GPU work + CPU
  sync wait (essentially **serial**: CPU launches → CPU sync waits
  for GPU → CPU does next-step prep → CPU launches again). vLLM
  step is **pipelined**: CPU prep for step N+1 happens during step
  N's GPU execution. vLLM's CUDA graph captures argmax + token-
  embedding inside the same graph, so its CPU never blocks on
  per-step sync.

The CPU/GPU serialization is the dominant fixable component: even
if the per-row kernel work stayed at ARLE's current speed, fully
pipelining CPU work against GPU work would cut step_total from
~82 ms to ~76 ms-equivalent (the GPU work duration alone), which
is a 7% gain. The bigger leverage is from the kernel side. F4
addresses the scheduling axis only; kernel work belongs to a later
M3.5 / M_b milestone using `ncu` deep-dive.

**Why is sync 65 ms?** Each `cuStreamSynchronize` waits for the
graph replay to finish. The graph contains 36 transformer layers
× (attention + MLP + RMSNorm + projections) plus embedding +
final norm + logits GEMM. 76 ms of GPU work per tick. The CPU sits
idle the whole time. Eliminating the wait does NOT speed up the
GPU — it lets the **next tick's CPU work start during the current
tick's GPU work**. Throughput improvement from F4 = (CPU work per
tick) / (sync wait per tick) = ~6 ms / 65 ms = ~10% best case
WITHOUT batch growth, but in practice F4 lets the scheduler keep
the GPU pipeline saturated as batch grows, reaching ~700 out tok/s
per the Phase 2 upper bound.

## P0 survey findings (2026-05-07)

`grep -rn run_or_capture infer/src/` + reads of:
- `infer/src/model/cuda_graph.rs` (`CudaGraphState`)
- `infer/src/model/qwen3/decode.rs` (single-token decode)
- `infer/src/model/qwen3/batch_decode.rs::decode_batch` (batched decode hot path)
- `infer/src/model/qwen3/forward.rs` (sample_batch_greedy_* methods)
- `infer/src/scheduler/cuda/decode.rs` (scheduler split-launch+readback)

State of CUDA graph capture today:

| Component | Captured into graph? | File:line |
|---|---|---|
| Embedding lookup | yes (single-token + batched) | `batch_decode.rs:1281` (graph_body) |
| All transformer layers | yes | `batch_decode.rs:1675-1727` (decode_batch_graph_body) |
| RMS norm + final norm | yes | inline in graph body |
| Logits GEMM | yes | `batch_decode.rs:1281` "embedding + layers + final norm + logits GEMM" |
| **Argmax sampling kernel** | **NO** — fired separately AFTER graph replay | `forward.rs:531-549` (sample_batch_greedy_launch) |
| **Token ID readback to CPU** | **NO** — `cuStreamSynchronize` + `clone_dtoh`/`memcpy_dtoh` | `forward.rs:557-571` (sample_batch_greedy_readback) |
| **EOS check** | **CPU side** (post-readback) | scheduler-side after token IDs land |

The scheduler's split sampling pattern (`decode.rs:296` launch +
`decode.rs:690` readback) already overlaps the launch with other
scheduler work, but the readback step **still synchronizes per
decode tick** because the next step's embedding input depends on
the just-sampled tokens (read from CPU back to a `decode_meta`
buffer via `memcpy_htod`).

## Two candidate designs

### F4-Small — CPU/GPU work overlap via prefetch token feed (~1.5 days, partial fix)

> **Codex review (2026-05-07, P2)**: the original "defer readback
> by 1 step" framing is broken. `collect_decode_batch_inputs()`
> feeds each launch from `req.generated_tokens.last()`, and that
> value is only appended after `sample_batch_greedy_readback()`
> completes. Skipping the readback would feed the launch a token
> from two steps ago — corrupting autoregression. Reading back
> before the next launch keeps the per-tick sync, defeating the
> goal. **Revised design below.**

The deferral has to happen at the **GPU buffer feed** layer, not
the CPU `generated_tokens.last()` layer:

1. Keep `sample_batch_greedy_launch` writing argmax output to
   `decode_ctx.argmax_out` (GPU buffer, already exists).
2. **Stop** reading `argmax_out` back to CPU on the decode hot
   path. Instead, copy `argmax_out` → a stable
   `next_decode_meta_gpu` buffer that the next launch's embedding
   reads directly. CPU never sees the tokens during decode.
3. The async readback (for streaming + stop-sequence detection)
   moves to the existing `copy_stream` (already in
   `tensor.rs::DeviceContext`). Each tick records an event; CPU
   polls with `cudaEventQuery` (non-blocking) and drains tokens
   when ready. Drain happens during the next decode tick's
   admission/plan window — not on the critical path.
4. EOS / stop checks initially stay text-match-on-CPU but with a
   1-tick delay (ok per acceptance).

This is closer to F4-Big in spirit but smaller in scope: the
in-graph sampling + permutation handling is deferred. Just stop
syncing on the readback.

- **Pros**: Same per-tick sync removal as F4-Big without graph
  surgery. ~120 LOC.
- **Cons**: Still does argmax outside the graph (one extra
  kernel launch per tick — ~80 µs). Slot-permutation
  invariant is simpler because `argmax_out → next_decode_meta_gpu`
  is a same-tick D2D copy with the same row order.
- **Expected gain**: ~50% off ITL (66 ms → 33 ms), ~700 out
  tok/s at s48 high-conc. Closes the gap to vLLM 647 in
  measured Phase 2 numbers.

### F4-Big — full in-graph sampling + GPU-side embedding feed (~2-3 days, full fix)

Match vLLM's pattern: argmax kernel + GPU-side write of token IDs
to a stable buffer. Next step's embedding lookup reads from that
GPU buffer directly (no CPU round-trip). CPU only reads token IDs
periodically (every K steps, where K = streaming chunk size, ~1-4
tokens for SSE streaming; can be 1 if the readback uses a
non-blocking `cudaEventQuery` poll instead of `cudaStreamSynchronize`).

Design pieces:
1. **Argmax kernel into graph**: extend `decode_batch_graph_body`
   in `batch_decode.rs:1281+` to launch `argmax_batch_logprob`
   after the logits GEMM. Output goes to a **stable GPU buffer**
   (`bufs.sampled_tokens_gpu`, allocated lazily like
   `logits_batch`).
2. **Embedding reads from sampled_tokens_gpu**: change
   `embedding_decode_into` callsite in graph body so that step
   N+1's input is `sampled_tokens_gpu` (written by step N's
   in-graph argmax). Skip the `memcpy_htod` of `decode_meta` for
   the token-ID portion (positions can stay H2D since they are
   computed CPU-side from `kv_cache.len()`).
3. **Async readback path**: a dedicated stream
   (`copy_stream` already exists in `tensor.rs::DeviceContext`)
   does `memcpy_dtoh` of `sampled_tokens_gpu` into a host buffer.
   Each decode tick records an event; the streaming layer polls
   the event with `cudaEventQuery` (non-blocking). Tokens read
   back asynchronously feed the HTTP SSE stream.
4. **EOS check**: a small GPU kernel that ORs all rows' EOS
   sentinel flags into a single GPU `bool`. CPU polls with
   `cudaEventQuery` on that single byte; only when EOS is hit
   does the CPU do a full readback. Until then, decode loop
   proceeds with no per-step sync.
5. **CUDA graph re-capture invariant**: argmax kernel and embed
   read both have stable input pointers (the
   `sampled_tokens_gpu` buffer never moves). Graph captures and
   replays cleanly. The `decode_meta` H2D (positions only) still
   needs to happen each tick — but H2D is async on the compute
   stream and does not require sync.
6. **Streaming chunk size**: HTTP SSE traditionally streams every
   token; under F4 the readback latency is amortized by the event
   poll. Worst case for stream latency: one decode-step delay
   between token generation and arrival at the client (still <
   100 ms, well below user perception threshold).

- **Pros**: Eliminates per-step `cuStreamSynchronize` entirely.
  Matches vLLM's pipeline pattern. Worst-case ITL bound by graph
  replay duration only (~10-20 ms per step at batch ~38).
- **Cons**: Larger change (~300 LOC across model + scheduler +
  HTTP streaming). New invariant on `sampled_tokens_gpu` lifetime.
  Stop-sequence detection (text match) currently runs on CPU after
  readback — preserve by keeping the K=1 readback path under a
  feature flag, ship F4-Big with `--decode-pipeline=async` opt-in
  initially.
- **Expected gain**: ~3.0× off ITL (66 ms → 22 ms), ~700 out
  tok/s at s48 high-conc, **reaches parity with vLLM 647**.
  Matches the upper-bound estimate from the Phase 2 wins entry.

## Decision

Land **F4-Small first** (1-day, low risk, closes 75% of gap), then
**F4-Big** (2-3 days, closes remaining gap to parity). Do not
short-circuit to F4-Big — the small version validates the approach
and provides incremental bench data; if it underperforms the
expected ~600 out tok/s, the assumption that sync is the bottleneck
needs re-checking before committing to F4-Big's larger change.

## Tasks

| # | Task | File | LOC est. | Owner | Trigger |
|---|---|---|---|---|---|
| 1 | F4-Small: defer readback by 1 step; current tick reads previous tick's tokens | `infer/src/scheduler/cuda/decode.rs` (`step_decode_launch`, `dispatch_decode_emits`) + `model/qwen3/forward.rs::sample_batch_greedy_readback` | ~80 | Codex | This plan landing |
| 2 | F4-Small: bench against Phase 1 baseline (462 out tok/s) at s48 high-conc | `scripts/bench_guidellm.sh` invocation, no code change | 0 | Claude | Task 1 commit |
| 3 | F4-Small: nsys trace post-fix; assert `cuStreamSynchronize` ≤ 25% of CPU API (down from 48.6%) | `nsys profile` direct | 0 | Claude | Task 2 numbers green |
| 4 | F4-Big: extend `decode_batch_graph_body` to include argmax kernel | `infer/src/model/qwen3/batch_decode.rs:1281-1340` (graph body) | ~50 | Codex | F4-Small <600 tok/s |
| 5 | F4-Big: stable `sampled_tokens_gpu` buffer; embedding reads from it for next step | `model/qwen3/batch_decode.rs` (`BatchDecodeBuffers`) + `model/qwen3/decode.rs::decode_kernels` | ~80 | Codex | Task 4 |
| 6 | F4-Big: async readback via copy_stream + cudaEventQuery polling | `model/qwen3/forward.rs::sample_batch_greedy_readback` + `scheduler/cuda/decode.rs::dispatch_emits` | ~100 | Codex | Task 5 |
| 7 | F4-Big: GPU EOS check kernel; non-blocking poll integration | `crates/cuda-kernels/csrc/misc/` (new EOS reduce kernel) + scheduler integration | ~80 | Codex | Task 6 |
| 8 | F4-Big: bench + nsys trace post-fix; assert parity with vLLM | bench + trace | 0 | Claude | Task 7 |

**F4-Small total: ~80 LOC, ~1 day. F4-Big additional: ~310 LOC, ~3 days.**

## Acceptance

### F4-Small

- `cargo test --release -p infer --features cuda --test e2e` passes.
- `cargo test --release -p infer --features cuda --test greedy_consistency` passes.
- Bench: `scripts/bench_guidellm.sh m3.6-f4small-s48-highconc` ≥ 580 out
  tok/s (≥ 25% above Phase 1 baseline of 462). TTFT p99 ≤ baseline + 1
  step (≤ 5 s).
- nsys trace: `cuStreamSynchronize` ≤ 25% of CPU API time.

### F4-Big

- All F4-Small tests + new test for `sampled_tokens_gpu` lifetime
  (`infer/tests/decode_in_graph_sampling.rs`).
- Bench: ≥ 640 out tok/s (≥ 0.95× vLLM s14 control of 647).
- nsys trace: `cuStreamSynchronize` ≤ 5% of CPU API time.
- HTTP streaming: SSE token delivery latency ≤ baseline + 1 step (no
  user-perceptible regression).

## Out of scope

- Spec-decode interactions. F4 changes only the regular decode path;
  spec-verify already runs its own kernel chain via
  `step_spec_decode_launch_from_path`.
- Metal backend. M3.6 is CUDA-only.
- Penalty/sampling beyond greedy. The non-greedy path
  (`sample_batch_with_penalties`) has its own readback and is more
  complex; F4-Small's readback deferral applies but F4-Big requires
  separate analysis. Scope to greedy in this plan.

## Risks + retreat

- **R1 — F4-Small breaks first-token latency**: under bursty
  request arrival, the deferred readback shifts the first token by
  one step. Mitigation: detect cold-start and skip the deferral on
  the first 1-2 tokens per request (only kick in at steady state).
- **R2 — F4-Big slot/row ownership drift**: codex review (2026-05-07
  P2) flagged that `running_batch` can finish/retract a slot and
  admit another while keeping the same row count, so row `j` can
  refer to a different request on the next graph replay. Embedding
  would consume the previous occupant's sampled token for the new
  row, corrupting generation. Mitigation: maintain a
  `sampled_tokens_owner: [SlotId; max_batch_size]` companion array
  that the scheduler checks on every `slot_indices` change; if any
  row's owner changed, **invalidate that row's `sampled_tokens_gpu`
  entry** (overwrite with a freshly-encoded prompt-suffix token via
  H2D before the next launch) AND re-capture the relevant graph
  slot. Batch-size-only check is insufficient.
- **R3 — Stop sequence detection lag**: F4-Big's batched readback
  delays stop-sequence text-match by K steps. Mitigation: K=1 + GPU
  EOS check covers the common case; text-match falls back to per-K
  readback only when text-stop is enabled.

## References

- Discovery + falsified F1/F2/F3 evidence:
  [`2026-05-07-m3.6-phase1-nsys-arle-s48-highconc.md`](../experience/wins/2026-05-07-m3.6-phase1-nsys-arle-s48-highconc.md)
- vLLM control bench (per-row ITL ground truth):
  [`2026-05-07-m3.6-phase2-vllm-s14-bench.md`](../experience/wins/2026-05-07-m3.6-phase2-vllm-s14-bench.md)
- CUDA graph state: `infer/src/model/cuda_graph.rs`
- Batched decode graph body: `infer/src/model/qwen3/batch_decode.rs:1281-1340`
- Sampling launch + readback split: `infer/src/model/qwen3/forward.rs:531-571`
- Scheduler dispatch: `infer/src/scheduler/cuda/decode.rs:296, 600, 690`
- vLLM in-graph sampling reference (their pattern):
  `vllm/v1/worker/gpu_model_runner.py::execute_model` graph capture path
- WukLab MLSys "scheduling overhead can dominate" (the inverse
  problem of ARLE's GPU-completion-wait dominance):
  https://mlsys.wuklab.io/posts/scheduling_overhead/
