chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[NPU] Fix NPU CI (#13834)

Browse Source 
Co-authored-by: c30031083 <chenxu140@huawei.com>
This commit is contained in:
Even Zhou
2025-11-25 10:09:36 +08:00
committed by GitHub
parent eb1d885400
commit db0ffc09ef
8 changed files with 11 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/sglang/srt/mem_cache/allocator_ascend.py
View File

@@ -112,9 +112,9 @@ class AscendPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
device=self.device,
)
torch.ops.npu.alloc_extend(
prefix_lens,
seq_lens,
last_loc,
prefix_lens.to(torch.int64),
seq_lens.to(torch.int64),
last_loc.to(torch.int64),
self.free_pages,
self.page_size,
out_indices,

1
python/sglang/srt/model_executor/cuda_graph_runner.py
View File

@@ -329,6 +329,7 @@ class CudaGraphRunner:
seq_len_fill_value=self.seq_len_fill_value,
encoder_len_fill_value=self.encoder_len_fill_value,
num_tokens_per_bs=self.num_tokens_per_bs,
cache_loc_dtype=self._cache_loc_dtype(),
)
self.tbo_plugin = TboCudaGraphRunnerPlugin()

3
python/sglang/srt/model_executor/input_buffers.py
View File

@@ -43,13 +43,14 @@ class GraphInputBuffers:
seq_len_fill_value: int,
encoder_len_fill_value: int,
num_tokens_per_bs: int,
cache_loc_dtype: torch.dtype,
) -> "GraphInputBuffers":
with torch.device(device):
input_ids = torch.zeros((max_num_token,), dtype=torch.int64)
input_embeds = torch.zeros((max_num_token, hidden_size), dtype=dtype)
req_pool_indices = torch.zeros((max_bs,), dtype=torch.int32)
seq_lens = torch.full((max_bs,), seq_len_fill_value, dtype=torch.int32)
out_cache_loc = torch.zeros((max_num_token,), dtype=torch.int64)
out_cache_loc = torch.zeros((max_num_token,), dtype=cache_loc_dtype)
positions = torch.zeros((max_num_token,), dtype=torch.int64)
mrope_positions = torch.zeros((3, max_num_token), dtype=torch.int64)
num_token_non_padded = torch.zeros((1,), dtype=torch.int32)

4
python/sglang/srt/model_executor/npu_graph_runner.py
View File

@@ -107,8 +107,8 @@ class NPUGraphRunner(CudaGraphRunner):
self.replay_prepare(forward_batch, pp_proxy_tensors)
else:
# In speculative decoding, these two fields are still needed.
self.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids)
self.positions[: self.raw_num_token].copy_(forward_batch.positions)
self.buffers.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids)
self.buffers.positions[: self.raw_num_token].copy_(forward_batch.positions)
# Replay
if not is_deepseek_nsa(self.model_runner.model_config.hf_config):

4
python/sglang/srt/speculative/spec_utils.py
View File

@@ -133,8 +133,8 @@ def assign_req_to_token_pool_func(
torch.ops.npu.cache_loc_assign(
req_pool_indices,
req_to_token,
start_offset,
end_offset,
start_offset.to(torch.int64),
end_offset.to(torch.int64),
out_cache_loc,
)