Add treemask mode to build_eagle_tree & release sgl-kernel 0.2.3 (#7756)

Co-authored-by: Pranjal Shankhdhar <pranjal.ssh@gmail.com>
This commit is contained in:
Lianmin Zheng
2025-07-05 12:17:05 -07:00
committed by GitHub
parent c04a8a820b
commit 5589b75024
6 changed files with 101 additions and 36 deletions

View File

@@ -72,6 +72,7 @@ def build_tree_kernel_efficient(
topk: int,
depth: int,
draft_token_num: int,
tree_mask_mode: int,
) -> None:
torch.ops.sgl_kernel.build_tree_kernel_efficient.default(
parent_list,
@@ -85,6 +86,7 @@ def build_tree_kernel_efficient(
topk,
depth,
draft_token_num,
tree_mask_mode,
)