[RFC]: Reuse DeepSeek-V4 compressor partial-state workspace across Attention layers

[jiangkuaixue123]

Motivation

Primary objective: make 32K-class prefill batches schedulable

The immediate reason for this RFC is that the current DeepSeek-V4 compressor partial-state allocation can prevent operators from using a large max_num_batched_tokens, such as 32768, even when that batch size is otherwise desirable for prefill throughput and TTFT.

On affected DeepSeek-V4 Ascend deployments, configuring --max-num-batched-tokens=32768 does not mean the engine can actually admit a 32K-token prefill batch. Before the model forward starts, the KV allocator reserves raw-token partial-state pages for every C4/C128 compressor layer. The admission requirement therefore includes the full scheduled chunk multiplied by the compressor-layer count. A long request can remain in waiting=capacity with reported KV-cache usage still at 0%, because admission fails before any block is allocated or any layer executes.

Operationally, this forces users to reduce max_num_batched_tokens and split a large prefill into smaller chunks. That avoids the immediate admission failure but reduces prefill efficiency, increases the number of scheduling/model-forward rounds, and can regress TTFT. The primary goal of this proposal is therefore not merely to reduce an internal memory number: it is to remove the compressor-layer multiplier from transient partial-state capacity so that 32K-class prefill batches can be admitted and executed within the same HBM budget, subject to the persistent KV, activation, and true in-flight-workspace requirements.

The corresponding vLLM-Ascend reproduction and allocator analysis are tracked in vllm-project/vllm-ascend#15204. That report also shows why reducing manager-side accounting alone is not a valid fix: the current kernel still indexes raw-token state, so the allocation and execution contracts must change together.

DeepSeek-V4 C4/C128 compression produces persistent compressed KV, but the current vLLM/vLLM-Ascend execution path also materializes FP32 per-token compressor partial states before emitting compressed KV.

Today those partial states are represented as cache-owning layers and allocated through the paged KV-cache lifecycle. Before model execution, admission/allocation must reserve state pages for every compressor layer that may consume the scheduled raw tokens. For a large prefill chunk, the transient component therefore grows approximately as:

scheduled raw tokens × compressor layer count × partial-state bytes

This mixes two different lifetimes:

1. Persistent state: compressed KV plus the small per-request, per-layer C4 tail / C128 accumulator / prefix checkpoint needed across chunks.
2. Transient state: per-token intermediate state needed only while one compressor layer consumes the current device batch.

The second component is not inherently persistent across decoder layers. DeepSeek-V4 executes Attention layers in order, so a bounded workspace can be reused after the previous layer has finished consuming it.

Multiple requests in one forward do not prevent this reuse. Their scheduled tokens can occupy disjoint packed ranges in one batch workspace, identified by the existing request metadata, cu_seqlens, positions, and slot mappings. The actual concurrency bound is the number of overlapping device forwards, uBatches, speculative substeps, or streams—not the number of requests and not the number of compressor layers.

This matters to AFD because compressor state remains on the Attention role, while DBO/uBatch execution and CAMAsyncAFDConnector can have multiple device batches in flight. An AFD implementation must therefore define explicit workspace ownership and completion semantics instead of assuming a single synchronous forward.

This RFC is a focused follow-up to:

• AFD DeepSeek-V4 support: [RFC]: Add DeepSeek V4 support on GPU and NPU for vLLM 0.26.0 #169
• vLLM-Ascend partial-state memory architecture discussion: [RFC]: DeepSeek-V4 compressor partial-state memory architecture vllm-ascend#15204
• vLLM online C128 experiment: Support online C128 compression for DeepSeek V4 vllm#48119
• vLLM model-wide eager workspace attempt and its stream-safety revert:
  • [DSv4 Perf] Optimize workspace reuse for eager break vllm#49236
  • Revert DSv4 eager workspace reuse vllm#52836

The objective is not to reduce scheduler accounting while leaving the kernel contract unchanged. That would be unsafe: the current compressor still indexes raw-token state and may access scheduled-but-unprocessed entries. The runtime representation, kernel consumption, workspace lifetime, and admission accounting must change together.

Proposed change

1. Separate persistent compressor state from transient workspace

Represent the Attention-side compressor memory as three explicit classes:

persistent compressed KV
persistent per-request/per-layer tail or checkpoint
transient per-in-flight-forward workspace

Persistent C4/C128 state remains isolated by request and layer because different layers produce different values. Only the full-chunk or tiled transient storage is reused across layers.

Target memory model:

Persistent:
  O(layer_count × active_requests × tail_or_checkpoint_size)

Transient:
  O(active_device_forwards × workspace_tokens × one_layer_state_size)

instead of:

O(layer_count × scheduled_raw_tokens × partial_state_size)

2. Use a batch workspace, not one workspace per request

For one packed forward containing requests r = 0..R-1, allocate one workspace slot sized for the total scheduled tokens:

workspace_tokens = sum(scheduled_tokens[r] for r in batch)

Each request receives a disjoint logical slice derived from existing packed-batch metadata. The same physical slot is then reused in execution order:

Layer L:
  load each request's persistent Layer-L tail
  produce/consume partial states in workspace slot S
  emit compressed KV
  write back each request's Layer-L tail/checkpoint

Layer L+1:
  reuse workspace slot S only after Layer L's consumers complete

Request boundaries, uneven lengths, empty ranges, and chunk positions must remain explicit in metadata; sharing storage must not merge request state.

3. Allocate by in-flight execution slot

A single global buffer is insufficient for async execution. Introduce a bounded workspace arena/ring whose ownership is tied to an in-flight execution identity, for example:

(forward generation, uBatch, speculative substep, backend stream domain)

A slot cannot be reused until every producer/consumer on the relevant CUDA/NPU streams has completed. Safe implementations may use:

• same-stream ordering where it is guaranteed;
• recorded completion events and stream waits;
• a workspace lease released from a completion callback;
• graph-capture-specific fixed slots with stable addresses.

The slot count should follow the maximum simultaneously active device batches, not request count or layer count. DBO and async CAM must use different slots for overlapping uBatches. Merged MTP/DSpark substeps must not reuse stale metadata or storage from another substep.

The revert in vLLM #52836 demonstrates why Python forward return or layer order alone is not a sufficient lifetime boundary when auxiliary streams are active.

4. Keep C4 and C128 workspaces separate initially

The first implementation should use separate C4 and C128 arenas because their state shapes, overlap behavior, alignment, and kernel schedules differ. A later optimization may overlay them with:

max(C4_workspace_bytes, C128_workspace_bytes)

only after profiling proves their device lifetimes never overlap and graph/stream synchronization is explicit.

5. Stage the implementation

Phase 0: measurement and lifetime tracing

• Record per-forward request/token counts, compressor-layer count, allocated partial-state pages, active uBatches/substeps, stream/event ownership, and peak HBM.
• Separate persistent tail/checkpoint bytes, transient bytes, cache-page padding, and unrelated activation memory.
• Establish native colocated vLLM/vLLM-Ascend baselines before enabling AFD.

Phase 1: bounded full-chunk workspace reuse

• Preserve the existing mathematical compressor flow.
• Move only current-chunk transient partial states out of persistent paged KV allocation.
• Reuse one workspace slot across layers after device completion.
• Keep per-request/per-layer tail/checkpoint storage persistent.

This phase removes the layer-count multiplier even if workspace size still scales with the scheduled chunk.

Phase 2: tiled or online compression

• Consume a bounded tile of raw-token partial states.
• Update the C4 ring or C128 online accumulator.
• Emit compressed KV at valid boundaries.
• Discard consumed transient state immediately.

This phase additionally removes full-chunk scaling from the transient component. The CUDA online C128 work in vLLM #48119 is a useful reference, but C4 overlap, Ascend kernels, prefix checkpoints, and AFD async execution require separate validation.

6. Update admission accounting only with the new execution contract

Do not divide state-page accounting by the compression ratio or layer count in isolation.

Scheduler/KV admission may stop reserving full-chunk state pages for every compressor layer only after:

• the compressor no longer indexes those pages as persistent raw-token storage;
• the bounded workspace is reserved independently;
• all scheduled-but-unprocessed tokens have a safe workspace slot;
• persistent tails/checkpoints remain capacity-accounted;
• cancellation and failure paths release leases without early reuse.

7. Prefix caching, P/D, and AFD connector semantics

• Prefix caching should persist only a representation sufficient to resume compression at a valid boundary.
• P/D transfer should transfer compressed KV and required checkpoint state, not transient workspace contents.
• AFD connectors should never serialize or transfer the transient compressor arena; it is local Attention-worker execution state.
• CAMAsyncAFDConnector PCP snapshot/build/restore must include persistent compressor state where required, while workspace leases remain local to the active forward/uBatch.

8. Validation and acceptance criteria

Correctness coverage must include:

• one and multiple requests in the same packed forward;
• unequal request lengths and empty/padded ranges;
• chunk boundaries before, at, and after multiples of 4 and 128;
• multiple consecutive chunks for the same request;
• mixed prefill/decode batches where supported;
• MTP/DSpark speculative substeps;
• DBO with at least two overlapping uBatches;
• synchronous CAMP2P and asynchronous CAM execution;
• prefix-cache hit/resume and P/D checkpoint restore where supported;
• cancellation, exception, and connector-cleanup paths;
• eager and supported graph-capture/replay modes with stable workspace addresses.

Required evidence:

1. Token/output parity with the native non-reuse baseline.
2. No cross-request or cross-layer contamination under stress concurrency.
3. Device-event evidence that a slot is not overwritten before its previous consumers finish.
4. Peak-memory accounting showing transient compressor memory no longer scales with compressor layer count.
5. Performance results showing synchronization does not erase the memory-capacity benefit.
6. Non-AFD DeepSeek-V4 behavior remains unchanged when AFD is disabled.
7. On the agreed target Ascend stack and HBM budget, a long-prefill workload can actually schedule and execute a 32K-token device batch with --max-num-batched-tokens=32768, without remaining in waiting=capacity solely because partial-state capacity is multiplied by the compressor-layer count. The report must include the scheduled token count, admission block requirement, workspace bytes, peak HBM, number of in-flight workspace slots, and native-baseline comparison.

Plugin boundary

Plugin-owned:

• AFD Attention-worker integration and feature gating.
• Workspace lease/slot ownership for AFD forwards, DBO uBatches, async CAM work items, and speculative substeps.
• Connector-local lifecycle, cancellation, PCP coordination, instrumentation, tests, and AFD recipes.
• Explicit validation that transient workspace never crosses the Attention/FFN connector boundary.

Compat helper:

• Exact-version adapters that translate AFD forward/uBatch identities into the native workspace API for the pinned vLLM and vLLM-Ascend baseline.
• Signature and schema drift checks for the DeepSeek-V4 compressor and model-runner hooks used by the adapter.

Compat patch:

• Prefer none. The compressor state representation, KV-cache/admission contract, and CUDA/NPU kernel ABI should be implemented in vLLM/vLLM-Ascend and consumed by afd-plugin.
• If an experimental AFD-only patch is required to collect evidence, it must be narrowly scoped, exact-version gated, disabled for non-AFD execution, tested for import/registration isolation, and carry an explicit upstream/removal plan.

Explicit class path:

• The DeepSeek-V4 AFD Attention model/runner introduced under [RFC]: Add DeepSeek V4 support on GPU and NPU for vLLM 0.26.0 #169.
• AFDAttentionModelRunner integration for synchronous execution.
• The NPU async Attention runner / CAMAsyncAFDConnector work-item lifecycle for multi-uBatch execution.

Upstream-owned and unchanged by the final plugin design:

• DeepSeek-V4 compressor mathematics and persistent-state semantics.
• KVCacheSpec/cache-manager and scheduler admission contracts.
• CUDA and Ascend compressor kernels and their workspace ABI.
• Native non-AFD DeepSeek-V4 execution.

Risks and alternatives

Risks

• Cross-stream overwrite: a later layer or uBatch may overwrite a slot while an earlier kernel still reads it. This is the primary correctness risk.
• Incorrect lifetime classification: C4 overlap, unfinished C128 groups, prefix checkpoints, or speculative state may be mistaken for transient data.
• Cross-request contamination: incorrect packed offsets may make one request read another request's partial state.
• Graph incompatibility: dynamically leased storage or changing addresses may invalidate CUDA/ACL graph capture assumptions.
• Under-accounting: reducing KV admission before reserving bounded workspace can move failure from scheduling to execution-time OOM.
• Excess synchronization: coarse stream synchronization may preserve correctness but regress throughput or eliminate AFD overlap.
• Version drift: DeepSeek-V4 compressor and cache contracts are still evolving in vLLM and vLLM-Ascend.
• Checkpoint incompatibility: prefix-cache or P/D state may not contain enough information to resume a tiled/online compressor at arbitrary boundaries.

Alternatives

1. Keep the current paged state cache. Lowest implementation risk, but transient memory continues to scale with layer count.
2. Reduce max_num_batched_tokens. Avoids some admission failures but directly reduces prefill efficiency and does not fix the lifetime model.
3. Accounting-only reduction. Rejected because the current kernel still addresses raw-token partial state and can read scheduled-but-unprocessed entries.
4. Online C128 only. Valuable first step, but does not solve C4 overlap or provide a common AFD async lifetime contract.
5. One global workspace without leases/events. Rejected because asynchronous streams and overlapping uBatches make it unsafe.
6. Implement everything inside afd-plugin. Useful only as a temporary evidence lane; the durable cache/kernel contract belongs upstream and must not alter non-AFD execution through plugin installation.

Feedback requested

1. Should the first implementation target Phase 1 full-chunk reuse, or go directly to tiled/online C4 and C128?
2. What is the authoritative device-completion boundary for releasing a workspace slot on CUDA and Ascend?
3. Should slots be owned by model-runner forwards, uBatch IDs, or an upstream workspace manager shared by all DeepSeek-V4 attention consumers?
4. What is the minimal persistent state required for C4, C128, prefix caching, and P/D resume at arbitrary chunk boundaries?
5. Should C4 and C128 remain permanently separate, or may they share an arena after stream-lifetime validation?
6. Which parts should land first in vLLM, vLLM-Ascend, and afd-plugin respectively?
7. Which GPU/NPU topology and async/DBO configuration should gate the first implementation PR?

Feedback period

One week, through 2026-09-07. Implementation may begin with instrumentation and a native-baseline reproducer while the ownership and synchronization contract is under review; scheduler accounting changes should wait for agreement on the complete execution contract.

CC list

@hsliuustc0106 @specture724 @jiaran-king @yujuancao07 @ShwStone

[jiangkuaixue123]

cc @ShwStone