chore: import upstream snapshot with attribution

python/sglang/kernels/ops/speculative/__init__.py

@@ -0,0 +1,33 @@
+"""Speculative-decoding kernels (Triton).
+
+The Triton kernels migrated here live in this package
+(``sglang.kernels.ops.speculative.<module>``); import them from there. Their
+``KernelSpec`` metadata is registered below for inventory (backend = Triton).
+"""
+
+from sglang.kernels.registry import register_kernel
+from sglang.kernels.spec import KernelBackend, KernelSpec
+
+# (module, public_fn) migrated from speculative/triton_ops.
+_TRITON_KERNELS = [
+    ("cache_locs", "assign_req_to_token_pool_func"),
+    ("cache_locs", "assign_extend_cache_locs_func"),
+    ("cache_locs", "generate_draft_decode_kv_indices"),
+    ("eagle", "fill_bonus_tokens"),
+    ("eagle", "fill_accept_out_cache_loc"),
+    ("gather_spec_extras", "gather_spec_extras"),
+    ("multi_layer_eagle", "rotate_input_ids_triton"),
+    ("spec_tree", "sgl_build_tree_kernel_efficient_triton"),
+    ("spec_tree", "verify_tree_greedy_kernel_triton"),
+]
+for _mod, _fn in _TRITON_KERNELS:
+    register_kernel(
+        KernelSpec(
+            op=f"speculative.{_fn}",
+            backend=KernelBackend.TRITON,
+            target=f"sglang.kernels.ops.speculative.{_mod}:{_fn}",
+        )
+    )
+del _mod, _fn
+
+__all__ = []

python/sglang/kernels/ops/speculative/cache_locs.py

@@ -0,0 +1,418 @@
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.utils import (
+    is_cpu,
+    is_cuda,
+    is_hip,
+    is_musa,
+    is_npu,
+    is_xpu,
+    next_power_of_2,
+)
+
+_is_cpu = is_cpu()
+_is_cuda = is_cuda()
+_is_hip = is_hip()
+_is_npu = is_npu()
+_is_musa = is_musa()
+_is_xpu = is_xpu()
+
+if _is_cpu:
+    from sgl_kernel import assign_extend_cache_locs_cpu, assign_req_to_token_pool_cpu
+
+
+@triton.jit
+def assign_req_to_token_pool(
+    req_pool_indices,
+    req_to_token,
+    start_offset,
+    end_offset,
+    out_cache_loc,
+    pool_len: tl.constexpr,
+    bs_upper: tl.constexpr,
+):
+    BLOCK_SIZE: tl.constexpr = 32
+    pid = tl.program_id(axis=0)
+    kv_start = tl.load(start_offset + pid)
+    kv_end = tl.load(end_offset + pid)
+    token_pool = req_to_token + tl.load(req_pool_indices + pid) * pool_len
+
+    length_offset = tl.arange(0, bs_upper)
+    start = tl.load(start_offset + length_offset, mask=length_offset < pid, other=0)
+    end = tl.load(end_offset + length_offset, mask=length_offset < pid, other=0)
+    out_offset = tl.sum(end - start, axis=0)
+
+    out_cache_ptr = out_cache_loc + out_offset
+
+    save_offset = tl.arange(0, BLOCK_SIZE) + kv_start
+    load_offset = tl.arange(0, BLOCK_SIZE)
+
+    num_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE)
+    for _ in range(num_loop):
+        mask = save_offset < kv_end
+        data = tl.load(out_cache_ptr + load_offset, mask=mask)
+        tl.store(token_pool + save_offset, data, mask=mask)
+        save_offset += BLOCK_SIZE
+        load_offset += BLOCK_SIZE
+
+
+def assign_req_to_token_pool_func(
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    start_offset: torch.Tensor,
+    end_offset: torch.Tensor,
+    out_cache_loc: torch.Tensor,
+    batch_size: int,
+):
+    if _is_cpu:
+        assign_req_to_token_pool_cpu(
+            req_pool_indices,
+            req_to_token,
+            start_offset,
+            end_offset,
+            out_cache_loc,
+            req_to_token.shape[1],
+        )
+        return
+    assign_req_to_token_pool[(batch_size,)](
+        req_pool_indices,
+        req_to_token,
+        start_offset,
+        end_offset,
+        out_cache_loc,
+        req_to_token.shape[1],
+        next_power_of_2(batch_size),
+    )
+
+
+@triton.jit
+def assign_draft_cache_locs_contiguous(
+    req_pool_indices,
+    req_to_token,
+    seq_lens,
+    out_cache_loc,
+    pool_len: tl.constexpr,
+    topk: tl.constexpr,
+    speculative_num_steps: tl.constexpr,
+):
+    BLOCK_SIZE: tl.constexpr = 128
+    pid = tl.program_id(axis=0)
+
+    copy_len = topk * speculative_num_steps
+    out_cache_ptr = out_cache_loc + pid * topk * speculative_num_steps
+
+    # Copy from req_to_token to out_cache_loc
+    kv_start = tl.load(seq_lens + pid)
+    token_pool = req_to_token + tl.load(req_pool_indices + pid) * pool_len
+    num_loop = tl.cdiv(copy_len, BLOCK_SIZE)
+    for i in range(num_loop):
+        copy_offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
+        mask = copy_offset < copy_len
+        data = tl.load(token_pool + kv_start + copy_offset, mask=mask)
+        tl.store(out_cache_ptr + copy_offset, data, mask=mask)
+
+
+@triton.jit
+def generate_draft_decode_kv_indices(
+    req_pool_indices,
+    req_to_token,
+    paged_kernel_lens,
+    kv_indices,
+    kv_indptr,
+    positions,
+    pool_len: tl.constexpr,
+    kv_indices_stride: tl.constexpr,
+    kv_indptr_stride: tl.constexpr,
+    bs_upper: tl.constexpr,
+    iter_upper: tl.constexpr,
+    num_tokens_upper: tl.constexpr,
+    page_size: tl.constexpr,
+):
+    BLOCK_SIZE: tl.constexpr = 128
+    iters = tl.program_id(axis=0)
+    bid = tl.program_id(axis=1)
+    topk_id = tl.program_id(axis=2)
+
+    num_steps = tl.num_programs(axis=0)
+    num_seqs = tl.num_programs(axis=1)
+    topk = tl.num_programs(axis=2)
+
+    kv_indices += kv_indices_stride * iters
+    kv_indptr += kv_indptr_stride * iters
+    iters += 1
+
+    load_offset = tl.arange(0, bs_upper)
+    seq_lens = tl.load(paged_kernel_lens + load_offset, mask=load_offset < bid, other=0)
+    seq_len = tl.load(paged_kernel_lens + bid)
+    cum_seq_len = tl.sum(seq_lens)
+
+    # Update kv_indices
+    kv_offset = cum_seq_len * topk + bid * iters * topk + topk_id * (seq_len + iters)
+    kv_ptr = kv_indices + kv_offset
+    token_pool_ptr = req_to_token + tl.load(req_pool_indices + bid) * pool_len
+
+    kv_offset = tl.arange(0, BLOCK_SIZE)
+    num_loop = tl.cdiv(seq_len, BLOCK_SIZE)
+    for _ in range(num_loop):
+        mask = kv_offset < seq_len
+        data = tl.load(token_pool_ptr + kv_offset, mask=mask)
+        tl.store(kv_ptr + kv_offset, data, mask=mask)
+        kv_offset += BLOCK_SIZE
+
+    extend_offset = tl.arange(0, iter_upper)
+    if page_size == 1 or topk == 1:
+        extend_data = tl.load(
+            token_pool_ptr + seq_len + topk_id * num_steps + tl.arange(0, iter_upper),
+            mask=extend_offset < iters,
+        )
+    else:
+        prefix_len = seq_len
+        last_page_len = prefix_len % page_size
+        num_new_pages_per_topk = (
+            last_page_len + num_steps + page_size - 1
+        ) // page_size
+        prefix_base = seq_len // page_size * page_size
+        start = (
+            prefix_base + topk_id * num_new_pages_per_topk * page_size + last_page_len
+        )
+        extend_data = tl.load(
+            token_pool_ptr + start + extend_offset,
+            mask=extend_offset < iters,
+        )
+
+    tl.store(kv_ptr + seq_len + extend_offset, extend_data, mask=extend_offset < iters)
+
+    # Update kv_indptr
+    bs_offset = tl.arange(0, num_tokens_upper)
+
+    zid = bid * topk + topk_id
+    if zid == 0:
+        zid = num_seqs * topk
+    positions = tl.load(positions + bs_offset, mask=bs_offset < zid, other=0)
+    base = tl.sum(positions)
+    tl.store(kv_indptr + zid, base + zid * iters)
+
+
+@triton.jit
+def align_evict_mask_to_page_size(
+    seq_lens,
+    evict_mask,
+    page_size: tl.constexpr,
+    num_draft_tokens: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    t_range = tl.arange(0, BLOCK_SIZE)
+
+    bid = tl.program_id(axis=0)
+    seq_len = tl.load(seq_lens + bid)
+    io_mask = t_range < num_draft_tokens
+    mask_row = tl.load(
+        evict_mask + bid * num_draft_tokens + t_range, mask=io_mask, other=0
+    )
+
+    num_trues = tl.sum(mask_row)
+    num_false = num_draft_tokens - num_trues
+
+    start = (seq_len + num_false - 1) // page_size * page_size - seq_len
+    for i in range(max(start, 0), min(start + page_size, num_draft_tokens)):
+        tl.store(evict_mask + bid * num_draft_tokens + i, False)
+
+
+@torch.compile(dynamic=True, disable=_is_npu)
+def get_src_tgt_cache_loc(
+    seq_lens: torch.Tensor,
+    out_cache_loc: torch.Tensor,
+    accept_index: torch.Tensor,
+    num_correct_drafts: torch.Tensor,
+    draft_token_num: int,
+    page_size: int,
+):
+    src_cache_loc = out_cache_loc[accept_index]
+    # zeros_like, not empty_like: any uncovered tail stays at slot 0 (padding)
+    # instead of caching-allocator garbage.
+    tgt_cache_loc = torch.zeros_like(src_cache_loc)
+    extended_len = seq_lens + draft_token_num
+    keep_len = torch.minimum(
+        (seq_lens + num_correct_drafts + 1 + page_size - 1) // page_size * page_size,
+        extended_len,
+    )
+    to_free_num_slots = extended_len - keep_len
+    return src_cache_loc, tgt_cache_loc, to_free_num_slots
+
+
+@triton.jit
+def get_target_cache_loc(
+    tgt_cache_loc,
+    to_free_slots,
+    num_correct_drafts,
+    to_free_num_slots,
+    out_cache_loc,
+    num_verify_tokens: tl.constexpr,
+    num_verify_tokens_upper: tl.constexpr,
+    bs_upper: tl.constexpr,
+):
+    bid = tl.program_id(axis=0)
+    offset = tl.arange(0, num_verify_tokens_upper)
+    bs_offset = tl.arange(0, bs_upper)
+
+    # write the first part to tgt_cache_loc
+    accept_len_all = tl.load(num_correct_drafts + bs_offset, mask=bs_offset < bid)
+    tgt_cache_loc_start = tl.sum(accept_len_all) + bid
+    copy_len = tl.load(num_correct_drafts + bid) + 1
+    out_cache_loc_row = tl.load(
+        out_cache_loc + bid * num_verify_tokens + offset, mask=offset < copy_len
+    )
+    tl.store(
+        tgt_cache_loc + tgt_cache_loc_start + offset,
+        out_cache_loc_row,
+        mask=offset < copy_len,
+    )
+
+    # write the second part to to_free_num_pages
+    to_free_num_slots_all = tl.load(to_free_num_slots + bs_offset, mask=bs_offset < bid)
+    to_free_num_slots_cur = tl.load(to_free_num_slots + bid)
+    out_cache_loc_start = num_verify_tokens - to_free_num_slots_cur
+    to_free_slots_start = tl.sum(to_free_num_slots_all)
+
+    copy_len = to_free_num_slots_cur
+    out_cache_loc_row = tl.load(
+        out_cache_loc + bid * num_verify_tokens + out_cache_loc_start + offset,
+        mask=offset < copy_len,
+    )
+    tl.store(
+        to_free_slots + to_free_slots_start + offset,
+        out_cache_loc_row,
+        mask=offset < copy_len,
+    )
+
+
+@triton.jit
+def filter_finished_cache_loc_kernel(
+    out_cache_loc,
+    tgt_cache_loc,
+    num_correct_drafts,
+    num_accept_tokens_filter,
+    bs_upper: tl.constexpr,
+    num_verify_tokens_upper: tl.constexpr,
+):
+    bid = tl.program_id(0)
+    bs_offset = tl.arange(0, bs_upper)
+
+    num_correct_drafts_all = tl.load(
+        num_correct_drafts + bs_offset, mask=bs_offset < bid
+    )
+    old_start = tl.sum(num_correct_drafts_all) + bid
+
+    num_accept_tokens_filter_all = tl.load(
+        num_accept_tokens_filter + bs_offset, mask=bs_offset < bid
+    )
+    new_start = tl.sum(num_accept_tokens_filter_all)
+
+    copy_len = tl.load(num_accept_tokens_filter + bid)
+    copy_offset = tl.arange(0, num_verify_tokens_upper)
+    value = tl.load(
+        tgt_cache_loc + old_start + copy_offset, mask=copy_offset < copy_len
+    )
+    tl.store(
+        out_cache_loc + new_start + copy_offset, value, mask=copy_offset < copy_len
+    )
+
+
+@triton.jit
+def assign_extend_cache_locs(
+    req_pool_indices,
+    req_to_token,
+    start_offset,
+    end_offset,
+    out_cache_loc,
+    pool_len: tl.constexpr,
+    bs_upper: tl.constexpr,
+):
+    BLOCK_SIZE: tl.constexpr = 32
+    pid = tl.program_id(axis=0)
+    kv_start = tl.load(start_offset + pid)
+    kv_end = tl.load(end_offset + pid)
+    token_pool = req_to_token + tl.load(req_pool_indices + pid) * pool_len
+
+    length_offset = tl.arange(0, bs_upper)
+    start = tl.load(start_offset + length_offset, mask=length_offset < pid, other=0)
+    end = tl.load(end_offset + length_offset, mask=length_offset < pid, other=0)
+    out_offset = tl.sum(end - start, axis=0)
+
+    out_cache_ptr = out_cache_loc + out_offset
+
+    load_offset = tl.arange(0, BLOCK_SIZE) + kv_start
+    save_offset = tl.arange(0, BLOCK_SIZE)
+
+    num_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE)
+    for _ in range(num_loop):
+        mask = load_offset < kv_end
+        data = tl.load(token_pool + load_offset, mask=mask)
+        tl.store(out_cache_ptr + save_offset, data, mask=mask)
+        load_offset += BLOCK_SIZE
+        save_offset += BLOCK_SIZE
+
+
+def assign_extend_cache_locs_func(
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    start_offset: torch.Tensor,
+    end_offset: torch.Tensor,
+    batch_size: int,
+    draft_token_num: int,
+    device,
+) -> torch.Tensor:
+    if _is_cuda or _is_hip or _is_musa or _is_xpu:
+        out_cache_loc = torch.empty(
+            (batch_size * draft_token_num,),
+            dtype=torch.int64,
+            device=device,
+        )
+        assign_extend_cache_locs[(batch_size,)](
+            req_pool_indices,
+            req_to_token,
+            start_offset,
+            end_offset,
+            out_cache_loc,
+            req_to_token.shape[1],
+            next_power_of_2(batch_size),
+        )
+
+        return out_cache_loc
+
+    elif _is_npu:
+        out_cache_loc = torch.empty(
+            (batch_size * draft_token_num,),
+            dtype=torch.int32,
+            device=device,
+        )
+        torch.ops.npu.cache_loc_update(
+            req_pool_indices,
+            req_to_token,
+            start_offset,
+            end_offset,
+            out_cache_loc,
+        )
+
+        return out_cache_loc
+
+    elif _is_cpu:
+        out_cache_loc = torch.empty(
+            (batch_size * draft_token_num,),
+            dtype=torch.int64,
+            device=device,
+        )
+        assign_extend_cache_locs_cpu(
+            req_pool_indices,
+            req_to_token,
+            start_offset,
+            end_offset,
+            out_cache_loc,
+            req_to_token.shape[1],
+        )
+
+        return out_cache_loc

python/sglang/kernels/ops/speculative/dflash.py

@@ -0,0 +1,246 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _dflash_accept_bonus_contig_kernel(
+    candidates_ptr,
+    target_top1_ptr,
+    accept_lens_out_ptr,
+    commit_lens_out_ptr,
+    bonus_ids_out_ptr,
+    out_tokens_ptr,
+    prefix_lens_ptr,
+    new_seq_lens_out_ptr,
+    candidates_row_stride,
+    target_row_stride,
+    accept_stride,
+    commit_stride,
+    bonus_stride,
+    out_tokens_row_stride,
+    prefix_lens_stride,
+    new_seq_lens_stride,
+    block_size,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row = tl.program_id(0)
+    cols = tl.arange(0, BLOCK_SIZE)
+    row_mask = cols < block_size
+    draft_mask = cols < (block_size - 1)
+
+    candidate_row_ptr = candidates_ptr + row * candidates_row_stride
+    target_row_ptr = target_top1_ptr + row * target_row_stride
+    candidate_tail = tl.load(candidate_row_ptr + cols + 1, mask=draft_mask, other=0)
+
+    accept_len = tl.full((), 0, tl.int32)
+    prefix_live = tl.full((), 1, tl.int32)
+    for col in range(BLOCK_SIZE - 1):
+        in_range = col < (block_size - 1)
+        candidate_id = tl.load(candidate_row_ptr + (col + 1), mask=in_range, other=0)
+        target_id = tl.load(target_row_ptr + col, mask=in_range, other=0)
+        match_i32 = (candidate_id == target_id).to(tl.int32)
+        keep = in_range & (prefix_live != 0) & (match_i32 != 0)
+        accept_len += keep.to(tl.int32)
+        prefix_live = tl.where(in_range, prefix_live & match_i32, prefix_live)
+
+    commit_len = accept_len + 1
+    bonus_id = tl.load(target_row_ptr + accept_len.to(tl.int64))
+    new_seq_len = tl.load(prefix_lens_ptr + row * prefix_lens_stride) + commit_len
+
+    tl.store(accept_lens_out_ptr + row * accept_stride, accept_len)
+    tl.store(commit_lens_out_ptr + row * commit_stride, commit_len)
+    tl.store(bonus_ids_out_ptr + row * bonus_stride, bonus_id)
+    tl.store(new_seq_lens_out_ptr + row * new_seq_lens_stride, new_seq_len)
+
+    out_val = tl.where(draft_mask, candidate_tail, 0)
+    out_val = tl.where(cols == accept_len, bonus_id, out_val)
+    tl.store(
+        out_tokens_ptr + row * out_tokens_row_stride + cols, out_val, mask=row_mask
+    )
+
+
+def _pick_num_warps(block_size: int) -> int:
+    if block_size <= 16:
+        return 1
+    if block_size <= 32:
+        return 2
+    if block_size <= 64:
+        return 4
+    return 8
+
+
+def _is_row_major_contiguous_2d(x: torch.Tensor) -> bool:
+    return x.ndim == 2 and x.is_contiguous()
+
+
+def _compute_dflash_accept_bonus_triton_unchecked(
+    candidates: torch.Tensor,
+    target_top1: torch.Tensor,
+    accept_lens_out: torch.Tensor,
+    commit_lens_out: torch.Tensor,
+    bonus_ids_out: torch.Tensor,
+    out_tokens_out: torch.Tensor,
+    prefix_lens: torch.Tensor,
+    new_seq_lens_out: torch.Tensor,
+) -> None:
+    batch_size, block_size = candidates.shape
+    if batch_size == 0:
+        return
+
+    if not _is_row_major_contiguous_2d(candidates):
+        raise ValueError("DFLASH Triton accept_bonus requires contiguous candidates.")
+    if not _is_row_major_contiguous_2d(target_top1):
+        raise ValueError("DFLASH Triton accept_bonus requires contiguous target_top1.")
+    if not _is_row_major_contiguous_2d(out_tokens_out):
+        raise ValueError(
+            "DFLASH Triton accept_bonus requires contiguous out_tokens_out."
+        )
+    if not accept_lens_out.is_contiguous():
+        raise ValueError(
+            "DFLASH Triton accept_bonus requires contiguous accept_lens_out."
+        )
+    if not commit_lens_out.is_contiguous():
+        raise ValueError(
+            "DFLASH Triton accept_bonus requires contiguous commit_lens_out."
+        )
+    if not bonus_ids_out.is_contiguous():
+        raise ValueError(
+            "DFLASH Triton accept_bonus requires contiguous bonus_ids_out."
+        )
+    if prefix_lens.ndim != 1:
+        raise ValueError("DFLASH Triton accept_bonus requires 1D prefix_lens.")
+    if not new_seq_lens_out.is_contiguous():
+        raise ValueError(
+            "DFLASH Triton accept_bonus requires contiguous new_seq_lens_out."
+        )
+
+    block = triton.next_power_of_2(block_size)
+    num_warps = _pick_num_warps(block)
+    _dflash_accept_bonus_contig_kernel[(batch_size,)](
+        candidates,
+        target_top1,
+        accept_lens_out,
+        commit_lens_out,
+        bonus_ids_out,
+        out_tokens_out,
+        prefix_lens,
+        new_seq_lens_out,
+        candidates.stride(0),
+        target_top1.stride(0),
+        accept_lens_out.stride(0),
+        commit_lens_out.stride(0),
+        bonus_ids_out.stride(0),
+        out_tokens_out.stride(0),
+        prefix_lens.stride(0),
+        new_seq_lens_out.stride(0),
+        block_size,
+        BLOCK_SIZE=block,
+        num_warps=num_warps,
+    )
+
+
+@triton.jit
+def _prepare_dflash_draft_block_contig_kernel(
+    bonus_tokens_ptr,
+    prefix_lens_ptr,
+    req_pool_indices_ptr,
+    req_to_token_ptr,
+    block_ids_out_ptr,
+    positions_out_ptr,
+    cache_loc_out_ptr,
+    bonus_tokens_stride,
+    prefix_lens_stride,
+    req_pool_indices_stride,
+    req_to_token_row_stride,
+    block_ids_row_stride,
+    positions_row_stride,
+    cache_loc_row_stride,
+    req_to_token_width,
+    block_size,
+    mask_token_id,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row = tl.program_id(0)
+    cols = tl.arange(0, BLOCK_SIZE)
+    row_mask = cols < block_size
+
+    prefix_len = tl.load(prefix_lens_ptr + row * prefix_lens_stride)
+    req_idx = tl.load(req_pool_indices_ptr + row * req_pool_indices_stride)
+    bonus_token = tl.load(bonus_tokens_ptr + row * bonus_tokens_stride)
+
+    logical_pos = prefix_len.to(tl.int64) + cols
+    valid = row_mask & (logical_pos < req_to_token_width)
+    req_row_ptr = req_to_token_ptr + req_idx * req_to_token_row_stride
+    slot_ids = tl.load(req_row_ptr + logical_pos, mask=valid, other=0)
+
+    block_ids = tl.full((BLOCK_SIZE,), mask_token_id, tl.int64)
+    block_ids = tl.where(cols == 0, bonus_token.to(tl.int64), block_ids)
+    tl.store(
+        block_ids_out_ptr + row * block_ids_row_stride + cols, block_ids, mask=row_mask
+    )
+    tl.store(
+        positions_out_ptr + row * positions_row_stride + cols,
+        logical_pos,
+        mask=row_mask,
+    )
+    tl.store(
+        cache_loc_out_ptr + row * cache_loc_row_stride + cols,
+        slot_ids.to(tl.int64),
+        mask=row_mask,
+    )
+
+
+def _prepare_dflash_draft_block_unchecked(
+    bonus_tokens: torch.Tensor,
+    prefix_lens: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    block_ids_out: torch.Tensor,
+    positions_out: torch.Tensor,
+    cache_loc_out: torch.Tensor,
+    mask_token_id: int,
+) -> None:
+    batch_size = int(bonus_tokens.numel())
+    if batch_size == 0:
+        return
+
+    if req_to_token.ndim != 2 or req_to_token.stride(1) != 1:
+        raise ValueError("DFLASH Triton prepare_block requires row-major req_to_token.")
+    if not _is_row_major_contiguous_2d(block_ids_out):
+        raise ValueError(
+            "DFLASH Triton prepare_block requires contiguous block_ids_out."
+        )
+    if not _is_row_major_contiguous_2d(positions_out):
+        raise ValueError(
+            "DFLASH Triton prepare_block requires contiguous positions_out."
+        )
+    if not _is_row_major_contiguous_2d(cache_loc_out):
+        raise ValueError(
+            "DFLASH Triton prepare_block requires contiguous cache_loc_out."
+        )
+
+    block_size = int(block_ids_out.shape[1])
+    block = triton.next_power_of_2(block_size)
+    num_warps = _pick_num_warps(block)
+    _prepare_dflash_draft_block_contig_kernel[(batch_size,)](
+        bonus_tokens,
+        prefix_lens,
+        req_pool_indices,
+        req_to_token,
+        block_ids_out,
+        positions_out,
+        cache_loc_out,
+        bonus_tokens.stride(0),
+        prefix_lens.stride(0),
+        req_pool_indices.stride(0),
+        req_to_token.stride(0),
+        block_ids_out.stride(0),
+        positions_out.stride(0),
+        cache_loc_out.stride(0),
+        int(req_to_token.shape[1]),
+        block_size,
+        int(mask_token_id),
+        BLOCK_SIZE=block,
+        num_warps=num_warps,
+    )

python/sglang/kernels/ops/speculative/eagle.py

@@ -0,0 +1,91 @@
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.utils import is_cpu, next_power_of_2
+
+_is_cpu = is_cpu()
+
+if _is_cpu:
+    from sgl_kernel import fill_accept_out_cache_loc_cpu, fill_bonus_tokens_cpu
+
+
+@triton.jit
+def fill_bonus_tokens(
+    accept_tokens,
+    accept_lens,
+    bonus_tokens_ptr,
+    accept_stride: tl.constexpr,
+):
+    # NOTE: we cannot fuse any in-place operations of `accept_lens` inside this kernel
+    # because this kernel reads accept_lens
+    pid = tl.program_id(axis=0)
+    # `accept_lens` includes the bonus token; the last accepted slot is at -1.
+    accept_len = tl.load(accept_lens + pid)
+
+    # accept_stride = per-req width of accept_tokens (= accept_index.shape[1]).
+    bonus_token_idx = accept_stride * pid + accept_len - 1
+    bonus_token = tl.load(accept_tokens + bonus_token_idx)
+    tl.store(bonus_tokens_ptr + pid, bonus_token)
+
+
+def fill_bonus_tokens_func(
+    accept_tokens: torch.Tensor,
+    accept_lens: torch.Tensor,
+    bonus_tokens: torch.Tensor,  # mutable
+    accept_stride: int,
+    batch_size: int,
+):
+    if _is_cpu:
+        fill_bonus_tokens_cpu(
+            accept_tokens,
+            accept_lens,
+            bonus_tokens,
+            accept_stride,
+        )
+        return
+    fill_bonus_tokens[(batch_size,)](
+        accept_tokens,
+        accept_lens,
+        bonus_tokens,
+        accept_stride,
+    )
+
+
+@triton.jit
+def fill_accept_out_cache_loc(
+    accept_index,
+    out_cache_loc,
+    accept_out_cache_loc,
+    size_upper: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    offset = tl.arange(0, size_upper)
+
+    masks = (tl.load(accept_index + offset, offset < pid, other=-1) != -1).to(tl.int64)
+    dst = tl.sum(masks)
+    src = tl.load(accept_index + pid)
+    if src > -1:
+        value = tl.load(out_cache_loc + src)
+        tl.store(accept_out_cache_loc + dst, value)
+
+
+def fill_accept_out_cache_loc_func(
+    accept_index: torch.Tensor,
+    out_cache_loc: torch.Tensor,
+    accept_out_cache_loc: torch.Tensor,  # mutable
+    size: int,
+):
+    if _is_cpu:
+        fill_accept_out_cache_loc_cpu(
+            accept_index,
+            out_cache_loc,
+            accept_out_cache_loc,
+        )
+        return
+    fill_accept_out_cache_loc[(size,)](
+        accept_index,
+        out_cache_loc,
+        accept_out_cache_loc,
+        next_power_of_2(size),
+    )

python/sglang/kernels/ops/speculative/fused_kv_materialize.py

@@ -0,0 +1,457 @@
+# Copyright 2023-2024 SGLang Team
+# 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.
+# ==============================================================================
+"""Fused Triton kernel for DFlash KV materialization.
+
+Combines: KV projection + RMSNorm + RoPE, then pool-managed KV writes.
+"""
+
+from typing import Callable, List, Optional
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _fused_norm_rope_kernel_stacked(
+    kv_ptr,  # [total_ctx, n_layers, kv_size * 2]
+    k_norm_weight_ptr,  # [n_layers, head_dim]
+    eps_ptr,  # [n_layers]
+    cos_sin_cache_ptr,  # [max_pos, rotary_dim]
+    positions_ptr,  # [total_ctx]
+    k_out_ptr,  # [n_layers, total_ctx, num_kv_heads, head_dim]
+    v_out_ptr,  # [n_layers, total_ctx, num_kv_heads, head_dim]
+    kv_stride_ctx,
+    kv_stride_layer,
+    k_norm_weight_stride_layer,
+    cos_sin_stride_pos,
+    k_out_stride_layer,
+    k_out_stride_ctx,
+    k_out_stride_head,
+    v_out_stride_layer,
+    v_out_stride_ctx,
+    v_out_stride_head,
+    total_ctx,
+    n_layers: tl.constexpr,
+    num_kv_heads: tl.constexpr,
+    head_dim: tl.constexpr,
+    kv_size: tl.constexpr,
+    rotary_dim: tl.constexpr,
+    half_rotary_dim: tl.constexpr,
+    BLOCK_HD: tl.constexpr,
+):
+    """Fused RMSNorm(K) + RoPE(K) materialization. Grid: (total_ctx, num_kv_heads, n_layers)."""
+    ctx_id = tl.program_id(0)
+    head_id = tl.program_id(1)
+    layer_id = tl.program_id(2)
+    if ctx_id >= total_ctx or layer_id >= n_layers:
+        return
+
+    position = tl.load(positions_ptr + ctx_id)
+    eps = tl.load(eps_ptr + layer_id).to(tl.float32)
+    kv_base = kv_ptr + ctx_id * kv_stride_ctx + layer_id * kv_stride_layer
+    k_base = kv_base + head_id * head_dim
+    v_base = kv_base + kv_size + head_id * head_dim
+    k_write = (
+        k_out_ptr
+        + layer_id * k_out_stride_layer
+        + ctx_id * k_out_stride_ctx
+        + head_id * k_out_stride_head
+    )
+    v_write = (
+        v_out_ptr
+        + layer_id * v_out_stride_layer
+        + ctx_id * v_out_stride_ctx
+        + head_id * v_out_stride_head
+    )
+
+    offs = tl.arange(0, BLOCK_HD)
+    mask_hd = offs < head_dim
+    mask_half = offs < half_rotary_dim
+
+    k_raw = tl.load(k_base + offs, mask=mask_hd, other=0.0).to(tl.float32)
+    v_raw = tl.load(v_base + offs, mask=mask_hd, other=0.0)
+
+    inv_rms = tl.rsqrt(tl.sum(k_raw * k_raw) / head_dim + eps)
+    norm_w = tl.load(
+        k_norm_weight_ptr + layer_id * k_norm_weight_stride_layer + offs,
+        mask=mask_hd,
+        other=1.0,
+    ).to(tl.float32)
+    k_normed = k_raw * inv_rms * norm_w
+
+    cos_sin_base = cos_sin_cache_ptr + position * cos_sin_stride_pos
+    cos_v = tl.load(cos_sin_base + offs, mask=mask_half, other=1.0).to(tl.float32)
+    sin_v = tl.load(
+        cos_sin_base + half_rotary_dim + offs, mask=mask_half, other=0.0
+    ).to(tl.float32)
+
+    k_first = tl.where(mask_half, k_normed, 0.0)
+    k_second_raw = tl.load(
+        k_base + half_rotary_dim + offs, mask=mask_half, other=0.0
+    ).to(tl.float32)
+    norm_w_second = tl.load(
+        k_norm_weight_ptr
+        + layer_id * k_norm_weight_stride_layer
+        + half_rotary_dim
+        + offs,
+        mask=mask_half,
+        other=1.0,
+    ).to(tl.float32)
+    k_second = k_second_raw * inv_rms * norm_w_second
+
+    k_rot_first = k_first * cos_v - k_second * sin_v
+    k_rot_second = k_second * cos_v + k_first * sin_v
+
+    tl.store(v_write + offs, v_raw, mask=mask_hd)
+    tl.store(k_write + offs, k_rot_first.to(v_raw.dtype), mask=mask_half)
+    tl.store(
+        k_write + half_rotary_dim + offs, k_rot_second.to(v_raw.dtype), mask=mask_half
+    )
+    mask_pass = (offs >= rotary_dim) & (offs < head_dim)
+    tl.store(k_write + offs, k_normed.to(v_raw.dtype), mask=mask_pass)
+
+
+def _fused_norm_rope_stacked(
+    kv: torch.Tensor,  # [total_ctx, n_layers, kv_size*2]
+    k_norm_weight: torch.Tensor,  # [n_layers, head_dim]
+    eps: torch.Tensor,  # [n_layers]
+    cos_sin_cache: torch.Tensor,  # [max_pos, rotary_dim]
+    positions: torch.Tensor,  # [total_ctx]
+    num_kv_heads: int,
+    head_dim: int,
+    rotary_dim: int,
+    k_out: Optional[torch.Tensor] = None,
+    v_out: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Fused RMSNorm + RoPE materialization for all layers."""
+    if kv.ndim != 3:
+        raise ValueError(
+            "Invalid stacked fused KV projection shape: "
+            f"got {tuple(kv.shape)}, expected 3D [total_ctx, n_layers, kv_size*2]."
+        )
+
+    total_ctx, n_layers, kv_dim = kv.shape
+    if total_ctx == 0:
+        empty = torch.empty(
+            (n_layers, 0, num_kv_heads, head_dim), dtype=kv.dtype, device=kv.device
+        )
+        return empty, empty
+
+    kv_size = num_kv_heads * head_dim
+    if kv_dim != kv_size * 2:
+        raise ValueError(
+            "Invalid fused KV projection shape: "
+            f"got {tuple(kv.shape)}, expected trailing dim {kv_size * 2}."
+        )
+    if rotary_dim <= 0 or rotary_dim > head_dim or rotary_dim % 2 != 0:
+        raise ValueError(
+            "Invalid fused KV rotary/head dim pair: "
+            f"rotary_dim={rotary_dim}, head_dim={head_dim}."
+        )
+    if k_norm_weight.shape != (n_layers, head_dim):
+        raise ValueError(
+            "Invalid stacked k_norm_weight shape for fused KV materialization: "
+            f"got {tuple(k_norm_weight.shape)}, expected {(n_layers, head_dim)}."
+        )
+    if eps.shape != (n_layers,):
+        raise ValueError(
+            "Invalid stacked eps shape for fused KV materialization: "
+            f"got {tuple(eps.shape)}, expected {(n_layers,)}."
+        )
+
+    half_rotary_dim = rotary_dim // 2
+    BLOCK_HD = triton.next_power_of_2(head_dim)
+
+    if positions.device != kv.device:
+        positions = positions.to(device=kv.device, dtype=torch.int64)
+    elif positions.dtype != torch.int64:
+        positions = positions.to(torch.int64)
+
+    expected_shape = (n_layers, total_ctx, num_kv_heads, head_dim)
+    if k_out is None:
+        k_out = torch.empty(expected_shape, dtype=kv.dtype, device=kv.device)
+    else:
+        if k_out.shape != expected_shape:
+            raise ValueError(
+                "Invalid k_out shape for fused KV materialization: "
+                f"got {tuple(k_out.shape)}, expected {expected_shape}."
+            )
+        if k_out.device != kv.device or k_out.dtype != kv.dtype:
+            raise ValueError(
+                "Invalid k_out device/dtype for fused KV materialization: "
+                f"got device={k_out.device}, dtype={k_out.dtype}, "
+                f"expected device={kv.device}, dtype={kv.dtype}."
+            )
+    if v_out is None:
+        v_out = torch.empty_like(k_out)
+    else:
+        if v_out.shape != expected_shape:
+            raise ValueError(
+                "Invalid v_out shape for fused KV materialization: "
+                f"got {tuple(v_out.shape)}, expected {expected_shape}."
+            )
+        if v_out.device != kv.device or v_out.dtype != kv.dtype:
+            raise ValueError(
+                "Invalid v_out device/dtype for fused KV materialization: "
+                f"got device={v_out.device}, dtype={v_out.dtype}, "
+                f"expected device={kv.device}, dtype={kv.dtype}."
+            )
+
+    _fused_norm_rope_kernel_stacked[(total_ctx, num_kv_heads, n_layers)](
+        kv,
+        k_norm_weight,
+        eps,
+        cos_sin_cache,
+        positions,
+        k_out,
+        v_out,
+        kv.stride(0),
+        kv.stride(1),
+        k_norm_weight.stride(0),
+        cos_sin_cache.stride(0),
+        k_out.stride(0),
+        k_out.stride(1),
+        k_out.stride(2),
+        v_out.stride(0),
+        v_out.stride(1),
+        v_out.stride(2),
+        total_ctx,
+        n_layers,
+        num_kv_heads,
+        head_dim,
+        kv_size,
+        rotary_dim,
+        half_rotary_dim,
+        BLOCK_HD,
+    )
+    return k_out, v_out
+
+
+class FusedKVMaterializeHelper:
+    """Fused KV materialization helper using batched projection.
+
+    Uses a single large GEMM across all layers, then a Triton kernel for fused
+    RMSNorm + RoPE materialization across all layers.
+    """
+
+    def __init__(
+        self,
+        layers: List,
+        rotary_emb,
+        num_kv_heads: int,
+        head_dim: int,
+        device: torch.device,
+        max_position_hint: Optional[int] = None,
+    ):
+        self.num_kv_heads = num_kv_heads
+        self.head_dim = head_dim
+        self.rotary_emb = rotary_emb
+        self.n_layers = len(layers)
+        self.device = device
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.layer_out_dim = 2 * self.kv_size
+
+        self.rotary_dim = int(getattr(rotary_emb, "rotary_dim", head_dim))
+        self.is_neox_style = bool(getattr(rotary_emb, "is_neox_style", True))
+
+        if not self.is_neox_style:
+            raise NotImplementedError("Only neox-style RoPE is supported.")
+        if self.rotary_dim <= 0 or self.rotary_dim > self.head_dim:
+            raise ValueError(
+                "Invalid fused KV rotary/head dim pair: "
+                f"rotary_dim={self.rotary_dim}, head_dim={self.head_dim}."
+            )
+
+        self.max_position_hint = (
+            max(int(max_position_hint) - 1, 0)
+            if max_position_hint is not None
+            else None
+        )
+        self._reserved_rope_cache_len = int(
+            getattr(self.rotary_emb, "cos_sin_cache", torch.empty((0,))).shape[0]
+        )
+        self._mm_out_supported = True
+        self._workspace_capacity = 0
+        self._workspace_dtype: Optional[torch.dtype] = None
+        self._proj_workspace: Optional[torch.Tensor] = None
+        self._k_workspace: Optional[torch.Tensor] = None
+        self._v_workspace: Optional[torch.Tensor] = None
+
+        kv_weights = []
+        k_norm_weights = []
+        eps_values = []
+
+        for layer_id, layer in enumerate(layers):
+            attn = layer.self_attn
+            if int(attn.num_kv_heads) != self.num_kv_heads:
+                raise ValueError(
+                    "num_kv_heads mismatch across layers for fused KV path: "
+                    f"expected {self.num_kv_heads}, got {int(attn.num_kv_heads)} at layer {layer_id}."
+                )
+            if int(attn.head_dim) != self.head_dim:
+                raise ValueError(
+                    "head_dim mismatch across layers for fused KV path: "
+                    f"expected {self.head_dim}, got {int(attn.head_dim)} at layer {layer_id}."
+                )
+            layer_rotary_dim = int(
+                getattr(attn.rotary_emb, "rotary_dim", self.head_dim)
+            )
+            layer_is_neox = bool(getattr(attn.rotary_emb, "is_neox_style", True))
+            if (
+                layer_rotary_dim != self.rotary_dim
+                or layer_is_neox != self.is_neox_style
+            ):
+                raise ValueError(
+                    "RoPE config mismatch across layers for fused KV path: "
+                    f"expected (rotary_dim={self.rotary_dim}, neox={self.is_neox_style}), "
+                    f"got (rotary_dim={layer_rotary_dim}, neox={layer_is_neox}) at layer {layer_id}."
+                )
+
+            qkv_w = attn.qkv_proj.weight
+            kv_weight = qkv_w[attn.q_size : attn.q_size + 2 * attn.kv_size]
+            kv_weights.append(kv_weight)
+            k_norm_weights.append(attn.k_norm.weight)
+            eps_values.append(float(attn.k_norm.variance_epsilon))
+
+        flat_kv_weight = torch.stack(kv_weights).reshape(
+            self.n_layers * self.layer_out_dim, -1
+        )
+        self.flat_kv_weight_t = flat_kv_weight.transpose(0, 1).contiguous()
+        self.k_norm_weights = torch.stack(k_norm_weights).contiguous()
+        self.eps_values = torch.tensor(
+            eps_values, dtype=torch.float32, device=self.device
+        )
+
+        if self.max_position_hint is not None:
+            self._ensure_rope_cache(self.max_position_hint)
+
+    def _ensure_rope_cache(self, max_position: int) -> torch.Tensor:
+        if max_position + 1 > self._reserved_rope_cache_len:
+            ensure_cos_sin_cache_length = getattr(
+                self.rotary_emb, "_ensure_cos_sin_cache_length", None
+            )
+            if callable(ensure_cos_sin_cache_length):
+                ensure_cos_sin_cache_length(max_position)
+                self._reserved_rope_cache_len = int(
+                    self.rotary_emb.cos_sin_cache.shape[0]
+                )
+
+        cos_sin_cache = self.rotary_emb.cos_sin_cache
+        if max_position >= int(cos_sin_cache.shape[0]):
+            raise RuntimeError(
+                "RoPE cos/sin cache is too short for fused KV materialization: "
+                f"max_position={max_position}, cache_len={int(cos_sin_cache.shape[0])}."
+            )
+        if cos_sin_cache.device != self.device:
+            cos_sin_cache = cos_sin_cache.to(self.device)
+        return cos_sin_cache
+
+    def _ensure_workspace(self, total_ctx: int, dtype: torch.dtype) -> None:
+        if (
+            self._workspace_capacity >= total_ctx
+            and self._workspace_dtype == dtype
+            and self._proj_workspace is not None
+            and self._k_workspace is not None
+            and self._v_workspace is not None
+        ):
+            return
+
+        new_capacity = max(1, total_ctx)
+        if self._workspace_capacity > 0:
+            new_capacity = max(new_capacity, self._workspace_capacity * 2)
+
+        self._proj_workspace = torch.empty(
+            (new_capacity, self.n_layers * self.layer_out_dim),
+            dtype=dtype,
+            device=self.device,
+        )
+        self._k_workspace = torch.empty(
+            (self.n_layers, new_capacity, self.num_kv_heads, self.head_dim),
+            dtype=dtype,
+            device=self.device,
+        )
+        self._v_workspace = torch.empty_like(self._k_workspace)
+        self._workspace_capacity = new_capacity
+        self._workspace_dtype = dtype
+
+    def materialize(
+        self,
+        ctx_hidden: torch.Tensor,
+        positions: torch.Tensor,
+        write_layer_kv: Callable[[int, torch.Tensor, torch.Tensor], None],
+    ) -> None:
+        """Materialize KV cache for all layers using batched projection."""
+        total_ctx = ctx_hidden.shape[0]
+        if total_ctx == 0:
+            return
+
+        if positions.ndim != 1:
+            positions = positions.reshape(-1)
+        if positions.numel() != total_ctx:
+            raise ValueError(
+                "positions must match ctx_hidden token count for fused KV materialization: "
+                f"positions={positions.numel()}, total_ctx={total_ctx}."
+            )
+
+        if ctx_hidden.device != self.device:
+            ctx_hidden = ctx_hidden.to(self.device, non_blocking=True)
+        if ctx_hidden.dtype != self.flat_kv_weight_t.dtype:
+            ctx_hidden = ctx_hidden.to(self.flat_kv_weight_t.dtype)
+        if positions.device != self.device:
+            positions = positions.to(
+                device=self.device, dtype=torch.int64, non_blocking=True
+            )
+        elif positions.dtype != torch.int64:
+            positions = positions.to(torch.int64)
+
+        max_position = (
+            self.max_position_hint
+            if self.max_position_hint is not None
+            else int(positions.max().item())
+        )
+        cos_sin_cache = self._ensure_rope_cache(max_position)
+
+        self._ensure_workspace(total_ctx, ctx_hidden.dtype)
+        assert self._proj_workspace is not None
+        assert self._k_workspace is not None
+        assert self._v_workspace is not None
+
+        proj_out_2d = self._proj_workspace[:total_ctx]
+        if self._mm_out_supported:
+            try:
+                torch.mm(ctx_hidden, self.flat_kv_weight_t, out=proj_out_2d)
+            except Exception:
+                self._mm_out_supported = False
+                proj_out_2d = torch.mm(ctx_hidden, self.flat_kv_weight_t)
+        else:
+            proj_out_2d = torch.mm(ctx_hidden, self.flat_kv_weight_t)
+
+        proj_out = proj_out_2d.view(total_ctx, self.n_layers, self.layer_out_dim)
+        tmp_k = self._k_workspace[:, :total_ctx]
+        tmp_v = self._v_workspace[:, :total_ctx]
+        cache_k, cache_v = _fused_norm_rope_stacked(
+            proj_out,
+            self.k_norm_weights,
+            self.eps_values,
+            cos_sin_cache,
+            positions,
+            self.num_kv_heads,
+            self.head_dim,
+            self.rotary_dim,
+            k_out=tmp_k,
+            v_out=tmp_v,
+        )
+        for layer_idx in range(self.n_layers):
+            write_layer_kv(layer_idx, cache_k[layer_idx], cache_v[layer_idx])

python/sglang/kernels/ops/speculative/gather_spec_extras.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _gather_rows_kernel(
+    idx_ptr,
+    s0,
+    d0,
+    n0,
+    s1,
+    d1,
+    n1,
+    s2,
+    d2,
+    n2,
+    s3,
+    d3,
+    n3,
+    HAS3: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    # One program == one (output row, column block). All buffers share the
+    # same gather index, so a single launch copies every buffer's row and
+    # the per-kernel launch bubbles between the old separate gathers vanish.
+    row = tl.program_id(0)
+    cb = tl.program_id(1)
+    src = tl.load(idx_ptr + row).to(tl.int64)
+    cols = cb * BLOCK + tl.arange(0, BLOCK)
+
+    m0 = cols < n0
+    tl.store(d0 + row * n0 + cols, tl.load(s0 + src * n0 + cols, mask=m0), mask=m0)
+
+    m1 = cols < n1
+    tl.store(d1 + row * n1 + cols, tl.load(s1 + src * n1 + cols, mask=m1), mask=m1)
+
+    m2 = cols < n2
+    tl.store(d2 + row * n2 + cols, tl.load(s2 + src * n2 + cols, mask=m2), mask=m2)
+
+    if HAS3:
+        m3 = cols < n3
+        tl.store(d3 + row * n3 + cols, tl.load(s3 + src * n3 + cols, mask=m3), mask=m3)
+
+
+def _row_width(buf: torch.Tensor) -> int:
+    """Flattened per-row element count (trailing dims), 1 for a 1-D buffer."""
+    return buf[0].numel() if buf.dim() > 1 else 1
+
+
+def _empty_like_rows(buf: torch.Tensor, m: int) -> torch.Tensor:
+    """Output buffer for `m` gathered rows of `buf` (same trailing dims/dtype/device)."""
+    return torch.empty((m, *buf.shape[1:]), dtype=buf.dtype, device=buf.device)
+
+
+def gather_spec_extras(
+    indices: torch.Tensor,
+    topk_p_buf: torch.Tensor,
+    topk_index_buf: torch.Tensor,
+    output_tokens_buf: torch.Tensor,
+    hidden_states_buf: Optional[torch.Tensor],
+):
+    """Gather spec extras (topk_p / topk_index / bonus_tokens / optional hidden
+    states) by a shared row index in a single fused Triton launch (one kernel
+    for all buffers) instead of one advanced-index gather per buffer.
+    `hidden_states_buf` is None when the build does not capture hidden states."""
+    # Source buffers are allocated once (torch.empty/full) and only ever mutated
+    # in place, so they are guaranteed row-contiguous. `indices` flows from
+    # several producers (req_pool_indices, filtered/merged future_indices); the
+    # kernel addresses it linearly, so normalize layout here (no-op when already
+    # contiguous) to avoid a silent wrong-result on a strided index tensor.
+    indices = indices.contiguous()
+    m = indices.shape[0]
+    has_hidden = hidden_states_buf is not None
+
+    topk_p = _empty_like_rows(topk_p_buf, m)
+    topk_index = _empty_like_rows(topk_index_buf, m)
+    bonus_tokens = _empty_like_rows(output_tokens_buf, m)
+    hidden_states = _empty_like_rows(hidden_states_buf, m) if has_hidden else None
+    if m == 0:
+        return topk_p, topk_index, bonus_tokens, hidden_states
+
+    n0 = _row_width(topk_p_buf)
+    n1 = _row_width(topk_index_buf)
+    n2 = _row_width(output_tokens_buf)
+    n3 = _row_width(hidden_states_buf) if has_hidden else 1
+    max_n = max(n0, n1, n2, n3)
+
+    # Dummy operands for the disabled hidden-states slot: the pointers must be
+    # valid even though the kernel never dereferences them (gated off by HAS3).
+    s3 = hidden_states_buf if has_hidden else indices
+    d3 = hidden_states if has_hidden else indices
+
+    block = min(1024, triton.next_power_of_2(max_n))
+    grid = (m, triton.cdiv(max_n, block))
+    _gather_rows_kernel[grid](
+        indices,
+        topk_p_buf,
+        topk_p,
+        n0,
+        topk_index_buf,
+        topk_index,
+        n1,
+        output_tokens_buf,
+        bonus_tokens,
+        n2,
+        s3,
+        d3,
+        n3,
+        HAS3=has_hidden,
+        BLOCK=block,
+    )
+    return topk_p, topk_index, bonus_tokens, hidden_states

python/sglang/kernels/ops/speculative/multi_layer_eagle.py

@@ -0,0 +1,96 @@
+# Copyright 2023-2024 SGLang Team
+# 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.
+# ==============================================================================
+
+import triton
+import triton.language as tl
+
+from sglang.srt.utils import is_cpu, is_npu
+
+_is_cpu = is_cpu()
+_is_npu = is_npu()
+
+if _is_cpu:
+    from sgl_kernel import rotate_input_ids_cpu
+
+
+@triton.jit
+def rotate_input_ids_kernel(
+    input_ids_ptr,
+    extend_start_loc_ptr,
+    extend_seq_lens_ptr,
+    topk_index_ptr,
+    select_index_ptr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+
+    start_loc = tl.load(extend_start_loc_ptr + pid)
+    seq_len = tl.load(extend_seq_lens_ptr + pid)
+    new_token = tl.load(topk_index_ptr + pid)
+
+    num_elements_to_shift = seq_len - 1
+
+    for off in range(0, num_elements_to_shift, BLOCK_SIZE):
+        offsets = off + tl.arange(0, BLOCK_SIZE)
+        mask = offsets < num_elements_to_shift
+
+        read_ptr = input_ids_ptr + start_loc + offsets + 1
+        val = tl.load(read_ptr, mask=mask)
+        tl.debug_barrier()
+
+        write_ptr = input_ids_ptr + start_loc + offsets
+        tl.store(write_ptr, val, mask=mask)
+        tl.debug_barrier()
+
+    if seq_len > 0:
+        if select_index_ptr is not None:
+            last_pos_ptr = input_ids_ptr + tl.load(select_index_ptr + pid)
+        else:
+            last_pos_ptr = input_ids_ptr + start_loc + seq_len - 1
+        tl.store(last_pos_ptr, new_token)
+
+
+def rotate_input_ids(
+    input_ids, extend_start_loc, extend_seq_lens, topk_index, select_index=None
+):
+    if _is_cpu:
+        rotate_input_ids_cpu(
+            input_ids,
+            extend_start_loc,
+            extend_seq_lens,
+            topk_index,
+            select_index,
+        )
+        return input_ids
+
+    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
+    grid = (batch_size,)
+
+    rotate_input_ids_kernel[grid](
+        input_ids,
+        extend_start_loc,
+        extend_seq_lens,
+        topk_index,
+        select_index,
+        BLOCK_SIZE=BLOCK_SIZE,
+    )
+    return input_ids

python/sglang/kernels/ops/speculative/spec_tree.py

@@ -0,0 +1,281 @@
+# Copyright 2023-2026 SGLang Team
+# 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.
+# ==============================================================================
+import triton
+import triton.language as tl
+
+
+@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,
+            )
+
+
+@triton.jit
+def verify_tree_greedy_kernel_triton(
+    predicts_ptr,
+    accept_index_ptr,
+    accept_token_num_ptr,
+    candidates_ptr,
+    retrieve_index_ptr,
+    retrieve_next_token_ptr,
+    retrieve_next_sibling_ptr,
+    target_predict_ptr,
+    batch_size: tl.constexpr,
+    num_speculative_tokens: tl.constexpr,
+    num_draft_tokens: tl.constexpr,
+):
+    """
+    Triton kernel for verifying EAGLE tree in greedy mode.
+    Each program handles one batch item.
+    """
+    bx = tl.program_id(0)
+
+    # Initialize
+    last_accept_retrieve_idx = tl.load(retrieve_index_ptr + bx * num_draft_tokens)
+    tl.store(accept_index_ptr + bx * num_speculative_tokens, last_accept_retrieve_idx)
+    # Cast to match dtype of loaded tensors to avoid type inconsistency
+    num_accept_tokens = tl.cast(0, last_accept_retrieve_idx.dtype)
+    cur_index = tl.cast(0, last_accept_retrieve_idx.dtype)
+
+    # Tree traversal loop
+    should_continue = 1
+    for j in range(1, num_speculative_tokens):
+        if should_continue:  # Early exit guard
+            cur_index = tl.load(
+                retrieve_next_token_ptr + bx * num_draft_tokens + cur_index
+            )
+
+            # Load target token once per level (before sibling search)
+            # last_accept_retrieve_idx is constant during sibling traversal
+            target_row = last_accept_retrieve_idx // num_draft_tokens
+            target_col = last_accept_retrieve_idx % num_draft_tokens
+            target_token = tl.load(
+                target_predict_ptr + target_row * num_draft_tokens + target_col
+            )
+
+            # Traverse siblings
+            found_match = 0
+            for _ in range(num_draft_tokens):  # Max iterations = num_draft_tokens
+                if found_match == 0:  # Early exit guard
+                    # Check if we've reached end of sibling list
+                    is_valid = cur_index != -1
+
+                    # Use masked loads with safe address (0 when invalid)
+                    safe_cur_index = (
+                        cur_index * is_valid
+                    )  # 0 if invalid, cur_index if valid
+                    safe_index = bx * num_draft_tokens + safe_cur_index
+
+                    # Load draft token info (loads from index 0 when invalid, but we won't use it)
+                    draft_index = tl.load(retrieve_index_ptr + safe_index)
+                    draft_token = tl.load(candidates_ptr + safe_index)
+
+                    # Check for token match (only valid when is_valid is True)
+                    token_match = is_valid & (draft_token == target_token)
+
+                    # Accept token using predicated stores (only write if matched)
+                    tl.store(
+                        predicts_ptr + last_accept_retrieve_idx,
+                        target_token,
+                        mask=token_match,
+                    )
+                    next_num_accept_tokens = num_accept_tokens + 1
+                    tl.store(
+                        accept_index_ptr
+                        + bx * num_speculative_tokens
+                        + next_num_accept_tokens,
+                        draft_index,
+                        mask=token_match,
+                    )
+
+                    num_accept_tokens = num_accept_tokens + token_match
+                    last_accept_retrieve_idx = (
+                        token_match * draft_index
+                        + (~token_match) * last_accept_retrieve_idx
+                    )
+                    found_match = token_match * 1 + (~is_valid) * (-1)
+
+                    # Masked load: only load next sibling when no match (hardware predication)
+                    # When matched: returns cur_index (other); when not matched: loads sibling
+                    cur_index = tl.load(
+                        retrieve_next_sibling_ptr + safe_index,
+                        mask=~token_match
+                        & is_valid,  # Only load when valid and NOT matched
+                        other=cur_index,  # Keep cur_index when matched or invalid
+                    )
+
+            if found_match != 1:
+                should_continue = 0
+
+    # Store final results
+    tl.store(accept_token_num_ptr + bx, num_accept_tokens)
+
+    target_row = last_accept_retrieve_idx // num_draft_tokens
+    target_col = last_accept_retrieve_idx % num_draft_tokens
+    final_target = tl.load(
+        target_predict_ptr + target_row * num_draft_tokens + target_col
+    )
+    tl.store(predicts_ptr + last_accept_retrieve_idx, final_target)