[Feature]: Full cudagraph support for MLA attention backend with DeepSeek MTP(Speculative decode)

[qiaoning]

🚀 The feature, motivation and pitch

On cudagraph capture stage, the MLACommonMetadataBuilder build metadata with max_query_len =1, with this PR

    def build_for_cudagraph_capture(
            self, common_attn_metadata: CommonAttentionMetadata) -> M:
        """
        This method builds the metadata for full cudagraph capture.
        Currently, only decode is supported for full cudagraphs with MLA.
        """
        m = common_attn_metadata
        assert m.num_reqs == m.num_actual_tokens, \
            "MLA only supports decode-only full CUDAGraph capture. " \
            "Make sure all cudagraph capture sizes <= max_num_seq."

        m.max_query_len = 1  # decode-only

And when run DeepSeek-R1 with DeepSeek MTP with simple-cuda-graph and run eagle with full cudagraph support patches, the max_query_len may be 2 (one actual decode token and one Speculative token), which will give an error output like ok, ok, ok, ok....

And we find that, when run DeepSeek with MTP, the MLACommonImpl.forward will always call self._forward_prefill, and self._forward_decode is NOT called anymore, and the cudagraph capture is decode-only, requir max_query_len=1, which may cause conflict? I'm NOT sure.

class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
    ...

    def forward(
        ...
    ) -> torch.Tensor:

        ...

        if has_prefill:
            output[num_decode_tokens:] = self._forward_prefill(
                prefill_q, prefill_k_c_normed, prefill_k_pe, kv_cache,
                attn_metadata)

        if has_decode: # Run DeepSeek R1 with MTP, has_decode is always False
            ...
            output[:num_decode_tokens] = self._forward_decode(
                decode_ql_nope, decode_q_pe, kv_cache, attn_metadata)

So, Is there any way to support DeepSeek MTP with full cudagraph？

@ProExpertProg @LucasWilkinson @zixi-qi @YaoJiayi Looking forward to your reply.

Other Refs:
#17211
#18435

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[ProExpertProg]

Yes, we're planning to overhaul cudagraph capture, dispatching, and replay in #20059. See my latest comment for design and Lucas's comment for spec-decoding support. But yeah the infra in that PR will make it possible to capture the cudagraphs for 2 tokens per request.

[qiaoning]

Yes, we're planning to overhaul cudagraph capture, dispatching, and replay in #20059. See my latest comment for design and Lucas's comment for spec-decoding support. But yeah the infra in that PR will make it possible to capture the cudagraphs for 2 tokens per request.

Thanks for your reply @ProExpertProg ! I have followed this PR before, it's really a great job. I read it carefully again, especially your recent proposal, and @fhl2000 and @LucasWilkinson 's discuss about spec-decode.

I think the refact infra now has the ability to capture/replay graph according to various conditions, including spec-decode.

And I see some codes about uniform_decode in class GPUModelRunner:

class GPUModelRunner(LoRAModelRunnerMixin):
    def __init__(...)
    ... 
        self.uniform_decode_query_len = 1 if not self.speculative_config else \
            1 + self.speculative_config.num_speculative_tokens

    def execute_model(...)
    ...
        is_uniform = max_query_len == self.uniform_decode_query_len and \
            num_scheduled_tokens == self.input_batch.num_reqs*max_query_len
        cudagraph_runtime_mode, batch_descriptor = \
            self.cudagraph_dispatcher.dispatch(
                BatchDescriptor(num_tokens=num_input_tokens,
                                is_uniform=is_uniform))
        # sanity check, in case speculative decoding on uniform batch
        # is not supported with full cudagraph
        if cudagraph_runtime_mode == CUDAGraphMode.FULL:
            assert attention_cuda_graphs, \
                f"Current batch {batch_descriptor} is not compatible with " \
                f"full cudagraph for attention backend " \
                f"{self.attn_backends[0].__name__}"

    def _dummy_run(...)
    ...
        # If cudagraph_separate_routine is enabled when use full cudagraph,
        # we need to manually activate the correct routine of attention backend
        # for mixed prefill-decode batches and uniform decode batches
        # separately during capturing. Uniform batch means that all
        # requests have identical query length, except a potential virtual
        # request (shorter) in the batch account for padding.
        # Uniform batch could either be common pure decode, where
        # max_query_len == 1, or speculative decode, where
        # max_query_len == 1 + num_spec_decode_tokens.

        # When setting max_query_len = 1, we switch to and capture the optimized
        # routine of FA2 for pure decode, i.e., Flashdecode + an optimization
        # for GQA/MQA.
        max_query_len = self.uniform_decode_query_len if is_uniform else \
                                                                num_tokens

But, I noticed that, vllm/v1/attention/backends/mla/common.py still has code like,

def build_for_cudagraph_capture(
            self, common_attn_metadata: CommonAttentionMetadata) -> M:
        ...
        assert m.max_query_len == 1  # decode-only

...

def can_run_in_cudagraph(
            self, common_attn_metadata: CommonAttentionMetadata) -> bool:
        return common_attn_metadata.max_query_len == 1

which means when use mla attention, the cudagraph capture is PURE_DECODE_ONLY, cannot run with uniform_decode with full cudagraph, I'm right?

So my question is:

1. Will max_query_len > 1 full cudagraph capture support on another PR?
2. Or mla attention(TritonMLA or FlashMLA) will use piecewise cudagraph or no cudagraph when MTP enabled, because full cudagraph capture is support only for PURE_DECODE_ONLY situation?