[Refactor] spec: Move platform dispatch and Triton kernels into kernels/ops

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

@@ -47,6 +47,7 @@
     ("pad", "unpad_draft_extend_output"),
     ("pad", "seqlens_expand_triton"),
     ("position", "compute_position_triton"),
+    ("table_qk_norm_rope", "table_qk_norm_rope_"),
     ("dsv4.attn_metadata_kernels", "expand_prefill_causally"),
     ("dsv4.attn_metadata_kernels", "build_page_table_positions"),
     ("dsv4.attn_metadata_kernels", "build_causal_swa_page_indices"),

python/sglang/kernels/ops/attention/table_qk_norm_rope.py

@@ -0,0 +1,83 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _table_qk_norm_rope_kernel(
+    qkv_ptr,
+    q_weight_ptr,
+    k_weight_ptr,
+    cos_sin_ptr,
+    pos_ptr,
+    row_stride,
+    q_size,
+    NHQ: tl.constexpr,
+    D: tl.constexpr,
+    EPS: tl.constexpr,
+):
+    t = tl.program_id(0).to(tl.int64)
+    h = tl.program_id(1)
+    pos = tl.load(pos_ptr + t).to(tl.int64)
+
+    HALF: tl.constexpr = D // 2
+    half_ar = tl.arange(0, HALF)
+    d_ar = tl.arange(0, D)
+    cos = tl.load(cos_sin_ptr + pos * D + half_ar).to(tl.float32)
+    sin = tl.load(cos_sin_ptr + pos * D + HALF + half_ar).to(tl.float32)
+
+    is_q = h < NHQ
+    col0 = tl.where(is_q, h * D, q_size + (h - NHQ) * D).to(tl.int64)
+    w_ptr = tl.where(is_q, q_weight_ptr.to(tl.int64), k_weight_ptr.to(tl.int64)).to(
+        tl.pointer_type(tl.bfloat16)
+    )
+
+    row = qkv_ptr + t * row_stride + col0
+    x = tl.load(row + d_ar).to(tl.float32)
+    ms = tl.sum(x * x, 0) / D
+    inv = 1.0 / tl.sqrt(ms + EPS)
+    w1 = tl.load(w_ptr + half_ar).to(tl.float32)
+    w2 = tl.load(w_ptr + HALF + half_ar).to(tl.float32)
+    x1 = tl.load(row + half_ar).to(tl.float32) * inv * w1
+    x2 = tl.load(row + HALF + half_ar).to(tl.float32) * inv * w2
+    x1 = x1.to(tl.bfloat16).to(tl.float32)
+    x2 = x2.to(tl.bfloat16).to(tl.float32)
+    o1 = x1 * cos - x2 * sin
+    o2 = x2 * cos + x1 * sin
+    tl.store(row + half_ar, o1.to(tl.bfloat16))
+    tl.store(row + HALF + half_ar, o2.to(tl.bfloat16))
+
+
+def table_qk_norm_rope_(
+    qkv: torch.Tensor,
+    positions: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    num_q_heads: int,
+    num_k_heads: int,
+    head_dim: int,
+    eps: float,
+) -> None:
+    """In-place QK RMSNorm + table-lookup neox RoPE on the fused QKV tensor.
+
+    Reads cos/sin from the SAME rotary table as the unfused path, so there is
+    no large-position angle drift (unlike theta-recompute kernels). V columns
+    are untouched.
+    """
+    T = qkv.shape[0]
+    if T == 0:
+        return
+    grid = (T, num_q_heads + num_k_heads)
+    _table_qk_norm_rope_kernel[grid](
+        qkv,
+        q_weight,
+        k_weight,
+        cos_sin_cache,
+        positions,
+        qkv.stride(0),
+        num_q_heads * head_dim,
+        NHQ=num_q_heads,
+        D=head_dim,
+        EPS=eps,
+    )

python/sglang/kernels/ops/sampling/renorm.py

@@ -0,0 +1,37 @@
+"""Top-k / top-p probability renormalization, resolved once per platform.
+
+A name is ``None`` where the platform has no kernel for it.
+"""
+
+from sglang.srt.utils import is_cuda, is_hip, is_musa, is_npu
+from sglang.utils import LazyImport
+
+if is_cuda():
+    # Resolved on first call: CUDA builds without FlashInfer must still import
+    # this module.
+    top_k_renorm_prob = LazyImport("flashinfer.sampling", "top_k_renorm_probs")
+    top_p_renorm_prob = LazyImport("flashinfer.sampling", "top_p_renorm_probs")
+elif is_hip():
+    from sglang.kernels.ops.sampling.renorm_triton import (
+        top_k_renorm_probs_triton as top_k_renorm_prob,
+    )
+    from sglang.kernels.ops.sampling.renorm_triton import (
+        top_p_renorm_probs_triton as top_p_renorm_prob,
+    )
+elif is_npu():
+    try:
+        from sgl_kernel_npu.sample import top_k_renorm_prob, top_p_renorm_prob
+    except ImportError:
+        top_k_renorm_prob = None
+        top_p_renorm_prob = None
+elif is_musa():
+    try:
+        from sgl_kernel import top_k_renorm_prob, top_p_renorm_prob
+    except Exception:
+        top_k_renorm_prob = None
+        top_p_renorm_prob = None
+else:
+    top_k_renorm_prob = None
+    top_p_renorm_prob = None
+
+__all__ = ["top_k_renorm_prob", "top_p_renorm_prob"]

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

@@ -10,24 +10,31 @@
 
 # (module, public_fn) migrated from speculative/triton_ops.
 _TRITON_KERNELS = [
+    ("cache_locs", "assign_draft_cache_locs_contiguous_func"),
     ("cache_locs", "assign_extend_cache_locs_func"),
     ("cache_locs", "assign_extend_cache_locs_uniform_func"),
+    ("cache_locs", "fill_front_out_cache_loc_func"),
     ("cache_locs", "generate_draft_decode_kv_indices"),
     ("cache_locs", "generate_draft_decode_window_kv_indices"),
+    ("dflash", "compute_dflash_correct_drafts_and_bonus"),
     ("eagle", "fill_bonus_tokens"),
     ("eagle", "fill_accept_out_cache_loc"),
     ("gather_spec_extras", "gather_spec_extras"),
     ("multi_layer_eagle", "rotate_input_ids"),
     ("spec_tree", "sgl_build_tree_kernel_efficient_triton"),
     ("spec_tree", "verify_tree_greedy_kernel_triton"),
+    ("spec_tree", "build_tree_kernel_efficient_func"),
+    ("spec_tree", "verify_tree_greedy_func"),
     ("topk1", "draft_topk1_postprocess"),
+    ("uno_tree", "build_uno_tree_proposal"),
     ("ragged_verify_kernels", "pad_verify_lens_to_bucket"),
     ("ragged_verify_kernels", "build_qo_indptr"),
     ("lilicorr", "lilicorr_topk_lse"),
     ("lilicorr", "lilicorr_sample_path"),
     ("reject_sampling", "chain_speculative_sampling_triton"),
     ("row_argmax", "row_argmax"),
     ("reconstruct_tree", "reconstruct_indices_from_tree_mask_triton"),
+    ("reconstruct_tree", "reconstruct_indices_from_tree_mask_func"),
 ]
 for _mod, _fn in _TRITON_KERNELS:
     register_kernel(

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

@@ -25,7 +25,10 @@
 _is_xpu = is_xpu()
 
 if _is_cpu:
-    from sgl_kernel import assign_extend_cache_locs_cpu
+    from sgl_kernel import (
+        assign_draft_cache_locs_contiguous_cpu,
+        assign_extend_cache_locs_cpu,
+    )
 
 
 def _get_oot_speculative_cache_locs_fn() -> Callable[..., torch.Tensor] | None:
@@ -79,6 +82,47 @@ def assign_draft_cache_locs_contiguous(
                 )
 
 
+def assign_draft_cache_locs_contiguous_func(
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    seq_lens: torch.Tensor,
+    out_cache_loc: torch.Tensor,  # mutable
+    topk: int,
+    num_cols: int,
+    positions: Optional[torch.Tensor] = None,  # mutable
+    mrope: Optional[torch.Tensor] = None,  # mutable
+):
+    """Copy each request's ``topk * num_cols`` slots from ``seq_lens`` on, and
+    write the draft positions when ``positions`` is given."""
+    if _is_cpu:
+        assert positions is None and mrope is None, "CPU writes no positions"
+        assign_draft_cache_locs_contiguous_cpu(
+            req_pool_indices,
+            req_to_token,
+            seq_lens,
+            out_cache_loc,
+            req_to_token.shape[1],
+            topk,
+            num_cols,
+        )
+        return
+    bs = seq_lens.shape[0]
+    assign_draft_cache_locs_contiguous[(bs,)](
+        req_pool_indices,
+        req_to_token,
+        seq_lens,
+        out_cache_loc,
+        req_to_token.shape[1],
+        topk,
+        num_cols,
+        positions=positions,
+        mrope=mrope,
+        BS=bs,
+        WRITE_POSITIONS=positions is not None,
+        WRITE_MROPE=mrope is not None,
+    )
+
+
 @triton.jit
 def _translate_token_ids(ids, v2p, mask, page_size: tl.constexpr):
     """Virtual token ids -> physical ones through the page-level v2p table,
@@ -384,131 +428,6 @@ def generate_draft_decode_window_kv_indices(
     tl.store(window_kv_indptr + zid, base)
 
 
-@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 rebuild_compact_draft_req_to_token(
     draft_req_to_token,
@@ -800,3 +719,34 @@ def assign_extend_cache_locs_func(
         )
 
         return out_cache_loc
+
+
+def fill_front_out_cache_loc_func(
+    num_front: torch.Tensor,
+    out_cache_loc: torch.Tensor,
+    front_out_cache_loc: torch.Tensor,  # mutable
+):
+    """Pack the first ``num_front[i]`` slots of request i's window in ``out_cache_loc``."""
+    bs = num_front.numel()
+    window = out_cache_loc.view(bs, -1)
+    rows = torch.arange(bs, device=window.device)
+    starts = torch.zeros_like(num_front)
+    if _is_cpu:
+        assign_extend_cache_locs_cpu(
+            rows,
+            window.int(),
+            starts,
+            num_front,
+            front_out_cache_loc,
+            window.shape[1],
+        )
+        return
+    assign_extend_cache_locs[(bs,)](
+        rows,
+        window,
+        starts,
+        num_front,
+        front_out_cache_loc,
+        window.shape[1],
+        next_power_of_2(bs),
+    )