[Core][Refactor] scheduler: refactor prefill chunk alignment logic

[jzakrzew]

Purpose

This is a behavior-preserving refactor that makes it possible to add prefill chunk alignment logic without introducing model-specific hacks directly into the scheduler. It introduces a PrefillChunkAlignmentPolicy interface, which separates prefill chunk-shaping decisions from the scheduler control flow. As part of this change, the existing Mamba align-mode chunk-splitting logic has been fully moved out of the scheduler and into the alignment policy layer.

The immediate motivation for this PR is the lack of progress on #38561. This PR does not include any batch-invariance-related logic, but it may provide a cleaner foundation for that work.

Test Plan

Added new tests for scheduler alignment policy behavior in tests/v1/core/test_scheduler.py.
Added tests for the policies themselves in tests/v1/core/test_prefill_chunk_alignment.py.

pytest tests/v1/core/test_scheduler.py -k align
pytest tests/v1/core/test_prefill_chunk_alignment.py

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

[gemini-code-assist]

Code Review

This pull request refactors the Mamba block-aligned splitting logic into a more flexible PrefillChunkAlignmentPolicy system. It introduces a protocol-based approach with a default no-op policy and a specific Mamba policy, allowing the scheduler to handle prefill chunk alignment more cleanly across different model types. A critical issue was identified in the Mamba alignment logic where a small token budget (less than the block size) could cause num_scheduled_tokens to snap to zero, potentially stalling the scheduler; this contradicts the code's internal documentation which suggests such small chunks should be allowed through without caching.