[Bugfix][Qwen3-Next] fixes the varlen issue in qwen3-next's MTP implementation.

[sighingnow]

This PR fixes the corner cases where guided decoding backend rollbacks draft tokens, causing unaligned verify batches.

Fixes #24730.
Fixes #24881.

[gemini-code-assist]

Code Review

This pull request aims to fix an issue with variable-length sequences in Qwen3-Next's multi-token prediction implementation, particularly for speculative decoding rollbacks. The changes span across the causal convolution Triton kernel, the Qwen3-Next model file, and the GatedDeltaNet attention backend. While the core logic for handling varlen inputs seems correct, I've identified a critical issue in the attention backend related to CUDA graph batch size calculation and a high-severity issue in the Triton kernel concerning the use of constexpr for runtime variables, which could lead to severe performance degradation or compilation errors.

[gemini-code-assist]

Changing self.decode_cudagraph_max_bs to be a token count (by multiplying with self.num_spec + 1) is incorrect, as this variable is used as a sequence count (batch size) for tensor allocations. For example, self.spec_state_indices_tensor is allocated with this as its first dimension (line 80), which is indexed by sequence, not by token. This change will lead to incorrect tensor allocations (either too large, wasting memory, or too small, causing out-of-bounds errors) and likely runtime failures.

To fix this correctly, decode_cudagraph_max_bs should remain a sequence count. A new variable should be introduced for the maximum token count if needed for the check at line 221.

[gemini-code-assist]

The kernel parameters state_len and seqlen are declared as tl.constexpr in the function signature (lines 635 and 634 respectively), but they are being reassigned here. constexpr values are meant to be compile-time constants and should not be modified at runtime. Passing runtime values to tl.constexpr parameters causes Triton to recompile the kernel for each unique value, which can lead to significant performance degradation and long compilation times. This reassignment is also confusing and can lead to unexpected behavior.

To fix this, you should change their type hints in the kernel signature to int. Additionally, for better code clarity and to avoid modifying input parameters, it's recommended to use new local variables for the updated values.

[chaunceyjiang]

100%|█████████████████████████████████████████████████████████████████████████████| 100/100 [01:11<00:00,  1.39it/s]
============ Serving Benchmark Result ============
Successful requests:                     100       
Maximum request concurrency:             10        
Benchmark duration (s):                  71.97     
Total input tokens:                      204357    
Total generated tokens:                  97082     
Request throughput (req/s):              1.39      
Output token throughput (tok/s):         1348.89   
Total Token throughput (tok/s):          4188.30   
---------------Time to First Token----------------
Mean TTFT (ms):                          540.71    
Median TTFT (ms):                        170.88    
P99 TTFT (ms):                           4819.26   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          6.88      
Median TPOT (ms):                        6.29      
P99 TPOT (ms):                           13.97     
---------------Inter-token Latency----------------
Mean ITL (ms):                           19.42     
Median ITL (ms):                         15.50     
P99 ITL (ms):                            127.42    
==================================================

Nice!!!! @sighingnow

[sighingnow]

Nice!!!! @sighingnow

@chaunceyjiang Thanks for the help to verify and the feedback!

[youkaichao]

thanks for the fix!

[david6666666]

thanks for the fix!