From 8e2ac2e628b1e573a14898347ed910e0bb21da7e Mon Sep 17 00:00:00 2001 From: Makcum888e <79456407+Makcum888e@users.noreply.github.com> Date: Thu, 30 Oct 2025 06:18:36 +0300 Subject: [PATCH] [NPU] fix pp_size>1 (#12195) --- .../sglang/srt/distributed/parallel_state.py | 14 +++++++----- python/sglang/srt/mem_cache/memory_pool.py | 4 ++-- .../sglang/srt/model_executor/model_runner.py | 4 +++- python/sglang/srt/utils/common.py | 22 ++++++------------- 4 files changed, 21 insertions(+), 23 deletions(-) diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 28e47f069..8d1e903d3 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -68,7 +68,7 @@ REDUCE_OP_SUM = int(torch.distributed.ReduceOp.SUM) @dataclass class GraphCaptureContext: - stream: torch.cuda.Stream if not _is_npu else torch.npu.Stream + stream: torch.get_device_module().Stream @dataclass @@ -498,7 +498,7 @@ class GroupCoordinator: maybe_pynccl_context = nullcontext() else: maybe_pynccl_context = pynccl_comm.change_state( - enable=True, stream=torch.cuda.current_stream() + enable=True, stream=torch.get_device_module().current_stream() ) pymscclpp_comm = self.pymscclpp_comm @@ -555,7 +555,7 @@ class GroupCoordinator: and input_.symmetric_memory ): with self.pynccl_comm.change_state( - enable=True, stream=torch.cuda.current_stream() + enable=True, stream=torch.get_device_module().current_stream() ): self.pynccl_comm.all_reduce(input_) return input_ @@ -655,7 +655,9 @@ class GroupCoordinator: world_size = self.world_size pynccl_comm = self.pynccl_comm - with pynccl_comm.change_state(enable=True, stream=torch.cuda.current_stream()): + with pynccl_comm.change_state( + enable=True, stream=torch.get_device_module().current_stream() + ): assert ( pynccl_comm is not None and not pynccl_comm.disabled ), "pynccl is required for reduce_scatterv" @@ -779,7 +781,9 @@ class GroupCoordinator: world_size = self.world_size pynccl_comm = self.pynccl_comm - with pynccl_comm.change_state(enable=True, stream=torch.cuda.current_stream()): + with pynccl_comm.change_state( + enable=True, stream=torch.get_device_module().current_stream() + ): assert ( pynccl_comm is not None and not pynccl_comm.disabled ), "pynccl is required for all_gatherv" diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 46e2c7c57..5f08eb29c 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1137,10 +1137,10 @@ class AscendTokenToKVPool(MHATokenToKVPool): torch_npu._npu_reshape_and_cache( key=cache_k, value=cache_v, - key_cache=self.k_buffer[layer_id].view( + key_cache=self.k_buffer[layer_id - self.start_layer].view( -1, self.page_size, self.head_num, self.head_dim ), - value_cache=self.v_buffer[layer_id].view( + value_cache=self.v_buffer[layer_id - self.start_layer].view( -1, self.page_size, self.head_num, self.head_dim ), slot_indices=loc, diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 1ccea6cc3..7ba527157 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1659,9 +1659,11 @@ class ModelRunner: get_attention_tp_size() ), head_dim=self.model_config.head_dim, - layer_num=self.model_config.num_hidden_layers, + layer_num=self.num_effective_layers, device=self.device, enable_memory_saver=self.server_args.enable_memory_saver, + start_layer=self.start_layer, + end_layer=self.end_layer, ) elif self.use_mla_backend and is_nsa_model: self.token_to_kv_pool = NSATokenToKVPool( diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index b12474a2e..cdc0586e2 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -1239,42 +1239,34 @@ def point_to_point_pyobj( dst: int = 1, ): """Send data from src to dst in group using DeviceToDevice communication.""" - + device = torch.get_device_module().current_device() if rank == src: if len(data) == 0: - tensor_size = torch.tensor( - [0], dtype=torch.long, device=torch.cuda.current_device() - ) + tensor_size = torch.tensor([0], dtype=torch.long, device=device) dist.send(tensor_size, dst=dst, group=group) else: serialized_data = pickle.dumps(data) size = len(serialized_data) tensor_data = torch.ByteTensor( np.frombuffer(serialized_data, dtype=np.uint8) - ).cuda( - device=torch.cuda.current_device() + ).to( + device=device ) # Move to GPU - tensor_size = torch.tensor( - [size], dtype=torch.long, device=torch.cuda.current_device() - ) + tensor_size = torch.tensor([size], dtype=torch.long, device=device) dist.send(tensor_size, dst=dst, group=group) dist.send(tensor_data, dst=dst, group=group) return data elif rank == dst: - tensor_size = torch.tensor( - [0], dtype=torch.long, device=torch.cuda.current_device() - ) + tensor_size = torch.tensor([0], dtype=torch.long, device=device) dist.recv(tensor_size, src=src, group=group) size = tensor_size.item() if size == 0: return [] - tensor_data = torch.empty( - size, dtype=torch.uint8, device=torch.cuda.current_device() - ) + tensor_data = torch.empty(size, dtype=torch.uint8, device=device) dist.recv(tensor_data, src=src, group=group) serialized_data = bytes(