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

[Kissmetothemoon]

Implements the proposal in #1552.

Motivation

Strict greedy MTP verification (_fwd_kernel_mtp_verify) discards the whole draft suffix at
the first mismatched token. ASD (Approximate Speculative Decoding, arXiv:2608.03447,
Apache-2.0 reference implementation) relaxes this:
a draft token is accepted while its local regret against the target logits,
r_i = max_v z_i(v) - z_i(x_i), stays within an explicit per-request budget. This raises
accept length and throughput for greedy users in exchange for a bounded, user-configurable
deviation. Budget B=0 (or m=0) recovers strict verification exactly — enforced as a
unit-test gold standard — and the default (flag unset) keeps today's behavior bit-for-bit.

ASD is the acceptance-side counterpart of LightSpec: LightSpec decides how much to verify;
ASD decides how strictly to accept. The planner only consumes accept_len statistics, so
the two compose without interference.

Measured on the paper's research evaluator (pure PyTorch, Qwen3-14B + DSpark block-7 draft,
8xL20, greedy, 1319 GSM8K requests): strict 66.21 tok/s / 79.682% vs ASD (B=2.0625)
75.04 tok/s (+13.34%) / 79.454% (-0.227pp); the B=0 arm is token-identical to strict
on every request. Engine-level LightLLM numbers will be filled in below before marking this
PR ready.

Modifications

• New server args (StartArgs): --mtp_asd_regret_budget (B, default None = strict),
  --mtp_asd_local_regret_ratio (g, suffix-value-weighted local cap), and
  --mtp_asd_block_max_mismatch (m, per-verify-step relaxed-token cap); validated in
  _launch_subprocesses (requires an enabled mtp_mode, non-negative values).
• New _fwd_kernel_mtp_asd_verify + mtp_asd_verify() in
  lightllm/common/basemodel/triton_kernel/mtp_utils.py, mirroring mtp_verify's launch
  contract. Per-row regrets are precomputed on device (max + gather over the same logits
  that produced next_token_ids; no host synchronization); the kernel applies the three
  ASD gates with the same accept_len semantics as strict verification (first infeasible
  position stops, the target row is still committed). Rows accepted under a nonzero regret
  commit the draft token into next_token_ids in place, so the token counter,
  scatter_mtp_next_tokens, and response construction all work unchanged.
• Budget ledger: ReqSamplingParamsManager.req_to_asd_cum_regret (allocated only when
  ASD is enabled, zeroed per request at the prefill seam init_req_sampling_params,
  updated in-kernel after each verify step).
• Dispatch: verify_mtp_tokens routes to ASD only when the flag is set AND the whole
  batch is greedy (top_k == 1, same rule as sample); sampled/mixed batches take the
  strict path unchanged. The three decode call sites (chunked_prefill, dp_backend,
  dp_backend overlap) pass logits and run_reqs through; the overlap path refreshes
  the pinned next_token_ids copy after verify when ASD is enabled (the early copy races
  ahead of the in-place commit).
• Tests: unit_tests/server/router/model_infer/mtp_speculative/test_mtp_asd_verify.py
  (6 CUDA tests: B=0/m=0 strict equivalence vs mtp_verify, hand-computed budget
  deduction, local-ratio gate, mismatch cap, all-accepted bonus semantics, cross-step
  budget persistence) + ASD arg-validation cases in
  unit_tests/server/test_mtp_start_args.py.

Accuracy Tests

• pytest unit_tests/server/router/model_infer/mtp_speculative/test_mtp_asd_verify.py — 6 CUDA tests
  (B=0/m=0 strict equivalence vs mtp_verify, hand-computed budget deduction, local-ratio
  gate, mismatch cap, all-accepted bonus semantics, cross-step budget persistence).
  Written alongside the kernel; I do not have a CUDA box at hand right now, so these are
  intended to run on GPU CI — happy to post results once CI (or a maintainer's GPU run)
  executes them.
• pytest unit_tests/server/test_mtp_start_args.py — ASD arg-validation cases pass locally
  (CPU, monkeypatched launch path).
• 3-arm GSM8K (strict / ASD B=0 / ASD q25) with a B=0 token-identity spot check:
  to be run on GPU hardware during review; the algorithm-level numbers above come from the
  paper's research evaluator.

Speed Tests

• 3-arm throughput (strict / B=0 / q25) on GPU hardware: to be filled during review.
  The only added hot-path work when ASD is enabled is one max+gather reduction over
  the verify logits plus a per-request verify kernel with the same launch shape as
  mtp_verify; strict path is untouched (zero overhead when disabled).

Usage

python -m lightllm.server.api_server --model_dir <model> \
  --mtp_mode dspark --mtp_step 3 \
  --mtp_asd_regret_budget 2.0625 \
  --mtp_asd_local_regret_ratio 0.25 \
  --mtp_asd_block_max_mismatch 2

Unset --mtp_asd_regret_budget = strict verification, zero behavioral change.
--mtp_asd_regret_budget 0 = ASD code path with exact strict semantics (sanity arm).
ASD is approximate decoding: outputs may deviate from strict greedy within the configured
budget; v1 applies to all-greedy batches only (sampled batches fall back to strict).

[shihaobai]

Thank you for your contribution. We’ll also test the performance gains on the H100.

[Kissmetothemoon]

Implements the proposal in #1552. Closes #1552.
Thanks for the quick review and for catching both issues on H200. Both are fixed,
and the branch has just been rebased onto the latest main — the PR now merges
cleanly. Details and a fresh end-to-end validation report below.

Bug 1 — CLI flags not registered

Fixed in lightllm/server/api_cli.py: the three ASD flags
(--mtp_asd_regret_budget, --mtp_asd_local_regret_ratio,
--mtp_asd_block_max_mismatch) are now registered in add_cli_args, matching the
existing StartArgs fields. Verified end-to-end: api_server --help lists all
three, and a launched server receives them correctly
(regret_budget=4.0 / local_regret_ratio=0.25 / block_max_mismatch=2).
A regression test (unit_tests/server/test_mtp_start_args.py) covers the
registration path so this cannot silently break again.

Bug 2 — gather OOB on bonus-row padding (-1)

Fixed in lightllm/common/basemodel/triton_kernel/mtp_asd.py: draft token ids are
clamped (clamp(min=0)) before logits.gather(), so padded -1 rows read row 0
and their outputs are masked downstream as before. The kernel test fixture contains
bonus rows with -1 padding (not vacuous coverage), run with
CUDA_LAUNCH_BLOCKING=1.

Validation report (8×L20, CUDA 12.8, torch 2.9.1+cu128, FA3 backend)

Re-validated on the rebased branch (fresh offline venv, clean working copy):

• Unit/kernel regression: 11 passed in 6.41s
  (test_mtp_asd_verify.py + test_mtp_start_args.py)
• Wiring: CLI → StartArgs passthrough OK; CuPy/CUDA-graph probe OK
• B=0 lossless equivalence (real server): Qwen3-8B + DSpark draft
  (block_size=7), greedy — strict MTP vs ASD(B=0) produced identical
  generated_text. (/generate returns text rather than token ids, so this is
  text-level equality.)
• The speculative path is confirmed exercised:
  mtp_accepted_token_num=15, mtp_total_verify_steps=6,
  mtp_total_verify_tokens=48.

Quick throughput signal (GSM8K, 16 prompts, concurrency 16, fixed 256 output
tokens, ASD with B=g=2.0625, m=1):

Arm	TPS	out tok/step	accepted tok/step
strict-pre	955.5	4.357	3.385
ASD	1209.9	4.836	3.862
strict-post	1059.0	4.367	3.397

ASD vs mean(strict): TPS +20.1% (one-sided +14.2%…+26.6%), verify steps −9.8%,
accepted tokens/step +13.9%.