[II] Align hybrid recurrent checkpoints with cache-group boundaries

[voipmonitor]

Behavior

Hybrid Mamba/attention serving can retain an explicitly configured recurrent-state checkpoint cadence while physical Mamba pages remain padded to the attention page size.

The scheduler derives recurrent alignment from the resolved MambaSpec groups. For EAGLE models, it converts the last reusable annotated draft-attention boundary onto the recurrent-state grid. A prefill budget smaller than the recurrent interval may therefore advance through private intermediate states and publish a checkpoint only on a boundary shared with target attention.

Technical reason

Platform._align_hybrid_block_size previously replaced every align-mode --mamba-block-size with the physical attention block size. Scheduler._mamba_block_aligned_split also read the physical cache block directly. Those two behaviors coupled logical recurrent checkpoints to physical page geometry.

For Kimi-K3 TP16/DCP16, 768-token recurrent checkpoints combined with 4,096-token prefill budgets did not land on 12,288-token target-attention boundaries. Native offload could restore only the first common 61,440-token boundary. A 12,288-token recurrent cadence permits three 4,096-token prefill steps per target boundary without changing physical page bytes.

Compatibility

• Default configurations are unchanged: when --mamba-block-size is not explicit, align mode still uses the physical attention block size.
• Physical KV page sizing and allocation are unchanged.
• Configurations without annotated EAGLE groups retain the conservative one-recurrent-block backoff.
• The behavior is model-independent for hybrid cache groups; Kimi-K3 TP16/DCP16 supplies the motivating geometry.

Validation

• 21 focused scheduler and Mamba alignment tests passed.
• 45 broader Mamba, hybrid prefix-cache, and offloading tests passed.
• The scaled Kimi-K3 regression uses 12 recurrent groups, 4 target-attention groups, and one EAGLE draft group. Scheduler budgets representing 768 and 4,096 tokens both stored 11 checkpoints in every recurrent and target group and restored the same full-context prefix.
• Ruff and formatting checks passed.

Full-model TP16/DCP16 validation is tracked in the source-composed Kimi-K3 image qualification.

Upstream relationship

• vllm-project/vllm#46251 is a draft that separates logical block sizes from allocation size through a hierarchical allocator. This pull request does not import that unfinished allocator; Kimi-K3 already uses uniform physical page bytes.
• vllm-project/vllm#52971 exports Kimi checkpoints from FlashKDA. It requires the FlashKDA prefill backend and addresses internal checkpoint export rather than explicit align-mode checkpoint cadence.

[coderabbitai]

Warning

Review limit reached

@voipmonitor, you've reached your PR review limit, so we couldn't start this review.

Next review available in: 11 minutes

Limit details: You’ve used the included review currently available.

You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository.

How can I continue?

Wait for the limit to reset, then comment @coderabbitai review or push new commits to the PR.

An organization admin can change what happens after included review limits in Billing.

How do review limits work?

CodeRabbit enforces per-developer PR review limits within each organization.

For paid Pro and Pro+ reviews, CodeRabbit uses a developer's included PR review attempts over the past 7 days to set the current hourly allowance. At typical activity levels, the full plan allowance applies. Higher sustained activity can lower the allowance until earlier attempts leave the 7-day window.

Please refer docs for additional details.

Review details

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 7fcd80d5-fb36-455f-9fce-8cd34f3a22e6

📥 Commits

Reviewing files that changed from the base of the PR and between 337ef76 and 308dd56.

📒 Files selected for processing (4)

• tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
• tests/v1/kv_connector/unit/offloading_connector/utils.py
• vllm/platforms/interface.py
• vllm/v1/core/sched/scheduler.py

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}