Progress update — 2026-10-22

Baseline: main 4a350344027960a48384626aa646f1765fda6d4f.

• [FEAT] Implement Wavefront Chunked Prefill Batching(WCPB) #60 closed without merging; its prototype measurements do not establish WCPB support on main.
• [Architecture] Unify stage execution with loop-level scheduling and mixed batching #118 merged shared stage execution, split-prefill loop scheduling and compatible prefill/decode RECURRENT mixing. [Refactor] Move runner result application to Scheduler (M2/R1) #107 merged Scheduler-owned execution feedback.
• Multi-chunk WCPB remains unimplemented: [Architecture] Unify stage execution with loop-level scheduling and mixed batching #118 explicitly excludes simultaneous same-request chunk/depth wavefront scheduling.
• The remaining implementation must represent multiple ready/in-flight chunks per request, preserve per-depth causal KV readiness and contiguous completed prompt frontiers, and bound retained hidden states and cancellation lifetimes.
• Reuse the merged execution/metadata/scheduler contracts. Evaluate WCPB against current split-prefill and mixed modes, including continuous decode traffic, TTFT/ITL/TPOT, memory and scheduler costs.
• Existing claims are unchanged. The old baseline description and experimental numbers below are historical; no measurements were rerun for this refresh.

Motivation.

Motivation

Issue #2 defines the primary optimization target as faster Ouro inference on one GPU with loop-level batching and last-exited KV semantics. The current runtime already supports chunked prefill, refill scheduling, asynchronous execution, multiple attention backends, CUDA Graphs, prefix caching, and preemption. The scheduler walkthrough documents the existing separation between admission, stage selection, and batch construction, while the module-boundary RFC calls for explicit scheduling tasks and separate ownership of logical request state, KV resources, and device execution.

The current chunked prefill path leaves parallelism inside a request's recurrent-depth dimension unused:

1. The scheduler selects a token range for each request.
2. The runner embeds those tokens.
3. The runner executes all recurrent depths for the selected range before returning.
4. The engine advances the prompt frontier and requeues the request.

With LAST_EXITED, the KV written for c0@d0 is sufficient for c1@d0; c1@d0 does not depend on c0@d1. These dependencies allow a scheduler to overlap work across chunks and depths when the required state is ready.

This is useful when active prompt work is sparse or uneven. A conventional chunk batch may contain only a few rows even though several depth/chunk tasks are ready. WCPB can use those tasks to fill the same recurrent invocation. Smaller prompt tasks also produce more scheduler boundaries, allowing refill policies to admit new prompt work or schedule decode work earlier.

Benefits and Chunk-Size Trade-offs

WCPB is motivated by a real trade-off in ordinary chunked prefill. The chunk size controls both execution occupancy and scheduling granularity:

The small-chunk benefit for decode is a scheduling property rather than a claim that every small-chunk workload has lower decode latency. A smaller chunk gives the policy more opportunities to select a decode batch between prompt batches, but the actual ITL/TPOT improvement depends on stage policy, active decode work, and device execution overlap. The static experiments in this RFC did not include a concurrent decode stream, so they measure the occupancy/throughput side of this trade-off rather than directly qualifying decode latency.

What the baseline measurements show

The controlled A100 experiments make the ordinary-chunk trade-off visible. For four 8064-token requests with max_num_batched_tokens=1024, baseline end-to-end throughput fell as the chunk size became smaller:

The same pattern appears in the mixed workloads. For the 32-request workload, baseline throughput changed from 662.25 tok/s at chunk 128 to 337.51 tok/s at chunk 16, while baseline prefill rows per batch changed from 480.00 to 60.00. For the 128-request workload, baseline throughput changed from 758.72 tok/s to 471.95 tok/s, and rows per batch changed from 451.39 to 97.57. These are paired measurements within each workload; they are not a claim that one chunk size is universally optimal.

Large chunks therefore remain useful when prompt throughput is the primary objective and enough prompt rows are available to fill the recurrent kernels. Their cost is reduced scheduling flexibility: a scheduler sees fewer points at which it can admit work or give decode a turn. Small chunks preserve those boundaries, but ordinary chunking cannot use the recurrent-depth dimension to compensate for the smaller per-chunk batch.

What WCPB adds

WCPB keeps the small chunk as the scheduling unit while allowing ready tasks from different chunks and recurrent depths to share the recurrent-core batch. This directly addresses the underfill mechanism in ordinary small-chunk prefill:

• It does not require increasing prefill_chunk_size to obtain more tokens in a recurrent invocation.
• It can combine c0@d1 with c1@d0 and later combinations when their hidden-state and depth-specific KV dependencies are ready.
• It retains more scheduler boundaries than a large ordinary chunk, while using independent ready tasks to improve the effective batch size.
• It is most useful when only a few requests are active or prompt lengths are uneven; in those cases, request-level batching alone may not provide enough rows.

The measurements support this mechanism. In the four-request 8064-token workload, WCPB raised rows per batch from 64.00 to 254.49 at chunk 16 and improved measured throughput from 275.82 to 412.01 tok/s (+49.38%). In the 32-request mixed workload, it raised rows per batch from 60.00 to 238.60 and improved throughput from 337.51 to 628.45 tok/s (+86.20%). In the 128-request mixed workload, it raised rows per batch from 97.57 to 361.96 and improved throughput from 471.95 to 740.29 tok/s (+56.86%). These chunk-16 results used one warmup and one measured run per variant, so they are directional measurements rather than three-run medians.

WCPB does not guarantee that a small chunk reaches the same absolute throughput as a large chunk. In the four-request workload, WCPB measured 465.31 tok/s at chunk 128 and 412.01 tok/s at chunk 16. Task metadata, hidden-state retention, dependency checks, and additional scheduling work remain costs. The intended benefit is to recover part of the occupancy lost by choosing a small chunk, while retaining its scheduling granularity. The appropriate acceptance target is therefore an end-to-end trade-off: higher effective rows and useful throughput at a chosen chunk size, together with bounded TTFT, decode ITL/TPOT, memory, and scheduler overhead.

For the SHARED KV layout, c1@d0 depends on the final-depth KV of the preceding prompt positions. The LAST_EXITED WCPB schedule therefore cannot be directly applied without changing the dependency graph. The first implementation supports only LAST_EXITED and rejects SHARED.

Proposed Design

Scheduler-visible phases and execution boundaries

The implementation keeps the existing prompt-prefill path and adds staged phases for mixed scheduling and WCPB.

With wavefront_chunked_prefill=False, an ordinary PREFILL batch executes its selected prompt range through the model’s full recurrent depth. In mixed mode, the scheduler can stage prompt work at boundaries needed to schedule prompt and decode requests together.

With WCPB enabled, prompt work uses two scheduler-visible phases:

PREFILLING_PRELUDE
-> embed a scheduled prompt range and prepare its inputs

PREFILLING_RECURRENT
-> execute one ready recurrent depth for scheduled prompt ranges

The prelude is a separate submission; the current implementation does not fuse it with depth 0. A WCPB prelude batch has loop_count=0, while a WCPB recurrent batch has loop_count=1. The scheduler records the execution boundary in the batch plan rather than relying on a user-facing atomic-prefill switch.

The model still processes prompt tokens through the configured recurrent depth before the prompt can proceed to CODA. WCPB changes how ready work is grouped and scheduled, not the logical number of recurrent depths.

Request state and scheduled work

The implementation does not create a persistent PrefillTask table or a general ExecutionTask runtime. The request owns its logical prompt progress; ScheduledRequest snapshots ranges and execution metadata for each batch.

The main progress state is:

• num_prefilled_tokens: prompt progress completed at the final recurrent depth, including progress within the current chunk.
• num_preluded_tokens: end of the prompt range whose prelude has been staged.
• prefill_depth_ranges: the ready token frontier for each recurrent depth.
• prefill_reclaimable_tokens: the start of the oldest resident chunk. It advances only when that chunk completes its final depth and its ring slot can be reused.
• prefill_window_origin: the origin used for ring-slot assignment, including a reused cached prefix.

Each scheduled range carries its token start and count, phase, depth, chunk start, and optional prefill_hidden_row_offset. It does not carry a persistent task ID or task-owned lifecycle state. One request may contribute multiple ranges at different depths to the same WCPB batch.

Readiness and wavefront batch construction

WCPB applies only to LAST_EXITED. Its depth-separated KV planes allow depth-0 work for a later chunk to proceed once the preceding depth-0 KV prefix is ready, even while an earlier chunk is doing deeper recurrent work.

The scheduler uses per-depth frontiers to select ready ranges:

• Depth-0 work follows the contiguous depth-0 KV prefix.
• Work at depth d + 1 follows the available hidden state from depth d.
• A recurrent range is limited by its depth frontier, staged prelude frontier, prompt end, chunk boundary, and remaining token budget.

The scheduler round-robins eligible request owners and ready ranges under max_num_batched_tokens, max_num_seqs, and KV-capacity constraints. It can place work from different chunks, requests, and recurrence depths into one recurrent-core batch. These are snapshots of ready ranges, not independently persisted tasks.

In mixed mode, the scheduler can combine decode and staged-prefill work in a batch when their execution phases are compatible. In particular, decode recurrent work and prompt recurrent work can share a recurrent batch. This behavior is implemented; it is not a non-goal.

Hidden-state storage and runner execution

The runner allocates staged prompt hidden-state storage with prefill_chunk_size * total_ut_steps rows per request. WCPB assigns chunks to ring slots relative to prefill_window_origin:

slot = chunk_index % total_ut_steps
prefill_hidden_row_offset = slot * prefill_chunk_size + chunk_offset

The runner embeds prelude ranges, gathers retained hidden rows for recurrent work, and packs them into the recurrent input batch. It scatters intermediate outputs back to their request ring rows for successor depths. The explicit row offset addresses the full request buffer, including the ring-slot offset. When it is None, the runner derives the row from token_start - chunk_start.

The final-depth output is used for prompt completion and CODA as needed; the runner does not retain a chunk’s intermediate state after that chunk has completed its final depth. Asynchronous submissions retain buffers until dependent device work has finished.

Progress, completion, and recovery

Prelude completion advances num_preluded_tokens and makes recurrent work eligible; it does not advance the fully prefilled prompt frontier.

Recurrent completion updates the corresponding depth frontier. Final-depth progress updates num_prefilled_tokens, but prefill_reclaimable_tokens advances only when the oldest resident chunk reaches its chunk end at the final depth. This keeps partial progress within a chunk separate from ring-slot reuse.

Prefix publication and PD progress are tied to completed prompt chunks. Imported requests initialize their prompt frontiers from the imported prefix. On preemption, staged progress is rewound to the request’s completed prompt frontier; unfinished staged work is recomputed after resume.

Compatibility and scope

WCPB requires LAST_EXITED; the scheduler rejects the SHARED KV layout because its position-major dependencies do not support this wavefront schedule. Unsupported staged-prefill combinations are rejected by scheduling-policy validation.

The current implementation includes paths for mixed scheduling, preemption and recovery, PD chunk transfer, speculative execution under compatible policies, asynchronous execution, and CUDA graph capture/replay. Mixed scheduling with speculative work is rejected; speculative WCPB uses a compatible non-mixed policy.

A generic task table, task-owned hidden-state handles, and out-of-order final-depth completion tracking remain future design options, not part of this implementation. The implementation enforces readiness and completion through request-owned frontiers, scheduled range snapshots, and the bounded hidden-state ring.

Goals

• Increase recurrent-core batch occupancy for chunked prompt prefill.
• Preserve LAST_EXITED causal attention and depth-specific KV semantics.
• Keep prompt prefill at full model depth so the first generated token has the same full-depth contract.
• Work with synchronous and asynchronous execution paths.
• Work with the existing Triton/FlashAttention execution and CUDA Graph infrastructure where a stable batch shape is available.
• Keep scheduler, logical KV management, and device execution responsibilities explicit.
• Make correctness and performance claims reproducible under controlled workloads.
• Provide an incremental path from the current minimal implementation to a more general ready-task scheduler.
• Introducing PD prefill/decode separation.

Non-goals

• Fused prefilling and decoding in the same device batch.
• Supporting WCPB with SHARED KV in the initial implementation.
• Adding preemption or recomputation of partially completed WCPB tasks in the initial implementation.
• Changing Ouro exit policies, gate thresholds, sampling semantics, or decode loop counts.
• Reducing the model's logical recurrent work or claiming fewer FLOPs.
• Replacing the scheduler with a generic task runtime or reproducing the full vLLM scheduler hierarchy.

Feedback Period.

No response

CC List.

No response

Anything else.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues in the issue tracker, and looked for the answer in the documentation.