[NPU][bugfix] Fix NPU KernelLaunch Failure in rotate_input_ids_triton with Empty Batch (#30589)
This commit is contained in:
@@ -15,9 +15,10 @@
|
|||||||
import triton
|
import triton
|
||||||
import triton.language as tl
|
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_cpu = is_cpu()
|
||||||
|
_is_npu = is_npu()
|
||||||
|
|
||||||
if _is_cpu:
|
if _is_cpu:
|
||||||
from sgl_kernel import rotate_input_ids_cpu
|
from sgl_kernel import rotate_input_ids_cpu
|
||||||
@@ -74,6 +75,13 @@ def rotate_input_ids(
|
|||||||
return input_ids
|
return input_ids
|
||||||
|
|
||||||
batch_size = extend_seq_lens.shape[0]
|
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
|
BLOCK_SIZE = 4096 if select_index is not None else 8
|
||||||
grid = (batch_size,)
|
grid = (batch_size,)
|
||||||
|
|
||||||
|
|||||||
@@ -63,6 +63,7 @@ class NPUMHATokenToKVPool(MHATokenToKVPool):
|
|||||||
end_layer: Optional[int] = None,
|
end_layer: Optional[int] = None,
|
||||||
enable_alt_stream: bool = True,
|
enable_alt_stream: bool = True,
|
||||||
enable_kv_cache_copy: bool = False,
|
enable_kv_cache_copy: bool = False,
|
||||||
|
**kwargs,
|
||||||
):
|
):
|
||||||
self.use_fia = get_bool_env_var("ASCEND_USE_FIA", "False")
|
self.use_fia = get_bool_env_var("ASCEND_USE_FIA", "False")
|
||||||
super().__init__(
|
super().__init__(
|
||||||
@@ -82,6 +83,7 @@ class NPUMHATokenToKVPool(MHATokenToKVPool):
|
|||||||
end_layer=end_layer,
|
end_layer=end_layer,
|
||||||
enable_alt_stream=enable_alt_stream,
|
enable_alt_stream=enable_alt_stream,
|
||||||
enable_kv_cache_copy=enable_kv_cache_copy,
|
enable_kv_cache_copy=enable_kv_cache_copy,
|
||||||
|
**kwargs,
|
||||||
)
|
)
|
||||||
|
|
||||||
def _create_buffers(self):
|
def _create_buffers(self):
|
||||||
|
|||||||
Reference in New Issue
Block a user