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

[NPU] Support chunk prefill for Qwen3.5/Qwen3.6 models (#25839)

Browse Source
This commit is contained in:
Liwansi
2026-05-21 14:44:25 +08:00
committed by GitHub
parent 4ea8282cb7
commit 190488e9a8
2 changed files with 8 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py
+1 -8
View File
@@ -350,14 +350,7 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase):
ssm_states.dtype, copy=False ssm_states.dtype, copy=False
) )
ssm_states[cache_indices] = last_recurrent_state ssm_states[cache_indices] = last_recurrent_state
if not forward_batch.spec_algorithm.is_none(): last_recurrent_state = last_recurrent_state.to(ssm_states.dtype, copy=False)
last_recurrent_state = last_recurrent_state.transpose(-1, -2).to(
ssm_states.dtype, copy=False
)
else:
last_recurrent_state = last_recurrent_state.to(
ssm_states.dtype, copy=False
)
ssm_states[cache_indices] = last_recurrent_state ssm_states[cache_indices] = last_recurrent_state
if h is not None: if h is not None:
self._track_mamba_state_extend( self._track_mamba_state_extend(
python/sglang/srt/mem_cache/memory_pool.py
+7
View File
@@ -289,6 +289,13 @@ class MambaPool:
device=device, device=device,
) )
if speculative_num_draft_tokens is not None: if speculative_num_draft_tokens is not None:
if _is_npu:
temporal_state = temporal_state.transpose(-1, -2)
temporal_state_shape = (
*temporal_state_shape[:-2],
temporal_state_shape[-1],
temporal_state_shape[-2],
)
# Cache intermediate SSM states per draft token during target verify # Cache intermediate SSM states per draft token during target verify
# Shape: [num_layers, size + 1, speculative_num_draft_tokens, HV, K, V] # Shape: [num_layers, size + 1, speculative_num_draft_tokens, HV, K, V]
intermediate_ssm_state_cache = torch.zeros( intermediate_ssm_state_cache = torch.zeros(