[BugFix] Fix FI accuracy issue when used for MLA prefill

[LucasWilkinson]

PR

VLLM_LOGGING_LEVEL=DEBUG vllm serve deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct
...
(EngineCore_DP0 pid=3662669) DEBUG 10-01 17:41:11 [v1/attention/.../mla/common.py:980] Using FlashInfer prefill for MLA
...
========================================================================
(vllm) lwilkinson@dgxB200-09:~/code/vllm$ python tests/evals/gsm8k/gsm8k_eval.py
Running GSM8K evaluation: 1319 questions, 5-shot
Evaluating: 100%|█████████████████████████████████████████████████████████████████████| 1319/1319 [00:28<00:00, 46.50it/s]

Results:
Accuracy: 0.782
Invalid responses: 0.000
Total latency: 28.381 s
Questions per second: 46.475

Main

(vllm) lwilkinson@dgxB200-09:~/code/vllm$ python tests/evals/gsm8k/gsm8k_eval.py
Running GSM8K evaluation: 1319 questions, 5-shot
Evaluating: 100%|█████████████████████████████████████████████████████████████████████| 1319/1319 [00:26<00:00, 49.12it/s]

Results:
Accuracy: 0.208
Invalid responses: 0.006
Total latency: 26.864 s
Questions per second: 49.100

[gemini-code-assist]

Code Review

This pull request addresses an accuracy issue with the FlashInfer MLA prefill implementation. The root cause appears to be an inconsistent shape of the log-sum-exp (LSE) tensor returned by FlashInfer, which is (q_len, num_heads) instead of the expected (num_heads, q_len). The changes correctly transpose the LSE tensor in both _run_prefill_new_tokens_fi and _run_prefill_context_chunk_fi to align with other backends, which should resolve the accuracy problem. The fix is logical and well-targeted. I have one suggestion to improve the code's robustness by using isinstance() for type checking, in line with Python best practices.

[benchislett]

I ran this to validate on 8xB200:

VLLM_ATTENTION_BACKEND=CUTLASS_MLA vllm serve deepseek-ai/DeepSeek-R1-0528 -tp 8 --max-model-len 8192 --no-enable-prefix-caching --port 8049 --max-num-batched-tokens 256

where it completely failed before, it passes GSM8k now.

I also ran it with MTP on my development branch (#25984), which relies more heavily on this prefill functionality, and it passed as well:

VLLM_ATTENTION_BACKEND=CUTLASS_MLA vllm serve deepseek-ai/DeepSeek-R1-0528 -tp 8 --max-model-len 8192 --no-enable-prefix-caching --port 8049 --max-num-batched-tokens 256 --speculative-config '{"method": "mtp", "num_speculative_tokens": 3}'

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9538|±  |0.0058|
|     |       |strict-match    |     5|exact_match|↑  |0.9500|±  |0.0060|

[benchislett]

I too, am curious why this didn't / doesn't cause a shape mismatch. However, it clearly works well to solve the problem.

I would approve if I had any understanding as to why this solves the issue.

[benchislett]

FIX #26042