[Bugfix][DSv4] Skip zero-query-len prefill chunks in FlashMLA sparse prefill (complements #49059)

[fireflyc]

Why

[Fixes #51486] DeepSeek-V4 sparse prefill crashes with Assertion res == CUresult::CUDA_SUCCESS failed (.../phase1.cuh:614) when KVTransfer / SimpleCPUOffloadConnector schedules prefill requests with 0 tokens while their KV is being loaded.

When such a request sits in the batch's prefill region with query_len == 0, DeepseekSparseSWAMetadata.get_prefill_chunk_plan can emit a chunk covering only zero-query-len requests, so query_start == query_end. _forward_prefill then slices q/out to an empty tensor and passes it to flash_mla_sparse_fwd, which fails to build the output TMA descriptor (s_q == 0) and aborts.

This is not a FlashMLA kernel bug — it is a missing invariant in the DeepSeek-V4 sparse-prefill chunking: it assumed every prefill-region request has ≥ 1 query token, which the scheduler (KVTransfer async load) does not guarantee.

What

1. vllm/v1/attention/backends/mla/sparse_swa.py — get_prefill_chunk_plan now skips zero-query-len prefill requests when advancing chunk boundaries, so they never form their own (empty) chunk.
2. vllm/models/deepseek_v4/nvidia/flashmla.py — _forward_prefill computes the chunk's query range before any KV gather, skips chunks with query_end <= query_start (defense in depth), and returns early when the plan is empty (all placeholder requests).
3. tests/kernels/attention/test_flashmla_sparse.py — two unit tests covering the chunk planner with zero-query-len requests (skipped/never-empty) and the all-zero → empty-plan case.

Behavior of the existing chunk planner is unchanged when there are no zero-query requests (existing test test_deepseek_v4_prefill_chunk_planning_expands_for_short_sequences still passes).

Why not duplicating an existing PR

• [Bugfix][DSv4][SM120] Skip empty sparse-MLA prefill chunks #49059 "[Bugfix][DSv4][SM120] Skip empty sparse-MLA prefill chunks" addresses the same root cause (empty prefill chunk with query_start == query_end) but only for the FlashInfer / SM120 path (vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py). It does not touch the FlashMLA path (vllm/models/deepseek_v4/nvidia/flashmla.py) or the shared chunk planner (vllm/v1/attention/backends/mla/sparse_swa.py).
• This PR is the FlashMLA / SM90 complement: same query_start == query_end empty-chunk bug, manifesting as a FlashMLA TMA-descriptor assert at phase1.cuh:614 instead of FlashInfer's reshape error.
• Verified with gh pr list --repo vllm-project/vllm --state open --search "..." on 2026-08-08: no PR other than [Bugfix][DSv4][SM120] Skip empty sparse-MLA prefill chunks #49059 addresses this. If maintainers prefer, these two changes can be consolidated into [Bugfix][DSv4][SM120] Skip empty sparse-MLA prefill chunks #49059.

Test plan

• .venv/bin/python -m pytest tests/kernels/attention/test_flashmla_sparse.py -v -k "chunk_plan" — passed (3 tests: the existing chunk-plan test + the 2 new ones).
• pre-commit — passed (the mypy var-annotated error on covered from the first run was fixed with covered: list[int] = []).

Model evaluation

This change only skips chunks that produce zero output tokens (placeholder requests); it does not alter the computation or output of any real prefill token, so no numerical/accuracy change is expected.

Real smoke test: served DeepSeek-V4-Flash-0731 with the reproduction config below (the same one that previously crashed with the phase1.cuh:614 TMA-descriptor assert) and completed mixed decode/prefill requests without the crash:

  --trust-remote-code \
  --kv-cache-dtype fp8 --block-size 256 --tensor-parallel-size 4 --enable-expert-parallel \
  --enable-prefix-caching --data-parallel-size 2 --gpu-memory-utilization 0.92 \
  --enable-chunked-prefill --max-num-batched-tokens 8192 --max-num-seqs 80 --max-model-len 393216 \
  --tokenizer-mode deepseek_v4 \
  --compilation-config '{"cudagraph_mode":"FULL_AND_PIECEWISE","custom_ops":["all"]}' \
  --speculative-config '{"method":"dspark","num_speculative_tokens":7,"draft_sample_method":"probabilistic"}' \
  --kv-transfer-config '{"kv_connector":"SimpleCPUOffloadConnector","kv_role":"kv_both","kv_connector_extra_config":{"cpu_bytes_to_use_per_rank":236223201280,"lazy_offload":false}}' \
  --reasoning-parser deepseek_v4 --enable-auto-tool-choice --tool-call-parser deepseek_v4 \
  --max-cudagraph-capture-size 128

then load-tested with:

evalscope perf \
      --parallel 48 \
      --api openai \
      --min-tokens 1024 \
      --max-tokens 1024 \
      --prefix-length 0 \
      --min-prompt-length 32768 \
      --max-prompt-length 32768 \
      --number 100 \
      --tokenizer-path Qwen/Qwen2.5-VL-3B-Instruct \
      --name $result_name

[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. Reviewers with write access and configured trusted contributors can comment /ci run whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use /ci run or /ci retry. New commits do not start CI automatically.

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.

🚀

[zihanlin-ai]

SM90 validation from a different producer of the zero-query state: we semantically backported these guards onto stock v0.26.0 without #49302, so PIECEWISE replay padding still produces the #54179 phantom requests. On 4x H20 TP4, all previously crashing cases pass: 512-token, 2,400-token cold, 312K cold, and 312K prefix-hit resend.

Instrumentation: the planner skip fired 432 times; the backend guard fired 0 times; no qlen=0 chunk reached the kernel. Thus the planner skip is sufficient for this producer, while the backend check remains defense-in-depth.

Caveat: v0.26.0's layout required semantic adaptation, so this was not verbatim PR-head validation. We also did not run #51486's SimpleCPUOffloadConnector path; this validates the guards against PIECEWISE padding only.

[fireflyc]

Thanks @zihanlin-ai for this — it's genuinely useful. Your instrumentation actually gives the strongest evidence for the structure we chose: the planner skip firing 432× while the backend guard fired 0× confirms that avoiding the zero-token chunk at the source is the real fix, and the query_end <= query_start check in _forward_prefill is just defense-in-depth.

Thanks again — this moves the PR forward.