iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Perf] Keep non-speculative GDN decode on the standard path (#59735)

Browse Source 
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>
This commit is contained in:
DimensionSTP
2026-10-02 13:44:53 +00:00
committed by GitHub
co-authored by Thien Tran Kimi Code
parent 24df2d3fc7
commit e4294144a8
2 changed files with 14 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/mamba/test_gdn_fused_mtp.py
+1 -1
View File
@@ -194,7 +194,7 @@ def test_fused_forward_uses_packed_entrypoint() -> None:
ba = torch.randn(num_tokens, 2 * HV, dtype=torch.bfloat16, device=device)
layer = types.SimpleNamespace(
prefix=PREFIX,
enable_fused_gdn_decode=True,
enable_fused_gdn_spec_decode=True,
norm=types.SimpleNamespace(
weight=torch.empty(V, dtype=torch.bfloat16, device=device)
),
vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py
+13 -8
View File
@@ -520,17 +520,22 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
raise ValueError(
f"VLLM_GDN_DECODE_KERNEL=cuda is not supported: {reason}"
)
logger.info_once(
"Falling back to the Triton GDN decode path: %s", reason
)
if self.speculative_config is not None:
logger.info_once(
"Falling back to the Triton GDN spec decode path: %s",
reason,
)
self.gdn_decode_kernel = "triton"
elif current_platform.is_cpu():
self.gdn_decode_kernel = "CPU"
elif current_platform.is_xpu():
self.gdn_decode_kernel = "XPU"
self.enable_fused_gdn_decode = self.gdn_decode_kernel == "cuda"
logger.info_once("GDN decode kernel: %s", self.gdn_decode_kernel)
self.enable_fused_gdn_spec_decode = (
self.gdn_decode_kernel == "cuda" and self.speculative_config is not None
)
if self.speculative_config is not None:
logger.info_once("GDN spec decode kernel: %s", self.gdn_decode_kernel)
compilation_config = get_current_vllm_config().compilation_config
if prefix in compilation_config.static_forward_context:
@@ -908,12 +913,12 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
mixed_qkvz, _ = self.in_proj_qkvz(hidden_states)
ba, _ = self.in_proj_ba(hidden_states)
use_fused_gdn_decode = (
self.enable_fused_gdn_decode
use_fused_gdn_spec_decode = (
self.enable_fused_gdn_spec_decode
and hidden_states.dtype == torch.bfloat16
and self.norm.weight.dtype in (torch.bfloat16, torch.float32)
)
if use_fused_gdn_decode:
if use_fused_gdn_spec_decode:
core_attn_out = torch.zeros(
(num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim),
dtype=hidden_states.dtype,