[ZJY0516]

Overview

Each Mamba layer declares the token alignment of the prefill checkpoints it can export from inside a forward pass once, in MambaBase.get_prefill_checkpoint_alignment(). MambaBase.get_kv_cache_spec derives num_prefill_checkpoint_blocks and prefill_checkpoint_alignment from it. Mamba2 (#57329), GLM-5.3-Flash KDA (#56960) and Kimi-K3 KDA (NVIDIA and AMD) drop their get_kv_cache_spec overrides. Split out of #60050.

Claims

• Spec fields are unchanged when every Mamba layer of the model declares an alignment, which holds for all Mamba2 and KDA models: Mamba2 1, FlashKDA 16, ROCm fused KDA KDA_CHECKPOINT_ALIGNMENT. Mamba layers that declare none still get no checkpoint fields.
• Four copies of the same replace(spec, num_prefill_checkpoint_blocks=..., prefill_checkpoint_alignment=...) override become one, so the two fields can no longer disagree. A checkpoint block reserved without an alignment would cache a state that is never written.
• The platform gets one place to read a layer's checkpoint grid. [Perf][GDN] Internal prefill checkpoints for the FlashInfer GDN prefill backend #60050 (merged) reads it to keep align-mode blocks on FlashInfer GDN's 64-token grid, and with this PR the KDA layers' alignments feed the same rule. That has no effect for FlashKDA (16). ROCm fused KDA blocks become multiples of KDA_CHECKPOINT_ALIGNMENT, so its checkpoints are no longer declined off-grid.

Validation

Tests (after merging main, 1x GB200 node):

Test	Result
tests/v1/attention/test_mamba2_prefill_checkpoints.py, tests/v1/attention/test_gdn_metadata_builder.py, tests/models/kimi_k3/test_kda_metadata.py, tests/models/kimi_k3/test_amd_kda_checkpoint.py, tests/models/qwen4_exp/test_ple.py	233 passed, 29 skipped
tests/kernels/mamba/test_mamba2_checkpoint_state.py	8 passed
tests/kernels/mamba/test_mamba_mixer2.py	3 passed
tests/models/kimi_k3/test_kda.py, tests/models/glm5next/test_kda_recurrent.py	123 passed, 7 skipped
tests/v1/core/test_mamba_align_chunk_split.py	48 passed
tests/v1/core/test_prefix_caching.py -k "mamba or Mamba"	30 passed
tests/v1/e2e/general/test_mamba_prefix_cache.py (mrv1, mrv1_async, mrv2, mrv2_async, mrv2_async_flashinfer)	5 passed
tests/models/language/generation/test_hybrid.py -k apc	7 passed

• The new test_layer_checkpoint_alignment_reaches_the_kv_cache_spec checks that a declared alignment, or none, reaches the spec together with its checkpoint block.
• test_amd_kda_checkpoint.py needs a gfx950 build with the fused KDA chunk kernel, so it is skipped here. Its _spec_for helper now runs the layer's real get_kv_cache_spec over stubbed state shapes instead of stubbing the base spec, which would bypass the derivation.

GSM8K (tests/evals/gsm8k/gsm8k_eval.py, 20-shot, 1319 questions, --enable-prefix-caching, 1x GB200), so prompts take and resume from internal checkpoints:

Model	Layer	Accuracy	Prefix cache hit rate
granite-4.0-h-tiny	Mamba2	0.788	89.7%
Kimi-Linear-48B-A3B-Instruct	Kimi-K3 KDA (NVIDIA)	0.895	54.0%

Details

• get_prefill_checkpoint_alignment() returns None by default. Readers act on it only in align mode, as before.
• Merging main. [Perf][GDN] Internal prefill checkpoints for the FlashInfer GDN prefill backend #60050 added get_prefill_checkpoint_alignment and the rule that the scheduler takes internal checkpoints only when every Mamba layer declares an alignment. Its FlashInfer GDN layer and Qwen4-Exp's PLE override get_kv_cache_spec to put their offsets on the prefix-cache grid. Both return the base spec unchanged when another layer cannot checkpoint. MambaBase.get_kv_cache_spec therefore applies that rule itself: it fills the checkpoint fields only when every Mamba layer declares an alignment. Otherwise those layers would inherit a checkpoint block that is never written. test_no_checkpoint_block_unless_every_mamba_layer_declares_one covers this.
• Kimi-K3 AMD's "KDA prefill checkpoint enabled" log moves from get_kv_cache_spec to __init__.
• Not a duplicate: no open PR touches these overrides (searched 2026-10-13). [Perf][Core][Mamba] Export prefill checkpoints at or before the initial-state column #60888 changes the checkpoint validity rule in kv_cache_interface.py, which this PR does not touch.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

🤖 Generated with Claude Code

[mergify]

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

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

[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.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-14T09:23:53.652132Z	e22519e	Draft marked ready

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: e22519ef62

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Preserve the fallback for duck-typed spec callers

When test_replayssm_does_not_reserve_speculative_state_blocks runs, it calls MambaBase.get_kv_cache_spec unbound on the SimpleNamespace constructed in tests/v1/attention/test_attention_backends_selection.py:32-48; that object has no get_prefill_checkpoint_alignment, so this unconditional dynamic call raises AttributeError before either parameterized case reaches its assertions. Add the default hook to that fixture or otherwise retain a safe default for this existing caller.

Useful? React with 👍 / 👎.

[mergify]

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

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

[ZJY0516]

/ci run --allow-stale

[github-actions]

✅ Triggered Buildkite CI #97930 for commit 750f663184d1.

⚠️ This PR is 2 commits behind upstream main. Running CI at your own risk because --allow-stale was requested; outdated CI configuration may cause failures. Before merging, merge or rebase onto the latest main, then rerun /ci run on the latest PR commit.

[mergify]

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

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