[spec decoding] fix multi_layer_eagle rotate_input_ids kernel registration (#31614)

This commit is contained in:
Qiaolin Yu
2026-07-17 15:18:01 -07:00
committed by GitHub
parent ec6a3163b7
commit e2d2e8d07e
2 changed files with 2 additions and 2 deletions
@@ -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: