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: