minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[GDN] Honor configured linear-attn verify backend in the kernel dispatcher (#34592)

Browse Source 
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Xiaoyu Zhang
2026-08-14 08:59:52 +08:00
committed by GitHub
co-authored by Claude Fable 5
parent 34219ed9a7
commit 704e512836
1 changed files with 15 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/linear/gdn_backend.py
+15 -5
View File
@@ -15,6 +15,7 @@ from sglang.srt.layers.attention.linear.utils import (
build_verify_intermediate_state_indices, build_verify_intermediate_state_indices,
get_linear_attn_decode_backend, get_linear_attn_decode_backend,
get_linear_attn_prefill_backend, get_linear_attn_prefill_backend,
get_linear_attn_verify_backend,
) )
from sglang.srt.layers.radix_linear_attention import RadixLinearAttention from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
from sglang.srt.mem_cache.memory_pool import MambaPool from sglang.srt.mem_cache.memory_pool import MambaPool
@@ -110,6 +111,7 @@ class GDNKernelDispatcher:
self, self,
decode_backend: LinearAttnKernelBackend, decode_backend: LinearAttnKernelBackend,
prefill_backend: LinearAttnKernelBackend, prefill_backend: LinearAttnKernelBackend,
verify_backend: Optional[LinearAttnKernelBackend] = None,
): ):
triton_kernel = TritonGDNKernel() triton_kernel = TritonGDNKernel()
self.tree_verify_kernel = triton_kernel self.tree_verify_kernel = triton_kernel
@@ -177,10 +179,15 @@ class GDNKernelDispatcher:
else: else:
raise ValueError(f"Unsupported GDN prefill backend: {prefill_backend}") raise ValueError(f"Unsupported GDN prefill backend: {prefill_backend}")
# Verify kernel: use FlashInfer when the selected FlashInfer kernel # Verify kernel. An explicitly configured verify backend wins; the
# supports MTP verify. SM90 uses the fp32-state path; SM100 uses the # historical auto rule (FlashInfer when the selected FlashInfer kernel
# bf16-state adapter in FlashInferGDNKernel. # supports MTP verify) only applies when no explicit choice was made.
if ( # SM90 FlashInfer verify requires a fp32 SSM state, so e.g.
# --mamba-ssm-dtype bfloat16 setups must be able to force Triton here.
if verify_backend is not None and verify_backend.is_triton():
self.verify_kernel = triton_kernel
self.verify_kernel_is_flashinfer = False
elif (
decode_backend.is_flashinfer() or prefill_backend.is_flashinfer() decode_backend.is_flashinfer() or prefill_backend.is_flashinfer()
) and flashinfer_kernel.supports_target_verify: ) and flashinfer_kernel.supports_target_verify:
self.verify_kernel = flashinfer_kernel self.verify_kernel = flashinfer_kernel
@@ -346,7 +353,10 @@ class GDNAttnBackend(MambaAttnBackendBase):
decode_backend = get_linear_attn_decode_backend() decode_backend = get_linear_attn_decode_backend()
prefill_backend = get_linear_attn_prefill_backend() prefill_backend = get_linear_attn_prefill_backend()
self.kernel_dispatcher = GDNKernelDispatcher(decode_backend, prefill_backend) verify_backend = get_linear_attn_verify_backend()
self.kernel_dispatcher = GDNKernelDispatcher(
decode_backend, prefill_backend, verify_backend
)
# Sized past the pool for attn_tp-padded warmup/MLP-sync batches (see helper). # Sized past the pool for attn_tp-padded warmup/MLP-sync batches (see helper).
self.verify_intermediate_state_indices = ( self.verify_intermediate_state_indices = (
build_verify_intermediate_state_indices( build_verify_intermediate_state_indices(