chore: import upstream snapshot with attribution

python/sglang/srt/speculative/dspark_components/kernels/dispatch.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+import torch
+
+
+def inputs_on_cuda(*args, **kwargs) -> bool:
+    """Route kernel dispatch by input placement: the first tensor argument
+    decides. CUDA inputs take the fused triton kernel; CPU inputs take the
+    torch reference implementation (triton is CUDA-only, and CPU-side callers
+    such as unit tests exercise the reference path)."""
+    for value in (*args, *kwargs.values()):
+        if isinstance(value, torch.Tensor):
+            return value.is_cuda
+    raise AssertionError("kernel dispatch requires at least one tensor argument")

python/sglang/srt/speculative/dspark_components/kernels/dspark_accept.py

@@ -0,0 +1,862 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import msgspec
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2
+from sglang.srt.speculative.dflash_utils import (
+    _get_or_create_chain_verify_buffers,
+    build_dflash_verify_target_probs,
+    compute_dflash_correct_drafts_and_bonus,
+)
+from sglang.srt.speculative.dspark_components.kernels.dispatch import inputs_on_cuda
+from sglang.srt.speculative.reject_sampling import chain_speculative_sampling_triton
+
+
+class AcceptSampling:
+    @classmethod
+    def execute(
+        cls, *args, **kwargs
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        candidates: torch.Tensor,
+        target_logits: torch.Tensor,
+        draft_probs: torch.Tensor,
+        sampling_info,
+        draft_input: DFlashDraftInputV2,
+        gamma: int,
+        verify_num_draft_tokens: int,
+        cutoff_verify_lens: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return accept_sampling(
+            candidates=candidates,
+            target_logits=target_logits,
+            draft_probs=draft_probs,
+            sampling_info=sampling_info,
+            draft_input=draft_input,
+            gamma=gamma,
+            verify_num_draft_tokens=verify_num_draft_tokens,
+            cutoff_verify_lens=cutoff_verify_lens,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        candidates: torch.Tensor,
+        target_logits: torch.Tensor,
+        draft_probs: torch.Tensor,
+        sampling_info,
+        draft_input: DFlashDraftInputV2,
+        gamma: int,
+        verify_num_draft_tokens: int,
+        cutoff_verify_lens: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return accept_sampling_triton(
+            candidates=candidates,
+            target_logits=target_logits,
+            draft_probs=draft_probs,
+            sampling_info=sampling_info,
+            draft_input=draft_input,
+            gamma=gamma,
+            verify_num_draft_tokens=verify_num_draft_tokens,
+            cutoff_verify_lens=cutoff_verify_lens,
+        )
+
+
+def _accept_sampling_core(
+    *,
+    candidates: torch.Tensor,
+    target_logits: torch.Tensor,
+    draft_probs: torch.Tensor,
+    sampling_info,
+    draft_input: DFlashDraftInputV2,
+    gamma: int,
+    verify_num_draft_tokens: int,
+    cutoff_verify_lens: Optional[torch.Tensor],
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    bs = candidates.shape[0]
+    device = candidates.device
+    if not sampling_info.need_top_k_sampling and not sampling_info.need_top_p_sampling:
+        target_probs = SoftmaxTemp.execute(
+            logits=target_logits,
+            temperatures=sampling_info.temperatures,
+            rows_per_request=verify_num_draft_tokens,
+        ).view(bs, verify_num_draft_tokens, -1)
+    else:
+        target_probs = build_dflash_verify_target_probs(
+            next_token_logits=target_logits,
+            sampling_info=sampling_info,
+            draft_token_num=verify_num_draft_tokens,
+            bs=bs,
+            max_top_k=draft_input.max_top_k,
+            uniform_top_k_value=draft_input.uniform_top_k_value,
+        )
+    (
+        retrieve_index,
+        retrieve_next_token,
+        retrieve_next_sibling,
+        predicts,
+        accept_index,
+        accept_token_num,
+    ) = _get_or_create_chain_verify_buffers(
+        bs=bs,
+        draft_token_num=verify_num_draft_tokens,
+        device=device,
+    )
+    uniform_samples = torch.rand((bs, gamma), dtype=torch.float32, device=device)
+    uniform_samples_final = torch.rand((bs,), dtype=torch.float32, device=device)
+    chain_speculative_sampling_triton(
+        predicts=predicts,
+        accept_index=accept_index,
+        accept_token_num=accept_token_num,
+        candidates=candidates,
+        retrive_index=retrieve_index,
+        retrive_next_token=retrieve_next_token,
+        retrive_next_sibling=retrieve_next_sibling,
+        uniform_samples=uniform_samples,
+        uniform_samples_for_final_sampling=uniform_samples_final,
+        target_probs=target_probs,
+        draft_probs=draft_probs,
+        threshold_single=1.0,
+        threshold_acc=1.0,
+        deterministic=True,
+    )
+    correct_len = accept_token_num
+    if cutoff_verify_lens is not None:
+        correct_len, cap_trim_lens = CapCorrectLen.execute(
+            correct_len=correct_len, verify_lens=cutoff_verify_lens
+        )
+    else:
+        cap_trim_lens = torch.zeros_like(correct_len)
+    return correct_len, cap_trim_lens, accept_index, predicts
+
+
+def accept_sampling(
+    *,
+    candidates: torch.Tensor,
+    target_logits: torch.Tensor,
+    draft_probs: torch.Tensor,
+    sampling_info,
+    draft_input: DFlashDraftInputV2,
+    gamma: int,
+    verify_num_draft_tokens: int,
+    cutoff_verify_lens: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    bs = candidates.shape[0]
+    device = candidates.device
+    correct_len, cap_trim_lens, accept_index, predicts = _accept_sampling_core(
+        candidates=candidates,
+        target_logits=target_logits,
+        draft_probs=draft_probs,
+        sampling_info=sampling_info,
+        draft_input=draft_input,
+        gamma=gamma,
+        verify_num_draft_tokens=verify_num_draft_tokens,
+        cutoff_verify_lens=cutoff_verify_lens,
+    )
+    row_ids = torch.arange(bs, dtype=torch.long, device=device)
+    accept_pos = accept_index[row_ids, correct_len.to(torch.long)].to(torch.long)
+    bonus = predicts[accept_pos].to(torch.int64)
+    return correct_len, bonus, cap_trim_lens
+
+
+@triton.jit
+def _gather_two_level_bonus_kernel(
+    accept_index_ptr,
+    predicts_ptr,
+    correct_len_ptr,
+    out_ptr,
+    cols,
+    n,
+    BLOCK: tl.constexpr,
+):
+    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < n
+    cl = tl.load(correct_len_ptr + offs, mask=mask, other=0).to(tl.int64)
+    accept_pos = tl.load(accept_index_ptr + offs * cols + cl, mask=mask, other=0).to(
+        tl.int64
+    )
+    bonus = tl.load(predicts_ptr + accept_pos, mask=mask, other=0)
+    tl.store(out_ptr + offs, bonus.to(tl.int64), mask=mask)
+
+
+def gather_two_level_bonus_triton(
+    *,
+    accept_index: torch.Tensor,
+    predicts: torch.Tensor,
+    correct_len: torch.Tensor,
+) -> torch.Tensor:
+    bs, cols = accept_index.shape
+    accept_index = accept_index.contiguous()
+    predicts = predicts.contiguous()
+    correct_len = correct_len.contiguous()
+    out = torch.empty(bs, dtype=torch.int64, device=accept_index.device)
+    BLOCK = 256
+    grid = (triton.cdiv(bs, BLOCK),)
+    _gather_two_level_bonus_kernel[grid](
+        accept_index, predicts, correct_len, out, cols, bs, BLOCK=BLOCK
+    )
+    return out
+
+
+def accept_sampling_triton(
+    *,
+    candidates: torch.Tensor,
+    target_logits: torch.Tensor,
+    draft_probs: torch.Tensor,
+    sampling_info,
+    draft_input: DFlashDraftInputV2,
+    gamma: int,
+    verify_num_draft_tokens: int,
+    cutoff_verify_lens: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    correct_len, cap_trim_lens, accept_index, predicts = _accept_sampling_core(
+        candidates=candidates,
+        target_logits=target_logits,
+        draft_probs=draft_probs,
+        sampling_info=sampling_info,
+        draft_input=draft_input,
+        gamma=gamma,
+        verify_num_draft_tokens=verify_num_draft_tokens,
+        cutoff_verify_lens=cutoff_verify_lens,
+    )
+    bonus = gather_two_level_bonus_triton(
+        accept_index=accept_index, predicts=predicts, correct_len=correct_len
+    )
+    return correct_len, bonus, cap_trim_lens
+
+
+try:
+    from flashinfer.sampling import softmax as _flashinfer_softmax
+except ImportError:
+    _flashinfer_softmax = None
+
+
+class SoftmaxTemp:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> torch.Tensor:
+        if not inputs_on_cuda(*args, **kwargs):
+            return cls.torch(*args, **kwargs)
+        if _flashinfer_softmax is not None:
+            return cls.flashinfer(*args, **kwargs)
+        return cls.triton(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        logits: torch.Tensor,
+        temperatures: torch.Tensor,
+        rows_per_request: int,
+    ) -> torch.Tensor:
+        return softmax_temp(
+            logits=logits,
+            temperatures=temperatures,
+            rows_per_request=rows_per_request,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        logits: torch.Tensor,
+        temperatures: torch.Tensor,
+        rows_per_request: int,
+    ) -> torch.Tensor:
+        return softmax_temp_triton(
+            logits=logits,
+            temperatures=temperatures,
+            rows_per_request=rows_per_request,
+        )
+
+    @classmethod
+    def flashinfer(
+        cls,
+        *,
+        logits: torch.Tensor,
+        temperatures: torch.Tensor,
+        rows_per_request: int,
+    ) -> torch.Tensor:
+        return softmax_temp_flashinfer(
+            logits=logits,
+            temperatures=temperatures,
+            rows_per_request=rows_per_request,
+        )
+
+
+def softmax_temp(
+    *,
+    logits: torch.Tensor,
+    temperatures: torch.Tensor,
+    rows_per_request: int,
+) -> torch.Tensor:
+    num_rows = logits.shape[0]
+    bs = num_rows // rows_per_request
+    assert (
+        bs * rows_per_request == num_rows
+    ), f"num_rows {num_rows} not divisible by rows_per_request {rows_per_request}"
+    temp_per_row = torch.repeat_interleave(
+        temperatures.reshape(bs).to(torch.float32), rows_per_request, dim=0
+    )
+    scaled = logits.to(torch.float32) / temp_per_row[:, None]
+    return torch.softmax(scaled, dim=-1)
+
+
+@triton.jit
+def _softmax_temp_kernel(
+    logits_ptr,
+    temp_ptr,
+    out_ptr,
+    vocab,
+    rows_per_request,
+    logits_row_stride,
+    BLOCK_V: tl.constexpr,
+):
+    row = tl.program_id(0)
+    temp = tl.load(temp_ptr + row // rows_per_request).to(tl.float32)
+    base = logits_ptr + row.to(tl.int64) * logits_row_stride
+    out_base = out_ptr + row.to(tl.int64) * vocab
+
+    row_max = -float("inf")
+    for v0 in range(0, vocab, BLOCK_V):
+        offs = v0 + tl.arange(0, BLOCK_V)
+        vmask = offs < vocab
+        x = tl.load(base + offs, mask=vmask, other=-float("inf")).to(tl.float32)
+        x = x / temp
+        row_max = tl.maximum(row_max, tl.max(x, axis=0))
+
+    sum_exp = 0.0
+    for v0 in range(0, vocab, BLOCK_V):
+        offs = v0 + tl.arange(0, BLOCK_V)
+        vmask = offs < vocab
+        x = tl.load(base + offs, mask=vmask, other=-float("inf")).to(tl.float32)
+        x = x / temp
+        e = tl.exp(x - row_max)
+        e = tl.where(vmask, e, 0.0)
+        sum_exp += tl.sum(e, axis=0)
+
+    for v0 in range(0, vocab, BLOCK_V):
+        offs = v0 + tl.arange(0, BLOCK_V)
+        vmask = offs < vocab
+        x = tl.load(base + offs, mask=vmask, other=-float("inf")).to(tl.float32)
+        x = x / temp
+        e = tl.exp(x - row_max)
+        tl.store(out_base + offs, e / sum_exp, mask=vmask)
+
+
+def softmax_temp_triton(
+    *,
+    logits: torch.Tensor,
+    temperatures: torch.Tensor,
+    rows_per_request: int,
+) -> torch.Tensor:
+    num_rows, vocab = logits.shape[0], logits.shape[-1]
+    bs = num_rows // rows_per_request
+    assert (
+        bs * rows_per_request == num_rows
+    ), f"num_rows {num_rows} not divisible by rows_per_request {rows_per_request}"
+    temperatures = temperatures.reshape(bs).to(torch.float32).contiguous()
+    out = torch.empty((num_rows, vocab), dtype=torch.float32, device=logits.device)
+    BLOCK_V = 4096
+    _softmax_temp_kernel[(num_rows,)](
+        logits,
+        temperatures,
+        out,
+        vocab,
+        rows_per_request,
+        logits.stride(0),
+        BLOCK_V=BLOCK_V,
+    )
+    return out
+
+
+def softmax_temp_flashinfer(
+    *,
+    logits: torch.Tensor,
+    temperatures: torch.Tensor,
+    rows_per_request: int,
+) -> torch.Tensor:
+    if _flashinfer_softmax is None:
+        raise RuntimeError(
+            "softmax_temp_flashinfer requires flashinfer.sampling.softmax, "
+            "which is unavailable in this environment"
+        )
+    num_rows, vocab = logits.shape[0], logits.shape[-1]
+    bs = num_rows // rows_per_request
+    assert (
+        bs * rows_per_request == num_rows
+    ), f"num_rows {num_rows} not divisible by rows_per_request {rows_per_request}"
+    temp_per_row = torch.repeat_interleave(
+        temperatures.reshape(bs).to(torch.float32), rows_per_request, dim=0
+    ).contiguous()
+    logits_2d = logits.to(torch.float32).contiguous()
+    return _flashinfer_softmax(logits=logits_2d, temperature=temp_per_row)
+
+
+class MixedAcceptSelectResult(msgspec.Struct):
+    correct_len: torch.Tensor
+    bonus: torch.Tensor
+    cap_trim_lens: torch.Tensor
+
+
+class SelectMixedAccept:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> MixedAcceptSelectResult:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        greedy_mask: torch.Tensor,
+        greedy_len: torch.Tensor,
+        greedy_bonus: torch.Tensor,
+        greedy_trim: torch.Tensor,
+        sampling_len: torch.Tensor,
+        sampling_bonus: torch.Tensor,
+        sampling_trim: torch.Tensor,
+    ) -> MixedAcceptSelectResult:
+        return select_mixed_accept(
+            greedy_mask=greedy_mask,
+            greedy_len=greedy_len,
+            greedy_bonus=greedy_bonus,
+            greedy_trim=greedy_trim,
+            sampling_len=sampling_len,
+            sampling_bonus=sampling_bonus,
+            sampling_trim=sampling_trim,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        greedy_mask: torch.Tensor,
+        greedy_len: torch.Tensor,
+        greedy_bonus: torch.Tensor,
+        greedy_trim: torch.Tensor,
+        sampling_len: torch.Tensor,
+        sampling_bonus: torch.Tensor,
+        sampling_trim: torch.Tensor,
+    ) -> MixedAcceptSelectResult:
+        return select_mixed_accept_triton(
+            greedy_mask=greedy_mask,
+            greedy_len=greedy_len,
+            greedy_bonus=greedy_bonus,
+            greedy_trim=greedy_trim,
+            sampling_len=sampling_len,
+            sampling_bonus=sampling_bonus,
+            sampling_trim=sampling_trim,
+        )
+
+
+def select_mixed_accept(
+    *,
+    greedy_mask: torch.Tensor,
+    greedy_len: torch.Tensor,
+    greedy_bonus: torch.Tensor,
+    greedy_trim: torch.Tensor,
+    sampling_len: torch.Tensor,
+    sampling_bonus: torch.Tensor,
+    sampling_trim: torch.Tensor,
+) -> MixedAcceptSelectResult:
+    correct_len = torch.where(
+        greedy_mask, greedy_len.to(sampling_len.dtype), sampling_len
+    )
+    bonus = torch.where(greedy_mask, greedy_bonus, sampling_bonus)
+    cap_trim_lens = torch.where(
+        greedy_mask, greedy_trim.to(sampling_trim.dtype), sampling_trim
+    )
+    return MixedAcceptSelectResult(
+        correct_len=correct_len, bonus=bonus, cap_trim_lens=cap_trim_lens
+    )
+
+
+@triton.jit
+def _mixed_accept_select_kernel(
+    greedy_mask_ptr,
+    greedy_len_ptr,
+    greedy_bonus_ptr,
+    greedy_trim_ptr,
+    sampling_len_ptr,
+    sampling_bonus_ptr,
+    sampling_trim_ptr,
+    correct_len_ptr,
+    bonus_ptr,
+    cap_trim_ptr,
+    bs,
+    BLOCK: tl.constexpr,
+):
+    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < bs
+    is_greedy = tl.load(greedy_mask_ptr + offs, mask=mask, other=0) != 0
+
+    g_len = tl.load(greedy_len_ptr + offs, mask=mask, other=0)
+    s_len = tl.load(sampling_len_ptr + offs, mask=mask, other=0)
+    tl.store(correct_len_ptr + offs, tl.where(is_greedy, g_len, s_len), mask=mask)
+
+    g_bonus = tl.load(greedy_bonus_ptr + offs, mask=mask, other=0)
+    s_bonus = tl.load(sampling_bonus_ptr + offs, mask=mask, other=0)
+    tl.store(bonus_ptr + offs, tl.where(is_greedy, g_bonus, s_bonus), mask=mask)
+
+    g_trim = tl.load(greedy_trim_ptr + offs, mask=mask, other=0)
+    s_trim = tl.load(sampling_trim_ptr + offs, mask=mask, other=0)
+    tl.store(cap_trim_ptr + offs, tl.where(is_greedy, g_trim, s_trim), mask=mask)
+
+
+def select_mixed_accept_triton(
+    *,
+    greedy_mask: torch.Tensor,
+    greedy_len: torch.Tensor,
+    greedy_bonus: torch.Tensor,
+    greedy_trim: torch.Tensor,
+    sampling_len: torch.Tensor,
+    sampling_bonus: torch.Tensor,
+    sampling_trim: torch.Tensor,
+) -> MixedAcceptSelectResult:
+    bs = greedy_mask.shape[0]
+    device = greedy_mask.device
+
+    correct_len = torch.empty(bs, dtype=sampling_len.dtype, device=device)
+    bonus = torch.empty(bs, dtype=sampling_bonus.dtype, device=device)
+    cap_trim_lens = torch.empty(bs, dtype=sampling_trim.dtype, device=device)
+    BLOCK = 256
+    _mixed_accept_select_kernel[(triton.cdiv(bs, BLOCK),)](
+        greedy_mask,
+        greedy_len,
+        greedy_bonus,
+        greedy_trim,
+        sampling_len,
+        sampling_bonus,
+        sampling_trim,
+        correct_len,
+        bonus,
+        cap_trim_lens,
+        bs,
+        BLOCK=BLOCK,
+    )
+    return MixedAcceptSelectResult(
+        correct_len=correct_len, bonus=bonus, cap_trim_lens=cap_trim_lens
+    )
+
+
+class AcceptGreedy:
+    @classmethod
+    def execute(
+        cls, *args, **kwargs
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        candidates: torch.Tensor,
+        target_logits: torch.Tensor,
+        verify_num_draft_tokens: int,
+        cutoff_verify_lens: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return accept_greedy(
+            candidates=candidates,
+            target_logits=target_logits,
+            verify_num_draft_tokens=verify_num_draft_tokens,
+            cutoff_verify_lens=cutoff_verify_lens,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        candidates: torch.Tensor,
+        target_logits: torch.Tensor,
+        verify_num_draft_tokens: int,
+        cutoff_verify_lens: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return accept_greedy_triton(
+            candidates=candidates,
+            target_logits=target_logits,
+            verify_num_draft_tokens=verify_num_draft_tokens,
+            cutoff_verify_lens=cutoff_verify_lens,
+        )
+
+
+def accept_greedy(
+    *,
+    candidates: torch.Tensor,
+    target_logits: torch.Tensor,
+    verify_num_draft_tokens: int,
+    cutoff_verify_lens: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    bs = candidates.shape[0]
+    target_predict = torch.argmax(target_logits, dim=-1).view(
+        bs, verify_num_draft_tokens
+    )
+    correct_len, bonus = compute_dflash_correct_drafts_and_bonus(
+        candidates=candidates,
+        target_predict=target_predict,
+    )
+    cap_trim_lens = torch.zeros_like(correct_len)
+    if cutoff_verify_lens is not None:
+        correct_len, cap_trim_lens = CapCorrectLen.execute(
+            correct_len=correct_len, verify_lens=cutoff_verify_lens
+        )
+        row_ids = torch.arange(bs, device=target_predict.device)
+        bonus = target_predict[row_ids, correct_len.to(torch.long)].to(torch.int64)
+    return correct_len, bonus, cap_trim_lens
+
+
+@triton.jit
+def _gather_row_bonus_kernel(
+    table_ptr,
+    idx_ptr,
+    out_ptr,
+    cols,
+    n,
+    BLOCK: tl.constexpr,
+):
+    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < n
+    idx = tl.load(idx_ptr + offs, mask=mask, other=0).to(tl.int64)
+    val = tl.load(table_ptr + offs * cols + idx, mask=mask, other=0)
+    tl.store(out_ptr + offs, val.to(tl.int64), mask=mask)
+
+
+def gather_row_bonus_triton(*, table: torch.Tensor, idx: torch.Tensor) -> torch.Tensor:
+    bs, cols = table.shape
+    table = table.contiguous()
+    idx = idx.contiguous()
+    out = torch.empty(bs, dtype=torch.int64, device=table.device)
+    BLOCK = 256
+    grid = (triton.cdiv(bs, BLOCK),)
+    _gather_row_bonus_kernel[grid](table, idx, out, cols, bs, BLOCK=BLOCK)
+    return out
+
+
+def accept_greedy_triton(
+    *,
+    candidates: torch.Tensor,
+    target_logits: torch.Tensor,
+    verify_num_draft_tokens: int,
+    cutoff_verify_lens: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    bs = candidates.shape[0]
+    target_predict = torch.argmax(target_logits, dim=-1).view(
+        bs, verify_num_draft_tokens
+    )
+    correct_len, bonus = compute_dflash_correct_drafts_and_bonus(
+        candidates=candidates,
+        target_predict=target_predict,
+    )
+    cap_trim_lens = torch.zeros_like(correct_len)
+    if cutoff_verify_lens is not None:
+        correct_len, cap_trim_lens = CapCorrectLen.execute(
+            correct_len=correct_len, verify_lens=cutoff_verify_lens
+        )
+        bonus = gather_row_bonus_triton(table=target_predict, idx=correct_len)
+    return correct_len, bonus, cap_trim_lens
+
+
+class FinalizeAcceptLensResult(msgspec.Struct):
+    commit_lens: torch.Tensor
+    new_seq_lens: torch.Tensor
+    cap_trim_lens: torch.Tensor
+
+
+class FinalizeAcceptLens:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> FinalizeAcceptLensResult:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        correct_len: torch.Tensor,
+        cap_trim_lens: torch.Tensor,
+        prefix_lens: torch.Tensor,
+    ) -> FinalizeAcceptLensResult:
+        return finalize_accept_lens(
+            correct_len=correct_len,
+            cap_trim_lens=cap_trim_lens,
+            prefix_lens=prefix_lens,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        correct_len: torch.Tensor,
+        cap_trim_lens: torch.Tensor,
+        prefix_lens: torch.Tensor,
+    ) -> FinalizeAcceptLensResult:
+        return finalize_accept_lens_triton(
+            correct_len=correct_len,
+            cap_trim_lens=cap_trim_lens,
+            prefix_lens=prefix_lens,
+        )
+
+
+def finalize_accept_lens(
+    *,
+    correct_len: torch.Tensor,
+    cap_trim_lens: torch.Tensor,
+    prefix_lens: torch.Tensor,
+) -> FinalizeAcceptLensResult:
+    commit_lens = correct_len.to(torch.int32) + 1
+    new_seq_lens = prefix_lens + commit_lens.to(prefix_lens.dtype)
+    return FinalizeAcceptLensResult(
+        commit_lens=commit_lens,
+        new_seq_lens=new_seq_lens,
+        cap_trim_lens=cap_trim_lens.to(torch.int32),
+    )
+
+
+@triton.jit
+def _finalize_accept_lens_kernel(
+    correct_len_ptr,
+    cap_trim_ptr,
+    prefix_lens_ptr,
+    commit_lens_ptr,
+    new_seq_lens_ptr,
+    cap_trim_out_ptr,
+    bs,
+    BLOCK: tl.constexpr,
+):
+    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < bs
+    commit = tl.load(correct_len_ptr + offs, mask=mask, other=0).to(tl.int32) + 1
+    prefix = tl.load(prefix_lens_ptr + offs, mask=mask, other=0)
+    trim = tl.load(cap_trim_ptr + offs, mask=mask, other=0).to(tl.int32)
+    tl.store(commit_lens_ptr + offs, commit, mask=mask)
+    tl.store(new_seq_lens_ptr + offs, prefix + commit, mask=mask)
+    tl.store(cap_trim_out_ptr + offs, trim, mask=mask)
+
+
+def finalize_accept_lens_triton(
+    *,
+    correct_len: torch.Tensor,
+    cap_trim_lens: torch.Tensor,
+    prefix_lens: torch.Tensor,
+) -> FinalizeAcceptLensResult:
+    bs = correct_len.shape[0]
+    device = correct_len.device
+
+    commit_lens = torch.empty(bs, dtype=torch.int32, device=device)
+    new_seq_lens = torch.empty(bs, dtype=prefix_lens.dtype, device=device)
+    cap_trim_out = torch.empty(bs, dtype=torch.int32, device=device)
+    BLOCK = 256
+    _finalize_accept_lens_kernel[(triton.cdiv(bs, BLOCK),)](
+        correct_len,
+        cap_trim_lens,
+        prefix_lens,
+        commit_lens,
+        new_seq_lens,
+        cap_trim_out,
+        bs,
+        BLOCK=BLOCK,
+    )
+    return FinalizeAcceptLensResult(
+        commit_lens=commit_lens,
+        new_seq_lens=new_seq_lens,
+        cap_trim_lens=cap_trim_out,
+    )
+
+
+class CapCorrectLen:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> tuple[torch.Tensor, torch.Tensor]:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        correct_len: torch.Tensor,
+        verify_lens: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return cap_correct_len(
+            correct_len=correct_len,
+            verify_lens=verify_lens,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        correct_len: torch.Tensor,
+        verify_lens: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return cap_correct_len_triton(
+            correct_len=correct_len,
+            verify_lens=verify_lens,
+        )
+
+
+def cap_correct_len(
+    *,
+    correct_len: torch.Tensor,
+    verify_lens: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    ell_r = (verify_lens.to(device=correct_len.device) - 1).to(correct_len.dtype)
+    capped = torch.minimum(correct_len, ell_r)
+    cap_trim_lens = correct_len - capped
+    return capped, cap_trim_lens
+
+
+@triton.jit
+def _cap_correct_len_kernel(
+    correct_len_ptr,
+    verify_lens_ptr,
+    capped_ptr,
+    trim_ptr,
+    n,
+    BLOCK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < n
+    cl = tl.load(correct_len_ptr + offs, mask=mask, other=0).to(tl.int64)
+    vl = tl.load(verify_lens_ptr + offs, mask=mask, other=0).to(tl.int64)
+    ell = vl - 1
+    capped = tl.minimum(cl, ell)
+    trim = cl - capped
+    tl.store(capped_ptr + offs, capped, mask=mask)
+    tl.store(trim_ptr + offs, trim, mask=mask)
+
+
+def cap_correct_len_triton(
+    *,
+    correct_len: torch.Tensor,
+    verify_lens: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    device = correct_len.device
+    correct_len = correct_len.contiguous()
+    verify_lens = verify_lens.to(device=device).contiguous()
+    n = correct_len.shape[0]
+    capped = torch.empty_like(correct_len)
+    trim = torch.empty_like(correct_len)
+    BLOCK = 1024
+    grid = (triton.cdiv(n, BLOCK),)
+    _cap_correct_len_kernel[grid](
+        correct_len, verify_lens, capped, trim, n, BLOCK=BLOCK
+    )
+    return capped, trim

python/sglang/srt/speculative/dspark_components/kernels/dspark_attn_metadata.py

@@ -0,0 +1,491 @@
+from __future__ import annotations
+
+from typing import Tuple
+
+import msgspec
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.speculative.dspark_components.kernels.dispatch import inputs_on_cuda
+from sglang.srt.utils import ceil_align
+
+
+class DsparkWindowGather(msgspec.Struct, frozen=True):
+    num_q: int
+    bs: int
+    context_lens: torch.Tensor
+    req_pool_indices_per_request: torch.Tensor
+    offsets: torch.Tensor
+    invalid: torch.Tensor
+
+
+class ComputeDsparkWindowGather:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> DsparkWindowGather:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        seq_lens_casual: torch.Tensor,
+        req_pool_indices_repeated: torch.Tensor,
+        block_size: int,
+        swa_window: int,
+    ) -> DsparkWindowGather:
+        return compute_dspark_window_gather(
+            seq_lens_casual=seq_lens_casual,
+            req_pool_indices_repeated=req_pool_indices_repeated,
+            block_size=block_size,
+            swa_window=swa_window,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        seq_lens_casual: torch.Tensor,
+        req_pool_indices_repeated: torch.Tensor,
+        block_size: int,
+        swa_window: int,
+    ) -> DsparkWindowGather:
+        return compute_dspark_window_gather_triton(
+            seq_lens_casual=seq_lens_casual,
+            req_pool_indices_repeated=req_pool_indices_repeated,
+            block_size=block_size,
+            swa_window=swa_window,
+        )
+
+
+class BuildDsparkSwaPageIndices:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> Tuple[torch.Tensor, torch.Tensor]:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        req_to_token: torch.Tensor,
+        full_to_swa_mapping: torch.Tensor,
+        req_pool_indices_per_request: torch.Tensor,
+        offsets: torch.Tensor,
+        invalid: torch.Tensor,
+        out_loc: torch.Tensor,
+        context_lens: torch.Tensor,
+        block_size: int,
+        swa_window: int,
+        page_index_aligned_size: int,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        return build_dspark_swa_page_indices(
+            req_to_token=req_to_token,
+            full_to_swa_mapping=full_to_swa_mapping,
+            req_pool_indices_per_request=req_pool_indices_per_request,
+            offsets=offsets,
+            invalid=invalid,
+            out_loc=out_loc,
+            context_lens=context_lens,
+            block_size=block_size,
+            swa_window=swa_window,
+            page_index_aligned_size=page_index_aligned_size,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        req_to_token: torch.Tensor,
+        full_to_swa_mapping: torch.Tensor,
+        req_pool_indices_per_request: torch.Tensor,
+        offsets: torch.Tensor,
+        invalid: torch.Tensor,
+        out_loc: torch.Tensor,
+        context_lens: torch.Tensor,
+        block_size: int,
+        swa_window: int,
+        page_index_aligned_size: int,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        return build_dspark_swa_page_indices_triton(
+            req_to_token=req_to_token,
+            full_to_swa_mapping=full_to_swa_mapping,
+            req_pool_indices_per_request=req_pool_indices_per_request,
+            offsets=offsets,
+            out_loc=out_loc,
+            context_lens=context_lens,
+            block_size=block_size,
+            swa_window=swa_window,
+            page_index_aligned_size=page_index_aligned_size,
+        )
+
+
+def compute_dspark_window_gather(
+    *,
+    seq_lens_casual: torch.Tensor,
+    req_pool_indices_repeated: torch.Tensor,
+    block_size: int,
+    swa_window: int,
+) -> DsparkWindowGather:
+    seq_lens_casual = seq_lens_casual.to(torch.int32)
+    num_q = seq_lens_casual.size(0)
+    assert num_q % block_size == 0, (
+        f"DSpark draft block forward must be uniform-gamma: num_q={num_q} not "
+        f"divisible by block_size={block_size}."
+    )
+    bs = num_q // block_size
+    device = seq_lens_casual.device
+
+    first_token = torch.arange(bs, device=device, dtype=torch.int64) * block_size
+    prefix_lens = (seq_lens_casual[first_token] - 1).to(torch.int32)
+    context_lens = torch.clamp(prefix_lens, max=swa_window).to(torch.int32)
+    req_pool_indices_per_request = req_pool_indices_repeated[first_token]
+
+    offsets = (
+        prefix_lens.to(torch.int64).unsqueeze(1)
+        - swa_window
+        + torch.arange(swa_window, device=device, dtype=torch.int64).unsqueeze(0)
+    )
+    invalid = offsets < 0
+    offsets = offsets.clamp(min=0)
+
+    return DsparkWindowGather(
+        num_q=num_q,
+        bs=bs,
+        context_lens=context_lens,
+        req_pool_indices_per_request=req_pool_indices_per_request,
+        offsets=offsets,
+        invalid=invalid,
+    )
+
+
+@triton.jit
+def _window_gather_kernel(
+    seq_lens_casual_ptr,
+    req_pool_rep_ptr,
+    context_lens_ptr,
+    req_pool_out_ptr,
+    offsets_ptr,
+    invalid_ptr,
+    block_size,
+    swa_window,
+    W_BLOCK: tl.constexpr,
+):
+    i = tl.program_id(0)
+    ft = i * block_size
+    prefix = tl.load(seq_lens_casual_ptr + ft).to(tl.int64) - 1
+    tl.store(context_lens_ptr + i, tl.minimum(prefix, swa_window).to(tl.int32))
+    tl.store(req_pool_out_ptr + i, tl.load(req_pool_rep_ptr + ft))
+    col = tl.arange(0, W_BLOCK)
+    cmask = col < swa_window
+    off = prefix - swa_window + col
+    tl.store(invalid_ptr + i * swa_window + col, off < 0, mask=cmask)
+    tl.store(offsets_ptr + i * swa_window + col, tl.maximum(off, 0), mask=cmask)
+
+
+def compute_dspark_window_gather_triton(
+    *,
+    seq_lens_casual: torch.Tensor,
+    req_pool_indices_repeated: torch.Tensor,
+    block_size: int,
+    swa_window: int,
+) -> DsparkWindowGather:
+    seq_lens_casual = seq_lens_casual.to(torch.int32).contiguous()
+    num_q = seq_lens_casual.size(0)
+    assert num_q % block_size == 0, (
+        f"DSpark draft block forward must be uniform-gamma: num_q={num_q} not "
+        f"divisible by block_size={block_size}."
+    )
+    bs = num_q // block_size
+    device = seq_lens_casual.device
+    req_pool_indices_repeated = req_pool_indices_repeated.to(device=device).contiguous()
+    context_lens = torch.empty(bs, dtype=torch.int32, device=device)
+    req_pool_out = torch.empty(bs, dtype=req_pool_indices_repeated.dtype, device=device)
+    offsets = torch.empty((bs, swa_window), dtype=torch.int64, device=device)
+    invalid = torch.empty((bs, swa_window), dtype=torch.bool, device=device)
+    W_BLOCK = triton.next_power_of_2(swa_window)
+    _window_gather_kernel[(bs,)](
+        seq_lens_casual,
+        req_pool_indices_repeated,
+        context_lens,
+        req_pool_out,
+        offsets,
+        invalid,
+        block_size,
+        swa_window,
+        W_BLOCK=W_BLOCK,
+    )
+    return DsparkWindowGather(
+        num_q=num_q,
+        bs=bs,
+        context_lens=context_lens,
+        req_pool_indices_per_request=req_pool_out,
+        offsets=offsets,
+        invalid=invalid,
+    )
+
+
+def build_dspark_swa_page_indices(
+    *,
+    req_to_token: torch.Tensor,
+    full_to_swa_mapping: torch.Tensor,
+    req_pool_indices_per_request: torch.Tensor,
+    offsets: torch.Tensor,
+    invalid: torch.Tensor,
+    out_loc: torch.Tensor,
+    context_lens: torch.Tensor,
+    block_size: int,
+    swa_window: int,
+    page_index_aligned_size: int,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    if offsets.ndim != 2 or offsets.shape[1] != swa_window:
+        raise ValueError(
+            "offsets must be [bs, swa_window]; "
+            f"got shape={tuple(offsets.shape)} (swa_window={swa_window})."
+        )
+    bs = offsets.shape[0]
+    device = offsets.device
+    context_lens = context_lens.to(device=device, dtype=torch.int32)
+
+    window_full_locs = req_to_token[
+        req_pool_indices_per_request[:, None].to(torch.int64), offsets
+    ]
+    window_full_locs = window_full_locs.masked_fill(invalid, 0)
+    window_swa_locs = full_to_swa_mapping[window_full_locs].to(torch.int32)
+    window_swa_locs = window_swa_locs.masked_fill(invalid, -1)
+
+    block_full_locs = out_loc[: bs * block_size].view(bs, block_size)
+    block_swa_locs = full_to_swa_mapping[block_full_locs].to(torch.int32)
+
+    target_width = ceil_align(swa_window + block_size, page_index_aligned_size)
+
+    swa_page_indices = _compact_dspark_window_then_block(
+        window_swa_locs=window_swa_locs,
+        block_swa_locs=block_swa_locs,
+        context_lens=context_lens,
+        target_width=target_width,
+        block_size=block_size,
+        swa_window=swa_window,
+    )
+
+    swa_page_indices = (
+        swa_page_indices.view(bs, 1, target_width)
+        .expand(bs, block_size, target_width)
+        .reshape(bs * block_size, target_width)
+        .contiguous()
+    )
+    swa_topk_lengths = (
+        (context_lens + block_size)
+        .view(bs, 1)
+        .expand(bs, block_size)
+        .reshape(bs * block_size)
+        .contiguous()
+        .to(torch.int32)
+    )
+    return swa_page_indices, swa_topk_lengths
+
+
+def _compact_dspark_window_then_block(
+    *,
+    window_swa_locs: torch.Tensor,
+    block_swa_locs: torch.Tensor,
+    context_lens: torch.Tensor,
+    target_width: int,
+    block_size: int,
+    swa_window: int,
+) -> torch.Tensor:
+    bs = window_swa_locs.shape[0]
+    device = window_swa_locs.device
+    out = torch.full((bs, target_width), -1, dtype=torch.int32, device=device)
+
+    j = torch.arange(swa_window, device=device, dtype=torch.int32).view(1, -1)
+    shift = (swa_window - context_lens.view(-1, 1)).to(torch.int32)
+    src_col = (shift + j).clamp_(min=0, max=swa_window - 1).to(torch.int64)
+    gathered = torch.gather(window_swa_locs, dim=1, index=src_col)
+    valid = j < context_lens.view(-1, 1)
+    out[:, :swa_window] = torch.where(valid, gathered, -1)
+
+    block_col = context_lens.view(-1, 1) + torch.arange(
+        block_size, device=device, dtype=torch.int32
+    ).view(1, -1)
+    block_rows = torch.arange(bs, device=device).view(-1, 1).expand(-1, block_size)
+    out[block_rows, block_col] = block_swa_locs
+    return out
+
+
+@triton.jit
+def _swa_page_indices_kernel(
+    req_to_token_ptr,
+    full_to_swa_ptr,
+    req_pool_ptr,
+    offsets_ptr,
+    out_loc_ptr,
+    context_lens_ptr,
+    out_ptr,
+    topk_ptr,
+    rt_stride,
+    swa_window,
+    block_size,
+    target_width,
+    TW_BLOCK: tl.constexpr,
+):
+    q = tl.program_id(0)
+    i = q // block_size
+    cl = tl.load(context_lens_ptr + i)
+    rp = tl.load(req_pool_ptr + i).to(tl.int64)
+    k = tl.arange(0, TW_BLOCK)
+    kmask = k < target_width
+    in_window = k < cl
+    src_col = tl.minimum(tl.maximum((swa_window - cl) + k, 0), swa_window - 1)
+    wmask = kmask & in_window
+    off = tl.load(offsets_ptr + i * swa_window + src_col, mask=wmask, other=0).to(
+        tl.int64
+    )
+    win_full = tl.load(req_to_token_ptr + rp * rt_stride + off, mask=wmask, other=0).to(
+        tl.int64
+    )
+    win_swa = tl.load(full_to_swa_ptr + win_full, mask=wmask, other=-1).to(tl.int32)
+
+    in_block = (k >= cl) & (k < cl + block_size)
+    bmask = kmask & in_block
+    bcol = tl.maximum(k - cl, 0)
+    blk_full = tl.load(out_loc_ptr + i * block_size + bcol, mask=bmask, other=0).to(
+        tl.int64
+    )
+    blk_swa = tl.load(full_to_swa_ptr + blk_full, mask=bmask, other=-1).to(tl.int32)
+
+    val = tl.where(in_window, win_swa, tl.where(in_block, blk_swa, -1))
+    tl.store(out_ptr + q * target_width + k, val.to(tl.int32), mask=kmask)
+    tl.store(topk_ptr + q, (cl + block_size).to(tl.int32))
+
+
+def build_dspark_swa_page_indices_triton(
+    *,
+    req_to_token: torch.Tensor,
+    full_to_swa_mapping: torch.Tensor,
+    req_pool_indices_per_request: torch.Tensor,
+    offsets: torch.Tensor,
+    out_loc: torch.Tensor,
+    context_lens: torch.Tensor,
+    block_size: int,
+    swa_window: int,
+    page_index_aligned_size: int,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    if offsets.ndim != 2 or offsets.shape[1] != swa_window:
+        raise ValueError(
+            "offsets must be [bs, swa_window]; "
+            f"got shape={tuple(offsets.shape)} (swa_window={swa_window})."
+        )
+    bs = offsets.shape[0]
+    device = offsets.device
+    req_pool = req_pool_indices_per_request.to(device=device).contiguous()
+    offsets = offsets.to(torch.int64).contiguous()
+    out_loc = out_loc[: bs * block_size].contiguous()
+    context_lens = context_lens.to(device=device, dtype=torch.int32).contiguous()
+    rt_stride = req_to_token.stride(0)
+    target_width = ceil_align(swa_window + block_size, page_index_aligned_size)
+    n_q = bs * block_size
+    swa_page_indices = torch.empty(
+        (n_q, target_width), dtype=torch.int32, device=device
+    )
+    swa_topk_lengths = torch.empty(n_q, dtype=torch.int32, device=device)
+    TW_BLOCK = triton.next_power_of_2(target_width)
+    _swa_page_indices_kernel[(n_q,)](
+        req_to_token,
+        full_to_swa_mapping,
+        req_pool,
+        offsets,
+        out_loc,
+        context_lens,
+        swa_page_indices,
+        swa_topk_lengths,
+        rt_stride,
+        swa_window,
+        block_size,
+        target_width,
+        TW_BLOCK=TW_BLOCK,
+    )
+    return swa_page_indices, swa_topk_lengths
+
+
+class BuildBlockSeqLensCausal:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> torch.Tensor:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        seq_lens: torch.Tensor,
+        block_size: int,
+        device: torch.device,
+    ) -> torch.Tensor:
+        return build_block_seq_lens_causal(
+            seq_lens=seq_lens,
+            block_size=block_size,
+            device=device,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        seq_lens: torch.Tensor,
+        block_size: int,
+        device: torch.device,
+    ) -> torch.Tensor:
+        return build_block_seq_lens_causal_triton(
+            seq_lens=seq_lens,
+            block_size=block_size,
+            device=device,
+        )
+
+
+def build_block_seq_lens_causal(
+    *,
+    seq_lens: torch.Tensor,
+    block_size: int,
+    device: torch.device,
+) -> torch.Tensor:
+    prefix = seq_lens.to(torch.int32)
+    steps = torch.arange(1, block_size + 1, device=device, dtype=torch.int32)
+    return (prefix[:, None] + steps[None, :]).reshape(-1)
+
+
+@triton.jit
+def _block_seq_lens_casual_kernel(
+    seq_lens_ptr,
+    out_ptr,
+    block_size,
+    n_out,
+    BLOCK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < n_out
+    row = offs // block_size
+    col = offs % block_size
+    prefix = tl.load(seq_lens_ptr + row, mask=mask, other=0)
+    tl.store(out_ptr + offs, (prefix + col + 1).to(tl.int32), mask=mask)
+
+
+def build_block_seq_lens_causal_triton(
+    *,
+    seq_lens: torch.Tensor,
+    block_size: int,
+    device: torch.device,
+) -> torch.Tensor:
+    seq_lens = seq_lens.to(device=device, dtype=torch.int64).contiguous()
+    n_rows = seq_lens.shape[0]
+    n_out = n_rows * block_size
+    out = torch.empty(n_out, dtype=torch.int32, device=device)
+    BLOCK = 256
+    grid = (triton.cdiv(n_out, BLOCK),)
+    _block_seq_lens_casual_kernel[grid](seq_lens, out, block_size, n_out, BLOCK=BLOCK)
+    return out

python/sglang/srt/speculative/dspark_components/kernels/dspark_draft_model.py

@@ -0,0 +1,443 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import msgspec
+import torch
+import torch.nn.functional as F
+import triton
+import triton.language as tl
+
+from sglang.srt.speculative.dspark_components.kernels.dispatch import inputs_on_cuda
+
+_BLOCK_V = 1024
+_IDX_SENTINEL = tl.constexpr(2147483647)
+
+
+class SampleStepTokens:
+    @classmethod
+    def execute(
+        cls,
+        *,
+        step_logits: torch.Tensor,
+        temperatures: torch.Tensor,
+        greedy_mask: torch.Tensor,
+        exp_noise: torch.Tensor,
+    ) -> torch.Tensor:
+        if step_logits.is_cuda:
+            return cls.triton(
+                step_logits=step_logits,
+                temperatures=temperatures,
+                greedy_mask=greedy_mask,
+                exp_noise=exp_noise,
+            )
+        return cls.torch(
+            step_logits=step_logits,
+            temperatures=temperatures,
+            greedy_mask=greedy_mask,
+            exp_noise=exp_noise,
+        )
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        step_logits: torch.Tensor,
+        temperatures: torch.Tensor,
+        greedy_mask: torch.Tensor,
+        exp_noise: torch.Tensor,
+    ) -> torch.Tensor:
+        return sample_step_tokens(
+            step_logits=step_logits,
+            temperatures=temperatures,
+            greedy_mask=greedy_mask,
+            exp_noise=exp_noise,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        step_logits: torch.Tensor,
+        temperatures: torch.Tensor,
+        greedy_mask: torch.Tensor,
+        exp_noise: torch.Tensor,
+    ) -> torch.Tensor:
+        return sample_step_tokens_triton(
+            step_logits=step_logits,
+            temperatures=temperatures,
+            greedy_mask=greedy_mask,
+            exp_noise=exp_noise,
+        )
+
+
+def sample_step_tokens(
+    *,
+    step_logits: torch.Tensor,
+    temperatures: torch.Tensor,
+    greedy_mask: torch.Tensor,
+    exp_noise: torch.Tensor,
+) -> torch.Tensor:
+    probs = torch.softmax(step_logits.float() / temperatures[:, None], dim=-1)
+    noise = torch.where(greedy_mask[:, None], 1.0, exp_noise)
+    return probs.div_(noise).argmax(dim=-1)
+
+
+@triton.jit
+def _online_partial_kernel(
+    logits_ptr,
+    temperatures_ptr,
+    greedy_mask_ptr,
+    exp_noise_ptr,
+    tile_max_ptr,
+    partial_key_ptr,
+    partial_idx_ptr,
+    V,
+    stride_row,
+    n_tiles,
+    BLOCK_V: tl.constexpr,
+):
+    row = tl.program_id(0)
+    tile = tl.program_id(1)
+    offs = tile * BLOCK_V + tl.arange(0, BLOCK_V)
+    mask = offs < V
+    logits = tl.load(
+        logits_ptr + row * stride_row + offs, mask=mask, other=float("-inf")
+    ).to(tl.float32)
+    temperature = tl.load(temperatures_ptr + row)
+    s = logits / temperature
+    tile_max = tl.max(s, axis=0)
+    greedy = tl.load(greedy_mask_ptr + row) != 0
+    noise = tl.load(exp_noise_ptr + row * V + offs, mask=mask, other=1.0)
+    denom = tl.where(greedy, 1.0, noise)
+    key = tl.exp(s - tile_max) / denom
+    key = tl.where(mask, key, -1.0)
+    tile_best = tl.max(key, axis=0)
+    idx = tl.where(key == tile_best, offs, _IDX_SENTINEL)
+    tl.store(tile_max_ptr + row * n_tiles + tile, tile_max)
+    tl.store(partial_key_ptr + row * n_tiles + tile, tile_best)
+    tl.store(partial_idx_ptr + row * n_tiles + tile, tl.min(idx, axis=0))
+
+
+@triton.jit
+def _online_combine_kernel(
+    tile_max_ptr,
+    partial_key_ptr,
+    partial_idx_ptr,
+    next_tokens_ptr,
+    n_tiles,
+    BLOCK_TILES: tl.constexpr,
+):
+    row = tl.program_id(0)
+    offs = tl.arange(0, BLOCK_TILES)
+    mask = offs < n_tiles
+    tile_max = tl.load(
+        tile_max_ptr + row * n_tiles + offs, mask=mask, other=float("-inf")
+    )
+    keys = tl.load(partial_key_ptr + row * n_tiles + offs, mask=mask, other=-1.0)
+    idxs = tl.load(
+        partial_idx_ptr + row * n_tiles + offs, mask=mask, other=_IDX_SENTINEL
+    )
+    global_max = tl.max(tile_max, axis=0)
+    rescaled = keys * tl.exp(tile_max - global_max)
+    rescaled = tl.where(mask, rescaled, -1.0)
+    best = tl.max(rescaled, axis=0)
+    cand = tl.where(rescaled == best, idxs, _IDX_SENTINEL)
+    tl.store(next_tokens_ptr + row, tl.min(cand, axis=0).to(tl.int64))
+
+
+def sample_step_tokens_triton(
+    *,
+    step_logits: torch.Tensor,
+    temperatures: torch.Tensor,
+    greedy_mask: torch.Tensor,
+    exp_noise: torch.Tensor,
+) -> torch.Tensor:
+    bs, V = step_logits.shape
+    device = step_logits.device
+    assert step_logits.stride(1) == 1, "step_logits rows must be contiguous"
+    stride_row = step_logits.stride(0)
+    temperatures = temperatures.to(torch.float32).contiguous()
+    greedy_mask = greedy_mask.to(torch.int32).contiguous()
+    exp_noise = exp_noise.to(torch.float32).contiguous()
+
+    n_tiles = triton.cdiv(V, _BLOCK_V)
+    block_tiles = triton.next_power_of_2(n_tiles)
+
+    tile_max = torch.empty((bs, n_tiles), dtype=torch.float32, device=device)
+    partial_key = torch.empty((bs, n_tiles), dtype=torch.float32, device=device)
+    partial_idx = torch.empty((bs, n_tiles), dtype=torch.int32, device=device)
+    next_tokens = torch.empty((bs,), dtype=torch.int64, device=device)
+
+    tile_grid = (bs, n_tiles)
+    row_grid = (bs,)
+
+    _online_partial_kernel[tile_grid](
+        step_logits,
+        temperatures,
+        greedy_mask,
+        exp_noise,
+        tile_max,
+        partial_key,
+        partial_idx,
+        V,
+        stride_row,
+        n_tiles,
+        BLOCK_V=_BLOCK_V,
+    )
+    _online_combine_kernel[row_grid](
+        tile_max,
+        partial_key,
+        partial_idx,
+        next_tokens,
+        n_tiles,
+        BLOCK_TILES=block_tiles,
+    )
+    return next_tokens
+
+
+_STACKED_WEIGHT_CACHE: dict[int, _StackedWkvWeight] = {}
+
+
+class CommitKvProj:
+    @classmethod
+    def execute(
+        cls,
+        *,
+        main_x: torch.Tensor,
+        wkv_linears: list[torch.nn.Module],
+    ) -> list[torch.Tensor]:
+        if main_x.is_cuda and _fused_commit_kv_proj_supported(wkv_linears=wkv_linears):
+            return cls.triton(main_x=main_x, wkv_linears=wkv_linears)
+        return cls.torch(main_x=main_x, wkv_linears=wkv_linears)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        main_x: torch.Tensor,
+        wkv_linears: list[torch.nn.Module],
+    ) -> list[torch.Tensor]:
+        return commit_kv_proj(main_x=main_x, wkv_linears=wkv_linears)
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        main_x: torch.Tensor,
+        wkv_linears: list[torch.nn.Module],
+    ) -> list[torch.Tensor]:
+        return commit_kv_proj_fused(main_x=main_x, wkv_linears=wkv_linears)
+
+
+def commit_kv_proj(
+    *,
+    main_x: torch.Tensor,
+    wkv_linears: list[torch.nn.Module],
+) -> list[torch.Tensor]:
+    return [linear(main_x)[0] for linear in wkv_linears]
+
+
+def commit_kv_proj_fused(
+    *,
+    main_x: torch.Tensor,
+    wkv_linears: list[torch.nn.Module],
+) -> list[torch.Tensor]:
+    num_stages = len(wkv_linears)
+    stacked = _stacked_wkv_weight(wkv_linears=wkv_linears)
+
+    if stacked.fp8_scale is not None:
+        quant_method = wkv_linears[0].quant_method
+        kv_all = quant_method.w8a8_block_fp8_linear(
+            input=main_x,
+            weight=stacked.weight,
+            block_size=quant_method.quant_config.weight_block_size,
+            weight_scale=stacked.fp8_scale,
+            input_scale=None,
+            bias=None,
+        )
+    else:
+        kv_all = torch.nn.functional.linear(main_x, stacked.weight)
+
+    head_dim = kv_all.shape[-1] // num_stages
+    return [
+        kv_all[:, i * head_dim : (i + 1) * head_dim].contiguous()
+        for i in range(num_stages)
+    ]
+
+
+class _StackedWkvWeight(msgspec.Struct):
+    weight: torch.Tensor
+    fp8_scale: Optional[torch.Tensor]
+
+
+def _stacked_wkv_weight(*, wkv_linears: list[torch.nn.Module]) -> _StackedWkvWeight:
+    key = id(wkv_linears[0])
+    cached = _STACKED_WEIGHT_CACHE.get(key)
+    if cached is None:
+        cached = _build_stacked_wkv_weight(wkv_linears=wkv_linears)
+        _STACKED_WEIGHT_CACHE[key] = cached
+    return cached
+
+
+def _block_quant_stack_applies(*, wkv_linears: list[torch.nn.Module]) -> bool:
+    quant_method = wkv_linears[0].quant_method
+    block_quant = hasattr(quant_method, "block_quant") and quant_method.block_quant
+    if not (block_quant and hasattr(quant_method, "w8a8_block_fp8_linear")):
+        return False
+    block_out = quant_method.quant_config.weight_block_size[0]
+    return all(
+        linear.weight.dtype == torch.float8_e4m3fn
+        and linear.weight.shape[0] % block_out == 0
+        for linear in wkv_linears
+    )
+
+
+def _dequant_supported(linear: torch.nn.Module) -> bool:
+    """Mirrors the preconditions asserted in _dequant_linear_weight."""
+    weight = linear.weight
+    if weight.dtype in (torch.bfloat16, torch.float16, torch.float32):
+        return True
+    if weight.dtype != torch.float8_e4m3fn:
+        return False
+    block = 128
+    out_dim, in_dim = weight.shape
+    expected_scale_shape = (
+        (out_dim + block - 1) // block,
+        (in_dim + block - 1) // block,
+    )
+    return tuple(linear.weight_scale_inv.shape) == expected_scale_shape
+
+
+def _fused_commit_kv_proj_supported(*, wkv_linears: list[torch.nn.Module]) -> bool:
+    """Whether _build_stacked_wkv_weight can handle these weights; unsupported
+    quant schemes fall back to the per-linear torch path in execute()."""
+    if _block_quant_stack_applies(wkv_linears=wkv_linears):
+        return True
+    return all(_dequant_supported(linear) for linear in wkv_linears)
+
+
+def _build_stacked_wkv_weight(
+    *, wkv_linears: list[torch.nn.Module]
+) -> _StackedWkvWeight:
+    if _block_quant_stack_applies(wkv_linears=wkv_linears):
+        weight = torch.cat([linear.weight for linear in wkv_linears], dim=0)
+        if wkv_linears[0].weight_scale_inv.dtype == torch.int32:
+            from sglang.srt.layers.quantization.fp8_utils import (
+                inverse_transform_scale_ue8m0,
+                transform_scale_ue8m0,
+            )
+
+            sf_fp32 = torch.cat(
+                [
+                    inverse_transform_scale_ue8m0(
+                        linear.weight_scale_inv, mn=linear.weight.shape[0]
+                    )
+                    for linear in wkv_linears
+                ],
+                dim=0,
+            )
+            scale = transform_scale_ue8m0(sf_fp32, mn=weight.shape[0])
+            return _StackedWkvWeight(weight=weight, fp8_scale=scale)
+        scale = torch.cat([linear.weight_scale_inv for linear in wkv_linears], dim=0)
+        if scale.dim() >= 2 and scale.stride(-2) != 1:
+            scale = scale.transpose(-2, -1).contiguous().transpose(-2, -1)
+        return _StackedWkvWeight(weight=weight, fp8_scale=scale)
+    weight = torch.cat(
+        [_dequant_linear_weight(linear) for linear in wkv_linears], dim=0
+    )
+    return _StackedWkvWeight(weight=weight, fp8_scale=None)
+
+
+def _dequant_linear_weight(linear: torch.nn.Module) -> torch.Tensor:
+    weight = linear.weight
+    if weight.dtype in (torch.bfloat16, torch.float16, torch.float32):
+        return weight.to(torch.bfloat16)
+    assert weight.dtype == torch.float8_e4m3fn, (
+        f"unsupported wkv weight dtype {weight.dtype} for the fused commit kv proj; "
+        f"execute() should have routed this to the torch path "
+        f"(_fused_commit_kv_proj_supported)"
+    )
+    block = 128
+    scale = linear.weight_scale_inv
+    out_dim, in_dim = weight.shape
+    expected_scale_shape = (
+        (out_dim + block - 1) // block,
+        (in_dim + block - 1) // block,
+    )
+    assert tuple(scale.shape) == expected_scale_shape, (
+        f"wkv weight_scale_inv shape {tuple(scale.shape)} does not match the "
+        f"128x128 block grid {expected_scale_shape} for weight {tuple(weight.shape)}; "
+        f"execute() should have routed this to the torch path "
+        f"(_fused_commit_kv_proj_supported)"
+    )
+    scale_full = scale.repeat_interleave(block, dim=0)[:out_dim]
+    scale_full = scale_full.repeat_interleave(block, dim=1)[:, :in_dim]
+    return (weight.to(torch.float32) * scale_full.to(torch.float32)).to(torch.bfloat16)
+
+
+_BLOCK = 1024
+
+
+class BuildStepLocal:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> torch.Tensor:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(cls, *, bias: torch.Tensor, base_local: torch.Tensor) -> torch.Tensor:
+        return build_step_local(bias=bias, base_local=base_local)
+
+    @classmethod
+    def triton(cls, *, bias: torch.Tensor, base_local: torch.Tensor) -> torch.Tensor:
+        return build_step_local_triton(bias=bias, base_local=base_local)
+
+
+def build_step_local(*, bias: torch.Tensor, base_local: torch.Tensor) -> torch.Tensor:
+    per_partition = base_local.shape[-1]
+    pad = per_partition - bias.shape[-1]
+    padded = (
+        F.pad(bias.to(torch.float32), (0, pad)) if pad > 0 else bias.to(torch.float32)
+    )
+    return base_local + padded
+
+
+@triton.jit
+def _build_step_local_kernel(
+    bias_ptr,
+    base_ptr,
+    out_ptr,
+    org_width,
+    per_partition,
+    BLOCK: tl.constexpr,
+):
+    row = tl.program_id(0)
+    tile = tl.program_id(1)
+    offs = tile * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < per_partition
+    base = tl.load(base_ptr + row * per_partition + offs, mask=mask, other=0.0).to(
+        tl.float32
+    )
+    bias = tl.load(
+        bias_ptr + row * org_width + offs, mask=offs < org_width, other=0.0
+    ).to(tl.float32)
+    tl.store(out_ptr + row * per_partition + offs, base + bias, mask=mask)
+
+
+def build_step_local_triton(
+    *, bias: torch.Tensor, base_local: torch.Tensor
+) -> torch.Tensor:
+    bs, per_partition = base_local.shape
+    org_width = bias.shape[-1]
+    base_local = base_local.contiguous()
+    bias = bias.contiguous()
+    out = torch.empty(
+        (bs, per_partition), dtype=torch.float32, device=base_local.device
+    )
+    grid = (bs, triton.cdiv(per_partition, _BLOCK))
+    _build_step_local_kernel[grid](
+        bias, base_local, out, org_width, per_partition, BLOCK=_BLOCK
+    )
+    return out

python/sglang/srt/speculative/dspark_components/kernels/dspark_schedule.py

@@ -0,0 +1,260 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.speculative.dspark_components.kernels.dispatch import (
+    inputs_on_cuda,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.speculative.dspark_components.dspark_planner import (
+        DSparkScheduleConfig,
+    )
+
+
+class ScheduleVerifyLensTopk:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> torch.Tensor:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        confidence: torch.Tensor,
+        budget: int,
+        cfg: DSparkScheduleConfig,
+    ) -> torch.Tensor:
+        return schedule_verify_lens_topk(confidence=confidence, budget=budget, cfg=cfg)
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        confidence: torch.Tensor,
+        budget: int,
+        cfg: DSparkScheduleConfig,
+    ) -> torch.Tensor:
+        return schedule_verify_lens_topk_triton(
+            confidence=confidence, budget=budget, cfg=cfg
+        )
+
+
+def compute_sort_survival(confidence: torch.Tensor) -> torch.Tensor:
+    return torch.cumprod(confidence.to(torch.float32), dim=1)
+
+
+def schedule_verify_lens_topk(
+    *,
+    confidence: torch.Tensor,
+    budget: int,
+    cfg: DSparkScheduleConfig,
+) -> torch.Tensor:
+    return schedule_verify_lens_topk_from_survival(
+        survival_probs=compute_sort_survival(confidence), budget=budget, cfg=cfg
+    )
+
+
+def schedule_verify_lens_topk_from_survival(
+    *,
+    survival_probs: torch.Tensor,
+    budget: int,
+    cfg: DSparkScheduleConfig,
+) -> torch.Tensor:
+    num_requests, _gamma = survival_probs.shape
+    max_len = cfg.resolved_max_verify_len()
+    device = survival_probs.device
+
+    selected_extra = torch.zeros(num_requests, dtype=torch.int64, device=device)
+    if budget > 0:
+        candidate_window = survival_probs[:, :max_len]
+        num_candidates = candidate_window.numel()
+        if num_candidates > 0:
+            request_index = (
+                torch.arange(num_requests, device=device)
+                .view(num_requests, 1)
+                .expand_as(candidate_window)
+            )
+            position_index = (
+                torch.arange(candidate_window.shape[1], device=device)
+                .view(1, candidate_window.shape[1])
+                .expand_as(candidate_window)
+            )
+            valid = candidate_window >= cfg.survival_eps
+
+            flat_prob = candidate_window.reshape(-1).to(torch.float64)
+            flat_request = request_index.reshape(-1)
+            flat_position = position_index.reshape(-1)
+            flat_valid = valid.reshape(-1)
+
+            order = _value_independent_descending_order(
+                probs=flat_prob,
+                positions=flat_position,
+                requests=flat_request,
+                valid=flat_valid,
+            )
+
+            take = min(int(budget), num_candidates)
+            chosen = order[:take]
+            chosen_requests = flat_request[chosen]
+            chosen_valid = flat_valid[chosen].to(torch.int64)
+            selected_extra.scatter_add_(0, chosen_requests, chosen_valid)
+
+    min_len = torch.full(
+        (num_requests,), cfg.min_verify_len, dtype=torch.int64, device=device
+    )
+    verify_lens = min_len + selected_extra
+    lower_bound = max(cfg.min_verify_len, 1)
+    verify_lens = torch.clamp(verify_lens, min=lower_bound, max=max_len)
+    return verify_lens.to(torch.int32)
+
+
+def _value_independent_descending_order(
+    *,
+    probs: torch.Tensor,
+    positions: torch.Tensor,
+    requests: torch.Tensor,
+    valid: torch.Tensor,
+) -> torch.Tensor:
+    masked_prob = torch.where(valid, probs, torch.full_like(probs, float("-inf")))
+    num_candidates = masked_prob.numel()
+    order = torch.arange(num_candidates, device=probs.device)
+    order = order[torch.argsort(requests[order], stable=True)]
+    order = order[torch.argsort(positions[order], stable=True)]
+    order = order[torch.argsort(-masked_prob[order], stable=True)]
+    return order
+
+
+@triton.jit
+def _schedule_topk_prep_kernel(
+    confidence_ptr,
+    survival_ptr,
+    selected_extra_ptr,
+    gamma,
+    cols,
+    G_P2: tl.constexpr,
+):
+    row = tl.program_id(0)
+    g = tl.arange(0, G_P2)
+    conf = tl.load(
+        confidence_ptr + row.to(tl.int64) * gamma + g, mask=g < gamma, other=1.0
+    ).to(tl.float32)
+    surv = tl.cumprod(conf, axis=0)
+    tl.store(survival_ptr + row.to(tl.int64) * cols + g, surv, mask=g < cols)
+    tl.store(selected_extra_ptr + row, 0)
+
+
+@triton.jit
+def _schedule_topk_finalize_kernel(
+    selected_extra_ptr,
+    out_ptr,
+    min_verify_len,
+    lower_bound,
+    max_len,
+    bs,
+    BLOCK: tl.constexpr,
+):
+    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < bs
+    extra = tl.load(selected_extra_ptr + offs, mask=mask, other=0).to(tl.int32)
+    lens = min_verify_len + extra
+    lens = tl.maximum(lens, lower_bound)
+    lens = tl.minimum(lens, max_len)
+    tl.store(out_ptr + offs, lens, mask=mask)
+
+
+@triton.jit
+def _schedule_topk_selected_extra_kernel(
+    survival_ptr,
+    selected_extra_ptr,
+    budget,
+    cols,
+    n,
+    survival_eps,
+    BLOCK_C: tl.constexpr,
+    BLOCK_CP: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    c = pid * BLOCK_C + tl.arange(0, BLOCK_C)
+    cmask = c < n
+    r = c // cols
+    p = c % cols
+    sp = tl.load(survival_ptr + c, mask=cmask, other=0.0)
+    valid_c = sp >= survival_eps
+    mp = tl.where(valid_c, sp, float("-inf"))
+    rank = tl.zeros([BLOCK_C], dtype=tl.int32)
+    for cp0 in range(0, n, BLOCK_CP):
+        cp = cp0 + tl.arange(0, BLOCK_CP)
+        cpmask = cp < n
+        rp = cp // cols
+        pp = cp % cols
+        spp = tl.load(survival_ptr + cp, mask=cpmask, other=0.0)
+        validp = spp >= survival_eps
+        mpp = tl.where(validp, spp, float("-inf"))
+        gt = mpp[None, :] > mp[:, None]
+        eq = mpp[None, :] == mp[:, None]
+        pos_lt = pp[None, :] < p[:, None]
+        pos_eq = pp[None, :] == p[:, None]
+        req_lt = rp[None, :] < r[:, None]
+        before = gt | (eq & (pos_lt | (pos_eq & req_lt)))
+        before = before & cpmask[None, :]
+        rank += tl.sum(before.to(tl.int32), axis=1)
+    selected = valid_c & (rank < budget)
+    tl.atomic_add(selected_extra_ptr + r, selected.to(tl.int32), mask=cmask)
+
+
+def schedule_verify_lens_topk_triton(
+    *,
+    confidence: torch.Tensor,
+    budget: int,
+    cfg: DSparkScheduleConfig,
+) -> torch.Tensor:
+    num_requests, gamma = confidence.shape
+    max_len = cfg.resolved_max_verify_len()
+    device = confidence.device
+    cols = min(max_len, gamma)
+    n = num_requests * cols
+
+    selected_extra = torch.empty(num_requests, dtype=torch.int32, device=device)
+    survival = torch.empty((num_requests, cols), dtype=torch.float32, device=device)
+    _schedule_topk_prep_kernel[(num_requests,)](
+        confidence.contiguous(),
+        survival,
+        selected_extra,
+        gamma,
+        cols,
+        G_P2=triton.next_power_of_2(max(gamma, 1)),
+    )
+    if budget > 0 and n > 0:
+        BLOCK_C = 64
+        BLOCK_CP = 256
+        grid = (triton.cdiv(n, BLOCK_C),)
+        _schedule_topk_selected_extra_kernel[grid](
+            survival,
+            selected_extra,
+            int(budget),
+            cols,
+            n,
+            float(cfg.survival_eps),
+            BLOCK_C=BLOCK_C,
+            BLOCK_CP=BLOCK_CP,
+        )
+
+    verify_lens = torch.empty(num_requests, dtype=torch.int32, device=device)
+    BLOCK = 256
+    _schedule_topk_finalize_kernel[(triton.cdiv(num_requests, BLOCK),)](
+        selected_extra,
+        verify_lens,
+        int(cfg.min_verify_len),
+        max(cfg.min_verify_len, 1),
+        int(max_len),
+        num_requests,
+        BLOCK=BLOCK,
+    )
+    return verify_lens

python/sglang/srt/speculative/dspark_components/kernels/dspark_verify_window.py

@@ -0,0 +1,871 @@
+from __future__ import annotations
+
+import msgspec
+import torch
+import triton
+import triton.language as tl
+
+from sglang.kernels.ops.speculative.cache_locs import assign_extend_cache_locs_func
+from sglang.srt.managers.schedule_batch import ScheduleBatch
+from sglang.srt.speculative.dspark_components.kernels.dispatch import inputs_on_cuda
+from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout
+
+
+class RaggedVerifyWindow(msgspec.Struct, frozen=True):
+    positions: torch.Tensor
+    verify_cache_loc: torch.Tensor
+    verify_ids: torch.Tensor
+
+
+class BuildRaggedVerifyWindow:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> RaggedVerifyWindow:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        batch: ScheduleBatch,
+        layout: RaggedVerifyLayout,
+        draft_block_ids: torch.Tensor,
+        draft_tokens: torch.Tensor,
+        bs: int,
+        device: str,
+        verify_num_draft_tokens: int,
+        model_runner,
+    ) -> RaggedVerifyWindow:
+        return build_ragged_verify_window(
+            batch=batch,
+            layout=layout,
+            draft_block_ids=draft_block_ids,
+            draft_tokens=draft_tokens,
+            bs=bs,
+            device=device,
+            verify_num_draft_tokens=verify_num_draft_tokens,
+            model_runner=model_runner,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        batch: ScheduleBatch,
+        layout: RaggedVerifyLayout,
+        draft_block_ids: torch.Tensor,
+        draft_tokens: torch.Tensor,
+        bs: int,
+        device: str,
+        verify_num_draft_tokens: int,
+        model_runner,
+    ) -> RaggedVerifyWindow:
+        return build_ragged_verify_window_triton(
+            batch=batch,
+            layout=layout,
+            draft_block_ids=draft_block_ids,
+            draft_tokens=draft_tokens,
+            bs=bs,
+            device=device,
+            verify_num_draft_tokens=verify_num_draft_tokens,
+            model_runner=model_runner,
+        )
+
+
+def build_ragged_verify_window(
+    *,
+    batch: ScheduleBatch,
+    layout: RaggedVerifyLayout,
+    draft_block_ids: torch.Tensor,
+    draft_tokens: torch.Tensor,
+    bs: int,
+    device: str,
+    verify_num_draft_tokens: int,
+    model_runner,
+) -> RaggedVerifyWindow:
+    prefix_lens = batch.seq_lens
+    verify_lens = layout.verify_lens.to(device=device, dtype=torch.int32)
+    padded_total = layout.graph_num_tokens
+
+    req_id, within, valid = compact_row_index(
+        verify_lens=verify_lens, padded_total=padded_total, device=device
+    )
+    safe_req = req_id.clamp(max=bs - 1)
+    positions = torch.where(
+        valid,
+        prefix_lens.to(torch.int64)[safe_req] + within,
+        torch.zeros_like(within),
+    )
+    real_cache_loc = assign_extend_cache_locs_func(
+        req_pool_indices=batch.req_pool_indices,
+        req_to_token=model_runner.req_to_token_pool.req_to_token,
+        start_offset=prefix_lens,
+        end_offset=prefix_lens + verify_lens.to(prefix_lens.dtype),
+        batch_size=bs,
+        draft_token_num=verify_num_draft_tokens,
+        device=device,
+    )
+    verify_cache_loc = torch.nn.functional.pad(
+        real_cache_loc, (0, padded_total - real_cache_loc.shape[0])
+    )
+    verify_cache_loc = torch.where(
+        valid, verify_cache_loc, torch.zeros_like(verify_cache_loc)
+    )
+
+    verify_ids = compact_verify_ids(
+        draft_block_ids=draft_block_ids,
+        draft_tokens=draft_tokens,
+        layout=layout,
+        device=device,
+    )
+
+    return RaggedVerifyWindow(
+        positions=positions,
+        verify_cache_loc=verify_cache_loc,
+        verify_ids=verify_ids,
+    )
+
+
+@triton.jit
+def _ragged_finalize_kernel(
+    req_ptr,
+    within_ptr,
+    prefix_ptr,
+    cache_ptr,
+    pos_out_ptr,
+    cache_out_ptr,
+    bs,
+    n,
+    real_len,
+    BLOCK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < n
+    req = tl.load(req_ptr + offs, mask=mask, other=0)
+    within = tl.load(within_ptr + offs, mask=mask, other=0)
+    valid = req < bs
+    safe_req = tl.minimum(req, bs - 1)
+    prefix = tl.load(prefix_ptr + safe_req, mask=mask, other=0)
+    pos = tl.where(valid, prefix + within, 0)
+    lmask = mask & (offs < real_len)
+    cl = tl.load(cache_ptr + offs, mask=lmask, other=0)
+    cl = tl.where(valid, cl, 0)
+    tl.store(pos_out_ptr + offs, pos, mask=mask)
+    tl.store(cache_out_ptr + offs, cl, mask=mask)
+
+
+def build_ragged_verify_window_triton(
+    *,
+    batch: ScheduleBatch,
+    layout: RaggedVerifyLayout,
+    draft_block_ids: torch.Tensor,
+    draft_tokens: torch.Tensor,
+    bs: int,
+    device: str,
+    verify_num_draft_tokens: int,
+    model_runner,
+) -> RaggedVerifyWindow:
+    prefix_lens = batch.seq_lens
+    verify_lens = layout.verify_lens.to(device=device, dtype=torch.int32)
+    padded_total = layout.graph_num_tokens
+
+    req_id, within, _valid = compact_row_index_triton(
+        verify_lens=verify_lens, padded_total=padded_total, device=device
+    )
+    real_cache_loc = assign_extend_cache_locs_func(
+        req_pool_indices=batch.req_pool_indices,
+        req_to_token=model_runner.req_to_token_pool.req_to_token,
+        start_offset=prefix_lens,
+        end_offset=prefix_lens + verify_lens.to(prefix_lens.dtype),
+        batch_size=bs,
+        draft_token_num=verify_num_draft_tokens,
+        device=device,
+    )
+    prefix_i64 = prefix_lens.to(device=device, dtype=torch.int64).contiguous()
+    positions = torch.empty(padded_total, dtype=torch.int64, device=device)
+    verify_cache_loc = torch.empty(
+        padded_total, dtype=real_cache_loc.dtype, device=device
+    )
+    BLOCK = 256
+    grid = (triton.cdiv(padded_total, BLOCK),)
+    _ragged_finalize_kernel[grid](
+        req_id,
+        within,
+        prefix_i64,
+        real_cache_loc,
+        positions,
+        verify_cache_loc,
+        bs,
+        padded_total,
+        real_cache_loc.shape[0],
+        BLOCK=BLOCK,
+    )
+
+    verify_ids = compact_verify_ids_triton(
+        draft_block_ids=draft_block_ids,
+        draft_tokens=draft_tokens,
+        layout=layout,
+        device=device,
+    )
+    return RaggedVerifyWindow(
+        positions=positions,
+        verify_cache_loc=verify_cache_loc,
+        verify_ids=verify_ids,
+    )
+
+
+_SEARCH_NBITS = 11
+
+
+class CompactRowIndex:
+    @classmethod
+    def execute(
+        cls, *args, **kwargs
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        verify_lens: torch.Tensor,
+        padded_total: int,
+        device,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return compact_row_index(
+            verify_lens=verify_lens,
+            padded_total=padded_total,
+            device=device,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        verify_lens: torch.Tensor,
+        padded_total: int,
+        device,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return compact_row_index_triton(
+            verify_lens=verify_lens,
+            padded_total=padded_total,
+            device=device,
+        )
+
+
+class CompactVerifyIds:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> torch.Tensor:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        draft_block_ids: torch.Tensor,
+        draft_tokens: torch.Tensor,
+        layout: RaggedVerifyLayout,
+        device: str,
+    ) -> torch.Tensor:
+        return compact_verify_ids(
+            draft_block_ids=draft_block_ids,
+            draft_tokens=draft_tokens,
+            layout=layout,
+            device=device,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        draft_block_ids: torch.Tensor,
+        draft_tokens: torch.Tensor,
+        layout: RaggedVerifyLayout,
+        device: str,
+    ) -> torch.Tensor:
+        return compact_verify_ids_triton(
+            draft_block_ids=draft_block_ids,
+            draft_tokens=draft_tokens,
+            layout=layout,
+            device=device,
+        )
+
+
+def compact_verify_ids(
+    *,
+    draft_block_ids: torch.Tensor,
+    draft_tokens: torch.Tensor,
+    layout: RaggedVerifyLayout,
+    device: str,
+) -> torch.Tensor:
+    req_id, within, valid = compact_row_index(
+        verify_lens=layout.verify_lens,
+        padded_total=layout.graph_num_tokens,
+        device=device,
+    )
+    bs = layout.verify_lens.shape[0]
+    safe_req = req_id.clamp(max=bs - 1)
+    anchors = draft_block_ids[:, 0]
+    drafts = draft_tokens[safe_req, (within - 1).clamp_min(0)]
+    verify_ids = torch.where(within == 0, anchors[safe_req], drafts)
+    verify_ids = torch.where(valid, verify_ids, torch.zeros_like(verify_ids))
+    return verify_ids.to(torch.int64)
+
+
+def compact_row_index(
+    *,
+    verify_lens: torch.Tensor,
+    padded_total: int,
+    device,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    verify_lens = verify_lens.to(device=device, dtype=torch.int64)
+    bs = int(verify_lens.numel())
+    incl = torch.cumsum(verify_lens, dim=0)
+    start = incl - verify_lens
+    real_total = incl[-1]
+    row = torch.arange(padded_total, device=device, dtype=torch.int64)
+    valid = row < real_total
+    req_id = torch.searchsorted(incl, row, right=True)
+    req_id = torch.where(valid, req_id, torch.full_like(req_id, bs))
+    within = torch.where(
+        valid, row - start[req_id.clamp(max=bs - 1)], torch.zeros_like(row)
+    )
+    return req_id, within, valid
+
+
+@triton.jit
+def _compact_row_index_kernel(
+    incl_ptr,
+    req_out_ptr,
+    within_out_ptr,
+    valid_out_ptr,
+    bs,
+    n,
+    BLOCK: tl.constexpr,
+    NBITS: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < n
+    row = offs.to(tl.int64)
+    real_total = tl.load(incl_ptr + (bs - 1))
+    lo = tl.zeros([BLOCK], dtype=tl.int32)
+    hi = tl.full([BLOCK], bs, dtype=tl.int32)
+    for _ in range(NBITS):
+        mid = (lo + hi) // 2
+        active = lo < hi
+        val = tl.load(incl_ptr + tl.minimum(mid, bs - 1), mask=mask, other=0)
+        go_right = val <= row
+        lo = tl.where(active & go_right, mid + 1, lo)
+        hi = tl.where(active & (~go_right), mid, hi)
+    req = lo
+    gidx = tl.maximum(req - 1, 0)
+    start = tl.load(incl_ptr + gidx, mask=mask, other=0)
+    start = tl.where(req > 0, start, 0)
+    valid = row < real_total
+    within = tl.where(valid, row - start, 0)
+    req_final = tl.where(valid, req.to(tl.int64), bs)
+    tl.store(req_out_ptr + offs, req_final, mask=mask)
+    tl.store(within_out_ptr + offs, within, mask=mask)
+    tl.store(valid_out_ptr + offs, valid, mask=mask)
+
+
+def compact_row_index_triton(
+    *,
+    verify_lens: torch.Tensor,
+    padded_total: int,
+    device,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    verify_lens = verify_lens.to(device=device, dtype=torch.int64).contiguous()
+    bs = verify_lens.shape[0]
+    incl = torch.cumsum(verify_lens, dim=0).contiguous()
+    req = torch.empty(padded_total, dtype=torch.int64, device=device)
+    within = torch.empty(padded_total, dtype=torch.int64, device=device)
+    valid = torch.empty(padded_total, dtype=torch.bool, device=device)
+    BLOCK = 256
+    grid = (triton.cdiv(padded_total, BLOCK),)
+    _compact_row_index_kernel[grid](
+        incl, req, within, valid, bs, padded_total, BLOCK=BLOCK, NBITS=_SEARCH_NBITS
+    )
+    return req, within, valid
+
+
+@triton.jit
+def _compact_verify_ids_gather_kernel(
+    req_ptr,
+    within_ptr,
+    draft_block_ids_ptr,
+    draft_tokens_ptr,
+    out_ptr,
+    bs,
+    gamma,
+    n,
+    BLOCK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < n
+    req = tl.load(req_ptr + offs, mask=mask, other=0)
+    within = tl.load(within_ptr + offs, mask=mask, other=0)
+    valid = req < bs
+    safe_req = tl.minimum(req, bs - 1)
+    anchor = tl.load(draft_block_ids_ptr + safe_req * gamma, mask=mask, other=0)
+    wcol = tl.maximum(within - 1, 0)
+    draft = tl.load(draft_tokens_ptr + safe_req * gamma + wcol, mask=mask, other=0)
+    v = tl.where(within == 0, anchor, draft)
+    v = tl.where(valid, v, 0)
+    tl.store(out_ptr + offs, v.to(tl.int64), mask=mask)
+
+
+def compact_verify_ids_triton(
+    *,
+    draft_block_ids: torch.Tensor,
+    draft_tokens: torch.Tensor,
+    layout: RaggedVerifyLayout,
+    device: str,
+) -> torch.Tensor:
+    req, within, _valid = compact_row_index_triton(
+        verify_lens=layout.verify_lens,
+        padded_total=layout.graph_num_tokens,
+        device=device,
+    )
+    bs = layout.verify_lens.shape[0]
+    gamma = draft_tokens.shape[1]
+    draft_block_ids = draft_block_ids.to(device=device, dtype=torch.int64).contiguous()
+    draft_tokens = draft_tokens.to(device=device, dtype=torch.int64).contiguous()
+    n = layout.graph_num_tokens
+    out = torch.empty(n, dtype=torch.int64, device=device)
+    BLOCK = 256
+    grid = (triton.cdiv(n, BLOCK),)
+    _compact_verify_ids_gather_kernel[grid](
+        req, within, draft_block_ids, draft_tokens, out, bs, gamma, n, BLOCK=BLOCK
+    )
+    return out
+
+
+class ScatterCompactToStrided:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> torch.Tensor:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        compact: torch.Tensor,
+        layout: RaggedVerifyLayout,
+        fill_value: float,
+        verify_num_draft_tokens: int,
+    ) -> torch.Tensor:
+        return scatter_compact_to_strided(
+            compact=compact,
+            layout=layout,
+            fill_value=fill_value,
+            verify_num_draft_tokens=verify_num_draft_tokens,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        compact: torch.Tensor,
+        layout: RaggedVerifyLayout,
+        fill_value: float,
+        verify_num_draft_tokens: int,
+    ) -> torch.Tensor:
+        return scatter_compact_to_strided_triton(
+            compact=compact,
+            layout=layout,
+            fill_value=fill_value,
+            verify_num_draft_tokens=verify_num_draft_tokens,
+        )
+
+
+def scatter_compact_to_strided(
+    *,
+    compact: torch.Tensor,
+    layout: RaggedVerifyLayout,
+    fill_value: float,
+    verify_num_draft_tokens: int,
+) -> torch.Tensor:
+    stride = verify_num_draft_tokens
+    bs = layout.verify_lens.shape[0]
+    dim = compact.shape[1]
+    device = compact.device
+    compact = compact[: layout.graph_num_tokens]
+    strided = torch.full(
+        (bs * stride + 1, dim), fill_value, dtype=compact.dtype, device=device
+    )
+    req_id, within, valid = compact_row_index(
+        verify_lens=layout.verify_lens,
+        padded_total=layout.graph_num_tokens,
+        device=device,
+    )
+    sink = bs * stride
+    strided_pos = torch.where(
+        valid,
+        req_id.clamp(max=bs - 1) * stride + within,
+        torch.full_like(within, sink),
+    )
+    strided.index_copy_(0, strided_pos, compact)
+    return strided[: bs * stride]
+
+
+@triton.jit
+def _scatter_compact_to_strided_kernel(
+    compact_ptr,
+    verify_lens_ptr,
+    start_ptr,
+    out_ptr,
+    stride,
+    dim,
+    fill_value,
+    BLOCK_D: tl.constexpr,
+):
+    o = tl.program_id(0).to(tl.int64)
+    dblk = tl.program_id(1)
+    i = o // stride
+    w = o % stride
+    vl_i = tl.load(verify_lens_ptr + i)
+    start_i = tl.load(start_ptr + i)
+    d = dblk * BLOCK_D + tl.arange(0, BLOCK_D)
+    dmask = d < dim
+    in_range = w < vl_i
+    src = tl.where(in_range, start_i + w, 0)
+    val = tl.load(compact_ptr + src * dim + d, mask=dmask & in_range, other=0)
+    val = tl.where(in_range, val, fill_value)
+    tl.store(out_ptr + o * dim + d, val, mask=dmask)
+
+
+def scatter_compact_to_strided_into(
+    *,
+    compact: torch.Tensor,
+    verify_lens: torch.Tensor,
+    out: torch.Tensor,
+    stride: int,
+    fill_value: float,
+) -> torch.Tensor:
+    dim = compact.shape[1]
+    fill_value = float(fill_value) if out.dtype.is_floating_point else int(fill_value)
+    compact = compact.contiguous()
+    verify_lens = verify_lens.to(dtype=torch.int64).contiguous()
+    start = (torch.cumsum(verify_lens, dim=0) - verify_lens).contiguous()
+    n_out = out.shape[0]
+    BLOCK_D = 1024
+    grid = (n_out, triton.cdiv(dim, BLOCK_D))
+    _scatter_compact_to_strided_kernel[grid](
+        compact,
+        verify_lens,
+        start,
+        out,
+        stride,
+        dim,
+        fill_value,
+        BLOCK_D=BLOCK_D,
+    )
+    return out
+
+
+def scatter_compact_to_strided_triton(
+    *,
+    compact: torch.Tensor,
+    layout: RaggedVerifyLayout,
+    fill_value: float,
+    verify_num_draft_tokens: int,
+) -> torch.Tensor:
+    stride = verify_num_draft_tokens
+    bs = layout.verify_lens.shape[0]
+    dim = compact.shape[1]
+    device = compact.device
+    out = torch.empty((bs * stride, dim), dtype=compact.dtype, device=device)
+    return scatter_compact_to_strided_into(
+        compact=compact,
+        verify_lens=layout.verify_lens.to(device=device),
+        out=out,
+        stride=stride,
+        fill_value=fill_value,
+    )
+
+
+class CommitInjectLayoutResult(msgspec.Struct):
+    swa_loc: torch.Tensor
+    positions: torch.Tensor
+
+
+class BuildCommitInjectLayout:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> CommitInjectLayoutResult:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        req_pool_indices: torch.Tensor,
+        req_to_token: torch.Tensor,
+        prefix_lens: torch.Tensor,
+        block_pos_offsets: torch.Tensor,
+        full_to_swa_mapping: torch.Tensor,
+        commit_lens: torch.Tensor,
+        stride: int,
+    ) -> CommitInjectLayoutResult:
+        return build_commit_inject_layout(
+            req_pool_indices=req_pool_indices,
+            req_to_token=req_to_token,
+            prefix_lens=prefix_lens,
+            block_pos_offsets=block_pos_offsets,
+            full_to_swa_mapping=full_to_swa_mapping,
+            commit_lens=commit_lens,
+            stride=stride,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        req_pool_indices: torch.Tensor,
+        req_to_token: torch.Tensor,
+        prefix_lens: torch.Tensor,
+        block_pos_offsets: torch.Tensor,
+        full_to_swa_mapping: torch.Tensor,
+        commit_lens: torch.Tensor,
+        stride: int,
+    ) -> CommitInjectLayoutResult:
+        return build_commit_inject_layout_triton(
+            req_pool_indices=req_pool_indices,
+            req_to_token=req_to_token,
+            prefix_lens=prefix_lens,
+            block_pos_offsets=block_pos_offsets,
+            full_to_swa_mapping=full_to_swa_mapping,
+            commit_lens=commit_lens,
+            stride=stride,
+        )
+
+
+def build_commit_inject_layout(
+    *,
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    prefix_lens: torch.Tensor,
+    block_pos_offsets: torch.Tensor,
+    full_to_swa_mapping: torch.Tensor,
+    commit_lens: torch.Tensor,
+    stride: int,
+) -> CommitInjectLayoutResult:
+    from sglang.kernels.ops.speculative.cache_locs import (
+        assign_extend_cache_locs_func,
+    )
+
+    bs = req_pool_indices.shape[0]
+    device = req_pool_indices.device
+
+    positions_2d = prefix_lens.unsqueeze(1) + block_pos_offsets[:stride]
+    positions = positions_2d.reshape(-1).to(dtype=torch.int64)
+
+    cache_loc = assign_extend_cache_locs_func(
+        req_pool_indices=req_pool_indices,
+        req_to_token=req_to_token,
+        start_offset=prefix_lens,
+        end_offset=prefix_lens + stride,
+        batch_size=bs,
+        draft_token_num=stride,
+        device=device,
+    ).to(dtype=torch.int64)
+    swa_loc = full_to_swa_mapping[cache_loc].to(torch.int32)
+
+    col = torch.arange(stride, device=device).view(1, -1)
+    committed = (col < commit_lens.to(torch.long).view(-1, 1)).reshape(-1)
+    swa_loc = torch.where(committed, swa_loc, torch.full_like(swa_loc, -1))
+
+    return CommitInjectLayoutResult(swa_loc=swa_loc, positions=positions)
+
+
+@triton.jit
+def _commit_inject_layout_kernel(
+    req_pool_ptr,
+    req_to_token_ptr,
+    prefix_lens_ptr,
+    block_pos_offsets_ptr,
+    full_to_swa_ptr,
+    commit_lens_ptr,
+    swa_loc_ptr,
+    positions_ptr,
+    rt_stride,
+    stride,
+    n,
+    BLOCK: tl.constexpr,
+):
+    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < n
+    r = offs // stride
+    c = offs % stride
+
+    prefix = tl.load(prefix_lens_ptr + r, mask=mask, other=0).to(tl.int64)
+    pos_off = tl.load(block_pos_offsets_ptr + c, mask=mask, other=0).to(tl.int64)
+    rp = tl.load(req_pool_ptr + r, mask=mask, other=0).to(tl.int64)
+    full_loc = tl.load(
+        req_to_token_ptr + rp * rt_stride + prefix + pos_off, mask=mask, other=0
+    ).to(tl.int64)
+    swa = tl.load(full_to_swa_ptr + full_loc, mask=mask, other=-1).to(tl.int32)
+
+    commit_len = tl.load(commit_lens_ptr + r, mask=mask, other=0).to(tl.int64)
+    swa = tl.where(c.to(tl.int64) < commit_len, swa, -1)
+
+    tl.store(swa_loc_ptr + offs, swa, mask=mask)
+    tl.store(positions_ptr + offs, prefix + pos_off, mask=mask)
+
+
+def build_commit_inject_layout_triton(
+    *,
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    prefix_lens: torch.Tensor,
+    block_pos_offsets: torch.Tensor,
+    full_to_swa_mapping: torch.Tensor,
+    commit_lens: torch.Tensor,
+    stride: int,
+) -> CommitInjectLayoutResult:
+    bs = req_pool_indices.shape[0]
+    n = bs * stride
+    device = req_pool_indices.device
+
+    swa_loc = torch.empty(n, dtype=torch.int32, device=device)
+    positions = torch.empty(n, dtype=torch.int64, device=device)
+    BLOCK = 256
+    _commit_inject_layout_kernel[(triton.cdiv(n, BLOCK),)](
+        req_pool_indices,
+        req_to_token,
+        prefix_lens,
+        block_pos_offsets,
+        full_to_swa_mapping,
+        commit_lens,
+        swa_loc,
+        positions,
+        req_to_token.stride(0),
+        stride,
+        n,
+        BLOCK=BLOCK,
+    )
+    return CommitInjectLayoutResult(swa_loc=swa_loc, positions=positions)
+
+
+class BuildOutTokens:
+    @classmethod
+    def execute(cls, *args, **kwargs) -> torch.Tensor:
+        if inputs_on_cuda(*args, **kwargs):
+            return cls.triton(*args, **kwargs)
+        return cls.torch(*args, **kwargs)
+
+    @classmethod
+    def torch(
+        cls,
+        *,
+        draft_tokens: torch.Tensor,
+        correct_len: torch.Tensor,
+        bonus: torch.Tensor,
+        verify_num_draft_tokens: int,
+        gamma: int,
+    ) -> torch.Tensor:
+        return build_out_tokens(
+            draft_tokens=draft_tokens,
+            correct_len=correct_len,
+            bonus=bonus,
+            verify_num_draft_tokens=verify_num_draft_tokens,
+            gamma=gamma,
+        )
+
+    @classmethod
+    def triton(
+        cls,
+        *,
+        draft_tokens: torch.Tensor,
+        correct_len: torch.Tensor,
+        bonus: torch.Tensor,
+        verify_num_draft_tokens: int,
+        gamma: int,
+    ) -> torch.Tensor:
+        return build_out_tokens_triton(
+            draft_tokens=draft_tokens,
+            correct_len=correct_len,
+            bonus=bonus,
+            verify_num_draft_tokens=verify_num_draft_tokens,
+            gamma=gamma,
+        )
+
+
+def build_out_tokens(
+    *,
+    draft_tokens: torch.Tensor,
+    correct_len: torch.Tensor,
+    bonus: torch.Tensor,
+    verify_num_draft_tokens: int,
+    gamma: int,
+) -> torch.Tensor:
+    bs = draft_tokens.shape[0]
+    out_tokens = torch.empty(
+        (bs, verify_num_draft_tokens),
+        dtype=torch.int64,
+        device=draft_tokens.device,
+    )
+    out_tokens[:, :gamma].copy_(draft_tokens)
+    out_tokens[:, gamma].fill_(0)
+    out_tokens.scatter_(1, correct_len.to(torch.int64)[:, None], bonus[:, None])
+    return out_tokens
+
+
+@triton.jit
+def _build_out_tokens_kernel(
+    draft_tokens_ptr,
+    correct_len_ptr,
+    bonus_ptr,
+    out_ptr,
+    gamma,
+    T,
+    n_out,
+    BLOCK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < n_out
+    b = offs // T
+    k = offs % T
+    cl = tl.load(correct_len_ptr + b, mask=mask, other=0).to(tl.int32)
+    bonus = tl.load(bonus_ptr + b, mask=mask, other=0)
+    draft_mask = mask & (k < gamma)
+    draft = tl.load(draft_tokens_ptr + b * gamma + k, mask=draft_mask, other=0)
+    val = tl.where(k == cl, bonus, tl.where(k < gamma, draft, 0))
+    tl.store(out_ptr + offs, val.to(tl.int64), mask=mask)
+
+
+def build_out_tokens_triton(
+    *,
+    draft_tokens: torch.Tensor,
+    correct_len: torch.Tensor,
+    bonus: torch.Tensor,
+    verify_num_draft_tokens: int,
+    gamma: int,
+) -> torch.Tensor:
+    bs = draft_tokens.shape[0]
+    T = verify_num_draft_tokens
+    device = draft_tokens.device
+    draft_tokens = draft_tokens.to(torch.int64).contiguous()
+    correct_len_i = correct_len.to(torch.int64).contiguous()
+    bonus_i = bonus.to(torch.int64).contiguous()
+    out = torch.empty((bs, T), dtype=torch.int64, device=device)
+    n_out = bs * T
+    BLOCK = 256
+    grid = (triton.cdiv(n_out, BLOCK),)
+    _build_out_tokens_kernel[grid](
+        draft_tokens, correct_len_i, bonus_i, out, gamma, T, n_out, BLOCK=BLOCK
+    )
+    return out