[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 <hinata@idihore.id>
Signed-off-by: MOVZX <movzx@yahoo.com>
This commit is contained in:
2026-10-09 20:16:36 +07:00
co-authored by Hinata
parent 69aed12cf8
commit db2851e54a
@@ -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: