chore: import upstream snapshot with attribution

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