diff --git a/python/sglang/srt/mem_cache/allocator_ascend.py b/python/sglang/srt/mem_cache/allocator_ascend.py index 1b1b84fe0..f66136025 100644 --- a/python/sglang/srt/mem_cache/allocator_ascend.py +++ b/python/sglang/srt/mem_cache/allocator_ascend.py @@ -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, diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index 2fb40cdec..b775b4535 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/cuda_graph_runner.py @@ -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() diff --git a/python/sglang/srt/model_executor/input_buffers.py b/python/sglang/srt/model_executor/input_buffers.py index 586399474..ec8d23476 100644 --- a/python/sglang/srt/model_executor/input_buffers.py +++ b/python/sglang/srt/model_executor/input_buffers.py @@ -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) diff --git a/python/sglang/srt/model_executor/npu_graph_runner.py b/python/sglang/srt/model_executor/npu_graph_runner.py index c42496280..abc276ff1 100644 --- a/python/sglang/srt/model_executor/npu_graph_runner.py +++ b/python/sglang/srt/model_executor/npu_graph_runner.py @@ -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): diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index b2f2ba12b..f20ea377a 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -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, ) diff --git a/test/srt/ascend/test_ascend_deepseek_mtp.py b/test/manual/ascend/test_ascend_deepseek_mtp.py similarity index 100% rename from test/srt/ascend/test_ascend_deepseek_mtp.py rename to test/manual/ascend/test_ascend_deepseek_mtp.py diff --git a/test/srt/ascend/test_ascend_tp1_bf16.py b/test/srt/ascend/test_ascend_tp1_bf16.py index f854605ce..fd0f96e73 100644 --- a/test/srt/ascend/test_ascend_tp1_bf16.py +++ b/test/srt/ascend/test_ascend_tp1_bf16.py @@ -8,9 +8,7 @@ from sglang.test.test_utils import ( DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, DEFAULT_URL_FOR_TEST, CustomTestCase, - is_in_ci, popen_launch_server, - run_bench_offline_throughput, ) TEST_MODEL_MATRIX = { @@ -71,26 +69,6 @@ class TestAscendTp1Bf16(CustomTestCase): finally: kill_process_tree(process.pid) - def test_b_throughput(self): - for model in self.models: - with self.subTest(model=model): - print(f"##=== Testing throughput: {model} ===##") - - output_throughput = run_bench_offline_throughput( - model, - [ - *self.common_args, - ], - ) - - print(f"##=== {model} throughput: {output_throughput} ===##") - - if is_in_ci(): - self.assertGreater( - output_throughput, - TEST_MODEL_MATRIX[model]["output_throughput"], - ) - if __name__ == "__main__": unittest.main() diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index 64246aaba..cf5ef8faa 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -366,7 +366,7 @@ suite_ascend = { ], "per-commit-16-npu-a3": [ TestFile("ascend/test_ascend_deepep.py", 400), - TestFile("ascend/test_ascend_deepseek_mtp.py", 400), + # TestFile("ascend/test_ascend_deepseek_mtp.py", 400), ], }