[spec decoding] fix multi_layer_eagle rotate_input_ids kernel registration (#31614)
This commit is contained in:
@@ -15,7 +15,7 @@ _TRITON_KERNELS = [
|
||||
("eagle", "fill_bonus_tokens"),
|
||||
("eagle", "fill_accept_out_cache_loc"),
|
||||
("gather_spec_extras", "gather_spec_extras"),
|
||||
("multi_layer_eagle", "rotate_input_ids_triton"),
|
||||
("multi_layer_eagle", "rotate_input_ids"),
|
||||
("spec_tree", "sgl_build_tree_kernel_efficient_triton"),
|
||||
("spec_tree", "verify_tree_greedy_kernel_triton"),
|
||||
("topk1", "draft_topk1_postprocess"),
|
||||
|
||||
@@ -76,7 +76,7 @@ def rotate_input_ids(
|
||||
|
||||
batch_size = extend_seq_lens.shape[0]
|
||||
|
||||
# rotate_input_ids_triton skipped: batch_size=0 (empty extend_seq_lens).
|
||||
# rotate_input_ids kernel 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:
|
||||
|
||||
Reference in New Issue
Block a user