[zheyuf]

Description

Brings the MiniMax-M3.1-Flash DSpark acceptance length (AL) on this branch to the SemiAnalysis golden. Two commits.

1. Non-causal DSpark draft blocks on TRTLLM (default path)

DSpark drafters are trained with a non-causal block and a two-sided sliding window, but the drafter ran every windowed layer causally.

• GQA DSpark layers now resolve non-causal (W-1, W-1).
• trtllm-gen windows only causal attention, so the TRTLLM backend serves these layers as the causal windowed decode, which also returns its log-sum-exp, plus a compiled merge of the block keys after each query.
• The workspace now holds the softmax stats that flashinfer carves off for the log-sum-exp.
• The window validator that refused these layers is removed.

Based on Pranav Shrestha's efa53ad (feat/minimax-m3.1-flash in NVIDIA/TensorRT-LLM; the same change is in NVIDIA#20922), with the compiled merge from NVIDIA#21130.

2. Golden sampling semantics (opt-in, off by default)

With use_rejection_sampling: true, DSpark draws its drafts ancestrally.

• Position k is drawn from softmax((base_k + bias(d_{k-1})) / T), where d_{k-1} is the token drawn before it. Drafts use the temperature only, like vLLM's probabilistic drafting. Greedy requests still draft the argmax chain.
• The distribution each draft was drawn from is stored as its proposal.
• With rejection sampling on, this sampler takes the place of the greedy-chained refine and of the SM107 TP2 fused draft head.

A new prototype field rejection_sample_method: standard | block (default standard) selects the verification rule. block replaces the chain kernel of the one-model rejection path with block verification (Sun et al., 2024). Ragged (confidence-scheduled) steps keep token-level verification.

Nothing changes with the default use_rejection_sampling: false. With standard, verification is unchanged, including the branch's native chain-sampling override. Upstream counterpart: NVIDIA#21085.

speculative_config:
  decoding_type: DSpark
  max_draft_len: 7
  speculative_model: <checkpoint>/dspark
  use_rejection_sampling: true
  rejection_sample_method: block

Results

GB300, TP4 EP4: 80 SPEED-Bench coding prompts, /v1/completions, T=1, top_p 0.95, 4096 tokens, 64 concurrent requests. Each AL cell is one to three passes.

	branch head (1aa2dfc)	this PR
default K7 AL (use_rejection_sampling: false)	3.149	3.593 (+14.1%)
decode step at forced AL, B1 / B8 / B64 (ms)	7.97 / 13.98 / 27.97	7.82 / 14.22 / 27.85

The step times are device-step medians from the iteration log, with acceptance forced to AL 3.6 and use_rejection_sampling: false. They come from one server restart per arm, and restarts vary by about ±2% at B1/B8. Earlier runs of the same merge with 5-10 restarts per arm measured +0.9 / +0.2 / +0.4%.

K	1	2	3	4	5	6	7
SA golden (vLLM, probabilistic + block)	1.82	2.44	2.88	3.19	3.46	3.81	4.04
this PR, golden semantics	1.815	2.449	2.872	3.187	3.506	3.855	4.055

• GB300 target minimax_m31_nvfp4__sm_103__tp4 (modeling_v2: require, TP4 EP1), golden K7: 4.045.

Test Coverage

• test_dspark_noncausal_window.py (GPU, SM100 family):
  • The TRTLLM lowering matches dense fp32 attention for K 3 and 7, at 16/1 and 64/4 query/KV heads, with contexts around the window edge.
  • The workspace fits the log-sum-exp stats at 64 heads and 512 requests.
• test_dspark_sampling.py (CPU):
  • Block verification is lossless (chi-square).
  • Ancestral drafting conditions each position on the drawn token and stores what it drew from.
• On GB300, with the head's wheel plus this PR:
  • These two files and test_dflash2_semantics.py: 43 passed.
  • DFlash / DSpark worker, rejection-sampling, sampling-op and RNG tests: 525 passed. The GB300 TP4 target's spec / sampling tests and the spec_decode catalog: 690 passed. The 46 and 4 failures there are the head's own: the same tests fail on the bare wheel.
  • api_stability: 38 passed. The regenerated golden manifest matches, apart from the existing gvr_deterministic_topk drift.
• Release Checks: cleanup_stale_prs.test.js and every pre-commit hook pass on the changed files. scripts/release_check.py also runs bandit over the whole package. It reports 10 issues, all in files this PR does not touch (the branch head has the same 10). Bandit on the changed files reports none.

PR Checklist

• rejection_sample_method is a new prototype LLM-args field; the golden manifest is updated.
• No new dependencies.

🤖 Generated with Claude Code