diff --git a/python/sglang/kernels/ops/speculative/multi_layer_eagle.py b/python/sglang/kernels/ops/speculative/multi_layer_eagle.py index a41c9f1d7..ebdc10dc4 100644 --- a/python/sglang/kernels/ops/speculative/multi_layer_eagle.py +++ b/python/sglang/kernels/ops/speculative/multi_layer_eagle.py @@ -15,9 +15,10 @@ import triton import triton.language as tl -from sglang.srt.utils import is_cpu +from sglang.srt.utils import is_cpu, is_npu _is_cpu = is_cpu() +_is_npu = is_npu() if _is_cpu: from sgl_kernel import rotate_input_ids_cpu @@ -74,6 +75,13 @@ def rotate_input_ids( return input_ids batch_size = extend_seq_lens.shape[0] + + # rotate_input_ids_triton skipped: batch_size=0 (empty extend_seq_lens). + # This is expected when a DP rank has no requests. + # TODO: @iforgetmyname Remove NPU-specific guard after triton-ascend fixes zero-sized grid kernel launch abort + if batch_size == 0 and _is_npu: + return input_ids + BLOCK_SIZE = 4096 if select_index is not None else 8 grid = (batch_size,) diff --git a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py index 718b46220..a685406f2 100644 --- a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py +++ b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py @@ -63,6 +63,7 @@ class NPUMHATokenToKVPool(MHATokenToKVPool): end_layer: Optional[int] = None, enable_alt_stream: bool = True, enable_kv_cache_copy: bool = False, + **kwargs, ): self.use_fia = get_bool_env_var("ASCEND_USE_FIA", "False") super().__init__( @@ -82,6 +83,7 @@ class NPUMHATokenToKVPool(MHATokenToKVPool): end_layer=end_layer, enable_alt_stream=enable_alt_stream, enable_kv_cache_copy=enable_kv_cache_copy, + **kwargs, ) def _create_buffers(self):