[Bugfix] Fix FlashInfer GDN warmup ValueError on SM90 GPUs

[tdoublep]

Summary

• PR [Bugfix] Warm up Triton autotuner for GDN layers during V1 profiling #36599 added Triton autotuner warmup for GDN layers during V1 profiling, which also exercises the FlashInfer path on SM90 GPUs
• FlashInfer's chunk_gated_delta_rule returns a single tensor when output_final_state=False, but fi_chunk_gated_delta_rule always unpacked two values, causing a ValueError
• Fix: handle the return value based on output_final_state — unpack the tuple when True, use the single tensor when False

Error before fix

WARNING 03-12 10:34:06 [qwen3_next.py:724] GDN prefill kernel warmup (T=16) failed for layer model.layers.0.linear_attn. First inference may OOM due to autotuner.
WARNING 03-12 10:34:06 [qwen3_next.py:724] Traceback (most recent call last):
WARNING 03-12 10:34:06 [qwen3_next.py:724]   File "/workspace/vllm/vllm/model_executor/models/qwen3_next.py", line 712, in _warmup_prefill_kernels
WARNING 03-12 10:34:06 [qwen3_next.py:724]     self.chunk_gated_delta_rule(
WARNING 03-12 10:34:06 [qwen3_next.py:724]   File "/workspace/vllm/vllm/model_executor/models/qwen3_next.py", line 176, in forward_cuda
WARNING 03-12 10:34:06 [qwen3_next.py:724]     return fi_chunk_gated_delta_rule(
WARNING 03-12 10:34:06 [qwen3_next.py:724]   File "/workspace/vllm/vllm/model_executor/models/qwen3_next.py", line 138, in fi_chunk_gated_delta_rule
WARNING 03-12 10:34:06 [qwen3_next.py:724]     output, final_state = chunk_gated_delta_rule_fi(
WARNING 03-12 10:34:06 [qwen3_next.py:724] ValueError: too many values to unpack (expected 2)

This error repeats for T=16, T=32, and T=64 for each GDN layer.

Test plan

• Ran tests/v1/e2e/test_mamba_prefix_cache.py::test_mamba_prefix_cache on H100 (SM90) with all caches cleared — passes without the ValueError warning

🤖 Generated with Claude Code

[ZJY0516]

Sorry about this. I think we still need this because flashinfer kernel is a jit kernel

[tdoublep]

@ZJY0516 OK, then we need just need to fix the return types. Let me update the PR

[gemini-code-assist]

Code Review

This pull request correctly addresses a crash during the Gated Delta Net (GDN) layer warmup on SM90 GPUs. The fix involves skipping the warmup, which is intended for Triton autotuning, on SM90 architectures as they use the FlashInfer backend and do not require this step. The change is implemented by adding a conditional check for CUDA and SM90 device capability. While the fix is correct, I've added a comment regarding code duplication that could be addressed to improve long-term maintainability.

I am having trouble creating individual review comments. Click here to see my feedback.

vllm/model_executor/models/qwen3_next.py (682-683)

While this check correctly fixes the issue, it duplicates the logic from ChunkGatedDeltaRule.__init__ which is used to determine whether to use the FlashInfer backend. This could lead to future maintenance issues if the backend selection logic changes but this check is not updated in tandem.

To improve maintainability and avoid this duplication, consider centralizing the backend choice logic. For example, you could add a property to the ChunkGatedDeltaRule class to indicate which backend is in use:

# In ChunkGatedDeltaRule
@property
def uses_flashinfer(self) -> bool:
    return self._forward_method == self.forward_cuda

Then, you could use this property here to make the decision, ensuring the warmup logic always stays in sync with the actual backend being used:

# In _warmup_prefill_kernels
if self.chunk_gated_delta_rule.uses_flashinfer:
    return

This would make the code more robust to future changes.

[tdoublep]

@ZJY0516 Please take another look. We now fix the error more explicitly, and still allow the warmup phase to happen when using FI kernels.

[ZJY0516]

Thanks for fixing this

[mergify]

Hi @tdoublep, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[mergify]

Hi @tdoublep, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[ZJY0516]

Could we merge this asap?

[mergify]

Hi @tdoublep, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[xyang16]

Can this fix get merged?