bugfix: separate greedy and probabilistic draft proposals for rejection sampling.

[ustcfy]

Summary

• Model greedy and probabilistic draft proposals explicitly with DraftProposal.
• Preserve dense draft distributions for probabilistic rejection sampling and use a delta proposal for greedy drafts.
• Propagate draft_sampling_mode through runtime options and the spawned-worker protocol.
• Enable probabilistic rejection sampling for reduced-vocab EAGLE-3 (d2t) by scattering the draft distribution into the full target vocabulary in the draft worker.
• Simplify DFlash/DSpark proposal transport and clarify sampling semantics.

Validation

• clang-format --dry-run --Werror on changed C/C++ files
• git diff --check
• Focused changed-line clang-tidy on sampler, DFlash, and DSpark paths
• Rebuild (all touched TUs compile clean, 0 warnings)
• Before/after acceptance-length measurement from server /vars counters (see Benchmark)
• GSM8K accuracy comparison (see Benchmark)

Benchmark

Env: Qwen3.5-2B (target) + Qwen3.5-2B-mtp (draft), MTP num_speculative_tokens=2,
single Ascend 910 die (TP1), chunked prefill, max_memory_utilization=0.85, block_size=128.
Workload: gsm8k 128-prompt corpus, temperature=0.7, top_p=0.9, seed=42,
max_tokens=2048, streaming. Acceptance is read from xLLM /vars counters (ground truth);
accept_length = accepted/steps + 1 (max 3.0 at num_spec=2). before = upstream 3906f5a9
(default temperature-following draft, no draft_sampling_mode flag); after = this branch.

Server launch (per die; ASCEND_RT_VISIBLE_DEVICES selects the die; before omits
--draft_sampling_mode):

ASCEND_RT_VISIBLE_DEVICES=<die> xllm \
  --model <Qwen3.5-2B> --draft_model <Qwen3.5-2B-mtp> \
  --speculative_algorithm=MTP --num_speculative_tokens 2 \
  --draft_sampling_mode=<greedy|probabilistic> \
  --enable_chunked_prefill=true --block_size=128 \
  --max_memory_utilization=0.85 --max_tokens_per_batch=19968 --max_seqs_per_batch=256 \
  --port <P> --master_node_addr 127.0.0.1:<MP> --transfer_listen_port <XP> \
  --nnodes 1 --node_rank 0

Client (evalscope, per cell): evalscope perf --api openai --url .../v1/chat/completions --dataset line_by_line --dataset-path <gsm8k-128.jsonl> --number 128 --parallel <conc> --max-tokens 2048 --temperature 0.7 --top-p 0.9 --seed 42 --stream. Acceptance is read
from the server /vars counters around each run.

Performance (number=128 per cell, one Ascend 910 die, one server at a time,
sequential, warm-up discarded). conc=1 (no queueing) and conc=32 (saturated)
bracket the range:

conc	config	E2E (s)	TTFT (ms)	TPOT (ms)	throughput (tok/s)
1	before	5.58	45.3	11.0	88.7
1	greedy (new default)	3.55	40.3	7.5	129.6
1	probabilistic	5.57	45.2	11.1	88.6
32	before	21.10	184.6	52.2	511.9
32	greedy	15.42	142.7	33.0	780.7
32	probabilistic	23.77	187.2	52.2	513.1

• New greedy default vs pre-PR default: +52% throughput / −37% TPOT / −27% E2E / −23% TTFT
  at conc=32 (780.7 vs 511.9 tok/s; 33.0 vs 52.2 ms; 15.42 vs 21.10 s; 142.7 vs 184.6 ms),
  and +46% throughput / −32% TPOT / −36% E2E / −11% TTFT at conc=1. Greedy skips the
  per-step dense [B, n_spec, vocab=248320] draft_probs materialization and the
  min(1, p/q) + residual recovery; because the first streamed token already includes one
  speculative decode step, the lighter step lowers TTFT too (hence the modest but real
  conc=1 TTFT gain with no queueing involved).
• probabilistic ≈ before: conc=1 TTFT 45.2 vs 45.3 ms / TPOT 11.1 vs 11.0 ms (within
  noise); conc=32 187.2 vs 184.6 ms / 52.2 vs 52.2 ms. The refactor preserves the old
  temperature-following path on the sampler side; opt back in with
  --draft_sampling_mode=probabilistic.
• accept_length (from /vars) ≈ 2.54 (greedy) vs 2.63 (before /
  probabilistic) — greedy accepts marginally fewer draft tokens, but the per-step savings
  outweigh it. Acceptance is draft-model-determined and otherwise unchanged by this PR.

Accuracy note (gsm8k, temperature=0, mean_acc = 0.72 for all configs): at temperature=0
the target uses all-greedy verify, so speculative decoding emits the target model's greedy argmax
regardless of draft_sampling_mode / num_spec / spec on-off / binary version; the target model
is byte-identical before↔after, so accuracy is invariant. Measured identical (0.72) for
after-greedy and after-probabilistic.

Losslessness of the probabilistic fix (Monte-Carlo)

gsm8k accuracy cannot expose the rejection-sampling defect: at temperature=0 the target
verifies greedily (argmax) and bypasses rejection entirely, and at temperature>0 accuracy is
too noisy to resolve a residual-distribution error. The correct probe is a distribution-match
test — fix a target distribution p and a draft distribution q, draw the rejection sampler
N times, and compare the empirical output distribution to p (lossless ⇔ total-variation
distance → 0). Both variants share the same acceptance rule u < min(1, p[x]/q[x]); they
differ only in the recovery distribution sampled on a rejected token.

With p = [.10 .20 .30 .20 .20], q = [.50 .20 .10 .10 .10], num_spec=1, N=2e6
(reproducible, seed-fixed):

variant	recovery distribution	empirical output	TV to p
target p	—	[.100 .200 .300 .200 .200]	—
after (this PR)	(p − q)+, full draft dist	[.100 .1996 .3003 .1997 .2005]	0.0008 (≈ noise floor 0.0006)
before	p with only the drafted token reduced by q[x]	[.100 .2879 .2335 .1887 .1899]	0.0879 (116× noise)

after converges to the target to the sampling-noise floor (lossless); before deviates
systematically (token 1 over-weighted 0.29 vs 0.20, token 2 under-weighted 0.23 vs 0.30) because
its (p−q)+ residual only subtracts draft mass at the single drafted token, not the full draft
distribution. On the shipped code, the RejectionSamplerTest.Random and
RandomFusedRecoveryDistribution unit tests assert the real sampler's recovery distribution.

Reduced-vocab EAGLE-3 (d2t) validation

The new-capability path in this PR — probabilistic rejection sampling for a reduced-vocab
EAGLE-3 draft — is validated separately on a dense single-die setup: Qwen3-8B (target, TP1,
one Ascend 910 die) + Qwen3-8B-EAGLE3 (draft, draft_vocab_size=32000 → full target
vocab_size=151936), speculative_algorithm=Eagle3, num_speculative_tokens=3. The draft
worker logs Loaded d2t tensor from state_dict, hot_token_id size: 32000, confirming the
reduced→full scatter map is active. gsm8k corpus, temperature=0.7, top_p=0.9, seed=42,
max_tokens=1024, 32 requests at concurrency 8; acceptance from /vars
(accept_length = accepted/steps + 1, max 4.0 at num_spec=3).

ASCEND_RT_VISIBLE_DEVICES=<die> xllm \
  --model <Qwen3-8B> --draft_model <Qwen3-8B-EAGLE3> \
  --speculative_algorithm=Eagle3 --num_speculative_tokens 3 \
  --draft_sampling_mode=<greedy|probabilistic> \
  --enable_chunked_prefill=true --block_size=128 \
  --max_memory_utilization=0.85 --max_tokens_per_batch=19968 --max_seqs_per_batch=128 \
  --port <P> --master_node_addr 127.0.0.1:<MP> --transfer_listen_port <XP> \
  --nnodes 1 --node_rank 0

config	accept_length	TPOT (ms)	throughput (tok/s)
before (pre-PR default)	2.09	17.9	430–551
greedy (new default)	2.10	16.6	462–542
probabilistic (d2t scatter)	2.08	18.4	414–488

Both modes produce coherent output with no init/runtime error, and acceptance is invariant
(~2.1 at num_spec=3) across before/after — the draft model is unchanged, so the PR changes
per-step cost and rejection-sampling correctness, not accept length. The new greedy default
trims TPOT vs the pre-PR default (16.6 vs 17.9 ms) by dropping draft-probs materialization
entirely. probabilistic exercises the reduced→full scatter: pre-PR, the draft distribution was
compressed to the selected token's probability before rejection sampling (measured 2.08 /
18.1 ms), so the (p−q)+ recovery on a rejected token saw q mass only at the drafted token;
this PR carries the full draft distribution — for reduced-vocab EAGLE-3, scattered from the
32000 draft vocab into the full 151936 target vocab — so recovery uses the true q.

[ustcfy]

Benchmark: draft_sampling_mode=greedy (new default) vs. the old default

Single-variable A/B isolating exactly this PR.

Setup: Qwen3.5-27B-w8a8 + MTP head, single Ascend 910, --speculative_algorithm=MTP --num_speculative_tokens=1, launched with --enable_shm --enable_schedule_overlap --enable_graph --enable_graph_double_buffer, block_size=128. Load: temperature=0.7, ignore_eos, steady state.

• before = old default enable_opt_validate_probs=false (parent commit 3906f5a91)
• after = new default draft_sampling_mode=greedy (this PR)

Throughput @ KV-saturated operating point (256 req × 256 tok, conc 64)

metric	before	after	Δ
throughput	280.8 tok/s	308.8 tok/s	+10.0%
accept length	1.69	1.64	≈ same (run-to-run 1.63–1.69)

TPOT (inter-token latency, ms/token)

concurrency	before	after	Δ
8 (no queueing)	51	46	−9.8%
64 (saturated)	81	76	−6%

Where the gain comes from

The gain is not from acceptance quality — accept length is unchanged. It is per-step compute, normalized per decode step:

stage	before	after	Δ
draft exec / step	2.62 ms	1.99 ms	−24%
target (verify) exec / step	2.52 ms	2.51 ms	≈ flat

Greedy proposals skip the entire probabilistic-draft machinery: the full-vocab (248320) softmax that produces draft_probs, restoring them to a dense [B, n_spec, V] tensor, and resampling from the draft distribution. msprof op-level over a decode window confirms it — the sampling-path ops all shrink (device ms, directional; the before window contains more steps):

op (device ms, decode window)	before	after
ScatterElementsV2 (dense probs materialize)	1844	957
SoftmaxV2 (draft vocab softmax)	60	27
Cumsum	159	65
Sort	126	49
DSARandomUniform	91	62

Net: at the throughput-optimal operating point, the greedy default gives +10% throughput / −6–10% TPOT, entirely from eliminating draft-side probability computation, with no change in acceptance length. The gain scales with vocab size and narrows at very large (GEMM-bound) batch.