[spec-overlap] bugfix for pd disaggregation and npu (#14088)

Co-authored-by: Even Zhou <even.y.zhou@outlook.com>

.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 }}

.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 }}

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
                 ),

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(
+            max_num_extend_tokens = next_power_of_2(extend_num_tokens)
-                prefix_lens.to(torch.int64),
+            bs = prefix_lens.shape[0]
-                seq_lens.to(torch.int64),
+            alloc_extend_kernel[(bs,)](
-                last_loc.to(torch.int64),
+                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:

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(
+            logits_output, _ = self.draft_runner.forward(
-                forward_batch.input_ids, forward_batch.positions, forward_batch
+                forward_batch, skip_attn_backend_init=True
             )
             if self.server_args.enable_nan_detection:
                 detect_nan(logits_output)

python/sglang/srt/speculative/spec_utils.py

@@ -117,7 +117,6 @@ 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,
@@ -127,16 +126,6 @@ def assign_req_to_token_pool_func(
         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,
-        )
 
 
 @triton.jit

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)
 

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),
     ],
 }