[NPU][eagle3] support qwen eagle3 on NPU (#14820)

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -88,7 +88,8 @@ class EAGLEDraftCudaGraphRunner:
             self.input_ids = torch.zeros((self.max_num_token,), dtype=torch.int64)
             self.req_pool_indices = torch.zeros((self.max_bs,), dtype=torch.int32)
             self.out_cache_loc = torch.zeros(
-                (self.max_num_token * self.speculative_num_steps,), dtype=torch.int64
+                (self.max_num_token * self.speculative_num_steps,),
+                dtype=self._cache_loc_dtype(),
             )
             self.positions = torch.zeros((self.max_num_token,), dtype=torch.int64)
             self.mrope_positions = torch.zeros(
@@ -132,6 +133,9 @@ class EAGLEDraftCudaGraphRunner:
                 f"Capture cuda graph failed: {e}\n{CUDA_GRAPH_CAPTURE_FAILED_MSG}"
             )
 
+    def _cache_loc_dtype(self):
+        return torch.int64
+
     def can_run(self, forward_batch: ForwardBatch):
         if self.require_mlp_tp_gather:
             cuda_graph_bs = (

python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py

@@ -92,7 +92,9 @@ class EAGLEDraftExtendCudaGraphRunner:
         with torch.device(model_runner.device):
             self.input_ids = torch.zeros((self.max_num_token,), dtype=torch.int64)
             self.req_pool_indices = torch.zeros((self.max_bs,), dtype=torch.int32)
-            self.out_cache_loc = torch.ones((self.max_num_token,), dtype=torch.int64)
+            self.out_cache_loc = torch.ones(
+                (self.max_num_token,), dtype=self._cache_loc_dtype()
+            )
             self.positions = torch.zeros((self.max_num_token,), dtype=torch.int64)
             self.mrope_positions = torch.zeros(
                 (3, self.max_num_token), dtype=torch.int64
@@ -204,6 +206,9 @@ class EAGLEDraftExtendCudaGraphRunner:
     def _create_graph(self):
         return torch.cuda.CUDAGraph()
 
+    def _cache_loc_dtype(self):
+        return torch.int64
+
     def _capture_init(self, run_once_fn):
         for _ in range(2):
             torch.cuda.synchronize()

python/sglang/srt/speculative/eagle_info_v2.py

@@ -468,8 +468,6 @@ def assign_extend_cache_locs_func(
         return out_cache_loc
 
     elif _is_npu:
-        import sgl_kernel_npu  # noqa: F401
-
         out_cache_loc = torch.empty(
             (batch_size * draft_token_num,),
             dtype=torch.int32,
@@ -482,6 +480,5 @@ def assign_extend_cache_locs_func(
             end_offset,
             out_cache_loc,
         )
-        out_cache_loc = out_cache_loc.to(dtype=torch.int64)
 
         return out_cache_loc