[RFC]: Mixed prefill/decode batching (last_exited) #81

Progress update — 2026-10-22

Baseline: main 4a350344027960a48384626aa646f1765fda6d4f.

• Implemented: [Architecture] Unify stage execution with loop-level scheduling and mixed batching #118 merged split-prefill and compatible prompt/decode RECURRENT batching. Enable non-atomic prefill with mixed scheduling under LAST_EXITED. PRELUDE/CODA remain separate; speculative rounds are excluded.
• Result application: [Refactor] Move runner result application to Scheduler (M2/R1) #107 merged Scheduler-owned progress/result updates on top of that stage architecture.
• Still separate: multi-chunk wavefront scheduling ([RFC]: Wavefront Chunked Prefill Batching (WCPB) implementation #65), mixed speculative verification, and qualification beyond the measured model/backend/workload combinations.
• Validation: [Architecture] Unify stage execution with loop-level scheduling and mixed batching #118 records CPU/GPU tests and real-model quality/performance evidence with explicit source pins. Its full GSM8K FA4 ordinary/mixed runs pass the declared gate, while Triton runs do not; later final-head changes are not fully qualified by those earlier runs.
• Existing contributor claims remain unchanged. Remaining work should compare this issue's acceptance criteria with current main and record validation/optimization gaps rather than reimplement mixed batching.

The motivation and proposed design below describe the pre-#118 baseline. See current scheduling contracts for the merged implementation.

Motivation.

Currently, prefill and decode always run in separate core calls. A PREFILL batch runs its chunks through all D loops within one engine step (ModelRunner._prefill_tokens()), and decode makes no progress in the meantime. With the default max_prefill_batches_before_decode = 1, RefillPolicy alternates PREFILL and RECURRENT batches when both are ready (PRELUDE and CODA steps aside). So a decode token that runs all D loops can wait behind up to D PREFILL batches, i.e. D × D core calls (16 for Ouro).

#31 gives indirect evidence of the cost: 2P2D cut P99 ITL by 53.5% / 47.9% vs. four replicas at 4 / 8 req/s. That run enabled async scheduling, CUDA Graphs, FA4 and other new features together, so it doesn't isolate this effect.

Mixing the two could help. A decode RECURRENT batch has at most max_num_seqs rows, yet its GEMMs read all recurrent-core weights on every call (about 1.23B parameters, ~2.5 GB in BF16 for Ouro-1.4B). Until they become compute-bound, extra prompt rows add little weight traffic, although longer prefixes add attention work for each prompt row and can limit the benefit. Sarathi-Serve, vLLM V1 and SGLang (--enable-mixed-chunk) mix chunked prefill with decode for the same reason. The CDB paper also notes that CDB "does not require phase separation and could also be combined with piggybacking" (App. A.1).

For a looped model, mixing also requires running prompts one depth at a time. Under last_exited, a prompt row at depth d only reads the depth-d KV of earlier positions, so a chunk can advance one depth per step, next to decode rows at their own depths. shared KV is position-major, so this RFC covers last_exited only.

On #60, @LyxWxj and I agreed on a split. @LyxWxj adds an opt-in per-depth prompt path, where atomic prefill stays the default and prompt work keeps its own batches; WCPB (#65) builds on it. This RFC lets prompt and decode work share a RECURRENT core call, and covers the decode-side scheduling and result routing. #65 lists "Fused prefilling and decoding in the same device batch" as a non-goal, so I'm tracking it here separately. Both parts need a shared item/task contract, which §1 proposes for review before it's locked.

Proposed Change.

Terms used below: a prompt item is one chunk of a request's prompt, a decode item is a request's current decode token, depth is the zero-based loop index, and a mixed batch is one RECURRENT core call containing both kinds.

1. Shared item/task contract (proposal)

On #60, @LyxWxj suggested replacing PREFILL with PRELUDE and RECURRENT. As a variant for discussion, I'd like to propose keeping PREFILL as the lifecycle stage and queue until the whole prompt has completed all depths, while prompt execution reuses PRELUDE/RECURRENT. This follows #65's note that "PREFILL remains a scheduler/lifecycle label for a request whose prompt is not complete". Concretely, prompt items stay in the PREFILL queue at every depth, while mixed execution uses RECURRENT batches. @LyxWxj, this is really step 1's call, so please confirm or push back. My reasons:

• In the RECURRENT queue, prompt items would count toward RefillPolicy's CODA condition (... or not q[RECURRENT]). With min_coda_batch_size > 1, they can delay small CODA batches by keeping q[RECURRENT] nonempty.
• _reserved_growth_blocks() can keep treating Stage.PREFILL as "prompt in progress", which matters with incremental KV plus preemption.
• A mixed batch stays a RECURRENT batch, and the runner and _update() dispatch on item.kind inside it, not on batch.stage == Stage.PREFILL, which today selects the atomic path.

Per-depth state needs a persistent owner, because ScheduledItem is frozen and rebuilt on every _take(). I propose a scheduler-owned PromptItem, kept on the request for now (request.prompt_item) and moved to an item queue once WCPB allows several per request:

@dataclass
class PromptItem:
    token_start: int                   # fixed at creation
    token_count: int                   # min(prefill_chunk_size, max_num_batched_tokens, remaining prompt)
    depth: int                         # next depth to run, 0..D-1
    embedded: bool                     # False until the prompt prelude has run
    state: PromptStateHandle | None    # runner-owned [token_count, hidden] storage

• The chunk length is fixed at creation from the configured budget, so an item always fits an otherwise empty batch. Today _make_scheduled_item() clamps to each step's remaining budget, which can't work across depths.
• ScheduledItem gains kind: ItemKind (DECODE or PROMPT). A PROMPT item references its PromptItem, whose range and depth are authoritative: the runner must not derive them from request.position or request.loops_done. SPECULATIVE items are out of scope.

After depth D-1, the last row is copied into the request's one-row hidden state before the handle is released, num_prefilled_tokens advances, and the prefix is published (publish_prefix() already checks every plane). If the prompt is complete, loops_done is set to D and the request goes to CODA as today. Otherwise the next prompt item is created.

Following @LyxWxj's preference on #60, the scheduler only holds an opaque PromptStateHandle, and the runner owns the storage:

• The handle records its owner (request object and KV allocation, as in AsyncState), so late results and reused request IDs can't touch it.
• It is released after the final depth and in ModelRunner.release(), which covers finish, abort_request() and the batch-failure path in LLMEngine._step(). [REFACT] [M7 2/3] Unify the sampling module, move the sampling RNG into the engine, and split suspend from release #80 would keep release() as the termination path when it splits out suspend().
• It never shares the decode one-row slots, and it uses its own storage, never a view into a mixed batch's output.
• Persistent storage is bounded by max_num_seqs × prefill_chunk_size × hidden elements (8 MiB for 16 × 128 in BF16) and budgeted separately in plan_cache(). Mixed execution may also keep transient buffers alive at the same time (gathered inputs, copy-back, result splitting). Today's startup probe doesn't exercise them, so the implementation must account for them with a conservative bound or an extended probe.

RECURRENT results become typed:

@dataclass
class RecurrentResult:
    decode_signals: list[float] | None   # if any decode item needs a signal: one per DECODE item, in item order; else None

LLMEngine._update() then treats each decode request by its own settings: trace requests exit by trace depth, fixed-depth requests (exit_threshold == 1) exit by max_loops and ignore any signal, and adaptive requests consume their signal as today, where a missing signal is an interface error. Since _update() currently indexes result[index] for every non-trace request, this changes the engine too. Prompt items return no host data; their completion is applied from the submitted item list. Explicit prefill ranges and tagged signals, as #32 wants for ModelRunnerOutput, can follow with R1.

I'd also like a few execution rules to be part of the contract, so step 1 and mixing behave the same way:

• Prompt-only per-depth work returns no exit signals, has no readback, and doesn't compute the gate.
• In a mixed step, the gate runs only when a decode item needs a signal, and only decode signals are returned.
• A fused prompt prelude (§2) uploads token IDs from pinned memory without blocking, since a blocking upload would add a host wait to every mixed step that embeds a chunk. The source must not be reused or changed until the copy completes.

2. Scheduling

Mixing gets its own MixedPolicy, used only when mixed_prefill_decode=True. I propose requiring the fused prelude: an item with embedded=False is embedded in its first RECURRENT step (#65's composite_depth0), which avoids a decode-free PRELUDE step per chunk.

MixedPolicy.schedule():
    if q[PRELUDE]:     return take(PRELUDE)      # decode preludes, as today
    if CODA is ready:  return take(CODA)         # same condition as RefillPolicy
    _admit()
    return take_mixed()

take_mixed() first takes decode items from q[RECURRENT] exactly as _take(RECURRENT) does today. It then takes prompt items from q[PREFILL] in admission order, each whole, and stops at the first one that doesn't fit the remaining max_num_batched_tokens.

• A prompt item that just ran a depth goes back to its original position. Scheduler.enqueue() appends to the tail today, which would make prompts take turns depth by depth and delay the oldest one's TTFT.
• A prompt item's KV grows only when its depth-0 work is scheduled, up to token_start + token_count. ensure_capacity() claims every depth plane at once, and admission already covers the capacity (full reservation, or reserve_outputs with incremental allocation and no preemption).

To make sure prompt work always fits, mixing requires prefill_chunk_size + max_num_seqs - 1 <= max_num_batched_tokens. While a prompt item waits, at most max_num_seqs - 1 decode items are active, so the head prompt item always fits and decode can't starve prompt progress. The item count can't bind first, because each active request contributes at most one item. The defaults (chunk 128, 8 seqs, budget 128) fail this guard, so startup rejects them and suggests something like --max-num-batched-tokens 256 (open question 6).

Prompt items never displace a ready decode item, so MixedPolicy doesn't need decode_due, and admission runs on every core step. Decode can still slow down in two ways: a mixed call may take longer than a decode-only one (§6 measures this), and decode still waits behind PRELUDE and CODA steps.

3. Execution

Each mixed batch is one core call.

• Decode rows come first, one per item, followed by each prompt item's rows in one contiguous run. Decode signals are then just logits[:num_decode].
• One _prepare_batch() call covers all rows. With FA4 this uses packed_prefill=True: decode rows leave the [N, 1] mode and become length-1 sequences, and each prompt item is one sequence. Triton, torch, FA2 and FA3 use the per-row path that non-FA4 prefill already uses.
• Decode hidden states are gathered and saved as today. Prompt rows read from and write back to their prompt state, and are embedded first if embedded=False.
• The gate follows the execution rules in §1, and only decode rows are read back.

4. Synchronization and waiting (bjf-frz's first point on #69)

This section is about explicit host waits in the steady-state synchronous eager path on CUDA with Triton or FlashAttention, without static buffers, prefill UVA, preemption or profiling. It doesn't rule out blocking inside allocations or backend libraries. The torch reference backend reads context lengths back in every attention call, so it's not covered on CUDA.

Since #51, _prepare_batch() copies its metadata from pinned memory without blocking. The host still waits at the .cpu() readback in execute() (decode RECURRENT steps with a signal, and CODA), and at the blocking token upload that starts a decode PRELUDE step or an atomic PREFILL batch.

Compared with step 1's separate batches, and given the non-blocking prompt upload from §1, a mixed step only does the readback its decode rows needed anyway, so mixing adds no sync point. When decode needs a signal and the steps alternate, a prompt's next depth waits for that readback in both designs. Separate steps do allow some overlap, though: the CPU can prepare the decode step while the GPU runs a prompt-only step. Mixing may reduce that overlap, but it also removes a scheduling boundary and a core call, so the net effect has to be measured.

Compared with today's atomic prefill, a PREFILL batch waits once, on its token upload, and then runs all D loops without a host sync. The per-depth path adds a CPU schedule and update per depth.

Prompt progress also waits behind decode's PRELUDE and CODA steps.

For the first version:

• Skip the gate readback when no decode item needs a signal. Signals only matter when exit_threshold < 1, but at the default of 1.0 the sync path still reads back an unused signal on every loop in ouro, ouro_delayed and random_lookahead. Within the scope above, skipping it removes the last explicit host wait from a fixed-depth mixed step; adaptive requests keep the readback in sync mode. This might help plain decode too, and I'll measure that before claiming anything.
• Prompt rows skip decode post-processing (readback, exit check, finalize). They still add CPU work, though: metadata, per-row KV bookkeeping in each layer, and progress updates.

Async is the main follow-up. It needs a mixed-depth, packed metadata path, event-based prompt-state lifetimes, and an in-flight state on PromptItem (like PrefillTask.state in #65).

5. Configuration and scope

The flag is SchedulerConfig.mixed_prefill_decode: bool = False (--mixed-prefill-decode), wired into both cli.py and serve.py; pd_serve.py rejects it. It requires step 1's per-depth path. Default behavior doesn't change, and the new device memory (prompt state plus any transient peak) is budgeted as described in §1.

The first version supports last_exited, synchronous eager execution, refill mode and all four exit modes, prefix caching and incremental KV, FCFS or priority admission without preemption, and the torch, Triton and FA4 backends (FA4 through the packed path, on SM9x or SM10x; on SM90 it has to be selected with flash_attn_4, since flash_attn picks FA3 there). FA2 and FA3 use the per-row path, but as #69 notes, they still need target-device validation. Ouro-1.4B is the initial validation target.

These are rejected at startup, with the reason in the error message:

• shared: position-major prompts.
• Async scheduling and CUDA Graphs: see §4 and §7.
• Speculative decoding: SpeculativePolicy never schedules RECURRENT items.
• PD: the prefill worker calls _take(PREFILL) directly.
• Preemption: prompt-state snapshot and restore aren't defined yet.
• no_refill: its cohort semantics.
• prefill_uva: under mixing, no prompt takes the atomic FA4 path.
• static_buffers and padding: until validated.

6. Validation

On CPU with a tiny model, tokens and exit depths must exactly match step 1's separate batches and atomic prefill, as in test_chunked_packed_generation_matches_serial. This covers fixed depth, seeded sampling and all four exit modes (the adaptive ones with threshold < 1), with fixed-depth and adaptive decode requests sharing batches with prompt rows. The tests also check:

• per-depth prompt KV agreement within existing tolerances, with prompt positions never finalized and exited decode tokens finalized as today;
• prefix reuse and incremental allocation;
• prompt-state release on abort, batch failure, request-ID reuse, and finish of decode rows that shared a batch with prompt rows;
• decode-first selection, continued admission during long prompts, and prompt admission order across depths;
• that mixed batches actually form;
• the budget boundary, at both the selector and the config level;
• rejection of every unsupported configuration.

On GPU in BF16, I'll check greedy agreement with separate batches under concurrent load, reporting the first divergent token and, for adaptive exit, the first exit-depth divergence, since bitwise equality isn't guaranteed. GSM8K-87 also runs concurrently with mixing on. The harness runs one request at a time today, so it needs a concurrent mode.

For performance, I'll submit a prefill-interference benchmark as a separate PR. It injects the same long prompts at fixed points of decode progress, so every arm sees the same prompt load at the same points.

• Arms: A0 is today's atomic prefill (the default), A1 is step 1's separate batches, and A2 is mixed. All three use the same budget, max_num_seqs and chunk size. A1 and A2 also use the same gate/readback skip, so comparing them isolates batching, while A0 vs. A2 shows the combined change.
• A run counts only if injected prompts get a free slot and KV capacity, and the decode population, injection points and request set are the same across arms. Each run records mixed core calls and rows per kind; a mixed-arm run with no mixed core call is invalid.
• Metrics: decode ITL p50/p99/max, TPOT, TTFT, prompt token-loops/s, output tokens/s, prompt depths per step, and prompt work per decode inter-token gap. Host sync sites, counts and wait times come from [FEAT] Add PyTorch profiling with coordinated PD capture and per-rank artifacts #72 profiler traces in separate untimed runs over the same workload, reported per step, per stage and per generated token. torch.cuda.set_sync_debug_mode is only an extra detector, since it doesn't catch every sync.
• Sweeps over chunk size, decode concurrency, prompt length and exit threshold, with each point the median of three interleaved runs on one GPU.
• End to end, as a separate check: A0/A1/A2 replay the same ShareGPT prompts and arrival schedule through /v1/completions with stream=true, on the same single-GPU synchronous eager setup. This borrows [feat] add NIXL PD disaggregation and KV scheduling features #31's workload, not its PD, async or graph setup. It reports client-side TTFT, TPOT, ITL p99, output tokens/s and request failures, plus A2's mixed-call count, since equal arrivals don't guarantee equal overlap.
• The goal is lower decode ITL p99 at equal or better output throughput, with TTFT and prompt-throughput changes reported alongside.

7. Plan

1. Agree on the contract here. @LyxWxj reviews it before it's locked.
2. Step 1 (@LyxWxj): the opt-in per-depth prompt path with separate batches. WCPB proceeds on the same contract.
3. This work, in one PR: MixedPolicy, the row layout, typed result routing, the readback skip, config validation and tests. I'll develop it against the agreed contract in parallel with step 1, and rebase when step 1 lands.
4. Follow-ups:
   • Async (§4).
   • CUDA Graphs. [FEAT] Capture fixed-depth speculative decode with CUDA Graphs #48 added packed-batch capture support for fixed-depth speculative decode; mixed-depth packed batches would still need validation.
   • Composing with WCPB. Item kinds, ranges, depth, the state handle and routing carry over, but the guard and the memory bound assume one prompt item per request.

8. Alternatives considered

• Piggybacking decode onto the atomic prefill's depth loop. This complicates exit handling and async, and overlaps with the step-1/WCPB direction.
• Putting prompt items in the RECURRENT queue. See the reasons in §1.
• Mixing PRELUDE and CODA as well. There's little to gain: prelude is an embedding lookup, and CODA already batches prompt-completed and decoding requests together.
• Giving prompts a fixed share of the budget. Decode-first plus the guard is simpler, and a share can come later if TTFT needs it.

9. Open questions

1. Lifecycle and queues (§1): keep PREFILL, or move prompt items to PRELUDE/RECURRENT as suggested on [FEAT] Implement Wavefront Chunked Prefill Batching(WCPB) #60?
2. Prompt-state storage: per-item tensors or a preallocated pool? And how should lifetimes follow events under async?
3. Prompt-only shortcut: when no decode item is ready, should a prompt item run its remaining depths in one step?
4. FA4 call mode for decode rows: packed, or a second [N, 1] launch per layer? Neither guarantees bitwise-identical decode results.
5. Should max_prefill_batches_before_decode mean anything under MixedPolicy?
6. Defaults and the guard: reject every config that fails the guard, including the defaults, as proposed? Or, in mixed mode only, derive an effective chunk min(prefill_chunk_size, max_num_batched_tokens - max_num_seqs + 1), fixed at creation and logged? That keeps the guarantee but still rejects configs with max_num_batched_tokens < max_num_seqs. Which would you prefer?

Feedback Period.

One week.

CC List.

@LyxWxj @bjf-frz @hsliuustc0106

Anything else.

References:

• Sarathi-Serve: Agrawal et al., Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve, OSDI 2024
• vLLM V1 chunked prefill
• SGLang --enable-mixed-chunk, and sgl-project/sglang#2273 (launching separate prefill and decode kernels for mixed chunked prefill)
• CDB paper v2, App. A.1

Related: #60, #65, #69, #32, #31, #63. Overlapping work: #46, #50, #57, #68, #80, and #32 R1.

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.