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

[Perf] Remove two operations in gdn_backend extend verify path (#22444)

Browse Source
This commit is contained in:
Jincong Chen
2026-04-10 17:53:57 +08:00
committed by GitHub
parent 1c76f322df
commit 0668a7f51a
1 changed files with 4 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/linear/gdn_backend.py
+4 -8
View File
@@ -255,6 +255,9 @@ 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) self.kernel_dispatcher = GDNKernelDispatcher(decode_backend, prefill_backend)
self.verify_intermediate_state_indices = torch.arange(
self.req_to_token_pool.size, dtype=torch.int32, device=model_runner.device
)
def init_forward_metadata(self, forward_batch: ForwardBatch): def init_forward_metadata(self, forward_batch: ForwardBatch):
super().init_forward_metadata(forward_batch) super().init_forward_metadata(forward_batch)
@@ -373,14 +376,7 @@ class GDNAttnBackend(MambaAttnBackendBase):
intermediate_conv_window_cache = ( intermediate_conv_window_cache = (
mamba_cache_params.intermediate_conv_window[0] mamba_cache_params.intermediate_conv_window[0]
) )
has_initial_states = torch.ones( intermediate_state_indices = self.verify_intermediate_state_indices
seq_len // forward_batch.spec_info.draft_token_num,
dtype=torch.bool,
device=forward_batch.input_ids.device,
)
intermediate_state_indices = torch.arange(
cache_indices.shape[0], dtype=torch.int32, device=cache_indices.device
)
else: else:
has_initial_states = forward_batch.extend_prefix_lens > 0 has_initial_states = forward_batch.extend_prefix_lens > 0