[ZJY0516]

Purpose

Fix #59065

With DCP > 1 and speculative decoding, FlashInfer MLA decode crashes when the decode batch is ragged, i.e. requests carry different numbers of query tokens. That happens whenever a request has fewer drafts than the others, and also during warmup. For example, Kimi-K3 with -tp 2 -dcp 2 and DSpark fails at engine startup:

ValueError: causal_seqlens_kv_global must have shape (9,), got (2,)

Here a 1+7 request and a 1-token request make 9 decode tokens.

DCP cannot take the compact ragged path (cum_seq_lens_q), because it needs LSE. So ragged batches fall to the "uneven length" branch. That branch has two problems:

• Shape mismatch. It splits q into one row per token, but it keeps block_table, seq_lens and causal_seqlens_kv_global per request. FlashInfer then rejects the call.
• Silent wrong reshape. The branch is only taken when num_decode_tokens % num_decodes != 0. A ragged batch whose token count happens to divide evenly (e.g. 8 + 2 tokens) is instead viewed as a uniform 5 + 5 batch, and each request's rows attend to the wrong KV.

Fix

A batch is now treated as ragged when num_decode_tokens != num_decodes * max_query_len. Ragged causal batches are split into single-token rows:

• block_table and seq_lens are expanded per row with the existing _flattened_decode_metadata.
• Each row's DCP bound is T - (tokens of its request after it). FlashInfer defines causal_seqlens_kv_global as the global exclusive causal bound of each row's newest query token, and dcp_tot_seq_lens (T) already includes all of the request's query tokens.

The non-DCP causal path is unchanged; it already goes through cum_seq_lens_q.

Not a duplicate

No open PR touches this branch. #56784 edits the same file, but only for backend selection.

Test Plan

• Kernel check on GB200, run directly against trtllm_batch_decode_with_kv_cache_mla with the cute-dsl backend, cp_world=2, both ranks, and (T, q_len) in {(301, 8), (97, 8), (130, 4), (64, 1)}:
  • One request with q_len rows and bound T gives bit-identical output and LSE (max error 0.0) to q_len single-token rows bounded at T - (q_len - 1 - j).
  • Negative control: bounding every split row at T gives a max output error of 0.26.
• tests/v1/attention/test_flashinfer_mla_dcp.py: test_flashinfer_mla_forward_uses_native_dcp_api is extended with a ragged (3, 1) batch. The new ragged-causal case fails without this change and passes with it.

pytest tests/v1/attention/test_flashinfer_mla_dcp.py tests/kernels/attention/test_flashinfer_mla_decode.py

Test Result

• tests/v1/attention/test_flashinfer_mla_dcp.py + tests/kernels/attention/test_flashinfer_mla_decode.py: 23 passed, 1 skipped (GB200).
• Without the source change, only ragged-causal fails.
• An end-to-end run with Kimi-K3 was not done here. The dcp2-*-dspark cases in [CI][Kimi-K3] Test prefix cache reuse with KV offload, P/D and DCP #59229 are strict xfails for this bug and will flip once this lands.

No output changes for batches that worked before, so no model eval was run.

AI assistance (Claude Code) was used for this change.

🤖 Generated with Claude Code

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

[janbernloehr]

Thanks for the fix! We validated the flashinfer_mla.py change at 1a6def812d94 with a controlled A/B run and confirmed that it resolves the startup crash reported in #59065 for our configuration.

Both runs used the same image, model, server flags and environment variables; the only runtime change was this PR’s patch to flashinfer_mla.py.

• Hardware/configuration: 2 nodes × 4 NVIDIA GB300, Ray executor, Kimi-K3 with TP=8, DCP=8, FlashInfer MLA, FP8 KV cache, and DSpark with 7 speculative tokens using the V2 model runner.
• Base image: vllm/vllm-openai:nightly-ac68c3087215e0a4f3cdfa218508c6aada57235d
• Software: vLLM 0.30.1rc1.dev396+gac68c3087, FlashInfer 0.7.0.post1, PyTorch 2.13.0+cu130.

Results:

	Unpatched	With this patch
Startup	Fails during warmup_kernels with causal_seqlens_kv_global must have shape (9,), got (2,)	Completes successfully
Serving	Never reached	5/5 requests completed, 0 errors

The serving check used synthetic prompts with 1,000 input / 1,000 requested output tokens, concurrency 1 and ignore_eos=true.

This validates startup and basic serving for the reported configuration. There was one run per arm, and no accuracy evaluation was performed. Multi-request ragged batches were exercised during warmup; we did not separately target the evenly divisible ragged case described in this PR.

Drafted with assistance from Codex (GPT-6).

[khluu]

/ci run

[github-actions]

✅ Triggered Buildkite CI #94737 for commit 8412d35a8081.

ℹ️ This PR is 815 commits behind upstream main. Merged upstream main into this branch as 8412d35a8081 before starting CI. Pull the branch before pushing more changes.

[pavanimajety]

PR looks good to me, thanks for the fix.

[janbernloehr]

@ZJY0516 can you update with main again?

[ZJY0516]

/ci run

[github-actions]

✅ Triggered Buildkite CI #97549 for commit 0e68d8e06b9a.