From 2e8f54e61edd9d60248d15e2090752ac9b3dd25d Mon Sep 17 00:00:00 2001 From: liupeng374 Date: Mon, 1 Dec 2025 22:58:20 +0800 Subject: [PATCH] [spec-overlap] bugfix for pd disaggregation and npu (#14088) Co-authored-by: Even Zhou --- .../workflows/release-docker-npu-nightly.yml | 2 +- .github/workflows/release-docker-npu.yml | 2 +- python/sglang/srt/managers/scheduler.py | 10 ++++--- .../sglang/srt/mem_cache/allocator_ascend.py | 19 +++++++----- .../sglang/srt/speculative/eagle_worker_v2.py | 4 +-- python/sglang/srt/speculative/spec_utils.py | 29 ++++++------------- scripts/ci/npu_ci_install_dependency.sh | 2 +- .../ascend/test_ascend_deepseek_mtp.py | 0 test/srt/run_suite.py | 2 +- 9 files changed, 32 insertions(+), 38 deletions(-) rename test/{manual => srt}/ascend/test_ascend_deepseek_mtp.py (100%) diff --git a/.github/workflows/release-docker-npu-nightly.yml b/.github/workflows/release-docker-npu-nightly.yml index 1ede19a35..1873c2c36 100644 --- a/.github/workflows/release-docker-npu-nightly.yml +++ b/.github/workflows/release-docker-npu-nightly.yml @@ -73,6 +73,6 @@ jobs: push: ${{ github.repository == 'sgl-project/sglang' && github.event_name != 'pull_request' }} provenance: false build-args: | - SGLANG_KERNEL_NPU_TAG=20251120 + SGLANG_KERNEL_NPU_TAG=20251128 CANN_VERSION=${{ matrix.cann_version }} DEVICE_TYPE=${{ matrix.device_type }} diff --git a/.github/workflows/release-docker-npu.yml b/.github/workflows/release-docker-npu.yml index 2b2506a28..dd054c9c0 100644 --- a/.github/workflows/release-docker-npu.yml +++ b/.github/workflows/release-docker-npu.yml @@ -70,6 +70,6 @@ jobs: push: ${{ github.repository == 'sgl-project/sglang' && github.event_name != 'pull_request' }} provenance: false build-args: | - SGLANG_KERNEL_NPU_TAG=20251120 + SGLANG_KERNEL_NPU_TAG=20251128 CANN_VERSION=${{ matrix.cann_version }} DEVICE_TYPE=${{ matrix.device_type }} diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 38704541f..15e7063ee 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -824,8 +824,10 @@ class Scheduler( draft_token_to_kv_pool = ( self.draft_worker.draft_worker.draft_runner.token_to_kv_pool ) + model_config = self.draft_worker.draft_worker.draft_runner.model_config else: draft_token_to_kv_pool = self.draft_worker.model_runner.token_to_kv_pool + model_config = self.draft_worker.model_config if ( self.disaggregation_mode == DisaggregationMode.DECODE @@ -837,12 +839,12 @@ class Scheduler( self.disagg_metadata_buffers = MetadataBuffers( buffer_size, hidden_size=( - self.draft_worker.model_config.hidden_size + model_config.hidden_size if self.spec_algorithm.is_eagle() else 16 # minimal padding size for RDMA ), hidden_states_dtype=( - self.draft_worker.model_config.dtype + model_config.dtype if self.spec_algorithm.is_eagle() else torch.float32 ), @@ -890,12 +892,12 @@ class Scheduler( self.disagg_metadata_buffers = MetadataBuffers( buffer_size, hidden_size=( - self.draft_worker.model_config.hidden_size + model_config.hidden_size if self.spec_algorithm.is_eagle() else 16 # minimal padding size for RDMA ), hidden_states_dtype=( - self.draft_worker.model_config.dtype + model_config.dtype if self.spec_algorithm.is_eagle() else torch.float32 ), diff --git a/python/sglang/srt/mem_cache/allocator_ascend.py b/python/sglang/srt/mem_cache/allocator_ascend.py index f66136025..74ac8983d 100644 --- a/python/sglang/srt/mem_cache/allocator_ascend.py +++ b/python/sglang/srt/mem_cache/allocator_ascend.py @@ -8,7 +8,7 @@ if TYPE_CHECKING: from sglang.srt.mem_cache.memory_pool import KVCache from sglang.srt.mem_cache.allocator import PagedTokenToKVPoolAllocator -from sglang.srt.utils import get_num_new_pages +from sglang.srt.utils import get_num_new_pages, next_power_of_2 def alloc_extend_kernel_ascend( @@ -104,21 +104,24 @@ class AscendPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator): return None if num_new_pages_item < 200: - import sgl_kernel_npu # noqa: F401 + from sgl_kernel_npu.mem_cache.allocator import alloc_extend_kernel out_indices = torch.empty( (extend_num_tokens,), dtype=torch.int64, device=self.device, ) - torch.ops.npu.alloc_extend( - prefix_lens.to(torch.int64), - seq_lens.to(torch.int64), - last_loc.to(torch.int64), + max_num_extend_tokens = next_power_of_2(extend_num_tokens) + bs = prefix_lens.shape[0] + alloc_extend_kernel[(bs,)]( + prefix_lens, + seq_lens, + last_loc, self.free_pages, - self.page_size, out_indices, - num_new_pages, + next_power_of_2(bs), + self.page_size, + max_num_extend_tokens, ) else: diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 251922a2e..663c4f341 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -386,8 +386,8 @@ class EagleDraftWorker(BaseDraftWorker): spec_info.hidden_states = hidden_states # Run forward - logits_output = self.draft_runner.model.forward( - forward_batch.input_ids, forward_batch.positions, forward_batch + logits_output, _ = self.draft_runner.forward( + forward_batch, skip_attn_backend_init=True ) if self.server_args.enable_nan_detection: detect_nan(logits_output) diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 108b4f0f8..2fb41b407 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -117,26 +117,15 @@ def assign_req_to_token_pool_func( out_cache_loc: torch.Tensor, batch_size: int, ): - if _is_cuda or _is_hip: - assign_req_to_token_pool[(batch_size,)]( - req_pool_indices, - req_to_token, - start_offset, - end_offset, - out_cache_loc, - req_to_token.shape[1], - next_power_of_2(batch_size), - ) - elif _is_npu: - import sgl_kernel_npu # noqa: F401 - - torch.ops.npu.cache_loc_assign( - req_pool_indices, - req_to_token, - start_offset.to(torch.int64), - end_offset.to(torch.int64), - out_cache_loc, - ) + assign_req_to_token_pool[(batch_size,)]( + req_pool_indices, + req_to_token, + start_offset, + end_offset, + out_cache_loc, + req_to_token.shape[1], + next_power_of_2(batch_size), + ) @triton.jit diff --git a/scripts/ci/npu_ci_install_dependency.sh b/scripts/ci/npu_ci_install_dependency.sh index 24192d6df..7cfdaaee6 100755 --- a/scripts/ci/npu_ci_install_dependency.sh +++ b/scripts/ci/npu_ci_install_dependency.sh @@ -49,7 +49,7 @@ wget -O "${BISHENG_NAME}" "${BISHENG_URL}" && chmod a+x "${BISHENG_NAME}" && "./ ### Install sgl-kernel-npu -SGL_KERNEL_NPU_TAG="20251120" +SGL_KERNEL_NPU_TAG="20251128" git clone --depth 1 https://github.com/sgl-project/sgl-kernel-npu.git --branch ${SGL_KERNEL_NPU_TAG} (cd sgl-kernel-npu && bash ./build.sh && ${PIP_INSTALL} output/deep_ep*.whl output/sgl_kernel_npu*.whl && cd "$(python3 -m pip show deep-ep | grep -E '^Location:' | awk '{print $2}')" && ln -s deep_ep/deep_ep_cpp*.so) diff --git a/test/manual/ascend/test_ascend_deepseek_mtp.py b/test/srt/ascend/test_ascend_deepseek_mtp.py similarity index 100% rename from test/manual/ascend/test_ascend_deepseek_mtp.py rename to test/srt/ascend/test_ascend_deepseek_mtp.py diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index f427562f3..49655d220 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -375,7 +375,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), ], }