[Feature][Core] Prefix checkpoint coordination and same-step pairing in V1 scheduler

[nicholaskh-ai]

Purpose

Implements RFC #55697.

Part 2 of 3 for Application-Directed Mamba Prefix Checkpointing (depends on PR #55873).

This PR introduces scheduler coordination, checkpoint boundary chunk splitting, and a dependency state machine for hybrid Mamba/GDN models in the V1 engine:

• Producer Truncation: Enforces prefill chunk splitting exactly at mamba_checkpoint_position so that the recorded state snapshot is anchored precisely at the checkpoint boundary.
• Unready State Machine: When a Producer is scheduled, its checkpoint block hash is marked unready in BlockPool to prevent concurrent consumers from reading uncomputed states. Consumers arriving while the checkpoint is unready are deferred (waiting_for_mamba_checkpoint = True) and prepended to step_skipped_waiting without blocking unrelated requests.
• Ready State Wakeup: When the Producer's prefill completes, mark_checkpoint_ready() unlocks the block hash, resuming deferred consumers as normal cache hits.
• Same-Step Pairing: When a Producer and Consumer arrive in the same scheduling pass, the Consumer inherits the Producer's prefix blocks directly and registers mamba_prefix_producer_id, enabling them to execute in the same engine step.

Duplicate-work Check

• PR [Bugfix][Core] Retain align-mode mamba grid and decode-end checkpoints #53803 targets automatic grid retention, not application-defined checkpoint coordination.
• PR [Bugfix] Stop mamba align prefill at the replay boundary #50409 and [Bugfix][V1] Use the Mamba cache group's block size for align-mode chunk splitting #54076 address replay boundary chunk splitting, not explicit checkpoint stops.
• PR [Mamba] Unify FI ReplaySSM STP/MTP/Prefix Caching Lifecycle #55688 targets ReplaySSM unification.

Test Plan

• Unit test for scheduler checkpoint coordination, same-step pairing, and unready-wakeup state machine:
  pytest tests/v1/core/test_mamba_checkpoint_scheduler.py -v
• Unit test for Mamba aligned chunk splitting at checkpoint stops:
  pytest tests/v1/core/test_mamba_align_chunk_split.py -k "test_internal_checkpoint_split" -v

AI assistance was used to prepare this change. The human submitter is responsible for reviewing the changed code and test results.

Generated with Devin

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Signed-off-by: nizhang1 nizhang1@coupang.com

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

[mergify]

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

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

[nicholaskh-ai]

CC @WoosukKwon @yewentao256

Hi reviewers, this is Part 2 of the implementation stack for RFC #55697. It adds scheduler-level checkpoint truncation, same-step Producer/Consumer pairing, and the unready/ready state machine in V1 scheduler.

Whenever you have a moment, we would appreciate your feedback. Thanks!

[mergify]

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

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

[yewentao256]

Thanks for the work! Will this apply to MRv2? We will deprecate V1 soon

[nicholaskh-ai]

Hi @yewentao256, thanks for the review and great question!

Yes, absolutely! Here is the compatibility breakdown across the stack:

• PR [Feature][Config] Support application-directed Mamba checkpoint token and input processing #55873 (Part 1 - Config & InputProcessor) & PR [Feature][Core] Prefix checkpoint coordination and same-step pairing in V1 scheduler #55875 (Part 2 - Scheduler Coordination):
  These are engine-level (InputProcessor, Scheduler, and BlockPool), completely runner-agnostic, and already apply seamlessly to both MRv1 and MRv2 without any changes.
• PR [Feature][Mamba] Batched two-phase grouped prefill execution for hybrid GDN/Mamba (+7.58x) #55876 (Part 3 - Worker & Kernels):
  The two-phase batched GDN kernels and GDNAttentionMetadataBuilder are also runner-agnostic. We have implemented the MRv2 wiring in vllm/v1/worker/gpu/model_states/mamba_hybrid.py and input_batch.py (plumbing mamba_prefix_producer_indices, mamba_checkpoint_positions, and mamba_checkpoint_source_block_ids through CommonAttentionMetadata), with unit tests passing. We will include native MRv2 support in PR [Feature][Mamba] Batched two-phase grouped prefill execution for hybrid GDN/Mamba (+7.58x) #55876 directly!

[nicholaskh-ai]

Update: We have implemented native MRv2 compatibility in PR #55876!

1. Full MRv2 Wiring:
   The checkpoint metadata (mamba_prefix_producer_indices, mamba_checkpoint_positions, and mamba_checkpoint_source_block_ids) is now fully wired through vllm/v1/worker/gpu/model_states/mamba_hybrid.py and input_batch.py with passing unit tests (test_prepare_attn_forwards_mamba_checkpoint_metadata).

2. Feature-Gated Execution:
   To ensure 100% production stability and zero regression as MRv2 undergoes ongoing upstream refactoring, the batched grouped GDN kernel in MRv2 is gated behind VLLM_ENABLE_MRV2_GROUPED_GDN=1 (disabled by default in MRv2).

3. End-to-End Verification:
   In our standard MRv2 production deployment, PR 1 (input token alignment) + PR 2 (scheduler prefix coordination) achieve 99.97% classification consistency and Pearson correlation > 0.9988 on 3,428 real-world multimodal requests.

The entire 3-PR stack is rebased, clean, conflict-free, and fully verified for both MRv1 and MRv2!

[mergify]

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

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