The unified decode path

This document records the invariants of the decode execution path after the persistent-batch unification: eager decode and CUDA-graph decode share one metadata path, one padding contract, one sampling route, and one output-buffer discipline. A deviation from the rules here is a bug unless this document is updated in the same change.

The problem this solves

Before unification every attention backend carried three decode-metadata implementations: an eager arm inside init_forward_metadata that built fresh tensors per step, a capture arm that allocated persistent buffers, and a replay arm that refreshed them in place. Twelve backends times two live decode paths drifted continuously — replay grew clamps, padding scrubs, and PD guards the eager arm lacked (and vice versa), and graph-only bugs surfaced only in end-to-end runs. A second dark path hid behind the capture ladder: a decode batch above max_cudagraph_capture_size fell back to the eager arm, a code path nothing exercised routinely.

Invariants

One decode metadata path

AttentionBackend.refresh_decode_metadata(bs, actual_bs, req_pool_indices, seq_lens, *, forward_mode, block_tables, num_extends, for_graph_replay, **cache_kwargs) is the ONLY way decode metadata is prepared:

• capture (init_forward_metadata_capture_cuda_graph) is INHERITED: the base default runs the idle-refresh arm (actual_bs=0, for_graph_replay=True) against the runner-seeded seq_lens and the runner's placeholder tables (placeholder_block_tables) — never live tables. Only a genuine capture-only asymmetry overrides it (see "Capture is inherited");
• replay = refresh (for_graph_replay=True) + graph.replay();
• eager decode = refresh (for_graph_replay=False) + the same forward Python the graph recorded.

init_forward_metadata serves extend/mixed (and idle warmup) ONLY; a pure DECODE call raises. There is deliberately no fresh-allocation decode arm anywhere. init_forward_metadata_replay_cuda_graph no longer exists.

Its extend inputs are one required, keyword-only bundle on every node — runner-facing (backends/base.py: router, V4.1, Mamba/KDA, composites) and leaf (backends/paged/base.py) alike: extend_seq_lens, extend_seq_lens_cpu, extend_prefix_lens, extend_prefix_lens_cpu are plain torch.Tensor ([>= num_extends] entries; empty, never None, when there are no extend requests) and extend_with_prefix is a plain bool, and query_shard is the forward's QueryShardPlan (which rows of the span this rank computes under query context parallelism: host integers and the staged row-index tensors) or None when every rank computes every row. Runner-facing nodes additionally take two host-only facts the scheduler knows and only V4.1 plans from: extend_replay_lens_cpu (how many leading rows of each extend re-feed already-cached positions — bounded replay, docs/design/scheduler.md) and extend_prompt_lens_cpu (the whole prompt length, so the backend can tell a prompt-completing chunk from an open one). Leaves never see them: the attention prologue writes every input row unconditionally, so the router and every other runner-facing node call reject_bounded_replay and fail loud on a non-zero replay instead of rewriting rows the prefix hit already shares. Likewise every node without a gathered-history extend arm calls reject_query_shard; the router forwards the shard to its leaves together with page_table_cpu, the host mirror of the extend rows of each leaf's kernel page table (built for a sharded forward, and on a KVP page-sharded group for an extend with a prefix or a pure extend; None otherwise). The one leaf that attends a shard — GPU DSA — counts page ownership from it on the host, and an MLA leaf declaring supports_mla_dcp hands it to _prepare_prefix_history, which splits the prefix-chunk gathers by page owner the same way. On a pure KVP extend the GPU DSA leaf also bounds the index pages any rank owns from it (max_owned_index_pages), so its indexer tiles span that owned width with the same plan on every rank. No default values: the runner passes the [:num_extends] slices of its input buffers on every call (the idle replay passes the empty [:0] slices), so a node that reads a field can never see a silently-defaulted one. This is deliberate — a = False default once hid extend_with_prefix being swallowed by a composite's **kwargs, and FlashMLA planned a ragged prefill for a prefix-cached batch.

Extend planning uses the existing CPU length mirrors for host arithmetic: total cache length is the cached prefix plus the current input length. FlashInfer receives host CSR metadata derived from those mirrors, with fresh pinned sources for asynchronous uploads. Planning does not download device lengths or indptrs; mixed batches apply the same rule to their extend rows.

Prefill token inputs arrive as the scheduler's read-only, one-dimensional int32 NumPy views. Input fill copies them directly into fresh pinned storage for the round, then uploads that storage asynchronously. It neither wraps the read-only source in a writable tensor nor creates an intermediate token copy. The pinned allocation must remain per-round: overlap can prepare the next round while the previous upload is still reading its source.

NPU graph task updates and replay run in inference mode because they write the output tensors allocated during capture. This applies to active and idle DP ranks; callers need not establish an inference context.

Buffer sizing: the ladder is a performance subset, never a capacity limit

ForwardStepRunner distinguishes max_capture_bs (top of the capture ladder, bounded by max_cudagraph_capture_size) from max_decode_bs (max_num_seqs // dp_size, floored at max_capture_bs). Persistent decode buffers are sized by max_decode_bs — init_cuda_graph_state runs unconditionally at wrapper construction, enforce_eager included. A decode above the ladder runs the same refresh with no graph; it is a first-class path, not a fallback.

State paging finalizes immutable metadata tapes for every supported decode batch size when the buffers are built. Finalization uploads CPU descriptions, so it belongs to startup, including sizes outside the capture ladder. Refresh supplies runtime registers and runs a prepared tape. Rebuilding the buffers rebuilds tapes that contain their addresses.

Distributed-argmax eligibility and communication state are prepared after heads, draft sharing and sampling configuration are wired, before graph owners or any forward. Prefill-only, eager and captured execution use that same prepared processor. A prepared fallback is distinct from uninitialized state; decode warmup is not an initializer. The kernel state keys on the runtime logits dtype, which a quantized head does not declare on its (possibly transposed or packed) weight, so preparation builds the state for every supported logits dtype up front, keyed by the head's num_embeddings_per_partition; the forward and the sharded argmax only select the prepared state matching logits.dtype, with the portable packed-key all-reduce prepared as the fallback for any other dtype. They never construct, rendezvous or allocate communication state.

Ascend context-parallel decode

Ascend DSA context-parallel decode uses one prepared communicator for fused TopK selection and LSE merging. Both stages run on the model stream. Its context object owns local page metadata only; it creates no auxiliary process group or stream. Bind the communicator, producer outputs and IPC tables before capture, and retain their storage until graph destruction. Missing resources fail before a context-parallel kernel is launched. The backend takes its CP degree and ranks from the host topology and passes that degree to every fused collective. The installed operators' Meta implementations validate topology and head layout before execution or graph capture; the framework carries no CP degree whitelist. DSA native fusion remains specialized for DCP8; prepared Ascend MLA also admits the H32 geometry on TP8/DCP8. Its head count, Q latent width and borrowed producer views follow the model, with the same projection, owner-store and FP32 LSE merge path. Bindings must supply Meta implementations for this startup check; a missing Meta kernel reports a binding upgrade, separately from unsupported topology.

DSA and MLA share the Ascend CP metadata owner and fused page/length kernel. The metadata is refreshed once per step and shared by that step's layers; captured buffers keep stable addresses. MLA refreshes FIA's host length binding before capture or task update and maps authoritative router slots once before the first Prolog. Its normalized FP32 partials use the same direct LSE merge as DSA, with empty-shard masking and layout conversion fused at that boundary.

Layer-shared Ascend metadata compares retained Tensor storage views when reusing prepared results. Equivalent views reuse the result; replacement storage or an in-place storage rebind invalidates it. The retained views share allocations without copying data and are cleared at the forward boundary. Prefill keeps the same storage-view contract for each layer's prepared cache planes.

DSA CP prepares its metadata before the first fused Prolog. One device launch expands request lengths into causal query rows, compacts owned pages, counts visible/initial/window tokens, writes the Indexer's grouped layout, and maps write slots to local INT64 addresses (-1 for skipped writes). The kernel writes directly into persistent consumer buffers. Layers share this result within the forward; replay recomputes it from the live request inputs. Eager and Graph execution use this same path. This DSA-specific preparation retains the common DCP ownership semantics without changing the shared placement implementation. Prefill and decode shapes without a prepared fused projection use the same persistent metadata outputs. They exchange query and TopK candidates over the existing device group, then merge sparse-attention partials on the model stream. That path creates no auxiliary communicator or stream. The model adapter must pass fused=True after fused Prolog and fused=False after ordinary projections. A fused call without its packed query fails before selection; a live native error is never retried through the collective path. A model with fused decode prepared can still execute ordinary prefill or smaller decode buckets. Capturing collective buckets requires HCCL graph support and target device validation; eager or fused-only validation does not establish that.

The public AscendDSABackend.attention_kernels(is_prefill=...) hook returns the stage's SFA provider. Plugins may override it while preserving the AscendDSAKernels tensor contracts and the common cache/CP metadata. Partial outputs use the shared DCP combine: zero empty-shard O before exchange, then normalize all-empty LSE rows before the native merge.

For T query rows and H global heads, the BF16 gathered query occupies T * H * 576 * 2 bytes. Each FP32 partial-output buffer occupies T * H * 512 * 4 bytes; exchange needs both send and receive buffers, with additional storage for sanitization/layout conversion, TopK and LSE. At 16K rows and 64 heads these are 1.125 GiB for Q and 2 GiB per partial buffer. Plan the prefill chunk against this transient peak as well as resident weights and cache; the native fused bucket's memory plan does not bound this path. The persistent gather workspace also reserves the retained current rows used by attention after exchange, and budgets every transport and slot view with the workspace pool's 256-byte alignment.

Single-rank decode uses the same preparation with identity page ownership: verify still needs one causal length and page-table row per candidate. Its Indexer and attention receive matching per-query batch axes, even without a collective; ordinary single-rank prefill retains its request-span layout.

Ascend vocabulary-parallel greedy selection

Greedy selection keeps each row's logit ordering and token-id tie break in one INT64 key. The HCCL backend handles INT64 MAX payloads up to 4 KiB with AllGather over an INT32 bit view followed by an exact local MAX. This preserves every key bit while allowing AIV communication, whose AllReduce does not support INT64. Larger payloads and other reductions retain their ordinary HCCL implementation. Eager and graph capture use the same collective path and return the original in-place output buffer. Ranks choose the collective from the shared dtype, operation and element count; local layout differences only require packing and copying back.

Rebinding a cache pool

Prefill autotune and graph capture use distinct writable pages per fabricated request and disjoint physical parents per cache group. Groups can share byte planes, so reusing the same parent across groups corrupts their state even when the group-local page IDs differ. The probe arena reserves the sum of each group's rounded-up parent demand, using its virtual packing for sharded groups. Decode capture continues to use the reserved null page.

Memory-saver boots skip the disposable graph-memory probe: the installed allocator cannot destroy a tagged probe graph pool and then capture again. They bind the real cache pool directly and capture the serving graphs once, leaving GPU-utilization headroom for graph memory. The rebind contract below applies to boots that perform the probe.

set_cache_pool may run more than once on the same backend tree: a memory probe binds a small pool, captures into a throwaway graph pool, then binds the real pool. The contract is that a rebound backend is indistinguishable from one first bound to that pool:

• Every node first answers validate_cache_pool for the whole subtree, and only then do the children and the node publish, without a second validation, so a rejected rebind moves nothing (a router's leaves exist only from its first bind on, so the first bind builds and binds them inside its own publish); set_cache_pool is that sequence, shared by every node through CachePoolBinding and never overridden; a node does its own work in _publish_cache_pool (set_kv_pool on the state backends is a retained alias). Atomicity covers rejections only: a failure inside a node's own binding work propagates, and the caller rebuilds the tree. A node rejects a pool that changes the geometry it owns: the router its group geometry (granularities, attention kinds, windows, and the row layout the leaves' kernels read), the state backends the state group ids and the state layers' ids and shapes, DeepSeek V4.1 the group ids and row geometry, Inkling the ShortConv checkpoint groups and their transfer policy. Page counts and, elsewhere, transfer policy may change. Paged leaves own kernel geometry only; the router validates group geometry for them. Qwen4-Exp's PLE and QSA indexer children validate their local fields during this same pass. Publishing the same pool again preserves their verify buffers; a different pool drops them (PLE's commit pointer tables, QSA's verify state); the router's geometry check covers the indexer's tables.
• For nodes accepting pool replacement, binding drops every pool-derived latch: pointer tables, scratch and views, per-forward metadata, the paged leaves' graph buffers, Inkling's ShortConv ring and pending remote restores, and side-state verify caches. The target router recreates KVOffloadAdapter against the new arena's offload engine. The state backends keep their pool-independent index buffers, so a same-geometry replacement stays usable without reinitialization of those buffers (a router in the same tree still needs init_cuda_graph_state before any metadata call); the caller still runs configure_runtime (with the new pool's specs, page counts, and request-slot capacity), init_cuda_graph_state, init_prefill_graph_state and preallocate_verify_workspace again after a rebind, as after a first bind. A probe pool must still hold max_bs state rows: the KDA raw-gate verify scratch is the bound pool's own conv slab. The backend tree covers only itself: the executor's own pool references (token_to_kv_pool, its cache runtime contract, the drafter's pool) and the layer-to-group stamps bind_cache_groups writes on the model are the caller's to re-publish, as are the graph owners' own pool references and the placeholder tables the decode runner sizes from the arena. Of the sequence above only init_prefill_graph_state runs inside capture_graphs(): configure_runtime and init_cuda_graph_state ran before the executor was returned, so set_cache_pool re-publishes them. preallocate_verify_workspace is the factory's, re-issued when the rebind rebuilds through it -- a rebind that skipped the factory would strand it.
• A rebind is an operation between the probe's capture_graphs() and the serving one, owned by reserve_and_rebind; nothing in the backend tree guards against a rebind at another time. That orchestrator releases both graph owners' captures first (the captured graphs record the buffers a publish drops, and eager kernels cache pointers they allocated inside a capture, such as FlashInfer's trtllm-gen MoE runner and Qwen4-Exp's uniform index bundles), unfreezes the device- global workspace pool the executor froze before capturing, rebinds the trees, re-runs bind_cache_groups and the initialization sequence above, freezes the workspace again and captures again.
• The KV budget reserves what the graphs will cost: a probe binds the smallest arena the family can run on (refused, with the numbers, when the memory profile cannot hold it, since a serving arena would not fit either) and captures a few entries of each ladder -- the widest three, then one a third and one two thirds of the way down -- with a driver-memory delta around each capture. The widest samples form a window priced at its positive bytes, plus one granule the window may hide (readings move in 2 MiB, from the driver's graph memory and the allocator's segments alike), over every marginal; each sample further down anchors its width at its reading plus that granule, capped at the window's rate when the reading is within three granules of it (one reading is lumpy) and at the reading less those three granules further above (a dearer entry stays dearer, and a granule more in any reading never lowers the reserve), and a skipped entry is priced on the line between the anchors around its width, or at the narrowest anchor below it. Every ladder is sampled and priced the same way, whatever shape its cost takes down the ladder: flat, falling with the entry's width, or lumpy. That is not a bound: a cost that drops between two anchors is priced short over that stretch. Nor is it known to be tight: Kimi-K3.1 on the PD prefill role, with 60 prefill entries sharing one pool, was reserved 11.64 GB and booted with 6.64 GB more device memory free than the same launch without the ladder. The result is reduced across ranks with MAX. The orchestrator releases the probe's graphs and collects the cycles they sit in, then rebuilds on the memory profile the probe build took -- where a boot without a reserve takes it -- minus the projection. The reserve covers the bytes inside the capture windows as projected -- what a boot without a probe captures there, one-time bytes the first captures take included; the probe releases them and the serving capture pays them again. On CUDA, what executor init and kernel tuning keep resident between the probe build and the probe -- buffers, and on a cold tuning cache the kernels of every tactic tried, though not the stack limit they raised, which is restored after tuning -- is measured the same way, including the free space a kept block pins in an allocator segment, and this startup residue joins each rank's projection before the MAX; its net is floored at zero. Off CUDA it stays on the headroom, as on a boot without a reserve; on ROCm, ROCr keeps the scratch memory tuning grows assigned to its queues, and reclaims it when a device allocation fails. The utilization headroom covers everything else: activations, fragmentation, the warmups and workspaces a capture allocates around its windows, the local memory a kept kernel reserves when it raises the stack limit again after tuning (a driver allocation that first drains the device, and the launch fails if it does not fit), and any shortfall of the projection, as it covers every graph and all of startup on a boot without a reserve. Profiling again after the probe would charge the cache a second time for what startup and the probe left allocated. The deltas read the whole device, so the probe assumes no other process allocates on it during startup. Not covered: a ladder every one of whose sampled marginals was served from slack, which is priced at nothing and says so in the log; and what the probe build allocates after its profile, such as attention backend workspaces, which the headroom funds. The EPD receive pool, which a multimodal prefill node allocates after its cache is sized, is left out of the profile instead, by each rank before the cross-rank minimum.

DP projection communication

DPColumnParallelLinear and DPRowParallelLinear accept full input channels for each rank's own tokens and return complete outputs in the same local token order. Their parallel mapping describes projection weight sharding, independently of attention's token ownership. For example, attention DP4 can use one TP4 projection group with projection dp_size=1; this does not mean the attention inputs are replicated.

Both eager and CUDA-graph execution require explicit physical row counts indexed by global rank in ForwardContext: collective_global_num_tokens from report_collective_sizing takes precedence over global_num_tokens. The counts include any graph padding and match the input rows on each owner. Empty owners participate when another rank in their subgroup has work. Missing counts are an error, not an instruction to assume equal counts across ranks. Ordinary TP with replicated token rows uses the existing ColumnParallelLinear and RowParallelLinear contracts instead.

The model runner calls prepare_dp_linear_communication to prepare fixed-capacity communication workspaces before cache-memory profiling and graph capture. Distributed initialization also reserves the shared all-reduce workspace for the complete enabled fusion token window. A prefill-only process may never warm a decode graph; its first fused reduction must not resize or rendezvous for that workspace on the serving thread. Disabled fusion retains the ordinary one-shot reduction capacity. Preparation binds each Linear and its workspace to the selected communication backend; forward operations dispatch through that same backend. Generic projection operations use the backend's ordinary collectives. AutoBackend composes the optimized projection dispatcher and reuses those generic operations for fallback. Sequential layers share model-private scratch on one stream, sized for the largest projection; matching configurations also share native resources. Concurrent streams or models use separate workspaces. Intermediate tensors borrow storage only for the current projection, so consumers finish before another projection reuses it. Final outputs belong to the caller and remain valid across later forwards. Graphs referencing a workspace are destroyed before it is released.

Fused projection runs communication, quantization, and FP8 GEMM in one kernel. Column projection combines all-gather, GEMM, and all-to-all. Row projection combines all-to-all, GEMM, and reduce-scatter.

Use DPColumnParallelLinear or DPRowParallelLinear with AutoBackend. The backend selects fusion for supported TP2 and TP4 groups within one node. Fusion requires FP8 weights with FP32 scales for each 128 × 128 block. At least one rank must have 64 physical input rows. Smaller batches, batch-invariant collectives, and unsupported configurations use the ordinary projection operations. All ranks use their physical row counts to select the same path.

Both fused kernels accept BF16 inputs and return BF16 outputs. Row projection also accepts E4M3 inputs with FP32 scales per 128 channels. Pass the scales to DPRowParallelLinear.forward through its scale argument. The kernel skips input quantization in this case. All ranks must use the same input format. For FP8 inputs, empty ranks must supply empty E4M3 inputs and FP32 scales.

Prepared FP8 row buckets can transpose activation scales inside 128-row blocks so each K group's scales occupy consecutive shared-memory banks. The exchange and GEMM use the same layout. This is limited to aligned owner boundaries, preserving each owner's byte range when FP8 and BF16 calls share the workspace. Scale values, FP32 accumulation order, and output rounding stay unchanged.

The backend prepares kernels before CUDA graph capture. Prepared row buckets can use different resident CTA counts. CTA zero advances inactive CTA epochs with the active grid so later bucket changes preserve the workspace's collective completion protocol. Each GEMM tile acquires input readiness from every owner it intersects, including tiles that cross an owner boundary. Output stores also wait for those owners to finish consuming the preceding call. A grid that keeps each CTA on one M tile must still wait for both owners when that tile crosses a boundary. TP4 row buckets that clear receive storage during input staging preserve one common clean prefix after reduction. It covers every smaller bucket that skips this staging clear. Larger buckets clear the remainder before publishing their input epoch; all producers and reducers acquire that epoch before using it. This contract applies across bucket and input-format changes. Supply runtime row counts within the prepared workspace capacity. Use each workspace on one stream. Do not run other distributed kernels at the same time.

Row projection rounds partial results to BF16, sums them in fixed rank order in FP32, and returns BF16. This result can differ from an ordinary BF16 collective reduction.

Padding contract

bs is the request count being prepared (the padded graph batch under replay); actual_bs is the live-request count. Requests in [actual_bs, bs) are padding and must resolve to the null page 0 / dummy slot so they never touch a live request's cache. Eager passes bs == actual_bs (unpadded — no wasted FLOPs); actual_bs == 0 is the idle replay. Eager idle bypasses the wrapper entirely (execute_idle_forward calls model_runner.forward(IDLE) directly). With a drafter, the eager idle then asks the drafter for its round (idle_forward_global_num_tokens): one list of per-rank token counts per draft forward the active ranks run, and runs one IDLE draft forward per entry over an empty window, entry i with spec_step_idx=i and that entry's counts — the row shape the active ranks' step i runs (the Eagle chain: the target's rows at step 0, one row per request after; multi-depth MTP: the target's rows at every depth; block drafters: one forward). The executor never derives a drafter's step count or shape itself.

Padding rows still route through the MoE layers, and under attention DP every rank's filler is interleaved with the real rows in the all-gathered MoE input. The expert load counters (--expert-distribution-recorder-mode stat) therefore read a device-side live-row mask (ExpertLoadRowMask, one [max_rows] bool buffer reserved before the first forward) that the graph owners mark from the host's padded/live counts before a padded replay — the decode graph in ForwardStepRunner, the bucket in PrefillGraph.replay — and clear right after it, on the same stream. Eager forwards are unpadded and read the all-True mask; the router counts a route if and only if its row is marked.

Pointer-stable per-bs views from one builder

Per-bs metadata objects (each leaf's _decode_views_by_bs[bs], the router's decode_write_locations views) are views over the persistent buffers, built by a single per-bs builder shared by capture and refresh, cached per bs. A bs never captured (above-ladder decode, enforce-eager) builds its views lazily on first refresh — no new storage, one-time cost. Views must be pointer-stable: a captured graph holds their addresses forever.

Kernel page tables and router write-slot buffers retain int32. Cache binding bounds local and virtual token slots to the int32 range, including sharded groups. Ascend adapters accept int32 router slots and prepare int64 buffers only for native consumers that require them.

Helpers that memoize tensors created inside capture must not return those tensors to eager callers. Keeping a Python reference preserves the allocation, but an earlier graph sharing the same private pool can overwrite its contents on replay. PLE's uniform index bundles are reused during capture only; the eager n-gram kernel writes uniform request indices alongside hash IDs, while ragged batches construct their indices outside the capture pool.

GDN verify shares memoized scratch seed indices (i * (T + 1)) between conv and recurrent reads in eager and captured forwards. FlashInfer FP32 MTP may use uninitialized output and a placeholder for a disabled intermediate cache: live rows are fully written, while negative padding rows skip state access and leave output undefined. Consumers must ignore padded output; enabled intermediate caches always require real storage.

A state group's block table holds two slots per request: the input checkpoint (slot 0: a claimed prefix hit, the block the previous forward wrote before the scheduler rotated it, or a fresh zeroed block for a prompt with no history; null once consumed) and the live block (slot 1). A forward reads slot 0 when it holds a block and slot 1 otherwise, and always writes slot 1 (backends/state/checkpoint.py). State backends refresh every state group's decode blocks in one prep-tape launch, up to eight groups; the tape loops over rows, so it covers every captured batch size. Target verify writes each group's committed-state blocks straight into the captured state_in buffers. The commit enqueued after the replay reads those buffers before the next refresh rewrites them. Verify keeps state_out at pad_slot_id and refills it only after a decode refresh has written live blocks into the same per-bs buffer.

After verification, GDN, KDA, and PLE resolve the accepted state's destination with commit_state_pages, once per state group and only for live requests: the live slot, with the acceptance clamped, in one launch. state_verify_commit_rows maps those blocks to layers and computes request * (verify_width + 1) + accepted for batched copies and ReplaySSM. Its inputs and outputs are contiguous: blocks are [groups, batch_size] for grouped state or [batch_size] for PLE, so no explicit strides are needed. Non-positive blocks resolve to row -1 so copies skip the null page. Keep this arithmetic in the kernels, without eager casts, gathers, index_select or repeat. GDN and KDA share their backend block resolver; PLE uses its own group's block vector and copies the shared context once and local convolution states in one batched launch.

GDN (prefill, decode, and verify), QSA, and gated residual kernels follow pdl_enabled(), passed explicitly to QSA indexing kernels. Waits precede producer-owned reads and outgoing triggers. A trigger permits successor setup, never publishes results; each kernel may delay it for performance. Streaming top-k, for example, avoids delaying scoring waves with waiting merge CTAs. Graphs retain their captured PDL setting; recapture to change it.

fused_gate_sigmoid_mul_add, sigmoid_mul, silu_and_mul, swiglu_oai, situ_and_mul, add3, and split AttnRes launchers read pdl_enabled() themselves. They use that same value for ENABLE_PDL and launch_pdl; model layers do not pass the platform PDL setting through their calls.

AttnRes partial kernels may trigger their successors before writing partial scratch. attnres_combine may preload only weights known to be independent of its predecessor; it waits before loading the prefix and the partial scratch (m, s, acc). A PDL trigger permits early launch but does not publish stores, and a later wait cannot repair values already loaded into registers.

Gated RMSNorm preloads weights only with weights_independent; a contiguous copy disables this preload. At RSAG-to-AR boundaries the next combine-norm preloads the all-gathered residual before its wait, so that collective must not trigger early. FlashInfer adapters preserve the upstream CuTe body and keep PDL compilation caches separate.

The Rubin GDN decode adapter skips a CTA when either state index is negative, before entering the upstream body that would otherwise redirect padding to row 0. Its PDL wait precedes the index reads and padding CTAs still signal. Automatic GDN kernel selection includes the bound state dtype so the generic FlashInfer BF16 path cannot overwrite the arena's null state. Recurrent state remains in the shared cache arena; kernel selection adds no state pool.

Gated residual projection weights separate logical rows from stored rows. Model loading extends BF16/FP16 down/inject storage for GEMM output alignment; later shards and reloads write into the same parameter storage. Forward passes the logical count explicitly and never repacks weights. The GEMM implementation computes the stored width and consumes only logical columns, while CuTe gives TMA a logical view of that same storage. Eager and CUDA graph calls therefore share both the weight layout and its update semantics.

Startup uses the ordinary dummy model traversal with kernel_warmup to prepare CuTe QSA, streaming top-k and gated residual configurations. It follows normal dispatch overrides and scratch budgets, independently of optional FlashInfer tactic profiling. Pipeline parallelism skips this traversal. The dummy extend has no cached prefix, so the MHA backend, which sends such an extend to its prefill kernel, also runs its cached-prefix extend kernel on the same rows during this traversal; that kernel compiles every variant a later batch can pick.

CuTe QSA keeps query rows and selected-slot row strides dynamic; model geometry, selection budgets and launch configurations stay static. Warmup covers encountered static signatures, including aligned and partial stream top-k query groups; other variants compile on demand.

PrimTS prepares native group/scheduler variants only during FlashInfer autotuning, retaining its runtime heuristic. Preparation uses scratch state independent of serving KV bindings and graph plans. Each actual geometry still needs an eager plan before graph capture.

QSA logits scoring uses the same paged kernel for every query layout. A batch whose request lengths are all available on the host may shorten its compressed block-table view to the maximum prefix-plus-query length, rounded to the cache group's logical block granularity. This changes neither the allocation nor the page mapping. Mixed batches with device-only decode lengths, and persistent decode views, retain the capacity bound; decode graph shapes stay fixed. Uniform query runs may share a K tile, but groups must never cross requests. A single-request forward is uniform regardless of its forward mode. Long runs use larger query groups; ragged layouts retain independent rows. Scoring writes only each query's complete-block prefix. CUDA, Triton, and gfx950 Gluon top-k read that per-row frontier on the device, so the score tail may remain uninitialized even when the frontier shrinks during graph replay. A bounded set of scoring CTAs walks tiles at the grid stride, stopping at each query group's current device frontier and reusing Q across its tiles. CTAs beyond that frontier skip Q/K reads, the dot and score stores, including for padded graph requests. Allocated score width retains its capacity bound; the grid depends on SM count, query-group count, capacity, and the loaded kernel's occupancy, never live lengths. Its CTA budget is three waves of the kernel's resident capacity, cached per compiled specialization in the platform layer. Both grid dimensions fit within that budget; excess query groups and longer frontiers need more device-loop iterations, not a new graph or a device-to-host read. Query-group counts use two compilation buckets: small grids fold away the group loop, while larger grids retain it with a runtime group count. This preserves decode register use without specializing on individual batch sizes. Only the query-group count disables implicit integer specialization; table views keep their allocation stride and score widths keep page alignment.

QSA top-k uses separate CUDA, Triton, and Gluon solutions selected through select_kernel; CUDA launch details stay inside its solution.

On NVIDIA, QSA scoring and both CUDA and Triton top-k follow the captured PDL setting and wait before producer-owned reads when enabled. gfx950 Gluon has no dependency wait. DSA and DSv4 persistent decode selection also follow that setting, including cluster launches. Small non-cluster PDL grids use a one-block register budget when their row count fits the SM count; larger grids retain the two-block budget. These precompiled launch-bounds variants use the same algorithm and captured shapes, without reading device lengths. The compiler determines register allocation and occupancy. Register and streaming selectors retain both budgets: spreading waiting CTAs also benefits short-row selection while scoring finishes. Serving workers use one GPU or homogeneous full GPUs, so the SM count used for this policy is cached once per process.

for_graph_replay is for graph-mechanics asymmetries only

for_graph_replay=True means a graph is in play — live replay AND the base default capture (which runs the idle-refresh arm). Two sanctioned branches on it exist:

• FlashMLA's tile schedule: flash_mla freezes its schedule on the first kernel call against a FlashMLASchedMeta (a request that has since crossed a page boundary loses its newest page), so the object is bound to one seq_lens value: eager refresh and every drafter seq_lens edit (advance_draft_forward_metadata, fill_block_decode_seq_lens) bind a fresh one, while a replay refresh leaves the slot alone — the captured graph re-runs the recorded schedule-builds, one per edit, against the live seq_lens buffer. The object lives on the backend, not on the decode views.
• DFLASH block-arm seeding (not for_graph_replay or actual_bs == 0): the drafter's recorded fill_block_decode_seq_lens rewrites the block-end lengths inside every replay, so only eager steps and the capture-time seeding fill them from Python.

Do not branch on this flag for anything a shared in-place refresh can express.

Capture is inherited

init_forward_metadata_capture_cuda_graph has a base default — run the idle-refresh arm (actual_bs=0, for_graph_replay=True) over the same persistent buffers replay refreshes — at both tiers: AttentionBackend (runner-facing; the router's version idle-fills its table stacks, republishes the decode write-location views, then runs each leaf's capture hook) and PagedAttentionBackend (kernel-facing leaves). That default IS the capture for every backend except a closed list of sanctioned overrides, each tied to something the idle refresh cannot express:

• FlashMLA (leaf): installs the keepalive tile-schedule object whose schedule-build the graph records (flash_mla freezes its schedule on the first kernel call against a sched-meta);
• Mamba (MambaAttnBackend): the warmup kernels need the arange query-start-loc, which the idle refresh deliberately zeroes;
• Inkling: conv-state seeding (paged conv reads pos = seq_len - 1, so capture must seed real lengths);
• HybridLinearAttnBackend / Qwen4ExpBackend / MSAHybrid: pure fan-out to their children so the real captures above are reached.

A new backend implements refresh_decode_metadata and inherits both init_cuda_graph_state (the page-table / cache-seqlens pair, sized by block_decode_expansion; extend it for extra persistent state) and capture; a new override must name its kernel-imposed asymmetry here. Leaf capture/refresh signatures are pinned by test_unified_decode_path.py::TestLeafSignatureConformanceTest.

Graded CUDA-graph support

A backend's graph capability is a class attribute, cuda_graph_support: CudaGraphSupport(decode_graph, prefill_graph), which an instance may only narrow or widen from its construction arguments (the DSA leaf below), never a scattered executor-side arch check. ModelExecutor.__init__ AND-composes it over the target and draft child_backends() trees once (resolve_cuda_graph_support), logs every culprit class, and downgrades the two graph subsystems (ForwardStepRunner.disable, PrefillGraph.disable). DSABackend disables the prefill graph (rationale on the class) unless its consumer narrows a padded replay to the live rows in an eager break: a KPool instance (GLM-5.3-Flash), and Kimi-K3.1 on the prefill role and, without --enable-mixed-batch, on the aggregated role (DSAConfig.padded_prefill_replay). That instance declaration follows launch arguments, so the prefill role and Kimi-K3.1's aggregated role check at boot that every rank plans the same prefill graphs (check_rank_uniform_prefill_plan; keyed on the role and the checkpoint's ModelConfig.dsa_padded_prefill_replay, never on a rank's graph decision) under the same indexer prefill logits budget, which sets how many KVP exchanges those breaks make. Qwen4-Exp's PLE and QSA indexer support the breakable graph: their model methods are eager break points, where live request metadata controls cache writes and padded token rows are sliced before state updates. The root composes its actual children, so a restriction from any other child still applies when there is no GDN leaf.

Rules: declarations are static "never works" facts — a runtime prefill capture failure is FATAL (no silent eager degrade: a family that cannot capture must declare it, or the boot dies). Resolution is device-side at startup and driven by class attributes or by instance declarations derived from replicated configuration, so every DP rank derives the same answer (event-loop.md). disable_prefill_graph in the config carries user intent only. decode_graph=False still requires refresh_decode_metadata and init_cuda_graph_state — eager decode runs the same unified path.

Prefill graph warmup runs on the same side stream as capture. Some kernels cache occupancy by stream and reject a cold CUDA graph capture; warming on the default stream does not prepare their capture-stream state. The capture stream waits for the dummy inputs and metadata before warmup, and warmup completes before capture begins. This applies to ordinary, encoder, and decoder captures. Graph-memory observation wraps only capture; eager warmups, decoder rearming and smoke-test replay stay outside the measured region.

One output layout per forward

Every ForwardContext and grammar completion carries a required, immutable ForwardOutputLayout. Ordinary prefill, mixed, and decode batches use the same contract as compact outputs: each emitting prefill has one output row, and each decode has its fixed verify width. Ordinary models emit one row for every extend request; a backend that skips incomplete-prefill outputs shortens only the emitting prefill prefix. Missing layout is not an execution mode.

Sampling parameters, cache progress, and acceptance lengths remain indexed by the original requests. Logits and token storage use the output layout's prefill/decode slices and per-request offsets. Grammar masks retain a fixed width per original request: sampling selects the first mask of each emitting prefill and the full mask span of each decode. This rule applies equally to ordinary mixed batches and batches with zero-output prefills. Grammar candidate preparation reads only the decode suffix of the live input buffer.

Graph capture uses the captured batch size in its layout. Replay temporarily pairs the padded context with a padded layout, then restores the original live layout before output/state processing. Grammar completions retain their immutable per-step layout, independently of later context updates. Idle graph warmup also supplies an explicit layout.

Prefill graphs around a row narrowing

A prefill forward whose row count drops once, at a fixed layer, by an amount that is not a function of the token bucket cannot be one token-shaped breakable graph. DeepSeek-V4.1 is the case: its CED decoder (layer 20 on) runs on a per-request tail of the prefill rows (decoder_view() — a completing prompt's last window, no rows for an open chunk, every decode row), so the row count at layer 20 depends on which requests complete.

The model declares the split instead of opting out: it implements PrefillGraph's NarrowingPrefillModel contract — encoder_forward (the token-shaped layers), narrowing_forward (the candidate source layer, all rows in, the view's rows out), decoder_forward (the remaining layers and the final norm on whatever rows it is given) and finish_forward (the sampled-row gather, the DSpark row report); forward is their composition, so eager and graphed prefill are one path. PrefillGraph captures the encoder per token bucket and the decoder per decoder-row bucket, from a fixed-row static state the narrowing lands into (leading rows copied, tail zeroed). The decoder graphs depend only on their row count, so one ladder serves every token bucket; it is the token ladder clipped to max_decoder_rows_per_request × max_num_seqs (a request contributes at most its window). A replay is encoder graph → eager narrowing → decoder graph, all under the bucket-pinned ambient context. The narrowing and decoder stages size their own collectives from their row counts (report_collective_sizing). The decoder graph replays with the narrowed row count as its valid rows, so its breaks scrub the static tail. A forward whose narrowed rows exceed the largest decoder bucket runs its decoder stage eager. Layers read their row plan from the live context, never from a loose argument a captured break would freeze. Capture runs the narrowing before every decoder run, as serving does: the decoder consumes per-forward backend state its predecessor produces (V4.1's reuse layers read the index source's selection, which later sources overwrite). Under attention DP the split graph stays off: the narrowed row count is rank-local (which prompts complete on this rank), so the decoder bucket and the collective shapes its graph bakes would differ across ranks, and the stages size their collectives from their own rows, which the DP metadata gather does not carry. Eager TP1 attention/dense with fused EP may narrow: the fused dispatch exchanges the actual local row counts, while the original DP counts remain capacity bounds. The CPU metadata gather also carries per-rank decode request counts. V4.1 DSpark sizes its one draft forward at block width times these counts, including in aggregated mixed batches. Prefill writes context KV; if another rank has decodes, prefill-only and idle ranks join its draft collectives with empty rows. A round with no decodes performs no draft forward. Decode graph warmup and capture use the same uniform padded decode counts as their target batches.

V4.1 completes deferred HC posts and reductions before stage handoffs or Engram updates. The fixed graph input state contains only completed residuals. Prefill warmup runs on the eventual capture stream because MegaMHC initializes per-stream barriers. Allocation observers wrap capture alone, leaving eager warmup allocations in utilization headroom. Fused HC all-reduce admission respects both deterministic and batch-invariant collective settings. Empty consumer rows carry no deferred post. Cache-only narrowing completes the incoming post before the global KV producer and skips new HC mixes; idle ranks still participate in expert-owned dispatch/combine, which never defers another reduction.

Drafter prefill graphs

With --enable-drafter-prefill-graph the drafter's work inside an extend forward replays from graphs too. The drafter declares what it records (BaseDrafter.extend_graph_plan), and the capture and the CUDA-graph memory projection read that one plan: Eagle's step 0 per target token bucket at each exact request count of --prefill-graph-capture-batch-sizes (prefill:drafter-step0), and its decode steps per request count (prefill:drafter-decode). They are breakable graphs in one private pool, apart from the target's prefill and decode pools. Their static buffers (the target's hidden rows at the largest bucket) are allocated outside the measured windows, like the target graph's static outputs, so the utilization headroom funds them. A drafter's single prefill graph of its own (captures_prefill_graph, the prefill:drafter ladder) is a separate hook with its own contract.

Step 0 is token-shaped up to the narrowing and one row per request after it, so its graph pads the rows to the bucket and bakes the request count. The sequence mixers stay at their breaks (the QSA indexer, PagedAttention.attend), and so do the drafter steps that publish backend state a later break reads: the accepted-prefix publication, and per decode step the top-k reset, the seq_lens advance and the write window (run_eager). Recorded into a segment they would pin per-forward host state to the capture's batch (the indexer backend's active metadata, the router's published window). Step 0's context is one fixed object per (bucket, requests); the replay stages its inputs into static buffers (the target's hidden rows with a zero tail, the gather rows, the frontier), and its outputs land in static buffers the decode steps read. The draft prologue writes K/V inside a captured segment, which holds no per-layer Host-load wait, so a step-0 replay first waits for the draft pool's whole Host load, as the target's replay does. These captures run each segment as soon as it is captured (run_segments_while_capturing): their breaks index with values a segment computes (the top-k rows the narrowing gathers, the decode steps' positions), which a capture alone leaves unwritten when the break runs.

A round replays only what the graphs baked: a pure extend (MIXED rounds run eager), no attention DP or query shard, a fitting bucket at the exact request count (a round over the largest bucket still replays the decode steps). Only a chain drafter on one rank, without a draft sampler, token history or a top-k carried across steps, over a draft model that declares supports_breakable_draft_extend (its draft backend also attends in the mode a break passes it), captures at all. ServerArgs refuses the switch for the launches its arguments rule out, and ModelExecutor for a boot whose plan is empty.

A replay is not bit-identical to a serving-time eager round in general. The graphs keep the kernels startup dispatch chose for them, as the decode graphs do: small-M BF16 GEMMs route to FlashInfer only before serving (compile_monitor.is_serving). A padded replay also runs the token-shaped kernels at the bucket's row count, where a kernel may pick its implementation or tactic by row count, as the target's padded replays do.

Prefill requests without generated outputs

The original execution batch and cache metadata always contain every request. A backend may declare skips_incomplete_prefill_outputs only when its model still produces all cache state required by later chunks. DeepSeek V4.1 uses this after the candidate-source layer writes global KV. Other backends keep the original output contract.

The scheduler packs completing prefills first, at most one incomplete prefill last among prefills, then decode requests. An immutable ForwardOutputLayout records E original prefills, P output-bearing prefills, D decode requests and verify width K. Completing prefills must form a prefix (validated on the CPU). Every forward carries a layout; one with P=E is the identity layout, and the executor's fast paths (whole-batch sample, whole-batch verify) still apply to it.

Logits and tokens use P+D*K rows. Accept lengths still use E+D request rows, with zero in [P,E). Sample uses the parameter prefix [😛]; verify uses the original [E:] suffix, retaining its batch-row coin offset. Grammar masks have a separate token axis. Each consumer uses the same token offset: i for i<P, P for P<=i<E (no storage), and P+(i-E)*K for decode. V4.1 explicitly marks selected logits rows so the logits processor never re-gathers them using original input indices.

The layout owns host queries for the shared prefill prefix, the original decode request suffix, the compact decode output suffix, and each request's stored output width. Consumers use these queries instead of deriving the same offsets independently. The executor aligns sampling parameters and token-indexed grammar masks at the sampler boundary; sampler interfaces and result buffers stay unchanged. Grammar owns matcher advancement and rollback, and zero accepted lengths already suppress advancement. Existing consumers that only need token_offset keep that interface. Models without cropped outputs carry the identity layout, under which every slice and offset above reduces to the original request-indexed contract.

Completion remains a local comparison in the executor and V4.1 attention metadata. Both use the same scheduled prefix, input count (including replay), and current prefill target, which may include previously generated tokens after re-admission. Contract tests keep those decisions consistent without adding state or parameters to the shared metadata initialization interface.

Cache advancement still consumes all input lengths for prefill, independently of output lengths. Future input writes, NaN/OOV attribution, grammar advances and V4.1 DSpark anchors operate only on output-bearing requests. Grammar keeps one completion per forward, including zero-output rounds; both deferred hostfunc and host fallback consume the frozen layout. A zero-row decoder bypasses its graph, normalization, LM head, sampler, and draft-context writes; the encoder and global KV producer have already executed.

When the decoder view is nonempty, the candidate-source layer retains the original full-row mHC and QKV projection shapes, then gathers the selected rows. Moving that gather before the projections changes split-K or quantized GEMM arithmetic and can change the retained logits. Only an empty decoder view bypasses those projections; its global KV producer still runs on all encoder rows.

Compute rows and output rows are distinct. When a batch contains a completing prefill or a decode, each incomplete prefill retains its original single decoder compute row. Removing that row changes quantized attention and MoE batch shapes and can alter other requests' logits even with identical input tokens and chunk boundaries. The retained row is omitted from logits_rows, so it still has no sampled token. Batches containing only incomplete prefills keep zero decoder rows and skip the entire decoder consumer stack. This preserves the optimization on cache-only rounds without changing the numerical shape of rounds that produce outputs.

Final prefill windows, bootstrap tokens, and PD candidate/cache handoff remain unchanged. V4.1 PD still requires layerwise transfer interval zero. This does not enable a cache-only prefill role or reduce resident model weights. PP and attention-DP narrowing retain their existing restrictions.

One draft metadata contract

The draft backend's decode metadata comes from refresh_decode_metadata and NOWHERE else — the same two steps in every round:

• decode round: target refresh, then draft refresh over the drafter-owned draft_seq_lens_buf (freshly seeded from the batch seq_lens);
• extend/mixed round: draft prefill init reading the accepted-prefix seq_lens view (never the mutable draft buffer), then the same draft refresh with one token per request — deliberately NOT the packed verify width.

Backends' init_forward_metadata must NOT double-fill draft decode metadata as a side effect (the deleted is_extend() and self.is_draft arms); the mixed/idle decode arms that remain serve the target's decode requests only. Drafters republish their in-loop seq_lens edits explicitly each step through advance_draft_forward_metadata (Eagle) / update_draft_forward_metadata (vanilla MTP frontier re-anchor) — metadata never aliases a buffer the drafter mutates behind the backend's back. Those two hooks are deliberately seq-lens-only: Eagle's step-0 accepted-prefix publish fires advance_draft_forward_metadata BEFORE the step-0 attention has consumed the verify-shaped write window, so the write-window publication is a separate, explicit drafter-loop call (publish_draft_step_locations, see "Write locations have one owner"). The router hands each hook to the leaf's hook of the same name, because the two edits describe different row shapes: the Eagle chain runs one row per request after step 0, the multi-depth MTP window k rows per request at every depth. A leaf whose decode kernels derive each row's causal bound from the request's single cache length needs the same seq_lens edit for both (the PagedAttentionBackend default routes update_ to advance_); a leaf holding per-row decode metadata re-expands it in update_ — DSA rewrites its per-token indexer rows (_dsa_seq_lens_2d, [bs * k, 1]) and their plan to the frontier, in place, while its advance_ re-plans [bs, 1] rows and leaves the per-token rows as the round's refresh published them. Neither hook clears the layer-shared sparse selection: the depth loop is one forward's worth of top-k reuse.

Backends with sharded KV must refresh derived local visibility in the same draft length-update hook as the global lengths. While page allocation and request order stay unchanged, they reuse the compact tables and ownership prefixes from the full refresh and update local visibility in place. Eager execution and CUDA graph replay use the same hooks and persistent buffers.

One named exception: draft-tree lanes (docs/design/tree-speculation.md) read TreeDraftInputs, which the drafter writes inside the round -- the frontier and lane window lengths once, then each step's lane masks, plus active, a Python flag set around each lane forward. The buffers are bound once, live at fixed addresses, and are written by in-graph ops before each lane forward reads them; the draft leaf's decode metadata itself is still refreshed only as above.

Step 0 narrows rows; the drafter owns the lengths, the model names the moment. Eagle's step 0 runs over the target's verify window (N rows per decode request), writes KV for every row, and continues from one live row per request (gather_ids), whose context is the accepted frontier valid_cache_len + accept_len — not the vc + N the round's refresh published. The drafter computes that frontier once per round (it is also step 1's cache_start) and attaches an AcceptedPrefixPublisher to the step-0 context as ctx.draft_narrowing; the model calls publish_accepted_prefix() right before the first kernel that reads the live rows (the MLA/MHA drafts at attention start; the QSA indexer after its verify-window layout, since that layout is derived from the decode-slot lengths), and a draft whose step 0 attends the whole verify window (the GLM DSA NextN heads) never calls it — the step loop publishes for step 1+. The call is idempotent (a copy of a fixed tensor into the leaves' buffers), so it carries no single-layer restriction. ForwardContext carries no drafter tensors: accept_lengths and draft_seq_lens_buf are gone, the handle's presence is the step-0 discriminator, and no model computes or edits seq_lens.

Both steps run unconditionally — there is no per-drafter opt-out. What makes that safe is the slot discipline: init writes prefill-slot metadata, refresh writes decode-slot metadata, and forwards read the slot matching their mode (forward_prefill_metadata / forward_decode_metadata; Inkling's conv wrapper mirrors this with conv_prefill_metadata / conv_decode_metadata). A round that runs no decode steps (vanilla MTP re-runs prompt requests as EXTEND depths) leaves the refreshed decode slot unread; a block drafter (DFLASH) re-runs the same refresh inside each block-decode step, overwriting it. A backend that lets one call clobber the other slot's metadata is in breach — that, not drafter special-casing, is the invariant to fix.

PD decode nodes

A PD decode-only node never runs an extend forward, so latches set on the extend path (the pre-unification _cache_groups_bound) stay False there. Refresh must therefore bind the group tables whenever they are delivered — never gate on an extend-latched flag — otherwise the kernels read the null page instead of the transferred KV. This rule predates unification and now protects eager decode too. (_cache_contract_bound is gone: every LCM pool publishes a cache contract, so the target allocates its write-location buffer unconditionally and drafts are gated structurally on is_draft.)

K3 DSpark pipeline prefill distributes target-tap projection across stages, while the final stage owns the proposal network and draft cache. After target prefill sampling, that stage runs the ordinary drafter; its completed call publishes the final cache producer barrier. PD transfers the sampled anchor and real draft candidates with the target and draft caches. Decode installs that window before its first ordinary verify round. Stage ownership changes where context and proposals are produced; candidate handoff and verification follow the same path as other speculative prefills.

Known limitation, multi-depth MTP (Mtp): the drafter's cross-round stash — per request-pool slot, the last k-1 committed tokens and the target hiddens one position behind them, which seed the rows of the frontier-anchored decode window that lie before this round's verify window — is drafter-private state the prefill node fills during its extend catch-up and the bootstrap payload does not carry. After a PD landing the decode node's first rounds read the slot's stash as it stands (never filled for this request), so the depth loop rewrites up to k-1 draft-KV positions in the prompt tail from wrong inputs until the stash has rolled those entries out (at most k-1 rounds; the prompt-tail planes the prefill node transferred were correct). Draft quality only: verification is exact. Intended fix: ship the slot's stash rows with the bootstrap payload, the way K3 DSpark hands over its anchor and candidates.

PD prefill nodes

The prefill role is not an eager role; it is a role with no decode step. ModelExecutorConfig.prefill_only turns the decode graph off (ForwardStepRunner.disable) because there is nothing for it to capture — the role's attention is configured at verify width one and allocates no verify scratch for a DECODE-shaped dummy — while the prefill graph keeps the same gating as any server (--enforce-eager, --disable-prefill-graph, --prefill-graph-max-tokens, the backend's declared support). Its extend forwards, chunked or prefix-hit, replay the breakable prefill graph through the same _run_target_forward dispatch; the KV handoff to decode is ordered behind the forward exactly as behind an eager one (the plan's remote-decode batch is emitted only once the final chunk's result has landed). Layerwise transfer keeps working under replay because the cache-step record lives inside the eager attention break (record_pd_cache_step, record_layer_cache_ready), after the layer's KV write on the same stream.

Pipeline parallelism is the one prefill configuration that forces eager: each stage threads its boundary state through an eager stage forward (ModelExecutor._run_target_forward), so ServerArgs.resolve_disaggregation sets enforce_eager for --pipeline-parallel-size > 1, not for the role.

Sampling has no greedy branch

Greedy requests normalize to top_k=1 in SamplingParams.__post_init__; the pool-indexed sampling route serves them, which is exactly what the captured graph records. SamplingBatchInfo.is_all_greedy and the eager-only argmax branches were deleted. Equivalence (top_k=1 == argmax, ties excepted) is pinned by test/runtime/sampling/test_greedy_route_equivalence.py (for sonic, whose kernels rank candidates at bf16 resolution -- a penalized greedy pick lies within one bf16 ulp of the maximum -- and realize top_k=-1 as bounded top-128 truncation, by test_sonic_backend.py).

Admission uses the selected sampling backend's finite top-k bound. GPU backends retain the exclusive 128 bound. ascend_full accepts positive finite values and caps them at 4096 when populating its per-request scalar pool, before the int32 write. The normalized disabled sentinel (top_k=-1) keeps whole-vocabulary sampling. The captured graph reads that same pool for single-step sampling and chain verification. Ascend sampling consumes the logits already audited and sanitized by the executor's NanGuard; it does not repeat a host-side NaN truth test inside the captured graph. Per-request numerical error flags keep their existing executor and output-processor path.

Synthetic acceptance length shares executor scheduling, logit processing, persistent output buffers and downstream state updates with ordinary verification. Only the greedy sampler supports this benchmark mode. It runs its ordinary full-width verification kernel before overriding the result with a forced draft prefix and a target token at the selected cutoff. This retains verification cost in the benchmark; synchronization and downstream state updates consume only the final synthetic outputs. The override reuses the target argmax at the cutoff, without resampling or a second executor path. DP sampling and other sampling backends are rejected. The persistent length buffer is refreshed in prepare_step or prepare_capture, outside graph capture, using a private device generator. Preparation must remain asynchronous so it does not stall forward launches. The sampler precomputes the configured length's floor, fractional part and integer flag at initialization and reuses them for each buffer refresh.

Non-speculative serving is the N == 1 case, not a second path

One sampling rule for every batch: prefill requests sample, decode requests verify (ModelExecutor._run_sampling). The decode candidate window is always [num_decodes, output_length] (_decode_candidates, a persistent input_ids_buf view): column 0 the last verified token, columns 1.. the draft candidates. Without a drafter, output_length == 1 — a one-column window that accepts nothing and resolves to exactly one sampled token through the same pool kernels, accept_length == 1 (test_decode_verify_n1_equivalence.py; Triton and FlashInfer's per-request stream use the same draw in sample() and verify(N=1); FlashInfer's legacy batch stream retains the coin route; sonic resolves the width-1 window through its single-step kernel, bitwise equal to sample(), pinned by test_sonic_backend.py). future_input_map is [pool, output_length] for the same reason: single-token decode is a width-1 candidate window.

Backends express verify geometry as a floor, not a mode: seq_lens clamp to clamp_min(q_len) unconditionally (drafts and plain decode have floor 1, where the clamp is the identity). What legitimately remains conditional on the drafter is the draft model's existence — draft backend refresh and the drafter loop itself — not the sampling or metadata shape of the target.

Sampling positions are independent of the KV frontier

SamplingBatchInfo borrows the accepted frontier (valid_cache_lengths) and the leading prefill requests' input lengths (extend_lengths). Paths that consume sampling offsets require a frontier, including standalone warmup and graph capture; missing metadata raises rather than substituting offset zero. Samplers that do not consume positions may leave it absent. Slicing a mixed batch slices the prefill lengths along with its request rows; a decode suffix has no extend lengths. These are execution inputs, not RNG state.

The sampling backend owns the pool-indexed offsets scratch required by the existing Gumbel/scalar-gather kernels. _prepare_sampling_offsets prepares it immediately before consumption, on the same stream, in eager execution and CUDA graph replay. The executor and graph runner own no sampling-position buffer and run no position-preparation prologue. Non-emitting prefill chunks, intermediate PP stages, greedy-only backends, and coin-based multi-token verification do not prepare offsets they do not read.

Preparation writes only the current call's unique positive request slots in one kernel, O(batch rows), without scanning the request pool. Slot zero is the scheduler's reserved padding slot: it stays zero, so repeated padding rows do not race. Unreferenced scratch entries may be stale and must not be consumed. The scratch is regenerated from execution inputs before use after a restore; it is not part of the request's retraction image. CUDA graph capture records the same preparation inside sampling, with stable scratch addresses. Frontier advancement happens after forward/replay, not before sampling reads it.

A prefill chunk predicts at frontier + input_length; a decode/verify row's first target is at frontier + 1, regardless of candidate width. A chain adds its candidate index and a tree uses node depth. For example, a 2049-token prompt predicts position 2049 whether computed as 1024 + 1024 + 1, 130 cached + 1024 + 895, or 2048 cached tokens plus one decode input. Constructing the base position does not implement MTP acceptance/residual RNG or domain separation; multi-token FlashInfer still uses its existing coins. The Triton chain verify under --enable-speculative-sampling (DFLASH with a DFlash2 draft) keys its accept coins and draft noise by salted offsets disjoint from these positions. A row with a recorded draft q then emits an exact sample of the same distribution, not the position-keyed draw. Greedy rows and rows without a q keep the position-keyed token.

With sampling_stream=per-request, FlashInfer and FlashInferFull route both sample() and verify(N=1) through the existing Gumbel pool kernel, keyed by request seed and this absolute position. Neither consumes the CPU coin generator. Random-stream policy, verification geometry, and TP agreement are separate:

• per_request_stream selects the position-keyed direct sampler for a window without draft candidates. has_draft_candidates describes the window, not a model family: wider windows still require acceptance and residual sampling.
• _prepare_step_hook skips coin preparation only where no coins are consumed. Wider per-request windows currently fall through to _refill_legacy_coins; they do not yet have position-keyed accept/residual randomness. That refill owns only generator advancement and buffer writes, not sampler selection.
• Each selected sampling implementation declares needs_tp_sync locally. The common epilogue gathers DP-owned rows or broadcasts replicated TP outputs as required; it does not infer synchronization from random-stream policy or window width. These are the existing communication choices, not new guarantees that per-request randomness alone makes per-rank probabilities identical.

_pack_single_verify_result only writes caller-owned output buffers (including DP-local buffers): one emitted token per row and its identity output index. The chain path adds its bonus token to the accepted-draft count before the common communication/logprob epilogue. Both paths therefore hand that epilogue emitted-token lengths. Full sampling reuses sample() for penalties and count updates exactly once, then packs its result.

This fixes the previous coupling to the number of prepare_step calls, including open prefill chunks that emitted nothing. Batch order, pool-slot reuse, prefix hits, and graph replay do not enter the random key. The fixed-seed token trajectory changes relative to the old coin-based decode implementation. Logits must still be identical for this to imply token equality, and sampling parameters and any penalty/grammar state must also match.

Multi-token FlashInfer speculative verification retains rejection sampling and its per-slot coin state; this change does not establish its invariance to cache hits or speculative scheduling. The batch sampling stream is also outside this position-based equality guarantee.

test_sampling_positions.py exercises the metadata and slot-write contract on CPU and CUDA, including untouched slots and duplicate padding. test_flashinfer_position_invariance.py exercises real CUDA sampling, bitwise selected-token logprobs, seed/position sensitivity, chunk counts, both phases, batch/slot permutations, and graph replay; it fails if CUDA is absent. These are sampler tests with constructed logits, not full-model acceptance tests.

Outputs are persistent-buffer slices on both paths

Optional sampled-token cumulative probabilities (--enable-output-cumulative-probs) follow the sampled logprobs' rows and lifetime: the forward returns them beside output_logprobs, a decode graph captures them with the same forward, and replay slices both to the live rows. They are computed once, at the common sampling boundary (sampling/diagnostics.py), from a copy of the target scores; no backend computes them. docs/guides/sampling-calibration.md is the contract: ordering, masking, supported configurations and cost.

sample() and verify() land their outputs in each sampling backend's persistent output buffers, on eager and replay alike. The FlashInfer backend packs tokens and accept lengths into one region (_output_pack_buf), so its get_packed_output_d2h collapses the two device-to-host copies into one; the Triton backends and sonic return separate token and length buffers and take the executor's two-copy path (get_packed_output_d2h returns None). The stop conditions cut the returned accept lengths in place (0 for a stopped request), so a backend whose decode rows return a buffer of ones refills it on every call, as sonic and FlashInfer do.

Logits capture state follows the tensor's device

The logits processor queries capture on the tensor's accelerator module; CPU tensors have no capture. Lazy communication-state probes stay outside capture and use the TP group's configured device collective, including HCCL on NPU. The multicast argmax provider retains its NVIDIA eligibility vote; an ineligible group keeps the portable vocabulary-parallel argmax or the plain gather, without entering CUDA APIs on another device.

What stays graph-only

Enumerated residue in ForwardStepRunner.__call__, all tied to the mechanics of replaying a recorded graph: input-buffer padding to the ladder bs plus the DFLASH sentinel req-pool rows, _set_graph_state_write_indices, the DeepEP dispatch-mode restore (deepep_adapter.replay()), the sampler-variant graph_key lookup, the TOKENSPEED_GRAPH_DEBUG metadata verify, output-buffer re-slicing, the ctx.bs save/restore, and the replay itself through the drafter (replay_decode_graph: EAGLE-Harness replays it on its round stream and launches its device-chained round behind it).

Address-freezing bugs — a refresh that binds metadata views over storage the captured graph never recorded — are assertable: capture snapshots the tensor identities reachable from the decode-metadata slots (graph_ptr_guard), and TOKENSPEED_GRAPH_DEBUG=1 re-verifies them before every replay (production replays pay one bool check). The snapshot has no exemption list: every tensor a slot reaches is an address the refresh must keep. Per-step-mutable objects a kernel owns (FlashMLA's tile schedule, which the kernel builds and freezes on first use) therefore live on the backend, outside the slots, not on the views — and so does the per-forward memo the models' layers share: the sparse indexer's selection (AttentionBackend.sparse_topk, a SparseTopKShare: GLM DSA's "shared" layers and the DSA / QSA MTP heads reuse the last indexer layer's top-k; QSA also keeps its layer-invariant row geometry in that share so the fused preparation runs once per forward). Every runner-facing node clears it when it builds a forward's metadata (the router's extend init / decode refresh / capture seeding), so the first layer computes, the rest reuse, and nothing outlives its forward; the drafter's in-loop seq_lens edits are not a new forward and leave the share alone — the drafter itself hands each draft step the top-k it reuses (or clears it) through the draft backend, and starts from the target backend's. The MHA leaf's decode split schedule (decode_schedule_buf, read by gfx950 verify) is the same kind of memo: its decode refresh marks it stale, the first verify layer computes it from the forward's lengths (inside the graph under replay), and the rest reuse it. ForwardContext carries none of this. What unification still can NOT test: mempool reuse and hostfunc semantics — the e2e regression matrix keeps graph-on and graph-off configurations for this reason.

Backend package layout

layers/attention/backends/ is organized by the role a node plays in the tree, not by model: base.py (the runner-facing AttentionBackend contract and the per-forward SparseTopKShare), support.py (graded CUDA-graph support) and cache_metadata.py (the runner's block-table bridge) stay at the root; paged/ holds the block-table route — the CacheGroupRouter, its geometry / table-stack / write-location helpers, and every kernel-facing paged leaf (base.py is PagedAttentionBackend; MHA, MLA, FlashMLA, TRT-LLM, TRT-LLM MLA, TokenSpeed MLA, DSA, MSA, QSA); state/ holds the recurrent consumers (Mamba/GDN and KDA); hybrid/ the layer-routing composite (linear.py is HybridLinearAttnBackend); and specific/ the bespoke single-model backends (DeepSeek V4.1, Qwen4-Exp's composite and side-cache consumers, and Inkling's dense + conv-state wrapper). A new leaf goes under paged/, a new recurrent family under state/. A model-shaped backend earns specific/ only when the ordinary router and ordinary paged or recurrent leaves cannot express it; use by one model alone is not a reason to introduce a bespoke backend.

Qwen4ExpBackend composes one attention backend, optional Qwen4ExpPLEBackend and optional QSAIndexerBackend. The attention child is the ordinary router, wrapped by the existing HybridLinearAttnBackend only when this view owns GDN layers. Forward dispatch and PD step recording stay with that child; the root broadcasts cache and metadata lifecycle calls. Registry construction selects the attention child; the Qwen4-Exp composite selects its consumers from the pool view's local fields, regardless of whether this view has GDN layers. Binding stays with the common validation and publication path after construction. The root initializes the common AttentionBackend attributes from its own AttnConfig, including draft status, verify width, dtype, and head geometry; these attributes do not depend on an attention child's wrapper shape. Draft views have no GDN or PLE child. PLE and QSA remain available on targets without linear-attention layers; the model retains their computation order.

QSA's full-KV attention uses the ordinary router and an MHA-derived leaf. The attention prologue writes its KV like any MHA layer's. Sparse attention has no MXFP8 block-scale input. Its compressed and recent cache groups belong to QSAIndexerBackend, not to extra attention leaves. The indexer backend refreshes stable raw group tables with the shared GroupTableStacks fill at expansion ratio one: block ids remain unchanged, holes become zero, and padded requests and column tails are cleared. QSA metadata and top-k kernels consume these raw block ids directly; neither their APIs nor the layout carry expansion factors. The recipe rejects compressed fields whose row count or group's token span differs from the model's single-page geometry before cache allocation. The indexer owns its query/sequence metadata and borrows the full-KV table and kernel page size from router.group_view. Layer-shared layout and top-k still use SparseTopKShare with the existing forward and MTP reuse boundaries. The router clears this share before the root prepares its indexer child; the indexer does not clear it again.

QSA block selection carries an explicit uniform query width into its kernel API. Decode derives it from its uniform row contract; extend metadata checks the scheduler's host query lengths, including fixed-width decode rows in a mixed batch. Ragged runs carry None. This adds no device-to-host length read. Materialized scoring can group a divisor of that width to share K within a request; it must retain each query's complete-block frontier and selection. Group size one and larger groups use the same scoring kernel, in both eager and captured forwards. Grouping must not be inferred from the total row count or page-table batch size for a ragged layout.

QSA block selection enters through the public QSA API and select_kernel. The materialized path composes the registered qsa_block_scores and qsa_block_topk_select operators; the fused path uses qsa_block_topk_stream. Required logits/stream features keep backend overrides within the requested algorithm, including forced kernel names that bypass ordinary operator matching. CUDA selection owns its launch details and scratch. Automatic selection uses query rows, table capacity, and supported layouts, never the prefill/decode mode. The caller's logits budget forces streaming when the score matrix would exceed it. Explicit paths bypass that automatic budget.

Blackwell CuTe streaming supports four BF16 index heads of width 128, 64-block pages, block top-k 512, and up to 2048 page-table columns. Within this geometry, ragged ownership and uniform tails stay on CuTe for every nonempty query count. Grids of at most 128 rows first try bounded scratch for split scoring; if that exceeds the caller's partial-workspace budget, they use workspace-free grouped scoring. This includes small decode grids with a zero-byte budget; capacities above 4096 compressed blocks use the overlap algorithm. Other geometries, larger page tables, and other GPU architectures retain portable Triton.

CuTe grouped scoring handles up to eight rows from one request. Uniform runs resolve their groups arithmetically, including a masked final group. Ragged runs resolve group ownership from device request indices at eight-row boundaries and ownership changes, splitting at the next boundary. A ragged grid launches one CTA per query row; only those boundaries start scoring groups. Long runs leave most CTAs inactive, exiting before pipeline setup. This trades a larger launch grid for avoiding a global group table, metadata allocation, or device-to-host read; aligned uniform runs retain their original unmasked specialization. Invalid query rows are zero-filled, have zero complete blocks, and never write output, but their threads participate in the same pipeline synchronization. Small query grids can also score rows independently. Grouped streams keep exact running selections in selector registers, with double-buffered score candidates overlapping selection and MMA across tiles. Small grids can split the key range into bounded global scratch, including uniform runs whose final group is partial. The caller's partial-workspace budget bounds that scratch; grouped CuTe streaming allocates only its output. Ragged grouping waits for PDL producers before reading request indices. Uniform and aligned grouping retain the internal waits before frontier reads and asynchronous Q copies.

Portable Triton streaming also allocates only its output. Each program scans a group contained in an explicit uniform request run and keeps each row's running top-k locally, with no split/merge buffers. Triton ragged runs use group size one in that same kernel.

For both implementations, candidate pruning must prove that discarded scores cannot enter the top-k, or perform a full-score fallback. Each row retains its complete-block frontier. Query count, request count, page-table width, and Q strides remain runtime values; group layout and bucketed split geometry specialize the compiled signature. PDL waits precede asynchronous Q copies and producer-owned metadata reads. Trimmed batched page-table views retain a runtime row stride. Speculative context padding counts toward their capacity; compact candidate offsets must restore the full logical block id instead of wrapping at 65536 blocks.

Different query groups can produce slightly different FP32 scores because their dot/reduction layouts differ; cross-layout bitwise equality is not a contract. Tests check each layout against the FP32 reference with rtol=1e-5, atol=1e-4, and validate selection exactly against that layout's own scores and tie-breaking rule. Near-ties may select different block IDs across layouts. Graph replay is compared with eager execution of the same layout so metadata-refresh checks do not depend on cross-layout rounding.

Qwen4-Exp attention callers pass topk_indices explicitly, using None for dense attention. The prologue has written the full KV cache before either path runs. Draft step zero still preserves the dense decode-context and KV-recording override, while QSA keeps its original context and narrows the selected top-k rows with the queries.

QSA uses positive decode_query_lengths for uniform decode/verification and None for prefill or mixed/ragged queries, including one-token prefill. CuTe serves supported uniform decode on SM100/SM103/SM107 and unscaled BF16 prefill on SM107 through the TMA gather solution. The TMA gather prefill path takes consecutive query positions within each request span and uses the existing logical metadata. KV lengths are views of the existing forward layout, without per-layer reconstruction. The TMA gather prefill path keeps no request state or persistent scratch. PrimTS accepts supported BF16 prefill and decode with logical block candidates on SM100/SM103; groups stay within each request. Other supported NVIDIA inputs use FA2 over equivalent physical slots. All kernels share the cache writer and sparse-attention call. The indexer returns physical slots across graph breaks and keeps logical candidates in the existing per-forward share for both stages. Draft steps carrying only physical slots use the direct-slot kernels.

The QSA leaf owns PrimTS plans and per-device scratch for its current KV pool. Calls, including graph replays, are ordered across streams. Rebinding drops eager plans; capture sessions retain graph plans and allocations until release. Scratch may grow without moving graph inputs. Each run regenerates its data.

QSA padding writes land in slot 0 and may contain NaN/Inf. FA2 and the direct-slot CuTe kernel read invalid candidates from the zero-initialized slot 1 of the reserved null page. TMA gather masks the entire null page before scoring. Writers must leave that read slot untouched. The indexer retains its -1 sentinel, and consumers derive their masks from the original selection.

QSAIndexerBackend privately owns QSAVerifyState only for a speculative target. Registry construction binds the cache plan and preallocates its workspace before model forward or graph capture. Draft and non-speculative indexer backends keep metadata but allocate no target verify workspace. Indexers use the root's indexer_backend; execution carries no separate indexer object and the root has no QSA state registry or type lookup. The staging flag records whether forward or capture has ever used staging, not whether one round is pending. Commit must not clear it: graph replay updates staging tensors without re-running the Python assignment. Staged keys retain the model dtype; commit converts them to the fixed BF16 raw cache. QSA compression callers explicitly select target-verification and draft staging; the runtime wrapper and kernel API require both controls, including None and False when staging is disabled.

Qwen4ExpPLEBackend resolves its own input/output state blocks and query lengths from the PLE cache group. It validates and slices rollback scratch by batch size using its own verify width; layers consume these views directly. It shares the two-slot block resolve with recurrent consumers, but neither uses Mamba metadata nor depends on Mamba's verify context or auxiliary-state hooks. GDN claims only the recurrent groups that back its own state fields.

The runner calls commit_state_after_acceptance once on the target after decode/mixed execution or graph replay, with live acceptance and num_extends. KDA uses the same replay lifecycle for standard target decode: its verify width and accepted count are one when there are no draft tokens. Its cache recipe reserves the record workspace for this case before sizing the arena. Replay capability determines this reservation even without a drafter; prefill-only engines do not reserve decode records. NVIDIA target verification and commit use saved corrections for both BF16 and FP16 activations; the same batched recovery also handles a layer subset. AMD uses its registered batched recovery over saved records. Fused convolution verification retains original QKV in its records workspace. The registered backend owns the gate projection when the model does not supply a gate. Split convolution producers retain original QKV in the same launch that computes convolved rows. Both arrangements write the existing replay workspace; commit reads that workspace after acceptance. A missing batched binding is an error, not a reason to recompute each layer through a second commit path. Since forward mode is derived from the extend count, zero means decode at this entry. Hybrid delegates to the recurrent child. A mixed batch gives prefill and decode rows their existing metadata and forwards; only the decode suffix enters the acceptance commit. The Qwen4-Exp root invokes its attention child, then PLE for decode and QSA for decode/mixed, excluding leading extends from QSA acceptance. Mixed rounds retain PLE's direct state writes. Each consumer commits once; stateless backends inherit a no-op. Transient verify storage belongs to these consumers; LCM remains the owner of the persistent request caches.

QSA verify staging and PLE commit-row buffers are preallocated for full decode capacity and sliced per batch. Cache recipes reserve their bytes before sizing the arena. preallocate_verify_workspace is called on the backend root and returns its verify buffers' bytes. The hybrid delegates to its recurrent child; Qwen4-Exp invokes its attention child, PLE, and QSA. Registry retains the recipe's preparation conditions and budget check, without opening the recurrent child. Inkling ring accounting and QCP history-gather allocation and sharing remain separate from verify preparation. Draft roots allocate no target verify workspace. Qwen4-Exp reserves no verify workspace when the target width is one, even with a draft model attached; this includes the inherited GDN/PLE staging budget and PLE commit rows.

One block-table route: router + leaves

The layering between the scheduler's block vocabulary and the kernels' page vocabulary is fixed, with exactly one conversion point:

layer	sees	never sees
C++ scheduler	per-group BlockTables: rows in block_granularity logical index, entries are CacheBlock ids	kernel pages, backends
bridge (CacheBatchMetadata)	contract-ordered group ids; {gid: [bs, W_g]} views over one packed int32 table, uploaded once as each row's span from the scheduler (its released prefix's null holes and its padding stay behind) and rebuilt on the device	pages, backends
CacheGroupRouter	attention group geometry (CacheGroupGeometry), each leaf's kernel_page_size, expansion, padding, and KV write-location slot math	kernel calls
QSAIndexerBackend	its raw compressed/recent group tables, query lengths, full-KV address view, and private verify workspace	MHA leaf metadata, persistent cache allocation
Qwen4ExpPLEBackend	its PLE state table, input/output state blocks, and verify workspace	Mamba metadata and verify context, persistent cache allocation
paged leaf (PagedAttentionBackend)	page_table (kernel pages, batch-ordered, padded), seq_lens, out_cache_loc; under a query shard, or on an extend with a prefix or a pure extend when its group is KVP page-sharded, also page_table_cpu, the host mirror of its extend rows, for page-owner counts	groups, block tables, contracts, draft/target table provenance
state consumers (Mamba/KDA, Inkling conv, V4.1)	their own groups' raw block_tables[gid] (block vocabulary)	other groups' tables, runner padding

The runner (ForwardStepRunner) does one thing with tables: hand the bridge's block_tables dict to the top-level backend. Capture / idle / prefill-graph dummy forwards use the runner's placeholder_block_tables(bs) (full-width zero tables, null page 0, slices of one persistent allocation) — always-contract delivery: the dict is complete on every path, so no backend carries a "no tables" arm. Delivery is guarded at both dispatch points: the runner's inline live-delivery check and the router's _check_live_delivery fail a live batch whose dict omits any consumed group — the persistent decode buffers would otherwise serve stale pages. Consumers take their own groups by positive claim on the attention kinds they read (cache_consumer_attn_kinds); extra groups ride through untouched.

Inside the router, GroupTableStacks holds the [G, max_bs, stack_max_num_pages] kernel-page table stack (each group's table expanded to its leaf's kernel_page_size and padded to the leaf's max_num_pages; the stack's column count is the widest group's) and the [G, max_bs * N] decode write-location stack. Both are allocated once and refilled in place: leaves copy their view out, while the decode write-slot views and the block drafters' draft_history_view read the stack storage inside captured graphs. A leaf whose decode refresh is only that copy, seq_lens clamped to its decode_seq_lens_floor (decode_buffer_targets, then publish_decode_metadata), is filled by the router together with every other such leaf in one batched copy, since at bs=1 each separate stream copy costs a few microseconds. QSA's indexer owns separate stacks for its two raw groups using the same fill at ratio one. The fill is one expand launch for every group when the tables are row-major views of one storage, as the bridge uploads them, and one launch per group otherwise. Its arguments are plain scalars (each group's offset in the upload and block count, live requests) — no device-side metadata tensor, because the per-step pinned staging + H2D it would need lands on the bs=1 latency path; padding requests ([actual_bs, bs)) and each group's column tail resolve to null page 0. The bridge's {gid: view} dict is the router's input; the router does not require the views to share one storage, only the one-launch fill does. The executor pads stopped requests' state rows in place (RuntimeStates.masked_state_tables), so the views stay in that storage; the bridge packs the state groups last and end to end, so that padding is one launch for every state group. The slot math lives in paged/write_locations.py as pure functions with one invariant: slot = table[req, pos // P] * P + pos % P is page-size invariant, so locations computed over the kernel-page stack equal raw-table locations bit for bit.

CacheBatchMetadata travels no further than the runner; no backend receives it (cache_metadata / forward_batch kwargs are gone). V4.1 consumes the same block_tables dict through its bespoke metadata build.

Deleted, for the record: decode_buffers.py, group_write_locations.py, draft_page_staging.py, expand_history_table as a backend-side step, and the capability flags uses_cache_groups, needs_group_block_tables, tables_self_padding, cache_active_pages_must_be_real, engine_owned_group_ids, table_tail_pad. None carried information not already implied by the pool's published specs plus the always-contract delivery.

Single-table leaves

A paged softmax attention leaf (PagedAttentionBackend: MHA, MLA, FlashMLA, TRTLLM, TRTLLM-MLA, TokenSpeed-MLA, MSA, DSA-over-dense) consumes exactly the pre-cache-group interface — page_table (kernel pages, batch-ordered, padded to [bs, max_num_pages]), seq_lens, out_cache_loc — and never perceives cache groups. Leaves own their persistent decode buffers (page_table_buf, seq_lens_buf) and copy the router's stack slice in on each refresh; they do not alias router storage. A single-group model is a router with one leaf; there is no single-table special case anywhere.

The sanctioned per-leaf residue, all kernel-imposed: verify_floor / block_decode_active (spec verify geometry as a clamp floor, which the router applies once per distinct floor in its batched copy), block_decode_expansion (whether block decode materializes one metadata entry per block position, or the leaf repeats one per request at forward time — FlashMLA, TRT-LLM MLA), FlashMLA's for_graph_replay tile-schedule swap, and the MLA family's num_extends decode-request slicing (override_num_extends).

One side channel exists beyond the table: set_request_slots(req_pool_indices), a no-op by default, which the router calls on every leaf after each metadata build (extend init, decode refresh, capture seeding). It serves a leaf that owns per-request side state indexed by pool slot — DSA's KPool tails — and doubles as that state's per-forward reset point. Paged KV leaves ignore it; it carries no table or page vocabulary.

Write locations have one owner

write_locations(layer, forward_mode) on the top-level backend addresses authoritative history — models, drafters, and the runner neither compute nor thread location vectors. forward_write_locations(layer, forward_mode) supplies the attention prologue's writes, appending the decode window for a draft's first step over a MIXED round (docs/design/attention-prologue.md). PagedAttention.forward, AttentionBackend.forward, model_runner.forward and the model forward chains above the attention layers carry no out_cache_loc parameter; InputBuffers has no location buffer; fill_input_buffers takes no table. For sparse KV offloading, the model prepares canonical Top-K IDs and positions with prepare_sparse_kv_access before writing latent KV to the returned hot slots. KVOffloadAdapter retains these mappings for prologue and router leaf access, requiring preparation for the same step and forward family. Index-K uses history destinations through write_locations.

• Extend: init_forward_metadata computes each group's span over the stacks ([sum(extend_seq_lens)], request-major); write_locations(layer, EXTEND) returns exactly that span.
• Decode / verify: refresh_decode_metadata publishes the token-major [bs * N] window views (decode_write_locations, pointer-stable per bs — the graph records them through the prologue's KV writes, and the pointer guard walks this slot). A MIXED round's draft refresh sets _decode_request_offset = num_extends so DECODE reads skip the extend requests.
• Draft steps: the drafters declare each step's window, the router owns the math and the address-stable storage. publish_draft_step_locations( cache_start, n) computes the window over the location stack (the same fused launch the decode refresh records — in-graph safe) and points write_locations at it: Eagle publishes its one advancing slot per step, vanilla MTP its re-anchored k-window once per round, DFLASH its block window after each block refresh (order matters: the refresh republishes the verify-shaped window). draft_write_locations_uniform(out, start, n) is the side-write variant — scratch resolution over the full-history table (draft_history_view) that must not clobber the published window (DFLASH's target-KV injection, DSpark context windows).
• Cross-backend reads: decode_window_locations() / extend_span_locations() expose the full-history group's published windows; DFLASH reads the TARGET router's windows through them to copy target-aligned KV into the draft cache (the pools share one page-id space).
• Writes outside the backend (the attention prologue) fetch their slots immediately before the write. An offloaded latent projection uses the compute slots returned by its sparse prepare. When a draft step-0 forward over a MIXED round writes the round's full K/V rows, dispatched as MIXED or as DECODE (the MLA draft's whole-batch write, a GQA draft's narrowed first step), forward_write_locations concatenates the EXTEND span and the DECODE window — eager-only, MIXED rounds never run under a captured graph. Target MIXED decode halves and later draft steps keep their ordinary decode-only windows.

Target capture is configured once during model setup

create_model_runner calls execution.factory.configure_draft_target after both models load and before cache construction. DFLASH/DSPARK models must implement the explicit TargetCaptureConfigurator interface; missing implementations fail at setup. A method with the same name on an unrelated object is not treated as an implementation or as evidence of prior setup. DFlash, DFlash2, and generic DSpark use their model's ordinary capture setup; K3 owns its trained stream/projection contract; DeepSeek V4.1 DSpark owns its checkpoint tap selection. models/target_capture.py contains only the shared interface. DFlash checkpoint parsing and target configuration live in DFlashDraftModel.configure_target, inherited by DFlash2 and generic DSpark. Setup calls only the parent interface configure_target; each concrete draft directly adapts to its target family. There is no generic DSpark helper probing for a DeepSeek-specific setter, and K3 targets need not implement that setter. EAGLE3 selection remains in this same setup phase, including its explicit server-argument override. The EAGLE3 and DFLASH capture setters run only on a target that opts in with BaseCausalLM.supports_aux_capture, one whose forward returns the taps layers_to_capture names. Any other target fails setup instead of handing the draft its final hidden state. DeepSeek V4.1 DSpark selects its taps with its own set_dspark_layers_to_capture and needs no opt-in.

This runs on every PP stage even when that stage has no executing drafter. wire_target only binds embeddings, heads, and other execution resources; it never selects capture layers, changes streams, or replaces the output layout. This also applies to V4.1's dedicated drafter with scheduler-owned context windows. There is no configured flag or optional-method probe in resource binding. A last pipeline stage borrows its local draft embedding when the target embedding lives elsewhere; this is resource binding, not a different proposal algorithm. Per-forward capture hooks consume the established configuration.

Embed/head sharing (shares_target_embed_head) follows the same stage ownership. get_embed_and_head returns None for a side the stage does not hold -- the embedding lives on the first stage, the head on the last -- and never dereferences an absent module. Stages before the last bind nothing: a draft built there only produces context. The last stage binds what the target reports, which is the head alone (embed=None); a pipeline-capable draft then keeps the embed_tokens shard its checkpoint ships for that layer, its loader rejects a pipeline checkpoint without one, and the factory checks afterwards that the draft's own get_embed_and_head still reports an embedding -- a draft that aliases None into its embedding is named at construction rather than failing at its first forward (BaseCausalLM.set_embed_and_head refuses embed=None outright, since a generic draft keeps none). Off the pipeline both sides are shared and the draft's copies are dropped before the KV-cache budget is profiled, as before.

A drafter sharing a target head constructs its head and logits processor from mapping.lm_head, just like the target. Attention TP is a separate layout: under attention DP the head is normally replicated even when attention uses multiple ranks. The factory checks the two logits layouts before binding.

Checkpoint tap labels remain zero-based completed-layer IDs. Prefix tap L is produced after L. AttnRes tap L is produced at L+1's entry by that layer's mixer, before input-layer normalization or snapshot mutation; the final tap belongs to the output mixer. Capture execution and projection-weight placement use this same owner on both PP and non-PP. There is no boundary deferral or recovery operation. Capture's mixed stream must not be replaced by the fused attention input, which already includes input-layer normalization.

execution/dspark_context.py contains DSparkContextProducer and the model interface it consumes. K3 tap ownership and projection arithmetic live in models/kimi_k3_dspark.py; the producer does not interpret K3 layer IDs. This interface covers DSpark context production, not a requirement for all draft algorithms. K3 DSpark is currently the model using this production path.

Pipeline stages use DSparkContextProducer: each stage normalizes the taps it owns if configured, applies their projection columns, and sums in FP32; the accumulator travels with the chunk's PP state and the final stage applies context normalization once and writes native context KV. The executor selects the producer from the pipeline configuration and the draft model's class (select_dspark_context_producer): on a pipeline a draft implementing DSparkContextModel gets a producer on every stage, a block drafter whose model does not implement it is rejected (it would draft from one stage's taps alone), EAGLE3 is refused (its aux taps span stages), and any other draft gets none. Off the pipeline every tap is local, so the drafter keeps its concatenated projection and its own context writes -- including the quantization-aware path, since raw per-tap weight slicing is not a quantized linear operation. PP drafts require unquantized projection weights.

Pipeline speculation is not DSPARK-only

An MTP (NextN) draft runs on a prefill pipeline without any cross-stage production: it consumes the post-final-norm hidden states the last stage already computes, so only that stage drafts. EAGLE3 stays off the pipeline: its aux taps come from several stages and the stage boundary bundle does not carry them; ServerArgs rejects it and select_dspark_context_producer refuses it again at executor construction. The MTP shape is:

• Construction. ServerArgs accepts DSPARK and MTP with --pipeline-parallel-size > 1 on the prefill role only (the chunk pipeline has no decode token feedback to draft against anywhere else); the dense == attention TP / CP = 1 rule stays DSPARK's, whose draft reduces attention-TP embedding partials over the dense TP group. Which stages build the draft model at all is the factory's decision, not each model's (factory.pipeline_stage_builds_draft): a block drafter is built on every stage because it produces context from every stage's taps, while an MTP draft is built on the last stage alone -- the other stages construct no draft runner, load no draft shard, and wire nothing, so ModelExecutor and configure_draft_target see draft_model_runner=None there and the in-tree K3 NextN draft needs no stage-aware shell. The skip is safe because nothing in draft construction is a world collective the skipping stages would have to join: the KV-budget all-reduce runs on every stage with or without a draft, the DeepEP/MoE communicators a draft layer reuses are the target's and scoped to the stage, and the only world-scoped step -- the InstantTensor weight iterator, which synchronizes over group.WORLD when a model declares no checkpoint_load_group -- is bounded by the factory, which loads a pipeline draft with the stage's rank set (ModelRunner(checkpoint_load_group=...), surfaced as LoadConfig.checkpoint_load_group; a model's own declaration, such as K3 DSpark's stage-subset filter group, still wins). With no producer configured the last stage's target forward captures CaptureHiddenMode.FULL for its drafter.
• Cache. Nothing changes: CacheLayerOwnership already places the draft cache layers as the last stage's trailing producer step, the draft backend and pool exist only there, and the merged plan, bootstrap placement and transfer routes are the same as for DSPARK.
• Layerwise CachePD. supports_pd_layerwise_finalization is decided per stage (device._supports_pd_layerwise_finalization): a stage owning no draft fields has nothing to finalize and answers True; the owning stage answers for ModelExecutor.draft_field_writer -- the producer when configured, else the drafter -- exactly as a non-PP engine does, and register_draft_final_step_counter reads the same property. The last stage registers the draft-final step counter; the others count target layers only.
• Handoff. The last stage samples, runs the drafter over the completing chunk and writes the candidate block into the reserved decode slot; the control plane shares (output_tokens, output_lengths, next_input_ids) with the other stages at commit -- with the logprob vectors and the NaN guard's per-request flags, see the QCP section -- so every rank's scheduler stamps the same bootstrap token and candidate window onto the remote decode. The PD wire and the decode side are untouched.

Expect a larger last-stage bubble (NextN layer plus draft extend and multi-step drafting); rebalance with --pp-layer-partition.

The producer is stateless across forwards. Each chunk owns its accumulator; queued chunks cannot alias it. A configured ctx.dspark_context_producer owns native context writes during target forward; otherwise the drafter owns them. This responsibility is fixed at construction, not inferred from a per-round readiness flag. The producer enqueues writes before the drafter on the same stream; failures propagate instead of selecting a fallback writer. The common block drafter always updates accepted-prefix lengths, but only projects/writes context when no producer is configured. Its capture sink is not armed for a forward with this producer, avoiding a second writer. The final PD readiness barrier remains after the whole proposal call, since proposal execution can write the same draft fields after context injection.

Per-forward drafter work rides on the context

What a drafter wants done during the target forward is a property of that forward, so it travels on ForwardContext — never as mutable state on the target model that someone must remember to reset. The executor's only seam is BaseDrafter.prepare_target_forward(ctx), called right before the target runs: the drafter decides under its own gate whether this round qualifies and attaches what it needs; a fresh context per round means nothing outlives it, and a model that sees no attachment does nothing. DFLASH is the one user: on decode-only rounds it attaches ctx.target_capture_sink, the target hands each captured tap to on_target_capture as it is produced, and the sink copies it (rounding a wider residual once) into its columns of one tap block. The last tap runs the same context projection and KV write as the drafter's run would on the concatenated taps, inline on the forward's stream: fc is split by output columns across attention TP, so the projection reads 1/TP of the weight and gathers the columns in order with the target's own collectives. No auxiliary stream is involved, so a capture point must hold no borrowed collective output the gather could overwrite. run writes the context itself only for a round the sink did not take, and refuses one it armed but never saw complete. These hooks consume the target's capture configuration; they do not change the tap selection or output layout.

The reverse direction rides on the context as well: a target that captures its taps on a row subset reports it as ctx.captured_rows (CapturedRows(positions, prefill_spans)). V4.1's CED narrowing is the one producer: its taps in layers 37–39 contain no rows for incomplete prefill chunks, the last window of each completing prefill, and all decode rows. The reported prefill spans retain a zero-length entry for each incomplete request, preserving the original request order. DSpark's prefill seeding (_seed_prefill_windows) reads these spans and positions instead of the input-length mirror and skips zero-length spans. A target with one captured row per input row leaves ctx.captured_rows as None, and the drafter keeps its buffer-based layout.

Shared prefill convolution preparation

Mamba/KDA extend metadata owns one immutable CausalConv1dPrefillMetadata per forward. Its two int32 maps associate convolution programs with request rows and local token chunks. The builder sizes them from the existing host length mirror and fills both directly from device query boundaries in one Triton launch. Every layer reads the same tensors and block size; the conv wrapper neither rebuilds them nor initializes/uploads per-layer scratch. This is transient execution metadata, not a new cache group or model state.

The same extend/mixed metadata owns a device int64 mirror of the int32 query boundaries. KDA layers share it for scan ABIs instead of casting per layer; the host int64 mirror still supplies launch planning without D2H. Decode refresh/capture does not allocate this prefill-only mirror.

Each metadata build allocates fresh index storage, including when two forwards have the same total token count but different request partitions. No subsequent forward refills a buffer an earlier forward may still read. Mixed batches include the decode rows' verify-token lengths in this same builder. Decode-only refresh/capture remains unchanged and carries no prefill convolution schedule. Breakable prefill graphs consume the live metadata in the eager attention break, as ordinary eager forwards do.

Prefill convolution returns [channels, tokens] with contiguous channels per token, independent of the input strides. Transposing to [tokens, channels] and splitting the packed channels can therefore remain views.

Prefill state staging fuses resumed conv-window copying, recurrent-state gather/zero, and history flags in one kernel. KDA gathers directly into FP32 scan state, including from BF16 cache storage; final writeback casts to the cache dtype. It preserves scheduler-owned block ids and arbitrary cache strides. Fresh rows never read null or stale recurrent state; their conv working windows remain unchanged. Shared input snapshots are read-only, output blocks are unique, and a private in-place source/destination is legal. This changes neither scan arithmetic nor cache allocation, retention, or checkpoint identity.

The NVIDIA CuteDSL prefill adapter declares its native v_major ([N, H, V, K]) state layout. The dispatch facade alone adapts a caller with another layout; the wrapper must not round-trip native state through FLA's [N, H, K, V] convention. Direct wrapper callers use the native layout for both initial and final state. Exact-length gate conversion to FP32 and beta packing retain their ordinary PyTorch operations. The native wrapper allocates the scan output. Ordinary attention breaks copy it into a stable graph-owned handoff buffer; inline KDA keeps output restoration and padding cleanup inside the graph, without that handoff copy. No output-buffer extension to the native wrapper is required. These preparation changes modify neither the native scan, its gate math, nor GEMM arithmetic.

Recurrent prefill subgraphs (KDA, Mamba2, GDN)

Capturing recurrent layers in the outer graph

CapacityPrefillBackend (state/prefill_capacity.py) owns this contract for KDA, Mamba2, and GDN; each subclass only states which forwards it admits, whether uncaptured shapes also run the capacity layout, and whether that layout stores in-forward State checkpoints. Supported pure-extend forwards use prepare_prefill_metadata before eager execution, startup capture, and replay. This consumer-stream seam builds or refreshes CapacityPrefillMetadata with the selected token and request capacities. Eager execution uses the live count; replay may round up to a captured count. The same metadata contract controls scan capacity and padding in every case; there is no temporary metadata binding or mutable inline flag. Only startup capture retains the metadata's addresses. Uncaptured shapes use temporary storage through the same builder. Whether a shape can be captured is also asked on its own, through admits_prefill_graph, which reads no forward context and writes nothing; the seam must return the same answer, and startup capture raises when a backend admits a shape and then refuses to prepare it.

Mamba2 keeps the scheduler metadata for uncaptured shapes: its chunk plans need only live bounds, so the capacity layout would only add work to eager forwards. A retained shape owns one persistent chunk plan, sized extent // chunk_size + sequences and padded with empty chunks; the preparation seam rewrites it in place, in one pinned upload, before each use. Unused request slots scan one padding row each past the live tokens, so the live tokens keep the packing of the unpadded forward and a capture of any request count matches eager bit for bit.

GDN admits capacity prefills only when the selected chunk-prefill kernel sizes its launch from the sequence count and reads the bounds on device (gdn_chunk_prefill_capturable); otherwise its layers keep their breaks. Like Mamba2, uncaptured shapes keep the scheduler metadata. The scan chunks each sequence from its own start, so captures of any request count match eager bit for bit.

A forward writing in-forward State checkpoints (--enable-state-checkpoint-in-forward) runs inline too: GDN is told the flag at construction, and its capacity metadata then keeps one fixed checkpoint slot per request slot (a row writes at most one). The captured scan stores a state at every page of the bound pool's prefix granularity, which divides every checkpoint offset, into one float32 scratch that all shapes share (forwards run in stream order), with a row per page of the widest captured token capacity. Before each use the refresh rewrites, in one pinned upload and a few device copies, each sequence's state count and each live checkpoint's scratch row, token index and block; an unused slot keeps an in-range token index and block -1, which the masked writers of the conv window and the recurrent state skip. The stored states are the scan's own, so the checkpoint blocks match eager bit for bit; the cost is the extra scratch writes and the shared scratch. A backend without the flag, and every other family, keeps such a forward on its live metadata at the eager break.

Qwen4-Exp's composite passes these seams to its attention child, so its GDN layers run inline; PLE and the QSA indexer keep their breaks over their own live metadata.

For retained shapes, the hybrid wrapper can omit the KDA attention break and capture neighboring projections, KDA kernels, and post-attention compute together. Full-attention layers keep their breaks. Before execution, the common preparation step validates live lengths and refreshes boundaries, convolution maps and state-block indices. All KDA layers read this same storage. Token lengths may vary within the bucket; live request counts may fill part of a capture. Output padding is cleared by the recurrent forward, replacing the attention break's handoff copy and tail scrub when the layer is captured.

The ordinary outer capture is retained for mixed batches and other request counts beyond captured capacity. All variants share the outer pool and execute serially, as existing bucket captures do. Layerwise PD transfer and data parallelism retain the ordinary route: host cache-step callbacks must remain live, and DP admission must stay rank-uniform. Retained metadata rejects a replacement cache pool; graph release and recapture remain the orchestrator's responsibility.

Replay uses a captured request capacity; live counts, row identities and lengths can change. The outer owner captures request capacities from prefill_graph_capture_batch_sizes (unset: the minimum count per token bucket) with one variant per token bucket and request count. ModelExecutorConfig requires this field explicitly: factories forward the configured list or None for the minimum-count policy, so missing configuration wiring fails at construction. Token buckets still follow the shared prefill token ladder. Capture requests have positive lengths and fit the model context and request buffers. At replay, unused execution slots have zero convolution length, negative state-block indices and one padding row each in the scan, past the live tokens. Native scans still receive only positive-length sequences; padding owns no cache blocks and its scan result reaches none: the masked state writer (store_state_rows) drops rows whose block is negative, and the padded output rows are zeroed by a device-side bound the capture cannot freeze. The live context and scheduler/MLA request counts remain unchanged. Selection reserves live_tokens + padded_requests in the token bucket; a full bucket can use the next existing token bucket, otherwise the ordinary fallback remains. Startup autotuning uses the same dummy-batch builder with an explicit minimum request count, ceil(num_tokens / context_len), independent of the configured capture request counts. Its token budget also respects rank-local request capacity. Request counts exceeding captured capacity retain the ordinary attention break and eager KDA. Replay refresh includes scan_query_start_loc, which the recurrent dispatcher consumes, as well as the convolution boundary and existing int64 mirror. Unifying metadata does not imply zero padding cost.

Startup capture and eager fallback

Supported KDA prefill uses the merged captures owned by PrefillGraph by default when prefill graphs are enabled. --disable-kda-prefill-graph disables KDA capture without changing ordinary prefill or decode graph settings. The shared ServerArgs configuration passes the setting explicitly to each KDA backend. Startup creates the configured token-bucket and request-capacity variants. Serving forwards only select and replay these captures, never warm up or capture a separate per-layer graph.

If no compatible merged capture exists, the ordinary outer graph retains its attention break and calls the same eager KDA implementation. State writes, PD cache-step recording and break-output copy/padding keep their existing order. Inputs outside the outer graph's admission rules run eager. New request shapes do not grow a backend-owned graph cache. Metadata refresh and eager execution may still allocate temporary buffers.

The outer owner holds all captures and outputs in one table keyed by token capacity and request capacity; None in the request-count position selects the ordinary attention-break capture. The backend retains startup metadata for the exact shapes that need stable addresses, not graphs or request state. Serving forwards never grow this retained table. The outer owner's serial shared-pool discipline applies to all variants; there is no separate KDA graph pool. Before recapture, it releases the old captures and resets retained prefill metadata through init_prefill_graph_state. Publishing a cache pool also drops retained prefill metadata. Graph release and cache-pool rebind remain coordinated by the orchestrator.

Fixed-capacity execution metadata

The capacity metadata overrides only the packed execution extent; real host lengths and GPU boundaries still agree. Its PrefillCapacity validates the packed bounds it builds. For KDA, an explicit KdaPrefillCapacity passed to the kernel facade admits the live CPU lengths: each sequence may fill the bucket, but their combined tokens must also fit it. The CuTeDSL adapter converts this descriptor to native planning bounds; cuDNN plans from tensor capacities and reads live boundaries on the device. Convolution maps reserve ceil(token_capacity / block_m) + sequences - 1 programs, bounding the sum of per-request rounded lengths without reserving the entire token bucket for every request. One GPU metadata refresh per forward marks inactive programs with PAD_SLOT_ID before all layers run: the convolution kernel otherwise performs unmasked prior-token loads even for an excess chunk. Scan inputs are cleared past the live device boundary inside the graph, since a capacity descriptor makes padding addressable to native full-tile loads. Both conv and scan read live GPU boundaries; total packed tokens, including dummy slots, must fit the physical extent. Native sequence slots are never empty: request padding scans one row per unused slot. Convolution skips their zero-length spans. A capture can serve smaller live batch counts without another schedule. Other solutions retain exact live-length planning and reject capacity mode.

For the pinned token-major CuTeDSL ABI, a fused preparation kernel scrubs padding, converts gates to FP32, and builds the device chunk plan. Its total chunk capacity is ceil(token_capacity / 16) + sequences - 1, while each sequence retains the full per-sequence walk bound. The third-party adapter passes this explicit plan to the existing native launch without replacing global functions or changing scan arithmetic. Routing and workspace partition rules remain owned by the native host. Unsupported layouts retain the public wrapper's capacity preparation.

The CuTeDSL preparation copies every input with its capacity padding scrubbed: the gate projection can produce undefined padding, and the inline path has no handoff copy to clean it. Per-call plan and scratch tensors belong to the active graph pool or eager invocation; they are not a mutable process-global plan shared across replay streams.

The cuDNN adapter reuses FlashInfer's cached plans and workspace, warmed by the existing startup capture loop. Replay matches eager execution at the same capacity bitwise; different capacities can select different scan partitions and therefore differ by floating-point rounding.

Changes to this capacity contract require validation of full-model overlap, memory use and performance in addition to kernel correctness.

Query context parallelism

--query-context-parallel-size N (mapping.attn.qcp_*, N == attn.tp_size == attn.kvp_size) shards every extend forward of a PD prefill engine over the attention TP group, over the pages --kv-parallel-size N shards. One forward, one path: the scheduler plans the whole chunk on every rank, the executor builds a QueryShardPlan from the request lengths (execution/query_shard.py) and puts it on ForwardContext.query_shard. The shard is zigzag per request: every request's new rows are cut into 2N near-equal segments and rank r takes segments r and 2N - 1 - r of every request, so every rank holds an early and a late part of sufficiently long requests — under causal attention the two approximately balance. Short requests can leave ranks empty; those ranks still join every history collective. A contiguous split of the chunk would hand whole requests to single ranks and serialize the group-by-group history gathers behind whichever rank owns the group's rows. A rank's rows are ascending span rows but not one range: the model's rows are input_ids_buf[plan.local_row_ids], positions_buf[plan.local_row_ids] and the same rows of every per-row input; per-request inputs keep the whole batch; a pipeline stage's boundary bundle carries shard rows. The collectives over the shard group concatenate the ranks' rows rank by rank (gather order, plan.span_to_gather_cpu), with the per-rank counts plan.row_counts as the CommManager row tables: a leg that only moves or sums rows (the dense and MoE legs, the head exchanges, the embedding's id gather and row reduce-scatter) works in that order unchanged, and whoever reads gathered rows by their span position — the KV and index-K writes by slot, the n-gram hash, the sampled and prompt rows of the exit — puts them in span order first (plan.to_span_order, or the RowSplit of plan.split_rows), one index_select of bytes. ctx.input_num_tokens and global_num_tokens keep the scheduler's full-chunk meaning: only collective sizing (CommManager) reads them, a model reads its rows from its tensors or ctx.query_shard.

Plugin API 4 replaces the contiguous query-shard contract with per-request zigzag rows. Plugins must declare support for API 4 after migrating their row-layout and request-history consumers; API 3 adapters are incompatible. QueryShardPlan.from_forward requires an explicit device for its row indices (use device="cpu" for host-only plans). The CPU permutation tables are span_to_gather_cpu and gather_to_span_cpu; their target-device copies are span_to_gather and gather_to_span.

Adapter migration: a query shard no longer exposes a contiguous range (local_slice, local_start, local_end). Use local_row_ids for device row selection and local_row_ids_cpu for host planning. Consumers of RequestTokenHistoryView must use the full request layout together with ctx.query_shard; row_offset cannot represent a zigzag shard. Upgrade adapters that rely on those fields together with the runtime. A row gather returns rank-major rows, so position-dependent consumers must restore span order before reading them; span-order partials must be put in gather order before reduce-scatter.

The contract a model (in tree or a plugin) implements:

• CommManager(query_sharded=mapping.attn.has_qcp) — declares that the model slices its rows by ctx.query_shard; a model that does not slice passes False and is refused at construction under a sharding mapping. Attention returns complete rows under this mapping — its weights are head-replicated (the default, mapping.attn.head_tp_size == 1), or under head TP over the shard group its own tail returns this rank's rows (below) — so the attention legs (pre_attn_comm, gather_residual, post_attn_comm, post_final_norm_comm) are identity on every forward — a sharded extend and the drafter's replicated decode steps alike — and needs_pre_attn_all_gather / needs_final_all_gather are False: nothing is ever scattered by attention. The dense and MoE legs follow the forward: with a shard, the existing all-gather / reduce-scatter legs over plan.rows_for_collective(ctx.collective_num_tokens) (the shard rows, or the sampled rows after a draft's narrowing); without one (decode steps, idle), the replicated all-reduce legs — which is why dense TP and the MoE TP×EP group must each be 1 or the attention TP width (validate_qcp, and the manager refuses other mappings). The row-layout conversions (slice_scattered_rows, gather_scattered_rows) are identity on a sharded forward, whose rows are the scattered share already. Fusion off.
• VocabParallelEmbedding.forward(input_ids, query_shard=ctx.query_shard): the ids a rank embeds are its shard, so the layer all-gathers the ids to the span for its vocab shard, sums the shards, and reduce-scatters the rows back (the same bytes as the replicated all-reduce). A model passes the shard explicitly; the default (None) is the replicated lookup.
• PagedAttention.latent_prologue(..., key_rows=QueryShardGather(ctx.query_shard, mapping.attn.qcp_group)) (through DeepseekV3AttentionMLA.forward_absorb_qkv_proj automatically): the prologue rotates the local rows without a cache (mla_prologue(cache=None)), all-gathers the rotated latent with the plan's row counts, puts it in span order and stores it owner-masked (latent_store); slots is the whole span. A rank whose shard is empty rotates nothing but joins the gather and the store, so a model must reach the prologue (and every other QCP collective: the index-K gather, each group's history gathers) with zero rows rather than return early; the dense MLA path (DeepseekV3AttentionMLA.forward, the expanded prologue) refuses a shard up front. Index-K the model gathers the same way (token_all_gather of the local keys before quantization, then span order) and writes with the owner mask.
• GPU DSA (backends/paged/dsa.py): init_forward_metadata builds a DSAQueryShardMetadata — request groups whose summed history fits the gather workspace (one whole history, reserved from the cache budget by the recipe's workspace_bytes, allocated once on the target tree by registry._prepare_fixed_workspaces and shared with the draft tree, whose extend step never gathers concurrently), each with its row_base in the request-major history-row numbering, this rank's local_query slice (a rank's rows are in request order, so a request range is one run of them; the counts come from plan.local_rows_per_request), and a HistoryGatherPlan (per-owner row counts from page_table_cpu, dcp/placement.py: owned_history_rows). For the indexer the model calls backend.gather_history_index_k(layer_id, pool, group) per group, which returns the group's index keys in position order in the leaf's index_k_format (DSAConfig.index_k_format; the pool read gather_index_k_rows(..., index_k_format=) refuses a plane of another dtype): fp8_scaled gives (fp8 [rows, head_dim] uint8, scales [rows, head_dim / 128] fp32), bf16 gives (keys [rows, head_dim] bf16, None). The model hands them to dsa_prefill_topk(q_local, w_local, group.gather.virtual_slots, row_starts_local, row_ends_local, ...) as the rows in workspace-row order of that format -- index_k_fp8=, index_k_scale= or index_k_bf16= -- never with index_k_cache; the facade routes the rows by the index_k_format trait and requires the index_k_workspace_rows feature (dsa.INDEX_K_WORKSPACE_ROWS_FEATURE), so only a leaf whose launcher takes rows of that format is selected and handed the row keywords (the in-tree DeepGEMM leaf for the FP8 pair; a plugin's bf16 leaf declares the feature beside index_k_format={"bf16"} and takes the index_k_bf16 keyword), and an override cannot force a plane-only leaf onto rows. The leaf probes that selection at construction (DSABackend.__init__ under qcp_size > 1, dsa.select_dsa_prefill_topk_for_rows with the configured format, page size, indexer geometry, and the envelope's batch-invariance and solution pin), so a platform without a declaring leaf is a NoKernelFoundError at startup rather than in the first sharded prefill. The model adds group.row_base to the returned rows and hands forward_sparse_prefill(topk_slots=<workspace rows>) the local rows; the arm gathers every group's KV (gather_history_kv, a collective every rank joins even without rows in the group) and attends the local rows with every head, return_lse=False, no combine. Latent KV stays in the padded rank-major layout of the equal-share history gather: position p lives at gathered_index[p]. The arm maps each nonnegative top-k workspace row within topk_lens through gathered_index[row - group.row_base], retaining the emitted order and the -1 sentinel. Index-K remains in position order for causal scoring. The position-to-row map and padded send slots belong to HistoryGatherPlan; this requires the equal-share history-gather planner. The mapped indices must not be sorted by physical row: sparse cores preserve the supplied selection order, including when slot_order="sorted" denotes position order. The gathered buffer handed to dsa_prefill as a flat [slots, dim] cache is a whole number of kernel pages (the workspace rows are padded by dsa_history_gather_workspace_rows), so the paged solutions' view of it holds too and no solution is wrong at runtime; the padding rows are never selected. Its capacity covers degree * send_rows rounded to kernel pages; if owner padding exceeds the reserved workspace, the latent gather uses a temporary buffer of that size for the group. Latent rows move as uint8 bytes directly into that buffer, without a position-order read-back. The index-K history gather remains separate. Both gathers preserve row bytes, including FP8 latent values and packed index-K keys with their scales. The workspace (HistoryGatherWorkspace, index_k_format recorded, rows of index_k_row_bytes(head_dim, format)) is sized by the recipe and allocated by the leaf from the same formula; the recipe refuses a draft whose index_k_format differs from the target's, naming both, since the two share one workspace, and the adopting leaf checks the recorded format with the rest of the geometry. The sparse cores take their head count from the query, never from layer.tp_q_head_num or a mapping assumption (DSABackend._query_heads: the model's heads or the attention-TP slice, any other count refused), so one model layer serves a forward whose rows carry every head and one whose rows carry the slice. The sharded arm attends with every head and refuses a query carrying the slice. The decode arm (the drafter's steps) keeps the DCP combine over the KVP pages; its form follows the query's heads (keep_all_heads for every head — head-replicated weights — the gather-and-reduce-scatter form for the slice, as the drafter's steps under head TP over the shard group carry). The model threads the sequence after forward_absorb_qkv_proj through DeepseekV3AttentionMLA.sparse_prefill_attn_v_proj (the backend's forward_sparse_prefill with the model's selection, then project_attended_heads: under head TP the tokens-to-heads exchange, the local w_vc); GLM-5's sparse prefill runs it in tree. A rank whose shard is empty calls it like every other rank — the core's history gathers are collectives — and projects nothing.
• The model exit hands the logits processor the shard's rows with the plan on LogitsMetadata.query_shard (BaseCausalLM.exit_logits is that one call; a model with its own exit does the same). The shard is a parameter of the processor's row selection: it scores the planned prompt rows first, then selects the sampled rows — hidden_states[gather_ids] on whole rows, gather_sampled_rows over the group on a shard (one byte-preserving all-gather, token_all_gather_rows, with sampled_rows_per_rank, then plan.sampled_to_request back to request order) — and runs the LM head on the batch's [bs, hidden] rows, so the vocab all-gather is the TP one. The group of both gathers is the processor's TP group: the LM head is vocab-sharded over it, every rank of it must end with the same rows, and the mapping makes the query shard group exactly that group (qcp_size == attn_tp_size); the processor refuses a plan of another width. The processor is the only caller of gather_sampled_rows (CommManager has no sampled-row leg: needs_final_all_gather is False under a shard and nothing gathers the final norm's rows). A FULL hidden capture stays the shard; a LAST capture is the gathered [bs, hidden] rows, whole on every rank — the aux taps' (Eagle3) when the model has them, each tap gathered the same way, and only on a LAST capture, since no other mode reads them selected. A model that selects its rows before the processor (logits_rows_selected) keeps that contract, and such a model cannot serve prompt logprobs, sharded or not. ctx.gather_ids keeps the batch's full layout on every forward, the drafters' extend steps included (Eagle's step 0 and every depth of the multi-depth Mtp drafter read the shard's rows of the shifted prefill ids and positions, chain the shard's hidden rows, and carry the plan on their context; Mtp sums its cross-chunk stash of target hiddens over the group, one owner per row). The plan derives local_sampled_index from each request's last row; QCP consumers read it directly instead of passing full-layout gather ids into a shard-local selector. It is used by gather_sampled_rows and by any model that narrows to its live rows itself. A draft model's FULL capture under a shard is the shard's rows, which is what the next depth consumes.
• Prompt logprobs (InputLogprobPlan, --input-logprob-chunk-tokens): the plan's rows are full-layout rows. The forward thread stages the whole plan's targets and slots on every rank (the shifted ids are the whole span everywhere, so the target audit flags the same requests on every rank) and keeps as this rank's InputLogprobRows.rows the plan's rows inside its shard, as indices into it, with the per-rank counts (rows_per_rank) and the index that restores plan order after the gather (gathered_to_rows): host arithmetic over the shard's row tables (QueryShardPlan.split_rows — the one row split of the plan, the same that splits the sampled rows). Scoring a row needs its full-vocabulary logits, and the head is vocab-sharded over the group: every rank must hold every planned row, so compute_input_token_logprobs all-gathers the planned rows' activations with those counts ([plan rows, hidden], put back in row order — a rank without a planned row contributes none and still joins) and then runs the unsharded chunk loop over the whole plan on every rank; the chunk schedule is thereby the same on every rank, which the vocab all-gather inside each chunk, a collective, requires. Every rank ends with the whole plan's fp32 vector — the tensor-parallel path's bit for bit, since each row meets the same operands — so no result gather follows, the per-request NaN audit agrees across the group and ModelExecutionResult.input_token_logprobs is the full vector in plan order on every rank; the commit path and the P→D bootstrap-logprob frame are unchanged. (Scoring only the local rows against a vocab-sharded head is not possible: a row's log-sum-exp needs every rank's vocab slice of that row, and the vocab all-gather assumes replicated rows. The LM-head work for prompt logprobs is therefore that of the TP path plus the activation gather of the planned rows; a vocab-parallel cross-entropy that trades the [rows, vocab] all-gather for per-row partials is the deferred logprob.topology-invariant item of numerics.md.)
• Pipeline parallelism: the last stage scores the prompt logprobs (on its shard, under QCP) and the control plane's share of its result (adopt_last_stage) carries output_logprobs, input_token_logprobs, and the NaN guard's output_nan_flags to the other stages with the sampled tokens; every stage pairs the vector with its own mirrored plan, which ModelExecutionResult carries whether or not the stage scored the rows. The flags travel with the values they audit: only the last stage holds logits and prompt logprobs to flag (the other stages' guards see placeholder outputs and the rank-consistent target audit at most), and every stage's output processor must take the same abort-or-finish branch for a request, or the stages' schedulers disagree on it. Prompt logprobs are therefore no longer refused on a pipeline split (supports_prompt_logprobs depends on the narrowing-model check only).
• Communication buffers (prepare_communication_runtime(max_forward_tokens)) stay sized by the whole chunk under QCP, not ceil(chunk / qcp): the all-gather / reduce-scatter legs' gathered side is the whole chunk on every rank (the dense and MoE legs gather the shard rows to the group), so only buffers of local rows could shrink, and the one model that prepares such buffers today does not shard queries. Revisit with a per-buffer audit when a sharding model allocates them.

Head TP over the query shards (--attn-head-tp-size N with --query-context-parallel-size N). The query shards hold different rows, so the head group of the decode-side layout (docs/serving/parallelism.md, "Decode-side TP layouts under attention DP") applies to them unchanged: mapping.attn.head_tp_group == qcp_group (the attention TP group), q_b_proj / kv_b_proj / o_proj are head-sharded over it, and the sharded extend forward runs the same exchange — the normalized q latent token-all-gathered to the span, q_b_proj and the absorption on this rank's head slice of every row, the heads-to-tokens all-to-all back to the shard rows with every head, the prologue (RoPE on the shard's own positions, the gathered KV write as on every QCP forward; there is no positions collective), the sparse core over the gathered history with every head and no LSE merge, the tokens-to-heads all-to-all, the local w_vc, and the o_proj tail: row-parallel plus a token reduce-scatter to the shard rows, or under --tp-batch-invariant attn the all-gather of the heads, the column-parallel GEMM and an all-to-all back. One resolver (comm_manager.head_tp_row_counts, the module's head_tp_leg_row_counts(ctx, num_rows, collective=)) hands every leg its per-rank counts: the DP tables under attention DP, the shard plan under QCP (row_counts for the legs up to the core, rows_for_collective(ctx.collective_num_tokens) after it — identical on an extend that does not narrow), so no model code branches on the layout. The drafter's decode steps on this engine hold every row on every rank: they exchange nothing — head_tp_exchanges(ctx) is False without a shard, the one fork of the head-TP path, asked by every leg and switching off the exchanges and nothing else — so they attend this rank's head slice of every row through the same attn_mqa layer (which declares every head; the DSA core reads the head count from the query, so no second layer exists for the slice), the DCP arm as attention TP runs it, and the tail all-reduces the row-parallel partials (or all-gathers the batch-invariant hidden shards), so the layer's rows stay replicated. The two forward forms of this engine are thus the sharded extend and the replicated decode step, and nothing else runs on it: the startup autotune's dummy extend carries the shard plan a real extend of its rows would (PrefillGraph.make_dummy_batch, the model taking its slice as _run_target_forward does), rather than an unsharded replicated-row extend the contract does not list. The expanded (dense MLA) prefill keeps refusing head TP: only the absorbed sparse prefill can take the exchange. The alternative to this fork — sharding the drafter's decode rows by a plan too, so every forward exchanges — needs a sharded DSA decode arm (history gathers for decode rows) that does not exist; the fork is the contained form until it does. mapping.attn.head_tp_serves_decode_only, which is False over the query shards, requires the Decode role for head TP without QCP. The prefill role's own rules (validate_qcp) and the mapping (head TP == qcp) gate the sharded Prefill layout.

Startup chooses its forward from the engine role. Every disaggregated Decode worker warms the existing Decode path at its maximum Decode batch, including ordinary TP and TP1 attention DP. It reserves no Prefill graph buckets and never captures an extend-shaped graph, even if a Prefill graph limit was supplied. Remote admission and snapshot restore do not execute local Prefill. Prefill workers and aggregated engines retain their extend warmup; disabling autotune disables tactic profiling, not the role's warmup traversal.

Eager only: the history gather runs in the attention break, so the GPU DSA backend declares CudaGraphSupport(prefill_graph=False) under a query shard and the executor never captures one (it also refuses to replay a prefill graph for a sharded forward). A graph-capable form needs equal padding on every rank so the gathers are even collectives inside the captured segment. Each rule lives once: the prefill role and --kv-parallel-size equal to the shard in validate_context_parallel_roles; MIXED rounds, attention DP and dense / MoE groups that are neither 1 nor the attention TP width in validate_qcp; the shard width (N == attn.tp_size) in Mapping; and the attention family (GPU DSA), device, and KV cache (bf16: the gathered write is latent_store, native rows only) in AttnConfig, which sees them resolved.

Non-goals

Extend/mixed metadata keeps its dynamic-shape construction path (init_forward_metadata), with PrefillGraph as its own capture story. The write-location kernels stay pure functions (paged/write_locations.py). QSA reuses the shared table fill without expansion (cache-concepts.md Principle 5).

Regression gates

• test/runtime/execution/test_kda_prefill_graph_cache.py is registered in runtime-1gpu; its direct-script entry point runs pytest. It covers request padding, capacity selection, metadata refresh, and state isolation. Native CuTeDSL replay tests run on NVIDIA SM100/SM103 and skip other devices; missing native dependencies on a supported device are errors, not skips.
• test/runtime/test_unified_decode_path.py — eager refresh and padded replay refresh produce identical live-request contents over the same buffers; lazy above-ladder views are pointer-stable; the graph_ptr_guard walk reports a rebound tensor by path and pins every tensor under the slots; FlashMLA's tile schedule stays off the views (capture keeps its object alive, replay refresh leaves it alone, eager refresh and every drafter seq_lens edit bind a fresh one, in-graph edits keep theirs alive); leaf capture/refresh signature conformance.
• test/runtime/execution/test_draft_target_wiring.py — the drafter's target-forward hook: DFLASH arms its capture sink on the context only on decode-only, non-idle rounds, the sink concatenates the taps and writes the context once at the last tap (test_dflash_context_projection.py: bit for bit what the projection of the concatenated taps writes, with fc split across attention TP); the executor calls the hook before the target forward (test_model_executor_cache_state.py); the target hands taps to the forward's sink in concat order (test_dspark_config.py).
• test/runtime/test_cache_group_router.py — router slot math, expansion, padding, placeholder delivery, per-group dispatch, draft window publication and address stability.
• test/runtime/test_qsa_backend.py — independent QSA raw-group metadata, target-only verify workspace, and live cache writes across eager execution and CUDA graph replay; test_qsa_verify_lifecycle.py — the Qwen4-Exp root commits GDN/PLE on decode and QSA on decode/mixed, using real acceptance rows once after execution, including PLE without GDN and failure cases.
• test/runtime/test_qwen4_backend_composition.py — local consumer selection, workspace accounting, draft hooks through the attention composite, and one PD cache step per layer.
• test/runtime/test_cudagraph_per_group.py, test_group_write_locations.py — per-group padding wiring and the write-location edge cases (holes, overflow, MTP re-anchor) on the unified path.
• grep -rn "init_forward_metadata_replay_cuda_graph\|is_all_greedy" python/ must stay empty.
• grep -rn "ENABLE_CP\|CP_METADATA" python/ must stay empty — the module-global context-parallel switch and its process-wide metadata holder were deleted; a parallel layout is a Mapping fact from an explicit server argument and per-forward row metadata rides ForwardContext.
• Every init_forward_metadata signature carries query_shard with no default (test_unified_decode_path.py binds it on every runner-facing node and leaf; leaves also take page_table_cpu), and grep -rn "input_num_tokens" python/tokenspeed/runtime/models/longcat_flash.py python/tokenspeed/runtime/models/base/causal_lm.py stays empty — a model under query context parallelism reads its row count from its tensors or ctx.query_shard, never from the scheduler's chunk count (test/runtime/distributed/test_query_shard.py).
• grep -rn "ctx.accept_lengths\|ctx.draft_seq_lens_buf\|_apply_correction" python/ must stay empty — the step-0 accepted prefix is published through ctx.draft_narrowing.publish_accepted_prefix(), never computed in a model (test/runtime/test_draft_advance_seqlens.py).
• grep -rn "ctx.dsa_\|dsa_swa_slot_mapping\|dsa_compressor_slot_cache" python/ must stay empty — the layer-shared sparse top-k is backend scratch (sparse_topk), cleared by every metadata build (test_cache_group_router.py).
• grep -rnE '^\s+extend_(seq|prefix|replay|prompt)_lens(_cpu)?: torch\.Tensor \| None,| extend_with_prefix: bool = False' python/tokenspeed/runtime/layers/attention/backends/ must stay empty — no init_forward_metadata parameter in the extend bundle is optional or defaulted (test/runtime/test_unified_decode_path.py binds the runner call shape against every runner-facing node and every leaf). Metadata dataclasses may still hold None for fields a decode batch does not carry; the contract is about the call, not the record. The one permitted match is the private QSAIndexerBackend._metadata helper, which is not an init_forward_metadata.
• grep -rn "select_out_cache_loc\|DraftPageStaging\|tables_self_padding\| cache_active_pages_must_be_real\|engine_owned_group_ids" python/ must stay empty — history write locations have one accessor (write_locations), sparse prepare returns its compute writes, and table delivery has no capability flags.
• grep -rn "out_cache_loc" python/tokenspeed/runtime/models/ matches only write_locations(...) fetches and the helper parameters they feed — never a forward-chain parameter threaded from the runner.
• grep -rn "AttentionArch.DSA\|qwen4_exp_has_side_state" python/tokenspeed/runtime/execution/ must stay empty — backend-imposed graph restrictions are cuda_graph_support declarations (test/runtime/test_cudagraph_support_resolution.py).
• grep -rn "def init_forward_metadata_capture_cuda_graph" python/ matches only the defaults (backends/base.py, paged/base.py, paged/router.py) and the sanctioned overrides listed in "Capture is inherited".
• New backends implement refresh_decode_metadata + init_cuda_graph_state; capture is inherited from the base default (idle refresh). Only a kernel-imposed capture asymmetry justifies an override.

Sparse KV offloading

Eager and captured decode execute the same _forward_step: the executor calls prepare_cache_batch, runs the target forward and verification, then calls writeback_accepted_kv before publishing completion. Request slots, positions, canonical Top-K IDs and accepted target-input lengths are GPU inputs. The arena-owned SparseKVOffload engine keeps hot buffers, tags, LRU, and metadata at fixed addresses. KVOffloadAdapter binds router layers to fields and retains each step's compute mappings. See cache-concepts.md, "Sparse KV offloading and physical storage", for allocation and admission.

At each sparse layer, prepare_sparse_kv_access returns the projection's compute write slots and publishes mapped attention reads. A declared selection producer can call prefetch_sparse_kv for an offloaded consumer: current-slot preparation runs on the execution stream, the prefetch stream waits for it, and lookup plus swap-in record the field's ready event. The consumer's prepare joins that event. Inline resolve and prefetched resolve produce the same ordered selection mapping.

Verification supplies accepted input lengths directly to GPU writeback masking. The write stream waits for execution, copies accepted rows to Host history and records write_done; execution waits for that event before completion. Capture records these dependencies and replay uses the captured buffers and GPU inputs. Offloading requires no CPU read of accepted lengths or sequence lengths. Null and padding requests install no tags and commit no KV rows.

When PD landing resets valid history lengths for a recycled request slot, invalidate_cache_residency runs outside graph execution. It fences prefetch and writeback streams and resets request-slot tags, seed flags, and LRU while preserving authoritative history. PD transfer completion precedes seeding. Here residency means the hot-cache contents managed by the arena's offload engine; the router adapter holds the compute mappings for accessing them.

Offloading disables the snapshot pool and uses the scheduler's existing capacity-abort path. The request's logical history blocks and slot are released, and the client receives a CapacityAbort finish. The fixed hot partition is reused only after the new owner's PD initialization invalidates its metadata. No snapshot store, restore, recovery prefill, or extend staging is used by offloaded arenas.

DeepSeek V4.1 context parallelism

V4.1 uses operator-local prefill query partitions and page-sharded global cache groups. Its eager and captured decode share metadata refresh and DCP partial-output reduction; SWA and compressor tails remain replicated. See the V4.1 context-parallel contract for ownership, graph lifetime and supported configurations.

Distributed draft tuning

The startup draft decode uses one request per attention TP rank so that sequence-parallel router projections have nonempty, equal local row counts. A single-request TP warmup can skip a GEMM on an empty peer while another peer waits in distributed tactic selection. The draft remains in the same autotune context as the target. Serving batches and graph capture sizes are unchanged. On NVIDIA, a speculative decode service must reserve at least one request per attention TP peer: max_num_seqs // data_parallel_size >= attn_tp_size. Server-argument validation rejects smaller capacities before allocating buffers or running a forward, including with --disable-autotune (which disables profiling, not startup warmup). Increase --max-num-seqs to at least data_parallel_size * attn_tp_size. Prefill-only roles do not run this draft decode warmup and are exempt. Other platforms run no FlashInfer tuning, so startup tuning has no cross-rank collectives: the draft decode warmup uses min(attn_tp_size, max_decode_bs) requests, and an empty peer skips its GEMMs as it does when serving fewer requests than peers. Decode state stays sized by --max-num-seqs.