[Perf] Keep non-speculative GDN decode on the standard path (#59735)
Signed-off-by: DimensionSTP <65501090+DimensionSTP@users.noreply.github.com> Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: Kimi Code <noreply@moonshot.ai>
DimensionSTP committed
Oct 2, 2026 at 22:44 UTC
e4294144a8ca4e57da712934dca33f89bef8fe78
2 files changed
+14
-9
tests/kernels/mamba/test_gdn_fused_mtp.py
+1
-1
@@ -194,7 +194,7 @@ def test_fused_forward_uses_packed_entrypoint() -> None:
194
ba = torch.randn(num_tokens, 2 * HV, dtype=torch.bfloat16, device=device)
195
layer = types.SimpleNamespace(
196
prefix=PREFIX,
197
- enable_fused_gdn_decode=True,
197
+ enable_fused_gdn_spec_decode=True,
198
norm=types.SimpleNamespace(
199
weight=torch.empty(V, dtype=torch.bfloat16, device=device)
200
),
vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py
+13
-8
@@ -520,17 +520,22 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
520
raise ValueError(
521
f"VLLM_GDN_DECODE_KERNEL=cuda is not supported: {reason}"
522
)
523
- logger.info_once(
524
- "Falling back to the Triton GDN decode path: %s", reason
525
- )
523
+ if self.speculative_config is not None:
524
+ logger.info_once(
525
+ "Falling back to the Triton GDN spec decode path: %s",
526
+ reason,
527
+ )
528
self.gdn_decode_kernel = "triton"
529
elif current_platform.is_cpu():
530
self.gdn_decode_kernel = "CPU"
531
elif current_platform.is_xpu():
532
self.gdn_decode_kernel = "XPU"
533
532
- self.enable_fused_gdn_decode = self.gdn_decode_kernel == "cuda"
533
- logger.info_once("GDN decode kernel: %s", self.gdn_decode_kernel)
534
+ self.enable_fused_gdn_spec_decode = (
535
+ self.gdn_decode_kernel == "cuda" and self.speculative_config is not None
536
+ )
537
+ if self.speculative_config is not None:
538
+ logger.info_once("GDN spec decode kernel: %s", self.gdn_decode_kernel)
539
540
compilation_config = get_current_vllm_config().compilation_config
541
if prefix in compilation_config.static_forward_context:
@@ -908,12 +913,12 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
913
mixed_qkvz, _ = self.in_proj_qkvz(hidden_states)
914
ba, _ = self.in_proj_ba(hidden_states)
915
911
- use_fused_gdn_decode = (
912
- self.enable_fused_gdn_decode
916
+ use_fused_gdn_spec_decode = (
917
+ self.enable_fused_gdn_spec_decode
918
and hidden_states.dtype == torch.bfloat16
919
and self.norm.weight.dtype in (torch.bfloat16, torch.float32)
920
)
916
- if use_fused_gdn_decode:
921
+ if use_fused_gdn_spec_decode:
922
core_attn_out = torch.zeros(
923
(num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim),
924
dtype=hidden_states.dtype,