[Kimi K3][Pref] Reuse internal checkpoint blocks for partial prefix caching

[ZeldaHuang]

Purpose

This PR is stacked on #52789 and extends its Kimi-K3 prefill checkpoint path to fine-grained partial prefix caching.

Main changes:

• Export the causal-conv checkpoint directly from the existing causal-conv forward kernel.
• Store recurrent checkpoints through the single FlashKDA prefill launch provided by the base PR.
• Re-key [Perf] Support internal prefill checkpoints for Mamba prefix caching, 9%~25% TTFT improvement #52789's request-owned internal checkpoint block at the partial hash boundary (for example, 128 tokens), instead of introducing a separate checkpoint-block owner.
• Reuse the existing req_to_blocks allocation, prefix hash, eviction, and free lifecycle.
• Reject checkpoint allocation when the requested offset is not aligned for the backend, preventing an unwritten block from becoming cache-visible.

This is not a duplicate of #52789: that PR provides the internal full-block checkpoint infrastructure; this stacked PR makes the same block lifecycle work for partial prefix boundaries and removes the remaining separate causal-conv checkpoint reconstruction.

Speculative decoding remains governed by #52789. The checkpoint fast path is not enabled with EAGLE in this PR.

Usage

No environment variable is required. The optimization is selected automatically when Kimi-K3 uses FlashKDA, prefix caching, and aligned Mamba cache mode.

vllm serve <model> \
  --kda-prefill-backend flashkda \
  --enable-prefix-caching \
  --mamba-cache-mode align \
  --prefix-match-unit 128

The required FlashKDA checkpoint support was merged in FlashKDA#7.

Performance

Setup:

• TP8 on 2 × 4 GB300
• Growing prefix from 120K to 164K tokens
• +2K tokens per round
• 8 concurrent streams
• Three runs per variant

Variant	P50 TTFT	Improvement vs non-partial
#52789, non-partial	2,151.53 ms	Baseline
#52789, partial128	2,375.23 ms	-10.40%
This PR, non-partial	2,141.67 ms	Baseline
This PR, partial128	1,398.49 ms	34.70%

Correctness and tests

Three consecutive full OCRBench runs with checkpoint mode and partial128:

• 89.4%
• 89.6%
• 89.3%
• Mean: 89.43%

Local regression after rebasing:

.venv/bin/python -m pytest -q \
  tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py \
  tests/v1/core/test_mamba_align_chunk_split.py
# 70 passed

pre-commit run
# all hooks passed, including ruff, mypy, DCO, and config validation

CUDA-only causal-conv and KDA metadata tests are included for CI/GPU validation.

AI assistance was used to rebase the branch, implement the lifecycle refactor, and add tests. The human submitter must review and validate the final change.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[yewentao256]

Thanks for building on #52789. Since this PR currently overlaps with its checkpoint infrastructure, please keep it in draft until #52789 lands, then rebase it on top and remove the duplicated pieces.

[yewentao256]

We should only use main branch when it is fully tested and reviewed, please hold until vllm-project/FlashKDA#7 landed

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ZeldaHuang.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[yewentao256]

I think previous PR doesn't support spec yet, @ZeldaHuang @ZJY0516 shall we open a new PR for this?

[ZJY0516]

I think previous PR doesn't support spec yet, @ZeldaHuang @ZJY0516 shall we open a new PR for this?

We will do it