[Feature] DeepSeek V4 continuous chunked prefill for grouped KV cache

[vegetabledoww]

Summary

After #132, DeepSeek V4 supports prompts up to 8192 tokens, configurable prefill chunk length, and MTP state continuity across chunks. Extend that foundation with starvation-safe continuous scheduling when chunked prefill and decode workloads coexist under grouped KV cache.

Area

Batching and scheduling

Motivation / Use Case

Grouped-cache kernels require each external dispatch to remain homogeneous. Prefill and decode rows therefore cannot share one dispatch, but a fixed phase preference can let a long prefill delay ready decode work or let sustained decode traffic block queued prefill requests.

Depth-2 asynchronous scheduling adds another ordering constraint: a preempted request must not be re-admitted before its older in-flight result drains, otherwise stale output can be applied to restarted request state. Packed-prefill metadata should also be rejected before malformed offsets or lengths can corrupt token and KV placement.

Proposed Behavior

• Keep every grouped-cache dispatch homogeneous: prefill-only or decode-only.
• When both phases have serviceable work, rotate the preferred phase across scheduler steps, for example decode -> prefill -> decode -> prefill, so neither side starves.
• Count admissible waiting requests as prefill demand. If a prefill phase cannot allocate work, allow decode to make progress on the next step.
• Under async scheduling, defer same-step re-admission of a preempted request until older in-flight output has drained.
• Validate packed prefill metadata before input construction: per-request cardinality, positive chunk lengths, seq_len == chunk_start + chunk_len, contiguous packed offsets, one-dimensional token IDs, and matching token/embedding extents.

Chunk size continues to use the controls introduced by #132. For example, an 8192-token prompt can run as 4 x 2048 when the effective chunk size is 2048; 128 remains an internal kernel tile rather than a serving chunk restriction.

Acceptance Criteria

• With grouped cache and runnable prefill plus decode work, every scheduler output contains exactly one phase and successive steps provide fair progress to both phases in synchronous and asynchronous modes.
• A waiting prefill is eventually admitted during continuous decode traffic, and decode resumes after an unschedulable prefill phase.
• A request preempted in an async step is not re-admitted in that same schedule() call, and its older result cannot mutate restarted token, placeholder, computed-token, or KV-cache state.
• Invalid packed-prefill metadata fails early with an actionable ValueError; valid multi-request packed chunks are accepted.
• Existing feat(deepseek): preserve MTP state across scheduler prefill chunks #132 behavior remains intact, including configurable chunk length, 8192-token prompts, AR and MTP execution, and cross-chunk MTP state continuity.
• Unit coverage includes phase rotation, both anti-starvation directions, async preemption/stale-result handling, and packed-metadata validation.
• On-device validation covers an 8192-token prompt split across multiple serving chunks while a decode-ready request remains live; both requests must complete with semantically correct output and without garbled or abnormally repeated text.

Out of Scope

• Mixing prefill and decode rows in the same external dispatch. Interleaving occurs at scheduler-step boundaries because grouped-cache scratch and padding constraints require homogeneous dispatches.
• Reimplementing chunk-length control, 8192-token support, or cross-chunk MTP state, which are already provided by feat(deepseek): preserve MTP state across scheduler prefill chunks #132.

Prototype Evidence

The prototype branch was validated on eight NPUs with MTP enabled. An 8192-token prompt ran as four exact 2048-token chunks at offsets 0, 2048, 4096, and 6144. The isolated baseline and the run with concurrent decode load both returned:

ORBIT-7319 | MAPLE-6047 | QUARTZ-2851

The trace also showed decode/prefill scheduler-step interleaving while both requests remained live.