From f63fddf0163e900810d4d331e3fa8eab30764e98 Mon Sep 17 00:00:00 2001 From: Yuannuo Feng Date: Wed, 9 Sep 2026 16:25:13 +0800 Subject: [PATCH] feat(mtp): add optional ASD acceptance policy for greedy MTP verification Approximate Speculative Decoding (arXiv:2608.03447): relax strict greedy MTP verification by accepting draft tokens whose target-logit regret stays within a bounded per-request budget. Default off (strict, lossless); budget=0 or max_mismatch=0 recovers strict verification exactly. All ASD logic lives in the new plugin module triton_kernel/mtp_asd.py; existing files only carry config fields and the dispatch seam. Implements ModelTC/LightLLM#1552. --- .../common/basemodel/triton_kernel/mtp_asd.py | 162 ++++++++++++ .../common/req_manager/req_sampling_params.py | 9 + lightllm/server/api_start.py | 7 + lightllm/server/core/objs/start_args_type.py | 30 +++ .../mode_backend/chunked_prefill/impl.py | 2 + .../mode_backend/dp_backend/impl.py | 17 ++ .../model_infer/mtp_speculative/utils.py | 52 +++- .../mtp_speculative/test_mtp_asd_verify.py | 245 ++++++++++++++++++ unit_tests/server/test_mtp_start_args.py | 60 +++++ 9 files changed, 577 insertions(+), 7 deletions(-) create mode 100644 lightllm/common/basemodel/triton_kernel/mtp_asd.py create mode 100644 unit_tests/server/router/model_infer/mtp_speculative/test_mtp_asd_verify.py diff --git a/lightllm/common/basemodel/triton_kernel/mtp_asd.py b/lightllm/common/basemodel/triton_kernel/mtp_asd.py new file mode 100644 index 0000000000..047637dcbf --- /dev/null +++ b/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 diff --git a/lightllm/common/req_manager/req_sampling_params.py b/lightllm/common/req_manager/req_sampling_params.py index 87271b6145..45f2f99b91 100644 --- a/lightllm/common/req_manager/req_sampling_params.py +++ b/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) diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index 6f8973425f..d7f85cd275 100644 --- a/lightllm/server/api_start.py +++ b/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: diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 9c89975de7..dc4f39c308 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/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) diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index 4d09476849..c0ba115057 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py +++ b/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", diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py index 9a81927bc1..3e455488df 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py +++ b/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( diff --git a/lightllm/server/router/model_infer/mtp_speculative/utils.py b/lightllm/server/router/model_infer/mtp_speculative/utils.py index 935fe92d33..641c55a8f5 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/utils.py +++ b/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, diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_mtp_asd_verify.py b/unit_tests/server/router/model_infer/mtp_speculative/test_mtp_asd_verify.py new file mode 100644 index 0000000000..aca56e3bd7 --- /dev/null +++ b/unit_tests/server/router/model_infer/mtp_speculative/test_mtp_asd_verify.py @@ -0,0 +1,245 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.mtp_asd import mtp_asd_verify +from lightllm.common.basemodel.triton_kernel.mtp_utils import mtp_verify + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="ASD verify kernel requires CUDA") + + +def _build_case(device="cuda"): + """One request with 2 draft rows + 1 bonus row; known per-row regrets. + + Layout: committed token 7, draft tokens [3, 5]. + row0: argmax 9 (logit 10.0), draft 3 (logit 8.0) -> regret 2.0 + row1: argmax 9 (logit 10.0), draft 5 (logit 7.4) -> regret 2.6 + row2: bonus row, argmax 6 + """ + req_to_next_token_ids = torch.tensor([[7, 3, 5, -1, -1]], dtype=torch.int64, device=device) + b_req_idx = torch.tensor([0, 0, 0], dtype=torch.int32, device=device) + b_mtp_index = torch.tensor([0, 1, 2], dtype=torch.int32, device=device) + b_req_mtp_start_loc = torch.tensor([0], dtype=torch.int32, device=device) + vocab = 10 + logits = torch.full((3, vocab), -10.0, dtype=torch.float32, device=device) + logits[0, 9] = 10.0 + logits[0, 3] = 8.0 + logits[1, 9] = 10.0 + logits[1, 5] = 7.4 + logits[2, 6] = 10.0 + new_next_token_ids = torch.tensor([9, 9, 6], dtype=torch.int64, device=device) + req_to_asd_cum_regret = torch.zeros(1, dtype=torch.float32, device=device) + return ( + req_to_next_token_ids, + b_req_mtp_start_loc, + new_next_token_ids, + b_req_idx, + b_mtp_index, + logits, + req_to_asd_cum_regret, + ) + + +def test_asd_zero_budget_equals_strict(): + """Gold standard: B=0 (and m=0) must reproduce strict verification exactly.""" + ( + req_to_next_token_ids, + b_req_mtp_start_loc, + new_next_token_ids, + b_req_idx, + b_mtp_index, + logits, + req_to_asd_cum_regret, + ) = _build_case() + + strict_len, strict_index = mtp_verify( + req_to_next_token_ids, b_req_mtp_start_loc, new_next_token_ids.clone(), b_req_idx + ) + for budget, max_mismatch in [(0.0, 2), (100.0, 0)]: + asd_ids = new_next_token_ids.clone() + asd_len, asd_index = mtp_asd_verify( + req_to_next_token_ids=req_to_next_token_ids, + b_req_mtp_start_loc=b_req_mtp_start_loc, + new_next_token_ids=asd_ids, + b_req_idx=b_req_idx, + b_mtp_index=b_mtp_index, + logits=logits, + req_to_asd_cum_regret=req_to_asd_cum_regret, + asd_budget=budget, + asd_local_ratio=100.0, + asd_max_mismatch=max_mismatch, + ) + assert torch.equal(asd_len, strict_len) + assert torch.equal(asd_index, strict_index) + assert torch.equal(asd_ids, new_next_token_ids) # no relaxed commit under strict equivalence + assert req_to_asd_cum_regret.item() == 0.0 + + +def test_asd_relaxed_acceptance_hand_computed(): + """row0 (regret 2.0) fits the budget and is relaxed-accepted with the draft token; + row1 (cumulative 4.6) exceeds B=3.0 and stops acceptance.""" + ( + req_to_next_token_ids, + b_req_mtp_start_loc, + new_next_token_ids, + b_req_idx, + b_mtp_index, + logits, + req_to_asd_cum_regret, + ) = _build_case() + + asd_ids = new_next_token_ids.clone() + accept_len, accepted_index = mtp_asd_verify( + req_to_next_token_ids=req_to_next_token_ids, + b_req_mtp_start_loc=b_req_mtp_start_loc, + new_next_token_ids=asd_ids, + b_req_idx=b_req_idx, + b_mtp_index=b_mtp_index, + logits=logits, + req_to_asd_cum_regret=req_to_asd_cum_regret, + asd_budget=3.0, + asd_local_ratio=100.0, + asd_max_mismatch=2, + ) + assert accept_len.tolist() == [2] + assert accepted_index.tolist() == [1, 1, 0] + assert asd_ids.tolist() == [3, 9, 6] # relaxed row commits the draft token 3 + assert req_to_asd_cum_regret.item() == pytest.approx(2.0) + + +def test_asd_local_ratio_gate(): + """The local gate rejects a late position even when the budget is ample. + + row1 regret 2.6 with suffix value K - i = 2 - 1 = 1: 2.6 / 1 > g=1.5 -> rejected. + row0 regret 2.0 with suffix 2: 2.0 / 2 = 1.0 <= 1.5 -> accepted. + """ + ( + req_to_next_token_ids, + b_req_mtp_start_loc, + new_next_token_ids, + b_req_idx, + b_mtp_index, + logits, + req_to_asd_cum_regret, + ) = _build_case() + + asd_ids = new_next_token_ids.clone() + accept_len, _ = mtp_asd_verify( + req_to_next_token_ids=req_to_next_token_ids, + b_req_mtp_start_loc=b_req_mtp_start_loc, + new_next_token_ids=asd_ids, + b_req_idx=b_req_idx, + b_mtp_index=b_mtp_index, + logits=logits, + req_to_asd_cum_regret=req_to_asd_cum_regret, + asd_budget=100.0, + asd_local_ratio=1.5, + asd_max_mismatch=5, + ) + assert accept_len.tolist() == [2] + assert asd_ids.tolist() == [3, 9, 6] + + +def test_asd_max_mismatch_cap(): + """m=1 lets only the first relaxed token through even with ample budget.""" + ( + req_to_next_token_ids, + b_req_mtp_start_loc, + new_next_token_ids, + b_req_idx, + b_mtp_index, + logits, + req_to_asd_cum_regret, + ) = _build_case() + + asd_ids = new_next_token_ids.clone() + accept_len, _ = mtp_asd_verify( + req_to_next_token_ids=req_to_next_token_ids, + b_req_mtp_start_loc=b_req_mtp_start_loc, + new_next_token_ids=asd_ids, + b_req_idx=b_req_idx, + b_mtp_index=b_mtp_index, + logits=logits, + req_to_asd_cum_regret=req_to_asd_cum_regret, + asd_budget=100.0, + asd_local_ratio=100.0, + asd_max_mismatch=1, + ) + assert accept_len.tolist() == [2] + assert asd_ids.tolist() == [3, 9, 6] + + +def test_asd_all_accepted_commits_bonus(): + """All draft rows feasible: accept_len == req_mtp_num, bonus row keeps the target id.""" + ( + req_to_next_token_ids, + b_req_mtp_start_loc, + new_next_token_ids, + b_req_idx, + b_mtp_index, + logits, + req_to_asd_cum_regret, + ) = _build_case() + + asd_ids = new_next_token_ids.clone() + accept_len, accepted_index = mtp_asd_verify( + req_to_next_token_ids=req_to_next_token_ids, + b_req_mtp_start_loc=b_req_mtp_start_loc, + new_next_token_ids=asd_ids, + b_req_idx=b_req_idx, + b_mtp_index=b_mtp_index, + logits=logits, + req_to_asd_cum_regret=req_to_asd_cum_regret, + asd_budget=100.0, + asd_local_ratio=100.0, + asd_max_mismatch=2, + ) + assert accept_len.tolist() == [3] + assert accepted_index.tolist() == [1, 1, 1] + assert asd_ids.tolist() == [3, 5, 6] # both draft tokens committed, bonus row keeps argmax 6 + assert req_to_asd_cum_regret.item() == pytest.approx(4.6) + + +def test_asd_budget_persists_across_steps(): + """The ledger carries over: regret spent in step 1 shrinks the allowance of step 2.""" + ( + req_to_next_token_ids, + b_req_mtp_start_loc, + new_next_token_ids, + b_req_idx, + b_mtp_index, + logits, + req_to_asd_cum_regret, + ) = _build_case() + + mtp_asd_verify( + req_to_next_token_ids=req_to_next_token_ids, + b_req_mtp_start_loc=b_req_mtp_start_loc, + new_next_token_ids=new_next_token_ids.clone(), + b_req_idx=b_req_idx, + b_mtp_index=b_mtp_index, + logits=logits, + req_to_asd_cum_regret=req_to_asd_cum_regret, + asd_budget=3.0, + asd_local_ratio=100.0, + asd_max_mismatch=2, + ) + assert req_to_asd_cum_regret.item() == pytest.approx(2.0) + + # Step 2: budget_used=2.0, row0 needs 2.0 more -> 4.0 > B=3.0 -> rejected at position 0. + asd_ids = new_next_token_ids.clone() + accept_len, accepted_index = mtp_asd_verify( + req_to_next_token_ids=req_to_next_token_ids, + b_req_mtp_start_loc=b_req_mtp_start_loc, + new_next_token_ids=asd_ids, + b_req_idx=b_req_idx, + b_mtp_index=b_mtp_index, + logits=logits, + req_to_asd_cum_regret=req_to_asd_cum_regret, + asd_budget=3.0, + asd_local_ratio=100.0, + asd_max_mismatch=2, + ) + assert accept_len.tolist() == [1] + assert accepted_index.tolist() == [1, 0, 0] + assert asd_ids.tolist() == [9, 9, 6] # nothing relaxed + assert req_to_asd_cum_regret.item() == pytest.approx(2.0) # unchanged diff --git a/unit_tests/server/test_mtp_start_args.py b/unit_tests/server/test_mtp_start_args.py index 87e12479d4..5c43209062 100644 --- a/unit_tests/server/test_mtp_start_args.py +++ b/unit_tests/server/test_mtp_start_args.py @@ -30,3 +30,63 @@ def test_mtp_prefill_still_requires_positive_step(monkeypatch): with pytest.raises(AssertionError): _launch_subprocesses(args) + + +def _patch_common_setup(monkeypatch): + monkeypatch.setattr("lightllm.server.api_start._set_envs_and_config", lambda args: None) + monkeypatch.setattr("lightllm.server.api_start.auto_set_max_req_total_len", lambda args: None) + monkeypatch.setattr("lightllm.server.api_start.auto_set_fused_shared_experts", lambda args: None) + monkeypatch.setattr("lightllm.server.api_start.set_unique_server_name", lambda args: None) + + +def test_mtp_asd_requires_mtp_mode(monkeypatch): + _patch_common_setup(monkeypatch) + args = StartArgs( + run_mode="prefill", + mtp_asd_regret_budget=2.0, + disable_cudagraph=True, + disable_vision=True, + disable_audio=True, + disable_shm_warning=True, + ) + + with pytest.raises(AssertionError, match="requires an enabled mtp_mode"): + _launch_subprocesses(args) + + +def test_mtp_asd_rejects_negative_budget(monkeypatch): + _patch_common_setup(monkeypatch) + args = StartArgs( + run_mode="prefill", + mtp_mode="dspark", + mtp_draft_model_dir=["draft"], + mtp_step=1, + mtp_asd_regret_budget=-1.0, + disable_cudagraph=True, + disable_vision=True, + disable_audio=True, + disable_shm_warning=True, + ) + + with pytest.raises(AssertionError, match="must be >= 0"): + _launch_subprocesses(args) + + +def test_mtp_asd_rejects_negative_ratio_and_mismatch(monkeypatch): + _patch_common_setup(monkeypatch) + base_kwargs = dict( + run_mode="prefill", + mtp_mode="dspark", + mtp_draft_model_dir=["draft"], + mtp_step=1, + mtp_asd_regret_budget=2.0, + disable_cudagraph=True, + disable_vision=True, + disable_audio=True, + disable_shm_warning=True, + ) + + with pytest.raises(AssertionError, match="must be >= 0"): + _launch_subprocesses(StartArgs(mtp_asd_local_regret_ratio=-0.1, **base_kwargs)) + with pytest.raises(AssertionError, match="must be >= 0"): + _launch_subprocesses(StartArgs(mtp_asd_block_max_mismatch=-1, **base_kwargs))