From db2851e54a2d4ae1b8cc5e2132297320abd58afb Mon Sep 17 00:00:00 2001 From: MOVZX Date: Mon, 5 Oct 2026 06:36:02 -0400 Subject: [PATCH] [GDN] Keep the fused packed forward under RecoverSSM RecoverSSM's verify needs the Triton decode kernel for its speculative rows, and the previous code met that by turning off the fused CUDA decode for the whole layer. That also turned off the fused packed forward, which carries every prefill chunk, so prefill ran the unfused path. The decode kernel stays Triton and the packed recurrent decode stays off; the fused packed forward stays on. A small helper makes the choice, with CPU tests for each combination. Co-Authored-By: Hinata Signed-off-by: MOVZX --- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 41 +++++++++++++++---- 1 file changed, 33 insertions(+), 8 deletions(-) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 1cebc9bae9..86d6f0120e 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -562,16 +562,18 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): elif current_platform.is_xpu(): self.gdn_decode_kernel = "XPU" - self.enable_fused_gdn_spec_decode = ( - self.gdn_decode_kernel == "cuda" and self.speculative_config is not None - ) - # RecoverSSM verify replaces the per-draft snapshot path (Triton decode - # only) self.use_gdn_recoverssm = self._uses_gdn_recoverssm() + ( + self.gdn_decode_kernel, + self.enable_fused_gdn_spec_decode, + self.enable_packed_recurrent_decode, + ) = self._decode_paths( + self.gdn_decode_kernel, + self.use_gdn_recoverssm, + self.enable_packed_recurrent_decode, + self.speculative_config is not None, + ) if self.use_gdn_recoverssm: - self.gdn_decode_kernel = "triton" - self.enable_fused_gdn_spec_decode = False - self.enable_packed_recurrent_decode = False logger.info_once( "GDN RecoverSSM speculative verify active (spec_query_len %d)", self.num_spec + 1, @@ -584,6 +586,29 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): raise ValueError(f"Duplicate layer name: {prefix}") compilation_config.static_forward_context[prefix] = self + @staticmethod + def _decode_paths( + gdn_decode_kernel: str, + use_recoverssm: bool, + packed_recurrent_decode: bool, + has_speculative_config: bool, + ) -> tuple[str, bool, bool]: + """(decode kernel, fused spec-decode forward, packed recurrent decode). + + The fused forward is on whenever the CUDA decode kernel is eligible and + the engine runs speculation (#59735 keeps non-speculative decode on the + standard path). RecoverSSM's verify replaces the per-draft snapshot path + and needs the Triton decode kernel for its spec rows, so with RecoverSSM + the decode kernel is Triton and the packed recurrent decode is off. The + fused forward stays on: it calls _forward_core for prefill and verify + rows alike and only fuses the input split and the gated norm. Turning it + off sends every prefill chunk down the unfused path. + """ + enable_fused_forward = gdn_decode_kernel == "cuda" and has_speculative_config + if use_recoverssm: + return "triton", enable_fused_forward, False + return gdn_decode_kernel, enable_fused_forward, packed_recurrent_decode + def _fused_gdn_decode_unsupported_reason( self, vllm_config: VllmConfig ) -> str | None: