diff --git a/benchmark/bench_build_eagle_tree.py b/benchmark/bench_build_eagle_tree.py new file mode 100644 index 000000000..694d5c56a --- /dev/null +++ b/benchmark/bench_build_eagle_tree.py @@ -0,0 +1,291 @@ +"""Benchmark build_tree_kernel_efficient (EAGLE draft-tree metadata) on XPU. + +Compares the SYCL kernel against the upstream Triton kernel (sgl_kernel.eagle_utils, +ported from sglang/kernels/ops/speculative/spec_tree.py) so the baseline is +exactly what the XPU path runs today (eagle_utils dispatches _is_xpu -> +sgl_build_tree_kernel_triton). + +The Triton baseline times its cumsum too: the Triton kernel takes +seq_len_prefix_sum as an input, so that launch is part of its cost. The SYCL +kernel folds the prefix sum into the kernel via a group reduction. + +Both providers are timed on identical inputs and their outputs are checked +against each other once per config before timing, so a regression shows up as a +correctness failure rather than a misleading speedup. +""" + +import itertools + +import pandas as pd +import torch +import triton +from sgl_kernel import TreeMaskMode, build_tree_kernel_efficient +from sgl_kernel.eagle_utils import sgl_build_tree_kernel_triton + + +def run_triton( + parent_list, + selected_index, + seq_lens, + bufs, + topk, + depth, + draft_token_num, + mode, +): + tree_mask, positions, r_index, r_next_token, r_next_sibling = bufs + sgl_build_tree_kernel_triton( + parent_list, + selected_index, + seq_lens, + tree_mask, + positions, + r_index, + r_next_token, + r_next_sibling, + topk, + depth, + draft_token_num, + mode, + ) + + +def run_sycl( + parent_list, + selected_index, + seq_lens, + bufs, + topk, + depth, + draft_token_num, + mode, +): + tree_mask, positions, r_index, r_next_token, r_next_sibling = bufs + build_tree_kernel_efficient( + parent_list, + selected_index, + seq_lens, + tree_mask, + positions, + r_index, + r_next_token, + r_next_sibling, + topk, + depth, + draft_token_num, + mode, + ) + + +# --------------------------------------------------------------------------- +# Inputs +# --------------------------------------------------------------------------- +def gen_draft_tree(bs, topk, num_steps, draft_token_num, device): + """Simulate the EAGLE draft loop to get a valid (parent_list, selected_index).""" + scores = torch.rand(bs, topk, dtype=torch.float32, device=device) + score_chunks = [scores] + parents_chunks = [ + torch.arange(-1, topk, dtype=torch.int64, device=device).expand(bs, -1) + ] + cum_scores = scores + for i in range(1, num_steps): + step_p = torch.rand(bs, topk, topk, dtype=torch.float32, device=device) + expand_scores = cum_scores.unsqueeze(2) * step_p + cum_scores, topk_cs_index = torch.topk( + expand_scores.flatten(start_dim=1), topk, dim=-1 + ) + score_chunks.append(expand_scores.flatten(start_dim=1)) + parents_chunks.append(topk_cs_index + (topk * topk * (i - 1) + topk)) + score_flat = torch.cat(score_chunks, dim=1) + selected_index = torch.sort( + torch.topk(score_flat, draft_token_num - 1, dim=-1).indices, dim=-1 + ).values + parent_list = torch.cat(parents_chunks[:-1], dim=1).contiguous() + return parent_list, selected_index.contiguous() + + +def alloc_bufs(bs, draft_token_num, seq_lens_sum, mode, device): + if mode == TreeMaskMode.QLEN_ONLY: + numel = bs * draft_token_num * draft_token_num + else: + numel = seq_lens_sum * draft_token_num + bs * draft_token_num * draft_token_num + tree_mask = torch.full((numel,), True, dtype=torch.bool, device=device) + positions = torch.zeros(bs * draft_token_num, dtype=torch.int64, device=device) + retrieve_buf = torch.full( + (3, bs, draft_token_num), -1, dtype=torch.int64, device=device + ) + return (tree_mask, positions, *retrieve_buf) + + +# Realistic EAGLE / MTP serving shapes: (topk, spec_steps, draft_token_num). +TREE_SHAPES = [ + (1, 3, 4), # MTP chain + (4, 3, 8), + (4, 3, 16), + (8, 4, 32), + (8, 5, 64), +] +BATCH_SIZES = [1, 8, 32, 64, 128, 256] +SEQ_LEN = 2048 # committed context per request + +configs = list(itertools.product(BATCH_SIZES, TREE_SHAPES)) +all_results = [] + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["bs", "tree_shape"], + x_vals=configs, + line_arg="provider", + line_vals=["sycl", "triton"], + line_names=["SYCL", "Triton (upstream)"], + styles=[("green", "-"), ("blue", "--")], + ylabel="Time (us)", + plot_name="build-eagle-tree-performance", + args={}, + ) +) +def benchmark(bs, tree_shape, provider): + topk, spec_steps, draft_token_num = tree_shape + device = "xpu" + torch.manual_seed(42) + + mode = TreeMaskMode.FULL_MASK + seq_lens = torch.full((bs,), SEQ_LEN, dtype=torch.int64, device=device) + seq_lens_sum = int(seq_lens.sum()) + parent_list, selected_index = gen_draft_tree( + bs, topk, spec_steps, draft_token_num, device + ) + + runner = run_sycl if provider == "sycl" else run_triton + bufs = alloc_bufs(bs, draft_token_num, seq_lens_sum, mode, device) + + args = ( + parent_list, + selected_index, + seq_lens, + bufs, + topk, + spec_steps, + draft_token_num, + mode, + ) + + # Correctness gate: both providers must agree before we trust the timings. + ref_bufs = alloc_bufs(bs, draft_token_num, seq_lens_sum, mode, device) + run_sycl( + parent_list, + selected_index, + seq_lens, + ref_bufs, + topk, + spec_steps, + draft_token_num, + mode, + ) + runner(*args) + torch.xpu.synchronize() + for name, got, want in zip( + ( + "tree_mask", + "positions", + "retrieve_index", + "retrieve_next_token", + "retrieve_next_sibling", + ), + bufs, + ref_bufs, + ): + if not torch.equal(got, want): + raise AssertionError( + f"{provider} disagrees with the SYCL reference on {name} " + f"(bs={bs}, topk={topk}, steps={spec_steps}, N={draft_token_num})" + ) + + for _ in range(10): + runner(*args) + torch.xpu.synchronize() + + ms = triton.testing.do_bench( + lambda: runner(*args), quantiles=[0.5, 0.25, 0.75], return_mode="median" + ) + if isinstance(ms, (tuple, list)): + ms = ms[0] + us = ms * 1e3 + + torch.xpu.empty_cache() + + # Mask bytes dominate the traffic: one bool per (draft token, tree column), + # plus the prefix columns the FULL_MASK layout skips over. + mask_bytes = bs * draft_token_num * draft_token_num + _, positions, r_index, r_next_token, r_next_sibling = bufs + + # Memory traffic (int64 = 8B, bool = 1B): reads are the tree-shape inputs, + # writes are the mask plus the four per-node metadata outputs. + read_bytes = (parent_list.numel() + selected_index.numel() + seq_lens.numel()) * 8 + write_bytes = ( + mask_bytes + + ( + positions.numel() + + r_index.numel() + + r_next_token.numel() + + r_next_sibling.numel() + ) + * 8 + ) + total_bytes = read_bytes + write_bytes + bandwidth_gb_s = total_bytes / (ms / 1e3) / 1e9 + + all_results.append( + { + "provider": provider, + "bs": bs, + "topk": topk, + "steps": spec_steps, + "draft_tokens": draft_token_num, + "us": us, + "nodes_per_sec_M": bs * draft_token_num / (ms / 1e3) / 1e6, + "mask_cells_per_sec_M": mask_bytes / (ms / 1e3) / 1e6, + "bandwidth_gb_s": bandwidth_gb_s, + } + ) + return us + + +if __name__ == "__main__": + benchmark.run(print_data=False) + print("Benchmark finished!") + + df = pd.DataFrame(all_results) + print("\n" + "=" * 88) + print("BUILD_EAGLE_TREE BENCHMARK RESULTS") + print("=" * 88) + print(df.to_markdown(index=False, floatfmt=".2f")) + + pivot = df.pivot_table( + index=["bs", "topk", "steps", "draft_tokens"], + columns="provider", + values="us", + ) + if {"sycl", "triton"}.issubset(pivot.columns): + pivot["speedup_x"] = pivot["triton"] / pivot["sycl"] + print("\n" + "=" * 88) + print("SYCL vs Triton (speedup = triton_us / sycl_us; >1 means SYCL wins)") + print("=" * 88) + print(pivot.to_markdown(floatfmt=".2f")) + print( + f"\n geomean speedup: " + f"{pivot['speedup_x'].prod() ** (1 / len(pivot)):.2f}x" + f"\n min: {pivot['speedup_x'].min():.2f}x" + f" max: {pivot['speedup_x'].max():.2f}x" + ) + + print("\n" + "=" * 88) + print("BANDWIDTH SUMMARY") + print("=" * 88) + for provider in df["provider"].unique(): + df_provider = df[df["provider"] == provider] + print(f"\n{provider}:") + print(f" Mean bandwidth: {df_provider['bandwidth_gb_s'].mean():.2f} GB/s") + print(f" Best bandwidth: {df_provider['bandwidth_gb_s'].max():.2f} GB/s") + print("\n") diff --git a/include/sgl_kernel_ops.h b/include/sgl_kernel_ops.h index b5b75b235..20515187a 100644 --- a/include/sgl_kernel_ops.h +++ b/include/sgl_kernel_ops.h @@ -884,6 +884,8 @@ void verify_tree_greedy( at::Tensor target_predict, int64_t sycl_stream = 0); +// tree_mask_mode mirrors sglang.srt.speculative.eagle_utils.TreeMaskMode: +// 0 = FULL_MASK, 1 = QLEN_ONLY. QLEN_ONLY_BITPACKING (2) is not implemented. void build_tree_kernel_efficient( at::Tensor parent_list, at::Tensor selected_index, @@ -895,7 +897,8 @@ void build_tree_kernel_efficient( at::Tensor retrive_next_sibling, int64_t topk, int64_t depth, - int64_t draft_token_num); + int64_t draft_token_num, + int64_t tree_mask_mode = 0); void segment_packbits( at::Tensor x, at::Tensor input_indptr, at::Tensor output_indptr, at::Tensor y, int64_t sycl_stream); diff --git a/python/sgl_kernel/__init__.py b/python/sgl_kernel/__init__.py index a56662e62..fbf6b54f5 100644 --- a/python/sgl_kernel/__init__.py +++ b/python/sgl_kernel/__init__.py @@ -189,6 +189,7 @@ def _export_jit_toolchain_env() -> None: sparse_attn_varlen_func, ) from sgl_kernel.speculative import ( + TreeMaskMode, build_tree_kernel_efficient, segment_packbits, tree_speculative_sampling_target_only, diff --git a/python/sgl_kernel/eagle_utils.py b/python/sgl_kernel/eagle_utils.py new file mode 100644 index 000000000..9649b0239 --- /dev/null +++ b/python/sgl_kernel/eagle_utils.py @@ -0,0 +1,341 @@ +from __future__ import annotations + +import math +from typing import List, Optional + +import torch +import triton +import triton.language as tl +from sgl_kernel.speculative import TreeMaskMode +from sgl_kernel.speculative import ( + build_tree_kernel_efficient as sgl_build_tree_kernel_efficient, +) + + +def organize_draft_results( + score_list: List[torch.Tensor], + token_list: List[torch.Tensor], + parents_list: List[torch.Tensor], + num_draft_token: int, +): + score_list = torch.cat(score_list, dim=1).flatten(1) + ss_token_list = torch.cat(token_list, dim=1) + top_scores = torch.topk(score_list, num_draft_token - 1, dim=-1) + top_scores_index = torch.sort(top_scores.indices).values + draft_tokens = torch.gather(ss_token_list, index=top_scores_index, dim=1) + + if len(parents_list) > 1: + parent_list = torch.cat(parents_list[:-1], dim=1) + else: + batch_size = parents_list[0].shape[0] + parent_list = torch.empty( + batch_size, 0, dtype=torch.long, device=parents_list[0].device + ) + + return parent_list, top_scores_index, draft_tokens + + +def build_tree_kernel_efficient( + bonus_tokens: torch.Tensor, + parent_list: torch.Tensor, + top_scores_index: torch.Tensor, + draft_tokens: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_sum: int, + topk: int, + spec_steps: int, + num_verify_tokens: int, + tree_mask_mode: TreeMaskMode = TreeMaskMode.FULL_MASK, + tree_mask_buf: Optional[torch.Tensor] = None, + fill_prefix_mask: bool = True, +): + draft_tokens = torch.cat((bonus_tokens.unsqueeze(1), draft_tokens), dim=1).flatten() + + # seq_lens_sum == sum(seq_lens); seq_lens: sequence length without draft tokens + bs = seq_lens.numel() + device = seq_lens.device + # e.g. for bs=1, tree_mask: num_draft_token, seq_lens_sum + num_draft_token (flattened) + # where each row indicates the attending pattern of each draft token + # if use_partial_packed_tree_mask is True, tree_mask: num_draft_token (flattened, packed) + if tree_mask_buf is not None: + tree_mask = tree_mask_buf + if tree_mask_mode == TreeMaskMode.QLEN_ONLY: + tree_mask.fill_(True) + elif tree_mask_mode == TreeMaskMode.QLEN_ONLY_BITPACKING: + tree_mask.fill_(0) + elif tree_mask_mode == TreeMaskMode.FULL_MASK: + # Only the [0, seq_len) prefix columns depend on this fill; the + # kernel below writes every tree cell itself. Skip the (up to + # 100s of MB) per-step memset when nothing reads the mask. + if fill_prefix_mask: + tree_mask.fill_(True) + else: + raise NotImplementedError(f"Invalid tree mask: {tree_mask_mode=}") + elif tree_mask_mode == TreeMaskMode.QLEN_ONLY: + tree_mask = torch.full( + (num_verify_tokens * bs * num_verify_tokens,), + True, + dtype=torch.bool, + device=device, + ) + elif tree_mask_mode == TreeMaskMode.QLEN_ONLY_BITPACKING: + packed_dtypes = [torch.uint8, torch.uint16, torch.uint32] + packed_dtype_idx = int(math.ceil(math.log2((num_verify_tokens + 7) // 8))) + tree_mask = torch.zeros( + (num_verify_tokens * bs,), + dtype=packed_dtypes[packed_dtype_idx], + device=device, + ) + elif tree_mask_mode == TreeMaskMode.FULL_MASK: + mask_shape = ( + seq_lens_sum * num_verify_tokens + + num_verify_tokens * num_verify_tokens * bs, + ) + # Same reasoning as the preallocated branch above. + tree_mask = ( + torch.full(mask_shape, True, dtype=torch.bool, device=device) + if fill_prefix_mask + else torch.empty(mask_shape, dtype=torch.bool, device=device) + ) + else: + raise NotImplementedError(f"Invalid tree mask: {tree_mask_mode=}") + + # TODO: make them torch.empty and fuse them into `sgl_build_tree_kernel` + retrieve_buf = torch.full( + (3, bs, num_verify_tokens), -1, device=device, dtype=torch.long + ) + retrieve_index, retrieve_next_token, retrieve_next_sibling = retrieve_buf + # position: where each token belongs to + # e.g. if depth of each draft token is [0, 1, 1, 2] and the prompt length is 7 + # then, positions = [7, 8, 8, 9] + positions = torch.empty((bs * num_verify_tokens,), device=device, dtype=torch.long) + + sgl_build_tree_kernel_efficient( + parent_list, + top_scores_index, + seq_lens, + tree_mask, + positions, + retrieve_index, + retrieve_next_token, + retrieve_next_sibling, + topk, + spec_steps, + num_verify_tokens, + tree_mask_mode, + ) + return ( + tree_mask, + positions, + retrieve_index, + retrieve_next_token, + retrieve_next_sibling, + draft_tokens, + ) + + +@triton.jit +def sgl_build_tree_kernel_efficient_triton( + parent_list_ptr, + selected_index_ptr, + verified_seq_len_ptr, + seq_len_prefix_sum_ptr, + tree_mask_ptr, + positions_ptr, + retrieve_index_ptr, + retrieve_next_token_ptr, + retrieve_next_sibling_ptr, + topk: tl.constexpr, + depth: tl.constexpr, + draft_token_num: tl.constexpr, + tree_mask_mode: tl.constexpr, + batch_size: tl.constexpr, + parent_list_stride: tl.constexpr, + selected_index_stride: tl.constexpr, +): + """ + Triton kernel for building EAGLE tree structure. + Each program handles one batch item (batch_idx). + """ + batch_idx = tl.program_id(0) + + # Calculate seq_tree_idx + seq_len = tl.load(verified_seq_len_ptr + batch_idx) + seq_len_prefix_sum = tl.load(seq_len_prefix_sum_ptr + batch_idx) + + # Cast initial value to match the dtype of loaded tensors to avoid type inconsistency + seq_tree_idx = ( + tl.cast(draft_token_num * draft_token_num * batch_idx, seq_len.dtype) + + seq_len_prefix_sum * draft_token_num + ) + + positions_offset = batch_idx * draft_token_num + tl.store(positions_ptr + positions_offset, seq_len) + + retrieve_index_offset = batch_idx * draft_token_num + + # Build retrieval index structure (reverse loop from draft_token_num-1 to 1) + for i in range(draft_token_num - 1, 0, -1): + current_token_idx = retrieve_index_offset + i + tl.store( + retrieve_index_ptr + batch_idx * draft_token_num + i, + current_token_idx, + ) + + parent_tb_idx = ( + tl.load(selected_index_ptr + batch_idx * selected_index_stride + (i - 1)) + // topk + ) + parent_position = 0 + found = 0 + + if parent_tb_idx == 0: + found = 1 + else: + parent_token_idx = tl.load( + parent_list_ptr + batch_idx * parent_list_stride + parent_tb_idx + ) + + # Find parent position + for pp in range(draft_token_num - 1): + if found == 0: + sel_idx = tl.load( + selected_index_ptr + batch_idx * selected_index_stride + pp + ) + if sel_idx == parent_token_idx: + parent_position = pp + 1 + found = 1 + + if found == 1: + # Update next token links + next_tok_addr = ( + retrieve_next_token_ptr + batch_idx * draft_token_num + parent_position + ) + next_tok = tl.load(next_tok_addr) + + if next_tok == -1: + tl.store(next_tok_addr, i) + else: + tl.store(next_tok_addr, i) + tl.store( + retrieve_next_sibling_ptr + batch_idx * draft_token_num + i, + next_tok, + ) + + tl.store(retrieve_index_ptr + batch_idx * draft_token_num, retrieve_index_offset) + + # Process all draft token indices for tree mask + for draft_tokenx in range(draft_token_num): + if tree_mask_mode == 0: # FULL_MASK + token_tree_idx = ( + seq_tree_idx + (seq_len + draft_token_num) * draft_tokenx + seq_len + 1 + ) + else: + token_tree_idx = ( + draft_token_num * draft_token_num * batch_idx + + draft_token_num * draft_tokenx + + 1 + ) + + tl.store(tree_mask_ptr + token_tree_idx - 1, 1) + for i in range(draft_token_num - 1): + tl.store(tree_mask_ptr + token_tree_idx + i, 0) + + if draft_tokenx > 0: + # Build tree path for draft_tokenx > 0 + cur_position = draft_tokenx - 1 + position = 0 + should_continue = 1 + + for _ in range(depth): + if should_continue: + position += 1 + tl.store(tree_mask_ptr + token_tree_idx + cur_position, 1) + + parent_tb_idx = ( + tl.load( + selected_index_ptr + + batch_idx * selected_index_stride + + cur_position + ) + // topk + ) + if parent_tb_idx == 0: + should_continue = 0 + else: + parent_token_idx = tl.load( + parent_list_ptr + + batch_idx * parent_list_stride + + parent_tb_idx + ) + + # Find cur_position for next iteration + found = 0 + for cp in range(draft_token_num - 1): + if found == 0: + if ( + tl.load( + selected_index_ptr + + batch_idx * selected_index_stride + + cp + ) + == parent_token_idx + ): + cur_position = cp + found = 1 + if found == 0: + should_continue = 0 + + tl.store( + positions_ptr + batch_idx * draft_token_num + draft_tokenx, + position + seq_len, + ) + + +def sgl_build_tree_kernel_triton( + parent_list: torch.Tensor, + selected_index: torch.Tensor, + verified_seq_len: torch.Tensor, + tree_mask: torch.Tensor, + positions: torch.Tensor, + retrieve_index: torch.Tensor, + retrieve_next_token: torch.Tensor, + retrieve_next_sibling: torch.Tensor, + topk: int, + depth: int, + draft_token_num: int, + tree_mask_mode: TreeMaskMode = TreeMaskMode.FULL_MASK, +): + """Triton-based implementation.""" + # TODO: Add support for QLEN_ONLY_BITPACKING mode + if tree_mask_mode == TreeMaskMode.QLEN_ONLY_BITPACKING: + raise NotImplementedError( + "QLEN_ONLY_BITPACKING is not supported in Triton implementation" + ) + + batch_size = verified_seq_len.shape[0] + seq_len_prefix_sum = torch.cumsum(verified_seq_len, dim=0) - verified_seq_len + + # Launch kernel with one program per batch item + grid = (batch_size,) + + sgl_build_tree_kernel_efficient_triton[grid]( + parent_list, + selected_index, + verified_seq_len, + seq_len_prefix_sum, + tree_mask, + positions, + retrieve_index, + retrieve_next_token, + retrieve_next_sibling, + topk=topk, + depth=depth, + draft_token_num=draft_token_num, + tree_mask_mode=int(tree_mask_mode), + batch_size=batch_size, + parent_list_stride=( + parent_list.stride(0) if parent_list.dim() > 1 else parent_list.shape[0] + ), + selected_index_stride=selected_index.stride(0), + ) diff --git a/python/sgl_kernel/speculative.py b/python/sgl_kernel/speculative.py index 2a5480c9e..cb20a96a2 100644 --- a/python/sgl_kernel/speculative.py +++ b/python/sgl_kernel/speculative.py @@ -1,3 +1,5 @@ +from enum import IntEnum + import torch from sgl_kernel.utils import get_xpu_stream @@ -58,6 +60,14 @@ def verify_tree_greedy( ) +class TreeMaskMode(IntEnum): + """Mirrors sglang.srt.speculative.eagle_utils.TreeMaskMode.""" + + FULL_MASK = 0 + QLEN_ONLY = 1 + QLEN_ONLY_BITPACKING = 2 + + def build_tree_kernel_efficient( parent_list: torch.Tensor, selected_index: torch.Tensor, @@ -70,7 +80,21 @@ def build_tree_kernel_efficient( topk: int, depth: int, draft_token_num: int, + tree_mask_mode: int = TreeMaskMode.FULL_MASK, ) -> None: + """Build the EAGLE draft-tree metadata in place. + + ``tree_mask`` layout depends on ``tree_mask_mode``: + * ``FULL_MASK`` -- rows of ``seq_len + draft_token_num`` per request, packed + back to back; only the trailing ``draft_token_num`` tree columns are + written here, so the caller owns the ``[0, seq_len)`` prefix fill. + * ``QLEN_ONLY`` -- a dense ``draft_token_num x draft_token_num`` block per + request. + ``QLEN_ONLY_BITPACKING`` is not implemented on XPU. + """ + if tree_mask_mode == TreeMaskMode.QLEN_ONLY_BITPACKING: + raise NotImplementedError("QLEN_ONLY_BITPACKING is not implemented on XPU") + torch.ops.sgl_kernel.build_tree_kernel_efficient.default( parent_list, selected_index, @@ -83,6 +107,7 @@ def build_tree_kernel_efficient( topk, depth, draft_token_num, + int(tree_mask_mode), ) diff --git a/src/sycl/SpecBuildTree.cpp b/src/sycl/SpecBuildTree.cpp new file mode 100644 index 000000000..60e2131ff --- /dev/null +++ b/src/sycl/SpecBuildTree.cpp @@ -0,0 +1,613 @@ +/* Copyright 2025 SGLang Team. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +// EAGLE speculative-decoding tree builder (build_tree_kernel_efficient). +// +// Semantics are those of the upstream Triton kernel +// (sglang/kernels/ops/speculative/spec_tree.py:sgl_build_tree_kernel_efficient_triton), +// but the work is reorganized for XPU: +// +// * Triton runs one *serial* program per batch item. Here a whole work-group +// owns one batch item and every draft token is a work-item, so the per-node +// work (ancestor walk, mask row, sibling links) runs in parallel. +// * The parent of each node is resolved once into `parent_pos` in shared local +// memory. Triton instead re-runs the O(draft_token_num) linear search over +// `selected_index` at every step of every ancestor walk, i.e. it does +// O(N^2 * depth) global loads per batch item where this does O(N^2) local +// ones (N = draft_token_num). +// * The sibling links are computed directly rather than by the serial +// head-insertion Triton uses. Head-inserting nodes N-1..1 leaves, for each +// parent p with ascending children i1= 1 corresponds to +// slot i-1 of `selected_index`. `parent_pos[i]` is the *node* index of i's +// parent, 0 when the parent is the root, and kParentNotFound when the parent +// token is absent from `selected_index` (Triton reads out of bounds / stops the +// walk in that case; we stop the walk and drop the link, which is well defined). + +#include +#include +#include + +#include + +#include "MemoryAccess.h" +#include "SYCLHelpers.h" +#include "Utils.h" +#include "sgl_kernel_export.h" + +namespace { + +// Mirrors sglang.srt.speculative.eagle_utils.TreeMaskMode. +enum class TreeMaskMode : int64_t { + FULL_MASK = 0, + QLEN_ONLY = 1, + QLEN_ONLY_BITPACKING = 2, +}; + +constexpr int32_t kParentNotFound = -1; +// Above this, a node's ancestor/child set no longer fits one uint64_t bitmask; +// the sibling/child lookup and the mask write both fall back to their O(N) +// per-node scan below instead. +constexpr int32_t kBitmaskFastPathMaxNodes = 64; + +// The mask write packs several tree_mask columns into one store, which only +// works if a bool is one byte and `true` is 0x01 (Itanium ABI / SPIR-V). +static_assert(sizeof(bool) == 1, "tree_mask pack stores assume 1-byte bool"); + +// Widest mask pack the host will instantiate, in 32-bit words: 4 words is 16 +// bytes, one Intel Xe LSC OWord message. +constexpr int32_t kMaxMaskPackWords = 4; + +// Expand the low 4 bits of `bits` into 4 bool bytes -- byte j is bit j as 0x00 +// or 0x01, exactly what `static_cast` would have stored, so one 32-bit +// word of a pack is one nibble of the ancestor mask. +// +// Multiplying by 0x00204081 (= sum_k 2^(7k), k = 0..3) puts bit j at bit 8j: +// the copy of the constant shifted by j contributes 2^(j + 7j), and no two +// contributions collide, so the product is a plain OR that the 0x01010101 mask +// then trims down to the byte lanes. Kept per nibble so this stays a 32-bit +// multiply -- Xe has no native 64-bit multiply. +inline uint32_t expand_mask_nibble(uint32_t bits) { + return ((bits & 0xFu) * 0x00204081u) & 0x01010101u; +} + +template +struct BuildTreeKernel : public __SYCL_KER_CONFIG_CONVENTION__ { + // Same sycl::vec + reinterpret_cast + store(offset, ptr) idiom as + // src/sycl/KVCache.cpp:44 -- a native sycl::vec is not a legal vector + // element type, and offset being in units of the whole pack is exactly the + // indexing the mask write wants. The width comes from the host (see + // can_vectorize_up_to at the launch site). + using MaskPack = sycl::vec; + static constexpr int32_t kPackCols = static_cast(sizeof(MaskPack)); + + BuildTreeKernel( + const int64_t* parent_list, + const int64_t* selected_index, + const seq_t* verified_seq_len, + bool* tree_mask, + int64_t* positions, + int64_t* retrieve_index, + int64_t* retrieve_next_token, + int64_t* retrieve_next_sibling, + int32_t topk, + int32_t depth, + int32_t draft_token_num, + int32_t parent_list_width, + int64_t parent_list_stride, + int64_t selected_index_stride, + bool full_mask, + int32_t row_blocks) + : parent_list_(parent_list), + selected_index_(selected_index), + verified_seq_len_(verified_seq_len), + tree_mask_(tree_mask), + positions_(positions), + retrieve_index_(retrieve_index), + retrieve_next_token_(retrieve_next_token), + retrieve_next_sibling_(retrieve_next_sibling), + topk_(topk), + depth_(depth), + draft_token_num_(draft_token_num), + parent_list_width_(parent_list_width), + parent_list_stride_(parent_list_stride), + selected_index_stride_(selected_index_stride), + full_mask_(full_mask), + row_blocks_(row_blocks) {} + + void sycl_ker_config_convention(sycl::handler& cgh) { + parent_pos_ = sycl::local_accessor(sycl::range<1>(draft_token_num_), cgh); + sel_local_ = sycl::local_accessor(sycl::range<1>(std::max(draft_token_num_ - 1, 1)), cgh); + ancestor_mask_ = sycl::local_accessor(sycl::range<1>(draft_token_num_), cgh); + child_ = sycl::local_accessor(sycl::range<1>(draft_token_num_), cgh); + } + + // Node index of `node`'s parent, or kParentNotFound. `sel_local` is this + // group's `selected_index` row staged in SLM (see operator()); every lane + // scans it, so keeping it local avoids an O(N^2) global-memory scan per + // group. + inline int32_t resolve_parent(const int64_t* sel_local, int64_t bid, int32_t node) const { + const int32_t parent_tb_idx = static_cast(sel_local[node - 1] / topk_); + if (parent_tb_idx == 0) { + return 0; // child of the root + } + // Triton indexes parent_list unconditionally; an out-of-range table index + // would read past the end, so treat it as "no parent" instead. + if (parent_tb_idx < 0 || parent_tb_idx >= parent_list_width_) { + return kParentNotFound; + } + const int64_t parent_token_idx = parent_list_[bid * parent_list_stride_ + parent_tb_idx]; +#pragma unroll 2 + for (int32_t pp = 0; pp < draft_token_num_ - 1; ++pp) { + if (sel_local[pp] == parent_token_idx) { + return pp + 1; + } + } + return kParentNotFound; + } + + void operator()(sycl::nd_item<1> item) const { + // Splitting one request's row range across `row_blocks_` groups (see the + // host-side launch) keeps small batches from leaving most compute units + // idle: phases 1-3 are cheap enough to redo per block, and only the mask + // write -- the part actually worth parallelizing further -- is + // partitioned by row. + // Xe has no native integer divide, so skip it entirely for the common + // case (row_blocks_ == 1, i.e. every batch large enough to already fill + // the machine) rather than pay an emulated div/mod on every launch for a + // split that isn't happening. + const int64_t group_id = item.get_group(0); + const int64_t bid = row_blocks_ == 1 ? group_id : group_id / row_blocks_; + const int32_t blk = row_blocks_ == 1 ? 0 : static_cast(group_id % row_blocks_); + const int32_t tid = static_cast(item.get_local_id(0)); + const int32_t lrange = static_cast(item.get_local_range(0)); + const int32_t num_nodes = draft_token_num_; + const int64_t seq_len = static_cast(verified_seq_len_[bid]); + + const int32_t rows_per_block = row_blocks_ == 1 ? num_nodes : (num_nodes + row_blocks_ - 1) / row_blocks_; + const int32_t row_start = blk * rows_per_block; + const int32_t row_end = sycl::min(row_start + rows_per_block, num_nodes); + + int32_t* parent_pos = parent_pos_.get_multi_ptr().get(); + int64_t* sel_local = sel_local_.get_multi_ptr().get(); + uint64_t* ancestor_mask = ancestor_mask_.get_multi_ptr().get(); + uint64_t* child = child_.get_multi_ptr().get(); + + // Stage this request's selected_index row in SLM once; resolve_parent's + // O(N) scan then hits local memory from every lane instead of re-reading + // global memory O(N^2) times per group. + const int64_t* sel_global = selected_index_ + bid * selected_index_stride_; +#pragma unroll 2 + for (int32_t p = tid; p < num_nodes - 1; p += lrange) { + sel_local[p] = sel_global[p]; + } + + // FULL_MASK rows are (seq_len + draft_token_num) wide and packed per batch + // item, so the row base needs sum(seq_len[0:bid]). One work-group owns one + // batch item, so reduce the prefix across the group instead of paying for a + // separate cumsum launch. + int64_t mask_base = 0; + int64_t row_stride = num_nodes; + int64_t col_offset = 0; + if (full_mask_) { + int64_t partial = 0; +#pragma unroll 2 + for (int64_t b = tid; b < bid; b += lrange) { + partial += static_cast(verified_seq_len_[b]); + } + const int64_t seq_len_prefix_sum = sycl::reduce_over_group(item.get_group(), partial, sycl::plus()); + mask_base = static_cast(num_nodes) * num_nodes * bid + seq_len_prefix_sum * num_nodes; + row_stride = seq_len + num_nodes; + col_offset = seq_len; + } else { + mask_base = static_cast(num_nodes) * num_nodes * bid; + } + sycl::group_barrier(item.get_group()); + + // Phase 1: resolve every node's parent into SLM. +#pragma unroll 2 + for (int32_t i = tid; i < num_nodes; i += lrange) { + parent_pos[i] = (i == 0) ? kParentNotFound : resolve_parent(sel_local, bid, i); + } + sycl::group_barrier(item.get_group()); + + const int64_t out_base = bid * num_nodes; + + // Phase 2: retrieve_index is the identity over the flattened draft tokens. +#pragma unroll 2 + for (int32_t i = tid; i < num_nodes; i += lrange) { + retrieve_index_[out_base + i] = out_base + i; + } + + // Phase 3: first child of node i, and i's next sibling. + if (num_nodes <= kBitmaskFastPathMaxNodes) { + // Same trick as the mask write below: a node's children fit one + // uint64_t bitmask, so "first child" / "next sibling" become a masked + // ctz instead of an O(N) scan per node (O(N^2) per group total). +#pragma unroll 2 + for (int32_t i = tid; i < num_nodes; i += lrange) { + child[i] = 0ull; + } + sycl::group_barrier(item.get_group()); +#pragma unroll 2 + for (int32_t i = tid; i < num_nodes; i += lrange) { + if (i > 0 && parent_pos[i] != kParentNotFound) { + sycl::atomic_ref< + uint64_t, + sycl::memory_order::relaxed, + sycl::memory_scope::work_group, + sycl::access::address_space::local_space>(child[parent_pos[i]]) + .fetch_or(1ull << i); + } + } + sycl::group_barrier(item.get_group()); +#pragma unroll 2 + for (int32_t i = tid; i < num_nodes; i += lrange) { + const uint64_t kids = child[i]; + retrieve_next_token_[out_base + i] = kids ? static_cast(sycl::ctz(kids)) : -1; + + int64_t next_sibling = -1; + if (i > 0 && parent_pos[i] != kParentNotFound) { + // Bits above i among i's siblings (i itself is always set); i == 63 + // is the top of the bitmask, so it never has a sibling above it. + const uint64_t above = (i >= 63) ? 0ull : (child[parent_pos[i]] & ~((1ull << (i + 1)) - 1)); + if (above) { + next_sibling = static_cast(sycl::ctz(above)); + } + } + retrieve_next_sibling_[out_base + i] = next_sibling; + } + } else { +#pragma unroll 2 + for (int32_t i = tid; i < num_nodes; i += lrange) { + int64_t next_token = -1; +#pragma unroll 2 + for (int32_t j = 1; j < num_nodes; ++j) { + if (parent_pos[j] == i) { + next_token = j; + break; + } + } + retrieve_next_token_[out_base + i] = next_token; + + int64_t next_sibling = -1; + if (i > 0 && parent_pos[i] != kParentNotFound) { +#pragma unroll 2 + for (int32_t j = i + 1; j < num_nodes; ++j) { + if (parent_pos[j] == parent_pos[i]) { + next_sibling = j; + break; + } + } + } + retrieve_next_sibling_[out_base + i] = next_sibling; + } + } + + // Phase 4: mask row for node i, plus its position (= depth in the tree). + // Every tree cell is written here, so the caller's prefix fill (which only + // supplies the [0, seq_len) columns) never affects this block. The mask is + // `bs*N^2` bytes -- the single largest thing the kernel touches -- so its + // access pattern dominates the kernel's bandwidth. + if (num_nodes <= kBitmaskFastPathMaxNodes) { + // Each node's ancestor set fits one uint64_t: build it in a register, + // stash it to SLM, then sweep rows with consecutive lanes writing + // consecutive columns of the same row (coalesced), instead of each lane + // owning a whole row (row_stride apart -- the worst-case store pattern). +#pragma unroll 2 + for (int32_t i = row_start + tid; i < row_end; i += lrange) { + if (i == 0) { + positions_[out_base] = seq_len; + ancestor_mask[0] = 1ull; // every draft token attends the root + continue; + } + uint64_t m = 1ull; + int32_t node_depth = 0; + int32_t cur = i; +#pragma unroll 2 + for (int32_t d = 0; d < depth_; ++d) { + ++node_depth; + m |= 1ull << cur; + const int32_t parent = parent_pos[cur]; + if (parent <= 0) { + break; // 0 -> reached the root (already marked); -1 -> unresolved + } + cur = parent; + } + ancestor_mask[i] = m; + positions_[out_base + i] = seq_len + node_depth; + } + sycl::group_barrier(item.get_group()); + // Write kPackCols columns per store instead of one bool byte at a time: + // kPackCols consecutive cells of a row are just expanded nibbles of `m`. + // Two things have to line up for that to be a win. + // + // * Lane count. lrange >= num_nodes always (see the host launch), so the + // byte loop below already has every lane writing a different column of + // the same row -- fully coalesced, nothing idle. Widening the store + // without also handing each lane a *different row* would just idle + // lanes: only num_nodes/kPackCols of them would have a pack to write. + // So the packed path folds rows into the lane space -- lane tid owns + // pack (tid % packs_per_row) of row (tid / packs_per_row) and strides + // by lrange/packs_per_row rows. + // * Alignment. A row starts col_offset (= seq_len under FULL_MASK) into + // the row and vec::store needs natural alignment, so the pack path is + // only legal when every row this group touches is pack-aligned. + // mask_base + col_offset and row_stride are both group-uniform, so one + // check covers all of them and the branch never diverges; when it + // fails the byte loop writes the whole mask and correctness never + // depends on alignment. + const int32_t packs_per_row = num_nodes / kPackCols; + const bool pack_aligned = packs_per_row > 0 && ((mask_base + col_offset) % kPackCols == 0) && + (row_stride % kPackCols == 0) && + (reinterpret_cast(tree_mask_) % static_cast(kPackCols) == 0); + const int32_t packed_cols = pack_aligned ? packs_per_row * kPackCols : 0; + + if (pack_aligned) { + // packs_per_row <= num_nodes / 4 <= 16 < 32 <= lrange, so row_step >= 2 + // and every pack of every row is covered; the max() is just a guard. + const int32_t row_step = sycl::max(lrange / packs_per_row, 1); + const int32_t lane_row = tid / packs_per_row; + const int32_t lane_pack = tid - lane_row * packs_per_row; + const int32_t lane_col = lane_pack * kPackCols; + if (lane_row < row_step) { +#pragma unroll 2 + for (int32_t row = row_start + lane_row; row < row_end; row += row_step) { + const uint64_t m = ancestor_mask[row]; + auto* pack_ptr = reinterpret_cast(tree_mask_ + mask_base + row_stride * row + col_offset); + MaskPack pack; +#pragma unroll + for (int32_t w = 0; w < kPackWords; ++w) { + pack[w] = expand_mask_nibble(static_cast(m >> (lane_col + 4 * w))); + } + pack.store(static_cast(lane_pack), pack_ptr); + } + } + } + + // Columns the packs did not cover: all of them when the pack path is off, + // otherwise just the num_nodes % kPackCols tail. + if (packed_cols < num_nodes) { +#pragma unroll 2 + for (int32_t row = row_start; row < row_end; ++row) { + bool* row_ptr = tree_mask_ + mask_base + row_stride * row + col_offset; + const uint64_t m = ancestor_mask[row]; +#pragma unroll 2 + for (int32_t c = packed_cols + tid; c < num_nodes; c += lrange) { + row_ptr[c] = static_cast((m >> c) & 1ull); + } + } + } + } else { +#pragma unroll 2 + for (int32_t i = row_start + tid; i < row_end; i += lrange) { + bool* row = tree_mask_ + mask_base + row_stride * i + col_offset; +#pragma unroll 2 + for (int32_t c = 0; c < num_nodes; ++c) { + row[c] = false; + } + row[0] = true; // every draft token attends the root + + if (i == 0) { + positions_[out_base] = seq_len; + continue; + } + + int32_t node_depth = 0; + int32_t cur = i; +#pragma unroll 2 + for (int32_t d = 0; d < depth_; ++d) { + ++node_depth; + row[cur] = true; + const int32_t parent = parent_pos[cur]; + if (parent <= 0) { + break; // 0 -> reached the root (already marked); -1 -> unresolved + } + cur = parent; + } + positions_[out_base + i] = seq_len + node_depth; + } + } + } + + const int64_t* parent_list_; + const int64_t* selected_index_; + const seq_t* verified_seq_len_; + bool* tree_mask_; + int64_t* positions_; + int64_t* retrieve_index_; + int64_t* retrieve_next_token_; + int64_t* retrieve_next_sibling_; + int32_t topk_; + int32_t depth_; + int32_t draft_token_num_; + int32_t parent_list_width_; + int64_t parent_list_stride_; + int64_t selected_index_stride_; + bool full_mask_; + int32_t row_blocks_; + + sycl::local_accessor parent_pos_; + sycl::local_accessor sel_local_; + sycl::local_accessor ancestor_mask_; + sycl::local_accessor child_; +}; + +} // namespace + +SGL_KERNEL_EXPORT void build_tree_kernel_efficient( + at::Tensor parent_list, + at::Tensor selected_index, + at::Tensor verified_seq_len, + at::Tensor tree_mask, + at::Tensor positions, + at::Tensor retrive_index, + at::Tensor retrive_next_token, + at::Tensor retrive_next_sibling, + int64_t topk, + int64_t depth, + int64_t draft_token_num, + int64_t tree_mask_mode) { + CHECK_INPUT(parent_list); + CHECK_INPUT(selected_index); + CHECK_INPUT(verified_seq_len); + CHECK_INPUT(tree_mask); + CHECK_INPUT(positions); + CHECK_INPUT(retrive_index); + CHECK_INPUT(retrive_next_token); + CHECK_INPUT(retrive_next_sibling); + + const auto mode = static_cast(tree_mask_mode); + TORCH_CHECK( + mode == TreeMaskMode::FULL_MASK || mode == TreeMaskMode::QLEN_ONLY, + "build_tree_kernel_efficient: unsupported tree_mask_mode ", + tree_mask_mode, + " (QLEN_ONLY_BITPACKING is not implemented on XPU)"); + + TORCH_CHECK(topk > 0, "build_tree_kernel_efficient: topk must be positive, got ", topk); + TORCH_CHECK(depth > 0, "build_tree_kernel_efficient: depth must be positive, got ", depth); + TORCH_CHECK( + draft_token_num > 0, "build_tree_kernel_efficient: draft_token_num must be positive, got ", draft_token_num); + + TORCH_CHECK(parent_list.scalar_type() == at::kLong, "parent_list must be int64"); + TORCH_CHECK(selected_index.scalar_type() == at::kLong, "selected_index must be int64"); + TORCH_CHECK(tree_mask.scalar_type() == at::kBool, "tree_mask must be bool"); + TORCH_CHECK(positions.scalar_type() == at::kLong, "positions must be int64"); + TORCH_CHECK(retrive_index.scalar_type() == at::kLong, "retrive_index must be int64"); + TORCH_CHECK(retrive_next_token.scalar_type() == at::kLong, "retrive_next_token must be int64"); + TORCH_CHECK(retrive_next_sibling.scalar_type() == at::kLong, "retrive_next_sibling must be int64"); + + const int64_t bs = verified_seq_len.numel(); + TORCH_CHECK( + selected_index.dim() == 2 && selected_index.size(0) == bs, + "selected_index must be (batch_size, *), got ", + selected_index.sizes(), + " for batch_size ", + bs); + TORCH_CHECK( + selected_index.size(1) >= draft_token_num - 1, + "selected_index must hold at least draft_token_num - 1 = ", + draft_token_num - 1, + " entries per request, got ", + selected_index.size(1)); + + for (const auto& out : {positions, retrive_index, retrive_next_token, retrive_next_sibling}) { + TORCH_CHECK( + out.numel() == bs * draft_token_num, + "build_tree_kernel_efficient: output tensors must hold batch_size * draft_token_num = ", + bs * draft_token_num, + " elements, got ", + out.numel()); + } + + // parent_list is (bs, width); organize_draft_results emits (bs, 0) when there + // are no non-root parents (single-step MTP), which must stay legal. + int64_t parent_list_stride = 0; + int64_t parent_list_width = 0; + if (parent_list.dim() > 1) { + TORCH_CHECK( + parent_list.size(0) == bs, "parent_list must be (batch_size, *), got ", parent_list.sizes(), " for bs ", bs); + parent_list_stride = parent_list.stride(0); + parent_list_width = parent_list.size(1); + } else { + parent_list_stride = parent_list.numel(); + parent_list_width = parent_list.numel(); + } + + const bool full_mask = (mode == TreeMaskMode::FULL_MASK); + const int64_t expected_mask_numel = bs * draft_token_num * draft_token_num; + TORCH_CHECK( + tree_mask.numel() >= expected_mask_numel, + "build_tree_kernel_efficient: tree_mask needs at least ", + expected_mask_numel, + " elements for tree_mask_mode ", + tree_mask_mode, + ", got ", + tree_mask.numel()); + + if (bs == 0) { + return; + } + + auto& queue = dpcppGetCurrentQueue(); + // One work-group per request; one work-item per draft token (nodes beyond the + // work-group size are handled by the strided loops in the kernel). + const int64_t max_wg = dpcppMaxWorkGroupSize(); + const int64_t local_range = std::min(std::max((draft_token_num + 31) / 32 * 32, 32), max_wg); + + // A batch too small to fill the machine on its own leaves most subslices + // idle -- one work-group per request regardless of how much row work that + // request has. Split each request's rows across a few extra blocks so a + // small batch still spreads across multiple subslices; phases 1-3 get + // redone per block (cheap, O(draft_token_num)), only the mask write is + // actually partitioned. Subslice count, not EU count, bounds how many + // work-groups can genuinely run concurrently (a group can't span + // subslices), so using EU count here would over-split large batches too. + const int64_t num_subslices = queue.get_device().get_info() * + queue.get_device().get_info(); + const int64_t max_row_blocks = std::max((draft_token_num + 31) / 32, 1); + const int64_t row_blocks = bs < num_subslices ? std::min(max_row_blocks, (num_subslices + bs - 1) / bs) : 1; + + // How many 32-bit words the mask write packs per store. can_vectorize_up_to + // takes the device's preferred int vector width and caps it by what the + // tree_mask base pointer's alignment actually permits; the per-row offset is + // re-checked in the kernel, which falls back to byte stores when it does not + // line up. A pack also has to fit inside one row. + int pack_words = std::min( + can_vectorize_up_to( + dpcppGetDeviceIdOfCurrentQueue(), reinterpret_cast(tree_mask.data_ptr())), + kMaxMaskPackWords); + while (pack_words > 1 && draft_token_num < 4 * pack_words) { + pack_words >>= 1; + } + + AT_DISPATCH_INDEX_TYPES(verified_seq_len.scalar_type(), "build_tree_kernel_efficient", [&] { + auto launch = [&](auto pack_words_tag) { + using Kernel = BuildTreeKernel; + Kernel kernel( + parent_list.data_ptr(), + selected_index.data_ptr(), + verified_seq_len.data_ptr(), + tree_mask.data_ptr(), + positions.data_ptr(), + retrive_index.data_ptr(), + retrive_next_token.data_ptr(), + retrive_next_sibling.data_ptr(), + static_cast(topk), + static_cast(depth), + static_cast(draft_token_num), + static_cast(parent_list_width), + parent_list_stride, + selected_index.stride(0), + full_mask, + static_cast(row_blocks)); + sycl_kernel_submit(bs * row_blocks * local_range, local_range, queue, kernel); + }; + switch (pack_words) { + case 4: + launch(std::integral_constant{}); + break; + case 2: + launch(std::integral_constant{}); + break; + default: + launch(std::integral_constant{}); + break; + } + }); +} diff --git a/src/torch_extension_sycl.cc b/src/torch_extension_sycl.cc index 2c58255f4..32c6feb4a 100644 --- a/src/torch_extension_sycl.cc +++ b/src/torch_extension_sycl.cc @@ -79,6 +79,15 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { "maybe_min_p_arr, float min_p_val, bool deterministic, Generator? gen) -> ()"); m.impl("min_p_sampling_from_probs", torch::kXPU, &min_p_sampling_from_probs); + /* + * Speculative decoding (EAGLE) + */ + m.def( + "build_tree_kernel_efficient(Tensor parent_list, Tensor selected_index, Tensor verified_seq_len, " + "Tensor! tree_mask, Tensor! positions, Tensor! retrive_index, Tensor! retrive_next_token, " + "Tensor! retrive_next_sibling, int topk, int depth, int draft_token_num, int tree_mask_mode=0) -> ()"); + m.impl("build_tree_kernel_efficient", torch::kXPU, &build_tree_kernel_efficient); + /* * Fast radix top-k (DeepSeek V3.2 indexer) */ diff --git a/tests/speculative/test_build_eagle_tree.py b/tests/speculative/test_build_eagle_tree.py new file mode 100644 index 000000000..8d9e182e6 --- /dev/null +++ b/tests/speculative/test_build_eagle_tree.py @@ -0,0 +1,399 @@ +import unittest + +import pytest +import torch +import utils +from sgl_kernel.eagle_utils import build_tree_kernel_efficient, organize_draft_results + +device = utils.get_device() + + +class TestBuildEagleTree(unittest.TestCase): + """Unit tests for build_eagle_tree functionality.""" + + def test_build_tree_kernel_efficient(self): + """Test the build_tree_kernel_efficient function with known inputs and expected outputs.""" + bonus_tokens = torch.tensor([29974, 13], device=device, dtype=torch.int32) + score_list = [ + torch.tensor( + [ + [[7.1127e-01, 2.8292e-01, 2.2995e-03, 1.7357e-03]], + [[9.7476e-01, 2.2219e-02, 6.5031e-04, 1.3212e-04]], + ], + dtype=torch.float32, + device=device, + ), + torch.tensor( + [ + [ + [6.9142e-01, 1.2863e-02, 1.6873e-03, 1.1871e-03], + [2.4787e-01, 1.8818e-02, 1.4204e-02, 9.2235e-04], + [2.2971e-03, 1.6700e-06, 1.8737e-07, 8.3146e-08], + [1.2771e-03, 2.4374e-04, 1.7832e-04, 1.1947e-05], + ], + [ + [8.4832e-02, 6.6068e-02, 5.8304e-02, 5.7851e-02], + [2.3616e-03, 1.1243e-03, 5.4368e-04, 2.7768e-04], + [2.5286e-04, 1.5578e-04, 2.8817e-05, 1.2888e-05], + [1.2834e-04, 2.5417e-06, 1.1279e-06, 1.6088e-08], + ], + ], + dtype=torch.float32, + device=device, + ), + torch.tensor( + [ + [ + [6.6438e-01, 2.6997e-02, 2.4236e-05, 4.0821e-06], + [2.4402e-01, 2.8409e-03, 5.0935e-04, 2.9022e-04], + [1.6178e-02, 2.0567e-03, 4.5892e-04, 3.0034e-05], + [1.3023e-02, 5.0497e-04, 3.6371e-04, 8.7750e-05], + ], + [ + [2.3263e-02, 2.0054e-02, 9.3990e-03, 2.7783e-03], + [6.4156e-02, 5.5506e-04, 1.0429e-04, 9.7211e-05], + [4.9950e-02, 5.0630e-03, 9.0068e-04, 3.3656e-04], + [7.5817e-03, 8.5731e-04, 6.9972e-04, 6.0793e-04], + ], + ], + dtype=torch.float32, + device=device, + ), + torch.tensor( + [ + [ + [6.6420e-01, 1.0525e-04, 6.5864e-05, 1.2253e-06], + [1.3019e-01, 1.0461e-01, 5.2083e-03, 1.6777e-03], + [2.0103e-02, 6.7335e-03, 1.2625e-04, 1.0364e-05], + [1.5142e-02, 7.0819e-04, 9.6595e-05, 8.7951e-05], + ], + [ + [5.8608e-02, 1.8840e-03, 7.8535e-04, 4.4400e-04], + [1.2185e-02, 2.0684e-03, 1.7418e-03, 1.4327e-03], + [6.2455e-03, 6.1487e-03, 2.6862e-03, 1.8034e-03], + [1.8590e-03, 1.6151e-03, 1.2481e-03, 3.6038e-04], + ], + ], + dtype=torch.float32, + device=device, + ), + ] + token_list = [ + torch.tensor( + [[29896, 29906, 29900, 29945], [13, 2, 29871, 28956]], + dtype=torch.int64, + device=device, + ), + torch.tensor( + [ + [ + 29889, + 29974, + 29945, + 29900, + 29974, + 29922, + 29930, + 29958, + 29889, + 29974, + 29930, + 29945, + 29974, + 29922, + 29930, + 29958, + ], + [ + 22550, + 4136, + 16492, + 8439, + 29871, + 2, + 3001, + 13, + 2, + 13, + 29906, + 29946, + 2, + 13, + 29871, + 259, + ], + ], + device=device, + ), + torch.tensor( + [ + [ + 29946, + 29945, + 29953, + 29906, + 29896, + 29945, + 29900, + 29906, + 29896, + 29945, + 29906, + 29953, + 29896, + 29945, + 29906, + 29946, + ], + [ + 29871, + 2, + 29901, + 29889, + 29871, + 2, + 395, + 259, + 29901, + 29871, + 2, + 29889, + 3001, + 1234, + 7146, + 2186, + ], + ], + device=device, + ), + torch.tensor( + [ + [ + 29946, + 29974, + 29945, + 29930, + 29889, + 29922, + 29974, + 29930, + 29974, + 29946, + 29930, + 29922, + 29889, + 29974, + 29945, + 29922, + ], + [ + 29941, + 29906, + 2, + 29946, + 29871, + 450, + 319, + 14990, + 29946, + 29941, + 2, + 29906, + 29871, + 2, + 3001, + 13, + ], + ], + device=device, + ), + ] + parents_list = [ + torch.tensor( + [[-1, 0, 1, 2, 3], [-1, 0, 1, 2, 3]], + dtype=torch.int64, + device=device, + ), + torch.tensor( + [[4, 8, 9, 10], [4, 5, 6, 7]], dtype=torch.int64, device=device + ), + torch.tensor( + [[20, 24, 21, 28], [24, 28, 20, 21]], + dtype=torch.int64, + device=device, + ), + torch.tensor( + [[36, 40, 41, 44], [36, 40, 44, 45]], + dtype=torch.int64, + device=device, + ), + ] + seq_lens = torch.tensor([5, 10], dtype=torch.int64, device=device) + topk = 4 + depth = 4 + num_draft_token = 8 + + parent_list, top_scores_index, draft_tokens = organize_draft_results( + score_list, token_list, parents_list, num_draft_token + ) + + ( + tree_mask, + position, + retrieve_index, + retrieve_next_token, + retrieve_next_sibling, + draft_tokens, + ) = build_tree_kernel_efficient( + bonus_tokens=bonus_tokens, + parent_list=parent_list, + top_scores_index=top_scores_index, + draft_tokens=draft_tokens, + seq_lens=seq_lens, + seq_lens_sum=torch.sum(seq_lens).item(), + topk=topk, + spec_steps=depth, + num_verify_tokens=num_draft_token, + ) + + # Verify expected outputs + self.assertEqual( + position.tolist(), + [5, 6, 6, 7, 7, 8, 8, 9, 10, 11, 12, 12, 12, 12, 13, 14], + "Position tensor does not match expected values", + ) + self.assertEqual( + retrieve_index.tolist(), + [ + [0, 1, 2, 3, 4, 5, 6, 7], + [8, 9, 10, 11, 12, 13, 14, 15], + ], + "Retrieve index tensor does not match expected values", + ) + self.assertEqual( + retrieve_next_token.tolist(), + [ + [1, 3, 4, 5, 6, 7, -1, -1], + [1, 2, -1, 6, -1, -1, 7, -1], + ], + "Retrieve next token tensor does not match expected values", + ) + self.assertEqual( + retrieve_next_sibling.tolist(), + [ + [-1, 2, -1, -1, -1, -1, -1, -1], + [-1, -1, 3, 4, 5, -1, -1, -1], + ], + "Retrieve next sibling tensor does not match expected values", + ) + self.assertEqual( + draft_tokens.tolist(), + [ + 29974, + 29896, + 29906, + 29889, + 29974, + 29946, + 29896, + 29946, + 13, + 13, + 22550, + 4136, + 16492, + 8439, + 29871, + 29941, + ], + "Draft tokens tensor does not match expected values", + ) + + def test_skip_prefix_fill_preserves_tree_blocks(self): + """fill_prefix_mask=False must leave every kernel-written cell intact. + + The fill only supplies the [0, seq_len) prefix columns; the qlen x qlen + tree block comes from the kernel and must be identical either way. + """ + bs, topk, spec_steps, num_draft_token = 2, 1, 3, 4 + seq_lens = torch.tensor([5, 10], dtype=torch.int64, device=device) + seq_lens_sum = int(seq_lens.sum().item()) + # topk=1 chain: token i descends from i-1; index 0 is the root. + parent_list = torch.tensor([[0, 0, 1]] * bs, dtype=torch.int64, device=device) + top_scores_index = torch.tensor( + [[0, 1, 2]] * bs, dtype=torch.int64, device=device + ) + draft_tokens = torch.arange( + bs * (num_draft_token - 1), dtype=torch.int64, device=device + ).view(bs, -1) + bonus_tokens = torch.tensor([101, 102], dtype=torch.int32, device=device) + mask_numel = seq_lens_sum * num_draft_token + num_draft_token**2 * bs + + def build(fill_prefix_mask): + # All-False start matches the real preallocated scratch: a skipped + # fill leaves the prefix stale-False. + tree_mask_buf = torch.zeros((mask_numel,), dtype=torch.bool, device=device) + return build_tree_kernel_efficient( + bonus_tokens=bonus_tokens, + parent_list=parent_list, + top_scores_index=top_scores_index, + draft_tokens=draft_tokens, + seq_lens=seq_lens, + seq_lens_sum=seq_lens_sum, + topk=topk, + spec_steps=spec_steps, + num_verify_tokens=num_draft_token, + tree_mask_buf=tree_mask_buf, + fill_prefix_mask=fill_prefix_mask, + ) + + def split_rows(tree_mask): + """Flat mask -> (all prefix columns, all tree-block cells).""" + prefixes, blocks = [], [] + offset = 0 + for seq_len in seq_lens.tolist(): + row_len = seq_len + num_draft_token + for tid in range(num_draft_token): + row = tree_mask[ + offset + row_len * tid : offset + row_len * (tid + 1) + ] + prefixes.append(row[:seq_len]) + blocks.append(row[seq_len:]) + offset += row_len * num_draft_token + return torch.cat(prefixes), torch.cat(blocks) + + filled = build(fill_prefix_mask=True) + skipped = build(fill_prefix_mask=False) + + filled_prefix, filled_blocks = split_rows(filled[0]) + skipped_prefix, skipped_blocks = split_rows(skipped[0]) + + self.assertTrue( + torch.equal(filled_blocks, skipped_blocks), + "Tree blocks diverged: the kernel must write every tree cell " + "regardless of the prefix fill", + ) + # Anti-vacuous: proves the two runs really differ on the prefix. + self.assertTrue(filled_prefix.all(), "Fill did not mark the prefix columns") + self.assertFalse( + skipped_prefix.any(), "Skipped fill unexpectedly touched the prefix columns" + ) + + for idx, name in enumerate( + ( + "positions", + "retrieve_index", + "retrieve_next_token", + "retrieve_next_sibling", + "draft_tokens", + ), + start=1, + ): + self.assertTrue( + torch.equal(filled[idx], skipped[idx]), + f"{name} diverged between filled and skipped runs", + ) + + +if __name__ == "__main__": + unittest.main()