[Qwen3Next] Fixes the cuda graph capture conditions under large batch sizes (#24660)

[sighingnow]

No description provided.

[gemini-code-assist]

Code Review

This pull request attempts to fix a CUDA graph capture condition for large batch sizes in the GDNAttentionBackend. However, the added condition m.num_actual_tokens <= self.decode_cudagraph_max_bs introduces a unit mismatch by comparing a token count with a sequence limit. This makes the check overly restrictive and prevents CUDA graph usage in many valid scenarios. My review provides a detailed explanation of the issue and suggests a more accurate approach to fix the underlying problem.

[gemini-code-assist]

This condition m.num_actual_tokens <= self.decode_cudagraph_max_bs appears to have a unit mismatch. m.num_actual_tokens is the number of tokens, while self.decode_cudagraph_max_bs is used as a limit on the number of sequences for sizing tensors like spec_state_indices_tensor and spec_sequence_masks.

Comparing tokens to sequences is likely incorrect and makes this check overly restrictive. For instance, with num_spec=7 and decode_cudagraph_max_bs=32, this change limits num_spec_decodes to 4 (since 4 * 8 <= 32), whereas the original code allowed up to 32 sequences.

The underlying issue is that batch_size can exceed self.decode_cudagraph_max_bs due to token padding. The batch_size is calculated as self.vllm_config.pad_for_cudagraph(m.num_actual_tokens) // (self.num_spec + 1).

A more accurate check would be to compute this batch_size and compare it against self.decode_cudagraph_max_bs, while also ensuring m.num_actual_tokens does not exceed self.compilation_config.max_capture_size to prevent errors from pad_for_cudagraph.

[fhl2000]

I think keep m.num_actual_tokens <= self.decode_cudagraph_max_bs is good. Then and num_spec_decodes <= self.decode_cudagraph_max_bs seems unnecessary.

[pUnK2008]

around 20 lines below there is one more similar place

[LucasWilkinson]

This makes sense to me assuming @vadiklyutiy confirms the fix

longer term hopefully: #23789 / #24002 will resolve this more broadly but this makes sense as a temporary fix 👍

[LucasWilkinson]

around 20 lines below there is one more similar place

Are you referring to:

        if (self.use_full_cuda_graph and num_prefills == 0
                and num_spec_decodes == 0
                and num_decodes <= self.decode_cudagraph_max_bs):

I think thats safe due to num_spec_decodes == 0 so in that case num_decodes == num_tokens

[vadiklyutiy]

around 20 lines below there is one more similar place

Are you referring to:

        if (self.use_full_cuda_graph and num_prefills == 0
                and num_spec_decodes == 0
                and num_decodes <= self.decode_cudagraph_max_bs):

I think thats safe due to num_spec_decodes == 0 so in that case num_decodes == num_tokens

agree