diff --git a/python/sglang/srt/disaggregation/ascend/conn.py b/python/sglang/srt/disaggregation/ascend/conn.py index d323d64d0..2ad9da407 100644 --- a/python/sglang/srt/disaggregation/ascend/conn.py +++ b/python/sglang/srt/disaggregation/ascend/conn.py @@ -109,13 +109,17 @@ class AscendKVManager(MooncakeKVManager): executor: concurrent.futures.ThreadPoolExecutor, dst_layer_ids: Optional[List[int]] = None, dst_device_kv_indices: Optional[npt.NDArray[np.int32]] = None, + dst_kv_item_len: Optional[int] = None, + dst_attn_tp_size: Optional[int] = None, ): if dst_device_kv_indices is not None: raise NotImplementedError( "Ascend PD transfer does not support HiSparse " "destination device KV indices" ) - + self._validate_envelope_kv_layout( + dst_kv_ptrs, dst_kv_item_len, dst_attn_tp_size + ) # Group by indices prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous( prefill_kv_indices, dst_kv_indices diff --git a/python/sglang/srt/mem_cache/pool_host/mha.py b/python/sglang/srt/mem_cache/pool_host/mha.py index 2d71a8a08..4646fb255 100644 --- a/python/sglang/srt/mem_cache/pool_host/mha.py +++ b/python/sglang/srt/mem_cache/pool_host/mha.py @@ -384,13 +384,18 @@ class MHATokenToKVPoolHost(HostKVCache): def backup_from_device_all_layer( self, device_pool, host_indices, device_indices, io_backend ): - ( - device_k_data_ptrs, - device_v_data_ptrs, - device_k_buffers, - device_v_buffers, - ) = self._resolve_device_transfer_buffers(device_pool) - device_kv_buffers = device_k_buffers + device_v_buffers + if io_backend == "kernel_ascend": + # NPU pools use contiguous multi-layer tensors and intentionally do + # not build the CUDA-style k_data_ptrs/v_data_ptrs arrays. + device_kv_buffers = None + else: + ( + device_k_data_ptrs, + device_v_data_ptrs, + device_k_buffers, + device_v_buffers, + ) = self._resolve_device_transfer_buffers(device_pool) + device_kv_buffers = device_k_buffers + device_v_buffers if io_backend == "kernel": if self.layout == "layer_first": if self.can_use_jit: