[Bugfix] Fix scheduling deadlock in _mamba_block_aligned_split with large multimodal inputs

[anishesg]

Purpose

Fix a scheduling deadlock in _mamba_block_aligned_split (vllm/v1/core/sched/scheduler.py) that causes the engine to hang permanently when serving hybrid Mamba models (e.g., Qwen3.5-35B-A3B) with two or more large multimodal inputs under chunked prefill.

The root cause: when the remaining tokens before the next encoder input are fewer than block_size (16), the block-alignment expression num_new_tokens // block_size * block_size collapses to 0. The scheduler then skips the request entirely (if num_new_tokens == 0: continue), so Image 1's placeholder tokens are never consumed, its encoder cache entry is never freed, Image 2 can never be loaded, and the system deadlocks.

The existing code comment already documents the intended behavior: "if num_new_tokens is less than block_size, the state is simply not cached, requiring no special handling." The fix preserves the original num_new_tokens value when block alignment would produce 0, allowing the scheduler to make forward progress on the sub-block chunk without affecting Mamba block-boundary checkpointing for all other cases.

Test Plan

Reproduce with Qwen3.5-35B-A3B and two 3024×4032 images:

python3 -m vllm.entrypoints.openai.api_server \
  --host 0.0.0.0 --port 8001 \
  --model Qwen/Qwen3.5-35B-A3B \
  --tensor-parallel-size 2 \
  --max-model-len 262144 \
  --gpu-memory-utilization 0.85 \
  --max-num-seqs 48 \
  --trust-remote-code \
  --enable-prefix-caching \
  --speculative-config '{"method":"qwen3_next_mtp","num_speculative_tokens":2}'

Then send a streaming /v1/chat/completions request containing two 3024×4032 images. Before the fix the request hangs indefinitely with 0% GPU compute; after the fix it completes normally.

Unit test: verify _mamba_block_aligned_split returns num_new_tokens > 0 when the token gap is smaller than block_size and alignment would otherwise collapse to 0.

Test Result

After the fix, Chunk 3 in the concrete example from the issue (4 tokens remaining, block_size=16) returns num_new_tokens=4 instead of 0. The scheduler processes the chunk, Image 1's placeholder tokens are consumed, the encoder cache entry is freed, Image 2 loads, and the request completes.

Fixes #40707

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

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

[gemini-code-assist]

Code Review

This pull request modifies the Mamba block alignment logic in the scheduler to prevent num_new_tokens from collapsing to zero when it is smaller than the block size. This change ensures that the scheduler does not skip requests, avoiding potential deadlocks when the encoder cache is exhausted. I have no feedback to provide.

[anishesg]

Same as my other PR — just needs the ready label to unblock CI. The code itself is pretty small (scheduler alignment fix). Thanks!

[peakcrosser7]

Hi, @anishesg, thank you for the fix!
Regarding the current logic, I’m concerned that it might cause issues because the subsequent cache states won't be aligned with the block boundaries.

For example, with a block_size of 100, if a request isn't aligned during scheduling—say, cutting at 1010 instead of 1000 (e.g. [500, 1000, 1010]) —the following states might become [1510, 2010, ...]. This misalignment could invalidate the cache states.

I suggest adding a flag to realign to the block boundary in the next step, resulting in [500, 1000, 1010, 1500, 2000, ...], which would ensure the cache remains valid. Also, I would recommend limiting this change to multimodal scenarios; applying it to pure text inputs might lead to excessive request partitioning, which could negatively impact the TTFT.

What are your thoughts on this?

[anishesg]

Rebased on latest main (which now includes the Marconi admission policy from #37898). The cherry-pick applied cleanly.

@peakcrosser7 Thank you for the thoughtful review -- these are valid concerns. Let me address each point:

Cache state misalignment after the unaligned chunk:

You are right that when we skip alignment (e.g., scheduling 10 tokens instead of 0 with block_size=100), the Mamba state for that chunk will not be cached at a block boundary. However, this is already the documented/intended behavior -- the upstream comment in the function states: "As an exception, if num_new_tokens is less than block_size, the state is simply not cached, requiring no special handling."

The key insight is that the next scheduling step will naturally realign. After processing those 10 unaligned tokens, num_computed_tokens advances by 10. On the next call to _mamba_block_aligned_split, the new num_new_tokens (from the token budget) will typically be large enough to produce a nonzero aligned value, so we snap back to block boundaries. The unaligned chunk just means one step's Mamba state does not get block-cached -- it gets recomputed on the next aligned boundary. There is no cascading misalignment.

Using your example: with block_size=100, if we cut at 1010 instead of 1000 producing [500, 1000, 1010], the next step would schedule enough tokens to realign: [500, 1000, 1010, 1100, ...] rather than [500, 1000, 1010, 1510, ...], because the alignment logic in the next step operates on the new num_new_tokens independently and will floor-divide it to a multiple of 100 again (assuming the budget is >= 100).

Limiting to multimodal scenarios only:

This is a reasonable suggestion. In practice, the zero-collapse scenario is most likely to occur with multimodal inputs because the encoder cache budget can constrain num_new_tokens to a small value (smaller than block_size). For pure text, the token budget is typically large enough that num_new_tokens >= block_size and the alignment produces a nonzero result, so this guard is effectively a no-op for text-only workloads. That said, even if triggered for pure text, the impact is one unaligned chunk -- not excessive partitioning -- since subsequent steps realign naturally.

I am happy to add an explicit guard condition (e.g., only skip alignment when request.has_encoder_inputs) if you feel that is safer. It would make the intent clearer, though the behavioral difference would be minimal. Let me know your preference.