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

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

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(
-                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:

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)

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