[Feature] Chunk-boundary states for gated-delta prefix caching

[Tsundoku958]

Thanks for your contribution and we appreciate it a lot. The following instructions would make your pull request more healthy and more easily receiving feedbacks. If you do not understand some items, don't worry, just make the pull request and seek help from maintainers.

Motivation

Final goal: Marconi-style prefix caching for hybrid LLMs

Qwen3-Next / Qwen3.5 are hybrid (Attention + Recurrent) LLMs whose gated delta rule (linear attention) layer updates its recurrent state in place, so it cannot be sliced along the sequence or rolled back from the tail to represent an arbitrary prefix like KV cache. This forces SSM-state prefix reuse to require exact matches, while fine-grained checkpointing at every boundary floods the cache with low-hit-rate entries and thrashing. Marconi (arXiv:2411.19379, https://arxiv.org/html/2411.19379v1) is the system built for exactly this, with three core ideas:

1. Where to get the state: leverage chunked state passing and, during prefill, directly reuse the state at the second-to-last chunk boundary as the checkpoint (gated-delta is exactly this kind of chunk-based recurrence; this PR's chunk_states maps to this step);
2. Which boundaries to store (admission): not all — checkpoint only at radix-tree branch points (nodes where multiple requests first share a prefix and diverge) + the tail token, reducing cache entries at the source and avoiding the thrashing of storing everything;
3. Which entries to evict: score radix nodes by a FLOP-aware utility = recency + α·flop_efficiency and evict the lowest first — since SSM state size is fixed regardless of sequence length, longer sequences save more FLOPs per byte and are worth keeping.

This yields 4.5×–34.4× token hit-rate gains over fine-grained+LRU baselines on LMSys/ShareGPT/SWEBench. Our final goal is to land this mechanism on LMDeploy's gated-delta models; implementation details of each mechanism can be found in the paper.

This PR's role: filling the producer-side gap

main already ships the complete SSM prefix-cache consumer side (the block_trie/checkpoint_lifecycle.py reserve/publish/restore lifecycle, block_size=64 naturally aligned to chunk boundaries, reserve_save(step=) already step-aware, and the long-context path already per-chunk). However, gated-delta prefill currently runs FLA chunk_gated_delta_rule(..., output_final_state=True) and only keeps the terminal final_state — every intermediate 64-chunk-boundary state is discarded. This is precisely the producer-side gap that the first idea above ("capture state at chunk boundaries") needs: without it, the existing checkpoint slots have nothing to fill at intermediate boundaries.

This PR fills that producer side: the gated-delta chunk forward now emits, in one prefill pass, the recurrent state at every 64-chunk boundary (chunk_states) together with the paired conv state (chunk_conv_states), corresponding to Marconi's first step of capturing state at chunk boundaries, as the foundation for later wiring into main's lifecycle + Marconi admission/eviction. The local kernel is ported from the chunked gated-delta-rule forward of flash-linear-attention (FLA, https://github.com/fla-org/flash-linear-attention), and a restore-invariant test (chunk_states[:, c] → run suffix → reproduces the terminal state) verifies the boundary states are exactly reusable.

Modification

• New lmdeploy/pytorch/kernels/cuda/chunk_gated_delta_rule.py: a forward-only Triton chunk forward kernel chain ported from FLA that additionally materializes the recurrent state at every chunk boundary; the public chunk_gated_delta_rule returns (o, final_state, chunk_states), plus chunk_conv_states (the W conv-input tokens preceding each boundary). Includes local prepare_chunk_indices/offsets (replacing FLA's to avoid D2H sync) and autotune chunk-count bucketing.
• Framework integration: backends/cuda/gated_delta_rule.py makes FLA optional, dispatches between FLA / local by chunk_indices, and returns a unified 3-tuple; backends/attention.py, backends/cuda/step_metadata.py, backends/cuda/op_backend.py prepare chunk metadata once per step and reuse across layers; backends/causal_conv1d.py + backends/cuda/causal_conv1d.py expose chunk_conv_states through the backend abstraction (no direct CUDA import); nn/gated_delta.py threads the metadata and unpacks the 3-tuple; models/qwen3_next.py, models/qwen3_5.py adapt to the third return value (currently _chunk_states as a placeholder; the convergence target removes it once the checkpoint bank is passed down).
• Tests: tests/pytorch/kernel/test_chunk_gated_delta_rule.py (56 passed) covers FLA numerical consistency (multiple shapes / packed varlen / production TP4 / non-contiguous fused-QKV view), prefill→decode handoff, prepare-metadata-once, and the Marconi restore invariant (one for recurrent, one for conv — restore at a boundary and run the suffix to reproduce the terminal state, with an independent F.conv1d/FLA reference); tests/pytorch/kernel/test_causal_conv1d.py gains a conv-state backend-dispatch equivalence test.

BC-breaking (Optional)

No BC-breaking.

• GatedDelta.__call__ changes its return value from a 2-tuple to a 3-tuple; downstream qwen3_next / qwen3_5 are already adapted, and the interface only adds (the new chunk_states is not consumed by downstream yet, purely placeholder).
• FLA is downgraded from a hard dependency to an optional one: it automatically falls back to the local kernel when FLA is absent, with no impact on existing functionality.
• AttentionMetadata gains two optional fields gated_delta_chunk_indices/offsets, defaulting to None and not breaking existing construction paths.

Use cases (Optional)

For enabling SSM prefix-cache reuse on gated-delta models such as Qwen3-Next / Qwen3.5 in follow-ups. This PR only provides the boundary-state producer and does not change existing inference behavior; wiring into main's existing checkpoint lifecycle behind a CacheConfig flag is a follow-up PR.

Future plan

This PR is only the producer side; the follow-ups wire chunk_states / chunk_conv_states into main's existing lifecycle and progressively implement Marconi's admission/eviction, corresponding to the paper §4.1 (capturing states during prefill) and the admission/eviction policies:

1. Wire into the existing checkpoint lifecycle (Marconi's "capture state at chunk boundaries"): pass the chunk-state bank + slot map down to GatedDelta/the model layer (passed by reference like past_key_value), write boundary states in place into the slots reserved by reserve_save(step=) via index_copy_, and return to the 2-tuple (removing the _chunk_states placeholder). The chunk_states[:, c] here is exactly the per-boundary state the paper captures.
2. Multi-boundary save per prefill: extend _prepare_prefill_cache_save's save_steps to one-set-per-sequence, calling reserve_save(step=s) for each block-aligned step, so a normal prefill saves multiple chunk-boundary checkpoints at once.
3. conv_state persistence (net-new gap vs main): main only covers recurrent_state; conv state must be checkpointed/restored at the same step or the restored layer is inconsistent.
4. Marconi judicious admission (v2): not all boundaries — checkpoint only at radix-tree branch points (where a prefix first diverges) + the tail token, reducing cache entries at the source and avoiding the thrashing that fine-grained checkpointing causes (the paper notes an all-store hit rate of only 0.4%).
5. FLOP-aware eviction (v2): score radix nodes by utility(n) = recency(n) + α·flop_efficiency(n) and evict the lowest first, with α tuned via offline grid search; exploit that SSM state is fixed-size so longer sequences save more FLOPs per byte.
6. flag: a CacheConfig flag (e.g. --enable-marconi-prefix), shallow fork controlling storage/restore only, keeping the production path unified.

References

• Marconi: Prefix Caching for the Era of Hybrid LLMs, arXiv:2411.19379 — https://arxiv.org/html/2411.19379v1 (source of the final goal and admission/eviction policies)
• flash-linear-attention (FLA), from which the ported chunked gated-delta-rule forward originates — https://github.com/fla-org/flash-linear-attention
  @grimoire @lvhan028
  [Feature] 当前qwen-next等SSM类模型prefix cache缓存命中率低 #4879

[grimoire]

We have other models that require state cache, not all of them are GDN, could this feature support models with other state type?

[Tsundoku958]

We have other models that require state cache, not all of them are GDN, could this feature support models with other state type?

The lifecycle side is already state-type-agnostic; only the boundary-state producer is GDN-specific by necessity.

This PR splits into two layers:

• Consumer side (store/restore/admission/eviction) is generic — and is itself a follow-up PR. This PR ships only the producer (emitting per-chunk-boundary state). Wiring it into main's existing StateCacheSpec + checkpoint_lifecycle + Marconi admission/eviction is a separate follow-up PR (as noted in the PR description). That lifecycle describes arbitrary state types, not GDN — DeepSeek-V4 already registers 3 SSM/compress specs through it — so when the consumer side lands it will be open to other state types by construction, not as an afterthought.
• Producer side (emitting per-chunk-boundary state) is per state-type, by necessity. What a boundary state is and whether it can be resumed from depends on each model's own recurrence, so there's no single kernel for all types — only a shared contract (emit a state at a token-aligned boundary that, resumed, reproduces the terminal state).
  So "other state-type models" split into two cases:

(a) True linear RNNs (KDA / GLA) — reusable, with only minor kernel changes. KDA's chunk forward (fla.ops.kda) is built on the same chunk_gated_delta_rule_fwd_h recurrence kernel gated-delta uses, and already exposes return_intermediate_states=True returning h = [B, NT, HV, V, K] — exactly the per-boundary state this PR materializes for GDN. So for KDA the producer is essentially a minor tweak of the existing kernel: the recurrence + output kernels and the restore invariant (resume from h[:, c] → run suffix → reproduces the terminal state) carry over directly; only the gate cumsum (KDA's carries Mamba-style A_log/dt_bias) and the WY/intra step need a per-type port. GLA is similar. Each such model is one (small) ported producer on the shared lifecycle.

(b) Sliding-window / compressed-KV state (DeepSeek-V4's v4_window_kv / v4_compress_state) — not this mechanism. That state is an in-place overwritten ring buffer holding only the last ratio tokens: no accumulating recurrence, no boundary state to resume from, nothing to chunk-parallelize. It is KV, so it belongs on main's existing KV-cache prefix-hit path, not the SSM boundary-state lifecycle this PR adds.

In short: shared, state-type-agnostic lifecycle + one (small) ported boundary-state producer per true-linear-RNN state type; this PR ships the GDN producer, the lifecycle wiring is a separate follow-up, and a KDA/GLA producer would be another such follow-up reusing the same lifecycle and kernel.

[grimoire]

Sliding window support in dsv4 is treated as a state, at least in engine level.

The state cache and it's prefix caching is designed for models that require fixed size caches, not just mamba like ssm, it is model agnostic. Adding Marconi might requires updating to kernels/modeling that are technically not ssm. It would increase the maintainance pressure(you would not add kernel support for non-nvidia backend, right?).

We need some time to evaluate weither we can accept this.

[Tsundoku958]

Sliding window support in dsv4 is treated as a state, at least in engine level.

The state cache and it's prefix caching is designed for models that require fixed size caches, not just mamba like ssm, it is model agnostic. Adding Marconi might requires updating to kernels/modeling that are technically not ssm. It would increase the maintainance pressure(you would not add kernel support for non-nvidia backend, right?).

We need some time to evaluate weither we can accept this.

You're right — this would increase maintenance pressure as you described, and I won't be implementing the kernel for other backends. To address this, when I implement the follow-up, I'll gate the feature behind a flag: the chunk-boundary checkpoint path runs only when the flag is on, and with it off everything falls back to main's existing single final-state checkpoint. The chunk kernel itself is also NVIDIA-only, so on non-NVIDIA backends the feature simply won't engage — no errors, no broken paths, just no multi-boundary gain. I'm hoping this opt-in approach mitigates the concern to some extent.