feat(mtp): add optional ASD acceptance policy for greedy MTP verification

lightllm/common/basemodel/triton_kernel/mtp_asd.py

@@ -0,0 +1,162 @@
+"""ASD (Approximate Speculative Decoding, arxiv:2608.03447) acceptance for MTP verify.
+
+This module is a self-contained plugin: it implements the budgeted relaxed acceptance
+rule used by ``mtp_speculative.utils.verify_mtp_tokens`` when ``--mtp_asd_regret_budget``
+is set. It shares the launch contract of ``mtp_utils.mtp_verify`` so the verification
+seam and all downstream consumers stay unchanged.
+"""
+
+import triton
+import triton.language as tl
+import torch
+
+
+@triton.jit
+def _fwd_kernel_mtp_asd_verify(
+    row_regrets,
+    row_draft_token_ids,
+    new_next_token_ids,
+    req_to_asd_cum_regret,
+    mtp_accept_len,
+    b_req_mtp_start_loc,
+    b_req_idx,
+    accepted_index,
+    verify_batch_size,
+    asd_budget,
+    asd_local_ratio,
+    asd_max_mismatch,
+    BLOCK_SIZE: tl.constexpr,
+):
+    cur_index = tl.program_id(0)
+    req_nums = tl.num_programs(axis=0)
+
+    req_start_loc = tl.load(b_req_mtp_start_loc + cur_index)
+    req_start_end = tl.load(
+        b_req_mtp_start_loc + cur_index + 1,
+        mask=cur_index + 1 < req_nums,
+        other=verify_batch_size,
+    )
+    req_mtp_num = req_start_end - req_start_loc
+    draft_num = req_mtp_num - 1  # K draft rows; the final row is the target (bonus) row
+    cur_req_idx = tl.load(b_req_idx + req_start_loc)
+
+    offset = tl.arange(0, BLOCK_SIZE)
+    req_offset = req_start_loc + offset
+    draft_pos_mask = offset < draft_num
+
+    regrets = tl.load(row_regrets + req_offset, mask=draft_pos_mask, other=0.0)
+    cum_regrets = tl.cumsum(regrets, axis=0)
+    cum_mismatches = tl.cumsum((regrets > 0).to(tl.int32), axis=0)
+    suffix_values = draft_num - offset  # K, K-1, ..., 1 on draft rows
+    local_ratios = tl.where(suffix_values > 0, regrets / suffix_values, 0.0)
+
+    budget_used = tl.load(req_to_asd_cum_regret + cur_req_idx)
+    feasible = (
+        draft_pos_mask
+        & (budget_used + cum_regrets <= asd_budget)
+        & (local_ratios <= asd_local_ratio)
+        & (cum_mismatches <= asd_max_mismatch)
+    )
+
+    # Acceptance stops at the first infeasible draft position; the target row at that
+    # offset is still committed (correction/bonus), mirroring strict accept_len semantics.
+    infeasible_positions = tl.where(feasible, BLOCK_SIZE, offset)
+    accept_draft_len = tl.min(infeasible_positions, axis=0)
+    accept_len = accept_draft_len + 1
+    tl.store(mtp_accept_len + cur_index, accept_len)
+
+    accepted_draft_mask = offset < accept_draft_len
+    step_regret = tl.sum(tl.where(accepted_draft_mask & draft_pos_mask, regrets, 0.0), axis=0)
+    tl.store(req_to_asd_cum_regret + cur_req_idx, budget_used + step_regret)
+
+    # Relaxed rows commit the draft token instead of the target-selected token.
+    relaxed_mask = accepted_draft_mask & (regrets > 0)
+    draft_token_ids = tl.load(row_draft_token_ids + req_offset, mask=relaxed_mask, other=0)
+    tl.store(new_next_token_ids + req_offset, draft_token_ids, mask=relaxed_mask)
+
+    accepted_index_values = tl.where(offset < accept_len, 1, 0)
+    tl.store(accepted_index + req_offset, accepted_index_values, mask=offset < req_mtp_num)
+    return
+
+
+def mtp_asd_verify(
+    req_to_next_token_ids: torch.Tensor,
+    b_req_mtp_start_loc: torch.Tensor,
+    new_next_token_ids: torch.Tensor,
+    b_req_idx: torch.Tensor,
+    b_mtp_index: torch.Tensor,
+    logits: torch.Tensor,
+    req_to_asd_cum_regret: torch.Tensor,
+    asd_budget: float,
+    asd_local_ratio: float,
+    asd_max_mismatch: int,
+):
+    """ASD (Approximate Speculative Decoding, arxiv:2608.03447) variant of ``mtp_verify``.
+
+    A draft token x_i is accepted while (1) the request-level cumulative regret stays within
+    ``asd_budget`` (regret r_i = max_v z_i(v) - z_i(x_i) against the target logits z_i),
+    (2) r_i / (K - i) <= ``asd_local_ratio``, and (3) the block contains at most
+    ``asd_max_mismatch`` relaxed tokens. Acceptance stops at the first infeasible position
+    and the target row is committed as usual, so ``asd_budget = 0`` or
+    ``asd_max_mismatch = 0`` exactly recovers strict greedy verification.
+
+    Same return contract as ``mtp_verify``; rows accepted under a nonzero regret commit the
+    draft token into ``new_next_token_ids`` in place, so downstream consumers (token counter,
+    scatter, response building) need no changes.
+
+    Args:
+        req_to_next_token_ids: (max_req_num, verify_width)
+        b_req_mtp_start_loc: (num_reqs,)
+        new_next_token_ids: (verify_batch_size,) target-selected ids; modified in place.
+        b_req_idx: (verify_batch_size,)
+        b_mtp_index: (verify_batch_size,) local row index of each row within its request.
+        logits: (verify_batch_size, vocab) the same logits that produced new_next_token_ids
+            (post-penalty, post-temperature).
+        req_to_asd_cum_regret: (max_req_num + 1,) per-request cumulative regret ledger.
+        asd_budget: request-level cumulative regret budget B.
+        asd_local_ratio: per-token regret / suffix-value cap g.
+        asd_max_mismatch: max relaxed tokens per verify step m.
+    Returns:
+        mtp_accept_len: (num_reqs,)
+        accepted_index: (verify_batch_size,)
+    """
+    verify_width = req_to_next_token_ids.shape[1]
+    BLOCK_SIZE = 16
+    assert verify_width <= BLOCK_SIZE, f"verify_width must be less than {BLOCK_SIZE}"
+    num_reqs = b_req_mtp_start_loc.shape[0]
+    verify_batch_size = b_req_idx.shape[0]
+    assert new_next_token_ids.shape == b_req_idx.shape == b_mtp_index.shape
+    assert logits.shape[0] == verify_batch_size
+
+    # Row-level regret precompute (device-side only, no host synchronization).
+    # Row r verifies draft column b_mtp_index[r] + 1; the bonus row of each request does
+    # not verify any draft, so its column is clamped and its regret is never consumed.
+    draft_columns = torch.clamp(b_mtp_index + 1, max=verify_width - 1).to(torch.int64)
+    row_draft_token_ids = req_to_next_token_ids[b_req_idx.to(torch.int64), draft_columns]
+    row_max_logits = logits.max(dim=-1).values
+    row_draft_logits = logits.gather(dim=-1, index=row_draft_token_ids.view(-1, 1)).squeeze(-1)
+    row_regrets = (row_max_logits - row_draft_logits).to(torch.float32)
+
+    mtp_accept_len = torch.empty((num_reqs,), dtype=torch.int32, device=req_to_next_token_ids.device)
+    accepted_index = torch.empty((verify_batch_size,), dtype=torch.int32, device=req_to_next_token_ids.device)
+
+    grid = (num_reqs,)
+    num_warps = 1
+    _fwd_kernel_mtp_asd_verify[grid](
+        row_regrets=row_regrets,
+        row_draft_token_ids=row_draft_token_ids,
+        new_next_token_ids=new_next_token_ids,
+        req_to_asd_cum_regret=req_to_asd_cum_regret,
+        mtp_accept_len=mtp_accept_len,
+        b_req_mtp_start_loc=b_req_mtp_start_loc,
+        b_req_idx=b_req_idx,
+        accepted_index=accepted_index,
+        verify_batch_size=verify_batch_size,
+        asd_budget=asd_budget,
+        asd_local_ratio=asd_local_ratio,
+        asd_max_mismatch=asd_max_mismatch,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+        num_stages=1,
+    )
+    return mtp_accept_len, accepted_index

lightllm/common/req_manager/req_sampling_params.py

@@ -40,6 +40,13 @@ def __init__(self, max_request_num):
             if get_env_start_args().mtp_dynamic_verify
             else None
         )
+        # Per-request cumulative ASD regret budget ledger. Only allocated when ASD acceptance
+        # is enabled; zeroed per request at the prefill seam in init_req_sampling_params.
+        self.req_to_asd_cum_regret = (
+            torch.zeros(max_request_num + 1, dtype=torch.float32, device="cuda")
+            if get_env_start_args().mtp_asd_regret_budget is not None
+            else None
+        )
 
         self.req_to_exponential_decay_length_penalty = torch.zeros(
             max_request_num + 1, dtype=torch.float32, device="cuda"
@@ -60,6 +67,8 @@ def init_req_sampling_params(self, req: "InferReq"):
         if self.req_to_next_token_scores is not None:
             self.req_to_next_token_scores[req.req_idx].fill_(0.0)
             self.req_to_next_token_scores[req.req_idx][0:1].fill_(1.0)
+        if self.req_to_asd_cum_regret is not None:
+            self.req_to_asd_cum_regret[req.req_idx].fill_(0.0)
         self.req_to_presence_penalty[req.req_idx].fill_(shm_param.presence_penalty)
         self.req_to_frequency_penalty[req.req_idx].fill_(shm_param.frequency_penalty)
         self.req_to_repetition_penalty[req.req_idx].fill_(shm_param.repetition_penalty)

lightllm/server/api_start.py

@@ -177,6 +177,13 @@ def _launch_subprocesses(args: StartArgs):
         assert args.mtp_draft_model_dir is None
         assert args.mtp_step == 0
 
+    # ASD acceptance params check (see --mtp_asd_regret_budget help)
+    if args.mtp_asd_regret_budget is not None:
+        assert args.mtp_mode is not None, "--mtp_asd_regret_budget requires an enabled mtp_mode"
+        assert args.mtp_asd_regret_budget >= 0.0, "--mtp_asd_regret_budget must be >= 0"
+        assert args.mtp_asd_local_regret_ratio >= 0.0, "--mtp_asd_local_regret_ratio must be >= 0"
+        assert args.mtp_asd_block_max_mismatch >= 0, "--mtp_asd_block_max_mismatch must be >= 0"
+
     # automatically set visual_dp based on visual_tp and tp.
     # In visual proxy mode keep the caller-provided visual_dp / visual_tp.
     if not args.visual_use_proxy_mode and args.visual_tp < args.tp and args.tp % args.visual_tp == 0:

lightllm/server/core/objs/start_args_type.py

@@ -201,6 +201,36 @@ class StartArgs:
     mtp_draft_model_dir: Optional[List[str]] = field(default=None)
     mtp_step: int = field(default=0)
     mtp_dynamic_verify: bool = field(default=False)
+    mtp_asd_regret_budget: Optional[float] = field(
+        default=None,
+        metadata={
+            "help": (
+                "enable Approximate Speculative Decoding (ASD) acceptance for MTP greedy verification "
+                "with this per-request cumulative regret budget B (arxiv:2608.03447). A draft token whose "
+                "target-logit regret keeps the request-level cumulative regret within B is accepted. "
+                "None (default) keeps strict lossless verification; B=0 recovers it exactly. Regrets are "
+                "measured on the post-penalty, post-temperature logits used for token selection."
+            )
+        },
+    )
+    mtp_asd_local_regret_ratio: float = field(
+        default=0.25,
+        metadata={
+            "help": (
+                "ASD local gate g: a draft token at position i (of K draft tokens) is accepted only if "
+                "regret_i / (K - i) <= g. Later draft positions get less slack."
+            )
+        },
+    )
+    mtp_asd_block_max_mismatch: int = field(
+        default=2,
+        metadata={
+            "help": (
+                "ASD per-block cap m: at most m relaxed (non-argmax) draft tokens are accepted per "
+                "verify step. m=0 recovers strict greedy verification."
+            )
+        },
+    )
     kv_quant_calibration_config_path: Optional[str] = field(default=None)
     pd_kv_page_num: int = field(default=16)
     pd_kv_page_size: int = field(default=1024)

lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py

@@ -285,6 +285,8 @@ def decode_mtp(
                 b_req_idx=model_input.b_req_idx,
                 b_req_mtp_start_loc=b_req_mtp_start_loc,
                 b_mtp_index=model_input.b_mtp_index,
+                logits=model_output.logits,
+                run_reqs=run_reqs,
             )
             accepted_index_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor(
                 key="accepted_index",

lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py

@@ -516,6 +516,8 @@ def decode_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq]):
                     b_req_idx=model_input.b_req_idx,
                     b_req_mtp_start_loc=b_req_mtp_start_loc,
                     b_mtp_index=model_input.b_mtp_index,
+                    logits=model_output.logits,
+                    run_reqs=run_reqs,
                 )
                 accepted_index_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor(
                     key="accepted_index",
@@ -804,7 +806,22 @@ def decode_overlap_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[Inf
                     b_req_idx=b_req_idx,
                     b_req_mtp_start_loc=b_req_mtp_start_loc,
                     b_mtp_index=b_mtp_index,
+                    logits=logits,
+                    run_reqs=run_reqs,
                 )
+                if self.args.mtp_asd_regret_budget is not None:
+                    # ASD rewrites next_token_ids in place during verify (relaxed rows commit
+                    # draft tokens). The early overlap copy above raced ahead of verify, so
+                    # refresh the pinned buffer to surface the committed ids downstream.
+                    (
+                        next_token_ids_cpu,
+                        next_token_logprobs_cpu,
+                        next_token_ranks_cpu,
+                    ) = self._async_copy_next_token_infos_to_pin_mem(
+                        next_token_ids=next_token_ids,
+                        next_token_logprobs=next_token_logprobs,
+                        next_token_ranks=next_token_ranks,
+                    )
                 mtp_accept_len0 = mtp_accept_len[:real_request_num0]
                 mtp_accept_len1 = mtp_accept_len[real_request_num0:]
                 accepted_index_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor(

lightllm/server/router/model_infer/mtp_speculative/utils.py

@@ -1,15 +1,17 @@
 from __future__ import annotations
 
 from collections import Counter
-from typing import TYPE_CHECKING, List, Tuple
+from typing import TYPE_CHECKING, List, Optional, Tuple
 
 import torch
 
+from lightllm.common.basemodel.triton_kernel.mtp_asd import mtp_asd_verify
 from lightllm.common.basemodel.triton_kernel.mtp_utils import (
     linear_att_mtp_state_index_update,
     mtp_scatter_next_token_ids,
     mtp_verify,
 )
+from lightllm.utils.envs_utils import get_env_start_args
 
 if TYPE_CHECKING:
     from lightllm.server.router.model_infer.infer_batch import InferReq
@@ -34,21 +36,57 @@ def alloc_mem_indexes(token_count: int) -> torch.Tensor:
     return g_infer_context.req_manager.mem_manager.alloc(token_count)
 
 
+def _is_all_greedy(run_reqs: List[InferReq]) -> bool:
+    """Mirror the greedy rule used by generic_post_process.sample (top_k == 1)."""
+
+    return all(req.sampling_param.shm_param.top_k == 1 for req in run_reqs)
+
+
 def verify_mtp_tokens(
     backend: ModeBackend,
     next_token_ids: torch.Tensor,
     b_req_idx: torch.Tensor,
     b_req_mtp_start_loc: torch.Tensor,
     b_mtp_index: torch.Tensor,
+    logits: Optional[torch.Tensor] = None,
+    run_reqs: Optional[List[InferReq]] = None,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Verify target tokens and update recurrent MTP state when required."""
+    """Verify target tokens and update recurrent MTP state when required.
 
-    accept_lengths, accepted_index = mtp_verify(
-        req_to_next_token_ids=backend.model.req_manager.req_sampling_params_manager.req_to_next_token_ids,
-        b_req_mtp_start_loc=b_req_mtp_start_loc,
-        new_next_token_ids=next_token_ids,
-        b_req_idx=b_req_idx,
+    When ASD acceptance is enabled (``--mtp_asd_regret_budget``) and the whole batch is
+    greedy, verification relaxes strict token equality under a bounded per-request regret
+    budget; relaxed rows commit their draft token into ``next_token_ids`` in place, so all
+    downstream consumers keep the same contract as with strict verification.
+    """
+
+    start_args = get_env_start_args()
+    sampling_params_manager = backend.model.req_manager.req_sampling_params_manager
+    use_asd = (
+        start_args.mtp_asd_regret_budget is not None
+        and logits is not None
+        and run_reqs is not None
+        and _is_all_greedy(run_reqs)
     )
+    if use_asd:
+        accept_lengths, accepted_index = mtp_asd_verify(
+            req_to_next_token_ids=sampling_params_manager.req_to_next_token_ids,
+            b_req_mtp_start_loc=b_req_mtp_start_loc,
+            new_next_token_ids=next_token_ids,
+            b_req_idx=b_req_idx,
+            b_mtp_index=b_mtp_index,
+            logits=logits,
+            req_to_asd_cum_regret=sampling_params_manager.req_to_asd_cum_regret,
+            asd_budget=start_args.mtp_asd_regret_budget,
+            asd_local_ratio=start_args.mtp_asd_local_regret_ratio,
+            asd_max_mismatch=start_args.mtp_asd_block_max_mismatch,
+        )
+    else:
+        accept_lengths, accepted_index = mtp_verify(
+            req_to_next_token_ids=sampling_params_manager.req_to_next_token_ids,
+            b_req_mtp_start_loc=b_req_mtp_start_loc,
+            new_next_token_ids=next_token_ids,
+            b_req_idx=b_req_idx,
+        )
     if backend.is_linear_att_mixed_model:
         linear_att_mtp_state_index_update(
             req_to_mtp_state_index=backend.model.req_manager.req_to_mtp_state_index,