PR #41142: [Bugfix][Qwen3-Next] Force a/b contiguity in fix_query_key_value_ordering

• Repository: vllm-project/vllm
• Author: Yeuvoir
• State: open | merged: False
• Link: https://github.com/vllm-project/vllm/pull/41142

Description (problem / solution / changelog)

Summary

Fixes #41112.

In GatedDeltaNetAttention.fix_query_key_value_ordering, the final reshape of b and a:

b = b.reshape(b.size(0), self.num_v_heads // self.tp_size)
a = a.reshape(a.size(0), self.num_v_heads // self.tp_size)

returns a non-contiguous view when num_v_heads == num_k_heads (i.e. np/ng == 1): the size-1 inner dim from the prior torch.split is squeezed without copying, leaving strides like (2K, 2) on the [B, num_v_heads] output.

This breaks fused_post_conv_prep (introduced in #37813), whose Triton kernel indexes a/b as

a_offsets = offs_t * stride_a_tok + i_hv
b_offsets = offs_t * stride_b_tok + i_hv

— it implicitly assumes the head dim has stride 1, so on a non-contiguous input it reads the wrong memory and the assertion downstream fires.

Fix

Call .contiguous() on b and a after the final reshape. This matches the non-interleaved (Qwen3.5) and LoRA paths in the same file which already do b = b.contiguous(); a = a.contiguous() after ba.chunk(...).

For typical Qwen3-Next configs (np/ng > 1, e.g. 32 v-heads / 16 k-heads) the reshape already produces a contiguous tensor, so .contiguous() is a true no-op — no performance impact on the default path.

Test plan

I believe this change is correct based on code review of:

• the non-interleaved sibling paths in the same file that already explicitly .contiguous() after the analogous split,
• the fused_post_conv_prep kernel's stride assumption.

I do not currently have a GPU available to run the upstream tests/kernels/test_fused_gdn_post_conv.py or an end-to-end Qwen3-Next generate, so I'd appreciate maintainers / anyone with access running:

• pytest tests/kernels/test_fused_gdn_post_conv.py -v
• An end-to-end Qwen3-Next prefill smoke test exercising gqa_interleaved_layout=True (e.g. via tests/v1/e2e/test_hybrid_chunked_prefill.py -k qwen3_next)
• Reproducer from #41112 (the case where a/b end up non-contiguous after fix_query_key_value_ordering)

Changed files

• vllm/model_executor/layers/mamba/gdn_linear_attn.py (modified, +5/-2)