diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py index cdb1dadb7..e15d3d239 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py @@ -350,14 +350,7 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase): ssm_states.dtype, copy=False ) ssm_states[cache_indices] = last_recurrent_state - if not forward_batch.spec_algorithm.is_none(): - 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 - ) + last_recurrent_state = last_recurrent_state.to(ssm_states.dtype, copy=False) ssm_states[cache_indices] = last_recurrent_state if h is not None: self._track_mamba_state_extend( diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 9f06d99c1..b1353902b 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -289,6 +289,13 @@ class MambaPool: device=device, ) 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 # Shape: [num_layers, size + 1, speculative_num_draft_tokens, HV, K, V] intermediate_ssm_state_cache = torch.zeros(