[zheyuf]

Stacked on #25 and the agg unified PR #26: review only the last 3 commits.

Description

Two per-step costs of DSpark agg serving on this branch, three commits. Commits 1-2 make ancestral drafting cheaper
than the greedy chain it replaced. Commit 3 keeps a request's drafter slot for the request's lifetime, so a context-only
step under TRTLLM_ARIA_NO_MIX no longer makes the next decode step resume every generation request.

1. Temperature-only sampling runs no filter

_temperature_only_params passes None for top-k / top-p / min-p. Under advanced_sampling_mode: no_topk, the
FlashInfer chain then runs softmax and sampling only, not a keep-all top-p renorm. The fused op (mode full) receives
its disable sentinels as before. Block verification's emitted-token sampler also uses these parameters.

2. One fused kernel per draft position

With use_rejection_sampling: true, DSpark drafts each block ancestrally: position k is drawn from
softmax((base_k + bias(d_{k-1})) / T), where d_{k-1} is the token drawn before it, and that distribution is stored
as the proposal. Until now this ran a full sampling pipeline per position: bias add, softmax, top-p renorm, sampling,
RNG bookkeeping and a [B, vocab] float32 proposal scatter, K times per step. The greedy chain it replaced ran one
batched sampling over the B x K rows. The sampled distribution does not change.

• One Triton kernel per position adds the Markov bias, scales by 1/T, and takes the argmax of the scaled logits plus
  Gumbel noise (Gumbel-max sampling). The bias GEMM stays on cuBLAS.
• The noise is counter-based: a Philox stream keyed by the request's (seed, offset) draft slot of position k. It is
  graph-safe and the same on every TP rank.
• The kernel reads the drafter's logits as they are (float32 or bfloat16) and writes the biased logits in float32,
  plus per-chunk softmax statistics. A small kernel picks the token and gathers the next Markov embedding.
• After the block, one kernel writes q_k = softmax(logits_k / T) for all positions into the slot-indexed
  draft_probs, from those float32 logits. It recomputes logits_k / T bit for bit as the draw used it, so the
  stored proposal is the distribution each token was drawn from.
• CUDA-graph padding rows draw from the bias alone.

3. Drafter slots for the request lifetime

• DFlashSpecMetadata.prepare freed the drafter slot of every request missing from the batch.
• Under TRTLLM_ARIA_NO_MIX=1 every context step is context-only, so each one evicted every generation request. The
  next decode step then paid a slot clear and a parked-context resume per request.
• Now a slot stays with its request:
  • prepare keeps the slot map in batch order;
  • _assign_slot reclaims the least recently batched holder only when no slot is free, and parks it as the
    eviction did;
  • a slot that leaves its request (reclaim, or a restarted prefill) is reset at once.
• A context-only step therefore leaves every generation slot untouched, and the next decode step does no slot writes.
• A request that does lose its slot still resumes its parked context.

Results

GB300, agg TP2 with the config of the c44 regression report (tot-regression-c44-tp2-s6), except
indexer_kv_dtype: fp8 (the image has no MiniMax MSA package) and a 100 GB/rank host cache (two servers per node).

• G = this branch's head minus its four DSpark commits; head = 519d279.
• Every server starts from the same frozen autotuner cache and pins TRTLLM_MXFP8_GEMM_BACKEND=trtllm. Without the
  pin, the decode-graph MXFP8 backend is re-profiled at each start, which adds about ±0.7 ms of noise per step.

Commits 1-2: fixed-shape decode

Forced AL 3.04; 2048-token ignore_eos requests, T=1, top_p 0.95; GPU pairs rotated. S1 = commit 1 only.

Decode device step (ms; iteration-log median, mean over servers):

	c1	c8	c44	c88
G	10.12	22.65	42.13	61.91
head	10.16	23.78	42.25	60.86
S1	10.11	23.58	41.87	60.39
commits 1-2	9.73	22.93	40.93	58.48
commits 1-2 / head	0.957	0.964	0.969	0.961
commits 1-2 / G	0.961	1.012	0.971	0.945

• Two servers per arm, four for commits 1-2; servers of the same arm agree within about 1%.
• At c8, commit 2's drafting is 0.6 ms cheaper than G's (table below), yet its step is 0.3 ms above G's. The rest of
  the c8 gap, about 0.9 ms, comes from the branch's other DSpark commits, not from drafting.
• Without the MXFP8 pin, every arm is 0.8-6.4 ms faster per step, and commits 1-2 / head is
  0.944 / 0.951 / 0.959 / 0.953.
• These rows predate commit 2's float32 logits buffer. With it, the step time is within 0.3% at c1 / c44
  (9.770 / 40.662 vs 9.794 / 40.772 ms, two servers each).

Drafting alone (B x 7 x 200064 logits, rank-256 Markov head, CUDA graph, us per step):

B	1	8	44	88
G: greedy chain + one batched sampling	483	821	2567	4755
head: sampling per position	736	1029	1879	3292
S1	637	912	1572	2663
commits 1-2	148	228	516	841

SA AL gate, K7: 80 SPEED-Bench coding prompts, T=1, top_p 0.95, 4096 tokens, c64. AL comes from the server
counters. Four passes per cell, two per GPU pair:

	head	commits 1-2
rejection_sample_method: standard	3.954	3.945
rejection_sample_method: block	4.063	4.057

The acceptance rate at every draft position agrees within 0.003. With commits 1-2, block verification costs
+0.28 / +0.32 / +0.88 / +1.50 ms per step (+1.4-2.9%) at c1 / c8 / c44 / c88.

Commit 3: agentic c44

Measured with commit 3 alone on the head. Replay of 531 agentic requests (52 conversations) at c44, forced AL 3.04,
generation steps at batch >= 40. Two sets of runs:

• tuned: the frozen autotuner cache only;
• pinned: the frozen cache plus the MXFP8 pin.

Median generation step (ms), "first step after a context step / other steps (difference)":

	tuned	pinned
G	40.50 / 39.39 (+1.11)	43.76 / 42.54 (+1.22)
head	44.18 / 41.16 (+3.02)	47.28 / 44.34 (+2.94)
commit 3	41.27 / 41.09 (+0.18)	44.34 / 44.17 (+0.17)

• Head and commit 3 have 4 servers per set; G has 2.
• Commit 3 removes the head's +3 ms on each first decode step after a context step, about 300 such steps per replay.
  The other steps are unchanged.
• P90 tok/s/user, head vs commit 3: 100.7 vs 106.3 (tuned), 93.4 vs 95.9 (pinned).
• tok/s/GPU, head vs commit 3: 154.0K vs 155.5K (tuned), 139.4K vs 145.5K (pinned).
• c1 decode tok/s/user p50 is equal: 427.1 vs 429.2 (tuned), 381.8 vs 382.7 (pinned).
• AL is unchanged:
  • agentic c44 natural AL: head 3.695 / 3.701 vs commit 3 3.703 / 3.719;
  • SA-gate K7 (standard), 3 fresh servers each: head 3.937 vs commit 3 3.916, inside the head's own spread
    (3.878-3.971);
  • forced AL measured 4.036-4.040 in every forced server;
  • no "no free drafter slot" line and no errors in any server log.

End to end: InferenceX AgentX on Balaram's recipes

agentx-harness 1.0.12 replays semianalysisai/cc-traces-weka-062126 in full mode (3600 s, 10 warmup requests per lane) at synthetic AL 4.04 (TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS=3.04). Each recipe is used verbatim from aria/m31f/recipes/. "good" is this branch without its four DSpark commits; it fails the SA AL gate (K7 2.96, 0.74 of golden). P90 tok/s/user is 1/P90(ITL), as InferenceX reports it.

point	arm	P90 tok/s/user	tok/s/GPU	vs 519 (P90, tok/s/GPU)	vs good (P90)
c1-tp4-s3	519	694.5	15,907		0.915
c1-tp4-s3	this PR	749.6	16,049	1.079, 1.009	0.988
c16-tp4-s3-r2	519	325.0	58,704		0.898
c16-tp4-s3-r2	this PR	351.3	59,058	1.081, 1.006	0.970
c44-tp2-s6*	519	46.1	93,750		0.741
c44-tp2-s6*	this PR	47.5	94,517	1.030, 1.008	0.763
c52-tp2-s6-r2*	519	27.9	81,545		0.794
c52-tp2-s6-r2*	this PR	29.9	84,398	1.073, 1.035	0.852

• Two runs per arm (519 at c44: 7). TP4 points match the recipe exactly.
• *TP2 points are approximate: two servers per node, 89 GB host cache per rank, no MSA-v1 decode. Their remaining gap to "good" comes from a separate host-cache issue, fixed in its own PR. They used the commit before the bf16 fix; the TP2 path draws from fp32 logits, so it is unaffected.

Test Coverage

• tests/unittest/_torch/speculative/test_dspark_markov_sampling.py (GPU):
  • With float32 and with bfloat16 drafter logits, drafts follow the drawn tokens: chi-square at p = 1e-4 for
    position 0, and for position 1 given the drawn position 0, with one and three vocab chunks.
  • The stored proposal equals the softmax of the float32 biased logits, and the drafter's logits are left as they
    are.
  • A CUDA-graph replay draws what the eager call draws, and a new offset draws anew.
  • Padding rows (NaN logits) draw legal tokens with finite proposals.
• tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sampling.py (CPU): temperature-only sampling runs no
  filter on the FlashInfer chain and gives the fused op its sentinels. The CPU mock of the old per-position loop is
  removed; the GPU test covers it.
• tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py: a context-only step keeps the generation
  slots; the next decode step writes no context length.
• tests/unittest/_torch/speculative/hw_agnostic/test_dflash_slot_resume.py:
  • a request whose slot is reclaimed while it sits out resumes its parked context;
  • a reclaimed slot and a restarted prefill start from base 0.
• Also checked:
  • the same inputs on two GPUs give identical tokens and proposals (TP consistency);
  • a greedy-sentinel row draws the argmax with a one-hot proposal;
  • the joint (d0, d1) chi-square passes;
  • every TP2 server captured the advanced-sampling decode graphs at all 27 batch sizes;
  • with all three commits, the GB300 TP4 target (modeling_v2: require, bfloat16 draft logits) starts, captures its
    graphs, and serves a short SA run at AL 4.10.
• Units on GB300 (the speculative hw_agnostic suite, test_tail_bookkeep.py, test_dflash_forward_context_onehot.py,
  the GB300 TP4 target's spec tests and the DFlash-touching files): no new failures from any commit.
  • All three commits together vs the head: 1733 vs 1727 passed, 90 vs 91 failed, 10 skipped each. There is no new
    failure, and every new or changed test passes.
  • Commit 3 also makes test_gen_init_sizing_is_inert_on_a_context_only_worker pass.
  • The failures that remain are the head's own and identical on both sides.
  • At this revision, commits 1-2's test files and the TP4 target's spec tests pass (80).

PR Checklist

• No API or config change. No new dependencies.

🤖 Generated with Claude Code