[YuDeng0102]

Summary

This draft implements the mixed-depth loop-level speculative round proposed in
Issue #43 on the current
stage/runtime architecture. It is an opt-in replacement for the block round
inside SpeculativeRunner.submit; result collection, rejection sampling, KV
commit and rollback remain shared with the block path.

The loop round:

• plans the complete wave schedule from token counts before launching kernels;
• prepares KV metadata and gather/scatter indices once per round;
• advances active positions at different recurrent depths in mixed waves;
• keeps draft IDs and branch state on the device until collection;
• works with deferred collection, multiple in-flight submissions, preemption
  between rounds and CUDA Graphs.

The loop round does not reduce recurrent work or critical-path length. It changes
the row counts seen by recurrent GEMMs, so its performance is hardware and
library dependent.

Cost model

For K=4, draft_loops=2, target_loops=4, both block and loop rounds execute
12 recurrent calls over 20C rows:

Round	Recurrent calls
Block	10 calls with C rows, 2 calls with 5C rows
Loop	4 calls with C rows, 8 calls with 2C rows

The loop is faster only when the device cost function for those row counts favors
the mixed shapes. GEMM kernel-selection thresholds and CUDA Graph limits can
therefore look like schedule gains. Measurements must keep graph fallbacks at
zero in both arms.

Measurements

Ouro-1.4B, BF16, Triton attention, CUDA Graphs, K=4, d=2, D=4, 16 prompts
cycled over the requested concurrency, 128 generated tokens, seven alternating
paired repeats, fixed per-request seeds, and Graph limits
cuda_graph_max_batch_size=1024, cuda_graph_max_graphs=128. Values below are
block time divided by loop time; values above 1 mean the loop round is faster.

C	Block ms/round	Loop ms/round	Round-time ratio	Throughput ratio	Acceptance block / loop
1	44.80	52.97	0.846	0.846	0.884 / 0.884
4	56.24	58.98	0.953	0.926	0.839 / 0.800
8	59.26	61.28	0.967	1.024	0.804 / 0.828
16	63.39	59.97	1.057	1.028	0.824 / 0.822
32	63.13	63.14	1.000	1.088	0.800 / 0.836
64	75.45	80.56	0.937	0.910	0.823 / 0.829

Throughput also includes acceptance, so it can differ from the direct round-time
comparison when BF16 outputs diverge. C=128 did not fit with two resident engines
on the 32 GB test GPU.

Top-p used temperature=0.8, top_p=0.8:

C	Block ms/round	Loop ms/round	Round-time ratio	Throughput ratio
1	47.72	56.43	0.846	0.846
16	73.52	71.97	1.021	0.985
64	123.66	128.16	0.965	0.965

With Graph max 128 at C=32, the block round's 160-row verification falls back to
eager while the loop remains captured; that comparison is not a fair schedule
comparison and is excluded from the controlled measurements above.

Correctness and scope

The migrated tests cover schedule invariants, per-round block/loop equivalence
for greedy and sampling, verification logits, KV and statistics, head calls,
seeded sampling, deferred collection, multiple in-flight submissions,
preemption between rounds, rollback and CUDA Graph/eager parity. The filtered
top-p/top-k rejection-sampling distribution test is retained.

Both modes consume each request generator in the same order: drafts by offset,
then verification. CPU/FP32 same-seed sampled output therefore agrees. GPU/BF16
greedy or sampled output can still diverge when different GEMM row counts produce
different floating-point results; exact GPU/BF16 text equality is not claimed.

Greedy early-rejection pruning is not included. It requires a host readback at
each wave, which conflicts with the asynchronous submit/collect path, and its
previous benefit was limited to early rejection at very low concurrency.

The option is experimental and hardware-specific. It should be enabled only
after measuring round time and graph fallback counters on the target deployment.

[0z5a]

Thanks for the careful write-up. I rebased this onto current main and looked into why the measured gains move around so much. Short version: the schedule itself has no intrinsic advantage over the block round. Measured gains and losses come from GEMM kernel-selection steps and from CUDA graph limits. I have a branch on top of your head that ports the round to main's runner and documents this. You can fast-forward to it or cherry-pick from it.

1. Port to main's runner (needed to merge)

After #52/#128/#141, main replaced the synchronous SpeculativeRunner.execute with planned submit/collect rounds. _core, _prepare_batch, ScheduledItem and Stage.SPECULATIVE no longer exist there, so worker/speculative.py, tests/test_speculative.py and docs/README.md conflict and the old loop code cannot be resolved mechanically. The port:

• plans every wave's KV metadata and gather/scatter indices on the host and uploads them once per round. The schedule (loop_waves(token_counts, d, D)) depends only on token counts, so nothing is built per wave.
• keeps draft IDs on the device and reads back once at collection. The round now composes with deferred collection, several in-flight submits, preemption between rounds and CUDA graphs.
• runs the head only on draft-ready rows inside the loop, plus one head call over all K+1 verification rows after the last wave. This removes the per-wave target head calls, per-branch torch.stack and .clone().
• reuses the block round's _collect_round. Both modes now consume each request's generator in the same order: drafts by offset, then verification.

Two consequences for the PR description:

• RNG order. The description says the loop changes the order in which draft sampling and verification consume the request generator. With this port it does not. On CPU/FP32, same-seed sampled outputs are identical to the block round (test_loop_sampling_matches_block_round_for_the_same_seed).
• GPU equality. On GPU/BF16, outputs can diverge for greedy too, because the two rounds run GEMMs with different row counts (see below). Exact greedy agreement holds on CPU/FP32, not on GPU/BF16.

Greedy early-rejection pruning is dropped. It needs a host readback every wave, which the asynchronous submit path cannot provide. Its measured benefit was limited to C=1 with an early mismatch.

2. Why the gains are unstable

For a uniform round with K=4, d=2, D=4, both rounds launch the same 12 recurrent calls over the same 20·C rows:

Round	Recurrent calls
Block	10 × C rows, 2 × 5C rows
Loop	4 × C rows, 8 × 2C rows

The loop round wins only if c(5C) − c(C) > 4·(c(2C) − c(C)), where c(M) is the cost of one call over M rows. For any affine c (flat when weight-bound, linear when compute-bound) the two are equal. So any difference comes from where cuBLAS kernel selection steps c(M) between C, 2C and 5C, and that depends on the GPU, the library and the model shape.

On an RTX 5090 (Ouro-1.4B, BF16, Triton, graphs, torch 2.13/cu130), the linear layers of one recurrent pass cost c(1)=1.90, c(2..16)≈2.66, c(20..32)≈2.0, c(64)=2.14, c(80..192)≈3.0–3.3 and c(640)=7.9 ms. Nsight Compute shows the cause:

• M=1 uses a GEMV kernel.
• M=2–16 use a 16×16 WMMA GEMM that is about 48% slower than the GEMV.
• M=20–32 use a split-K 32×32 WMMA GEMM that is faster than the GEMV.
• M=80 moves to a larger tile.

Paired results: 7 alternating repeats × 128 tokens, 16 prompts, fixed per-request seeds, zero graph fallbacks in both arms. Ratios are block over loop.

C	ms/round block	ms/round loop	Round-time ratio	c(M) prediction	Throughput ratio
1	44.80	52.97	0.846	0.846	0.846
4	56.24	58.98	0.953	0.955	0.926
8	59.26	61.28	0.967	0.980	1.024
16	63.39	59.97	1.057	1.227	1.028
32	63.13	63.14	1.000	1.067	1.088
64	75.45	80.56	0.937	0.922	0.910

Per-round time is the schedule comparison. Throughput also includes acceptance, which differs once BF16 outputs diverge. For example, at C=32 acceptance is 0.800 vs 0.836, and that alone produces the +8.8% throughput with identical round time. Top-p behaves the same way: round-time ratios are 0.846, 1.021 and 0.965 at C=1, 16 and 64.

Two configuration effects can look like schedule gains:

• Graph batch limit. With cuda_graph_max_batch_size=128, the block round's 5C-row verification calls fall back to eager once C > 25, while the loop round's 2C-row waves still replay graphs. At C=32 on the 5090 this alone moves the block round from 63.13 to 66.0 ms/round (448 recurrent fallbacks) and leaves the loop round unchanged. That is a 4.7% "gain" that comes only from the graph limit. Your A100 table uses limit 128 for C=32 and C=64.
• Graph cache size. Ragged rounds at the end of generation create more distinct wave sizes in the loop round and can exhaust cuda_graph_max_graphs=16.

The benchmark now has --cuda-graph-max-graphs and reports graph counters for both arms.

For reference, your head at its own settings shows 0.857 / 1.177 / 1.202 throughput ratio at C=1 / 16 / 32 on the same GPU (3 × 32 tokens, graph limit 128, default 16-graph cache). These numbers are not comparable to the table above, for three reasons:

• that baseline arm is the old synchronous block round, not main's submit/collect round;
• graph fallbacks were neither controlled nor reported;
• the ratios include acceptance differences from BF16 divergence.

The closest controlled contrast: with the graph limit at 128 and C=32, the ported loop round measures 1.140 throughput ratio, against your head's 1.202, and its per-round ratio is 1.047 (66.0 vs 63.01 ms) because of the eager fallbacks.

3. Benchmark and tests

• Benchmark. It uses fixed per-request seeds, so sampled arms draw the same random numbers. It reports the median of per-repeat paired ratios with a bootstrap interval, quartiles, ms/round, acceptance and graph counters. Defaults are 128 tokens × 7 repeats; 3 × 32 cannot separate ±5% from noise.
• Tests. The existing speculative tests on main run in both modes. tests/test_loop_speculative.py is rewritten against main's runner: schedule invariants; per-round equivalence of results, verification logits, KV and stats for d=1/2/3 with greedy and sampling; head calls inside the loop; CLI and config handling; failure cleanup; rollback. Your filtered top-p/top-k rejection-sampling test is kept.
• Results. Full CPU suite: 877 passed, 378 skipped. GPU-marked suite on the RTX 5090 (Triton): all speculative and loop cases pass in both modes. The only failures are FA4 on SM120 (unsupported), PD cases that need a second GPU, and test_cuda_memory_plan_reclaims_allocator_residue_on_recreation. The same failure set appears on the same GPU for an unrelated branch that does not contain this change, so it is environmental.

Suggestion

The mode is correct and now mergeable, but I would present it as an experimental, hardware-specific option and not claim a throughput improvement. On a given deployment it is worth enabling only where measured round time favours it at the deployed concurrency. A design that actually removes work or critical-path length, such as cross-round overlap or retaining hidden state across rounds, is where a general speedup could come from.

Branch: 0z5a/vllm-rlt:loop-spec-port-20261024 @ 80a737d (6 commits on top of 68c724b, a fast-forward; diff). The first commit (merge of main) does not pass tests on its own, because the old loop code needs the port in the next commit. Please fast-forward or squash instead of bisecting through it.

[YuDeng0102]

Thanks for the detailed investigation and for porting the loop round to the current main runner.

The breakdown of the identical recurrent work, GEMM shape effects, and CUDA Graph limit effects is very helpful. We will use 80a737d (0z5a/vllm-rlt:loop-spec-port-20261024) as the new base, migrate the documentation and tests to the current submit/collect architecture, and update the PR description to present this as an experimental, hardware-specific option rather than a general throughput improvement.

We will also correct the RNG-order and GPU/BF16 equality statements and keep the greedy early-pruning path out of the migrated change.