[Bugfix][MTP][Sparse MLA] Allow sparse MLA with MTP to run with FULL cudagraphs

[MatthewBonanni]

Purpose

DeepseekV32IndexerMetadataBuilder currently only reports support for UNIFORM_SINGLE_TOKEN_DECODE. Therefore, when running a sparse MLA model with MTP, FULL cudagraphs are never captured.

In reality, the deepGEMM kernel fp8_paged_mqa_logits supports MTP with num_speculative_tokens=1 (i.e. next_n = 2).

This PR changes the reported support to UNIFORM_BATCH and adds an explicit error for num_speculative_tokens > 1, rather than letting the kernel itself crash.

Test Plan

vllm serve deepseek-ai/DeepSeek-V3.2 \
    -tp 8 -ep \
    --speculative-config '{"method": "mtp", "num_speculative_tokens": 1}' \
    --no-enable-prefix-caching

with

vllm bench serve \
    --dataset-name random \
    --input-len 128 \
    --output-len 2048 \
    --seed 42 \
    --ignore-eos \
    --temperature 0 \
    --num-prompts 200 \
    --request-rate inf \
    --max-concurrency 64

Test Result

Main:

============ Serving Benchmark Result ============
Successful requests:                     200       
Failed requests:                         0         
Maximum request concurrency:             64        
Benchmark duration (s):                  434.34    
Total input tokens:                      25600     
Total generated tokens:                  409600    
Request throughput (req/s):              0.46      
Output token throughput (tok/s):         943.04    
Peak output token throughput (tok/s):    640.00    
Peak concurrent requests:                109.00    
Total token throughput (tok/s):          1001.98   
---------------Time to First Token----------------
Mean TTFT (ms):                          1372.90   
Median TTFT (ms):                        1520.53   
P99 TTFT (ms):                           2167.85   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          54.12     
Median TPOT (ms):                        53.95     
P99 TPOT (ms):                           57.18     
---------------Inter-token Latency----------------
Mean ITL (ms):                           107.91    
Median ITL (ms):                         102.92    
P99 ITL (ms):                            165.42    
---------------Speculative Decoding---------------
Acceptance rate (%):                     99.43     
Acceptance length:                       1.99      
Drafts:                                  205340    
Draft tokens:                            205340    
Accepted tokens:                         204174    
Per-position acceptance (%):
  Position 0:                            99.43     
==================================================

PR:

============ Serving Benchmark Result ============
Successful requests:                     200       
Failed requests:                         0         
Maximum request concurrency:             64        
Benchmark duration (s):                  193.55    
Total input tokens:                      25600     
Total generated tokens:                  409600    
Request throughput (req/s):              1.03      
Output token throughput (tok/s):         2116.20   
Peak output token throughput (tok/s):    1472.00   
Peak concurrent requests:                120.00    
Total token throughput (tok/s):          2248.46   
---------------Time to First Token----------------
Mean TTFT (ms):                          1238.56   
Median TTFT (ms):                        1223.26   
P99 TTFT (ms):                           2081.36   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          24.73     
Median TPOT (ms):                        24.87     
P99 TPOT (ms):                           26.77     
---------------Inter-token Latency----------------
Mean ITL (ms):                           49.27     
Median ITL (ms):                         45.99     
P99 ITL (ms):                            59.30     
---------------Speculative Decoding---------------
Acceptance rate (%):                     99.26     
Acceptance length:                       1.99      
Drafts:                                  205513    
Draft tokens:                            205513    
Accepted tokens:                         203997    
Per-position acceptance (%):
  Position 0:                            99.26     
==================================================

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request correctly enables FULL cudagraph support for sparse MLA models with MTP by changing _cudagraph_support to UNIFORM_BATCH. It also adds a necessary safeguard to prevent crashes by raising a ValueError for unsupported num_speculative_tokens > 1, which is a limitation of the fp8_paged_mqa_logits kernel. The changes are well-implemented, improving both functionality and robustness. The code is clean and the logic is sound.

[jeejeelee]

cc @WoosukKwon @zyongye

[ElizaWszola]

Two questions:

1. What was the previous scenario for running with num_speculative_tokens > 1? Was it a kernel crash or did the limited support prevent it?
2. Does anything need to be updated in cuda_graphs.md for this PR?

[mergify]

Documentation preview: https://vllm--34457.org.readthedocs.build/en/34457/

[MatthewBonanni]

@ElizaWszola

1. Previously there was an assertion within DeepGEMM as in [Bug]: [H200] DeepSeek V3.2 MTP > 1 run into error (FLASHMLA_SPARSE backend) #31845. I was under the impression that this was due to a kernel limitation because comments throughout vLLM seem to indicate that the indexer kernel doesn't support num_speculative_tokens > 1. This may not be true, though, and it may have actually just been an issue with metadata building. @LucasWilkinson has a fix here: [BugFix] Add support for MTP num_speculative_tokens > 1 with sparse MLA #34552
2. Sparse MLA backends weren't previously discussed in cuda_graphs.md but I've added FlashInferMLASparse to it now - good catch