Skip to content

[FEAT] Support random sampling in async block speculation - #187

Draft
ycjcx123 wants to merge 10 commits into
InferMatrixLab:mainfrom
ycjcx123:feat/m7-async-spec-sampling
Draft

ycjcx123 wants to merge 10 commits into
InferMatrixLab:mainfrom
ycjcx123:feat/m7-async-spec-sampling

Conversation

@ycjcx123

Copy link
Copy Markdown
Contributor

Draft Status & Rebase Plan: GPU evidence for 740a097 includes the November 15 primary run and the November 7 independent-seed replication. Both used the same executable code and tests. I plan to rebase onto the latest upstream main in the next few days, rerun the affected tests, and mark this PR ready for review once they pass.

Required human contributor answers

1. What outcome are you trying to achieve?

Make asynchronous block speculation support random sampling (temperature > 0, top_k, top_p) whose committed tokens still follow the target distribution.

2. What is the highest-leverage variable for that outcome?

Correctness of the committed output distribution, together with not breaking the asynchronous pipeline or greedy execution.

3. How will this work move that variable rather than create busywork?

Device-side rejection sampling checked against analytically computed expectations (deterministic quadrature and Monte Carlo), CUDA tests, and end-to-end real-model comparisons, so the distribution is measured rather than argued from the construction.

4. Who is the customer or user, and what do they actually need?

Users of asynchronous speculative decoding who need temperature, top-k and top-p sampling to work there, and maintainers who need the shared result layout and RNG ownership to stay correct for greedy and synchronous decoding.

PR necessity and bugfix evidence

A PR is needed: this is a new capability, not a bugfix, so there is no observed failure to link. Async block speculation previously rejected every non-greedy request (StageEngine._add_request raised for temperature != 0); this PR lifts that restriction for async block mode only and keeps the rejection for pipeline and PD decode stage.

Purpose

Enable temperature > 0, top_k, and top_p sampling in asynchronous block speculation (SpeculativeConfig(num_speculative_tokens=3, mode="block") with async_scheduling=True), including batches containing both greedy and random requests. Previously these requests raised ValueError; they now use device-side speculative rejection sampling without host reads during submission.

Three contracts govern the implementation:

  • Committed sequence: random result rows contain accepted drafts followed by the correction or bonus token. The CPU output and next-round GPU input use that same sequence; an accepted random draft need not equal the target argmax. Greedy rows retain their existing values. Effective K=0 samples from the target.
  • Request-owned RNG: the existing RngRegistry supplies one fixed block of 2K+1 uniforms per random request per submitted round. Greedy sampling acquires no generator. Fixed-seed replay requires the same execution mode, round shapes and batch plan; equality with synchronous sampling or across eager/Graph modes is not guaranteed. Reusing the current synchronous sampler would introduce host reads, so this PR uses device-side sampling instead.
  • Numerical errors: an invalid distribution at the stopping position fails the whole round through existing fatal cleanup. Live-row checks precede EOS truncation; stale rows and positions after the stopping point are ignored. Failed device slots do not advance. Per-request error isolation is outside scope.

Local async block execution is validated in eager and CUDA Graph modes on one RTX 5090 (SM120) with the Triton attention backend. Pipeline speculation and PD async decode remain greedy-only; FA4 is unvalidated on SM120, and async block preemption remains rejected. Usage and the algorithm are documented in docs/async-speculative.md.

Part of #43 and #66#issuecomment-6078423491: this implements and qualifies stochastic sampling for the asynchronous block mode listed as open there (docs/runtime-history.md), building on the sampling/RNG refactor from umbrella RFC #32 (#80).

Related work: #157/#164 overlap the async runner's reference lifetime, and #83/#78 overlap result/RNG interfaces. #153's 43 CUDA tests were checked on an integration snapshot, reported below.

Issues encountered & questions for maintainers (shared with #153)

  1. SM120 Paged FlashAttention-4 rejection in existing tests (same issue encountered in #153):
    • Observed failure: In the existing CUDA regression selection (test_async_speculative_graphs.py [9], test_input_preparation.py [6], test_cuda_graph.py [12]), 27 CUDA cases hit the same paged FA4 rejection on SM120 during setup at vllm_rlt/kernels/flash_attention.py:21:
      ValueError: Paged FlashAttention-4 is not supported on SM120 by flash-attn-4==4.0.0b30; use --attention-backend triton until a compatible paged implementation is available.
    • Base verification: Re-running those exact 27 node IDs on unmodified base 444638b fails 27/27 with the same failing node IDs and rejection message. Excluding unavailable FA parameterizations yields 121 passed, 2 skipped, 0 failed.
    • Question for maintainers (applies to both #153 and this PR):
      • Is SM120 + --attention-backend triton validation with flash_attn_4 and multi-GPU PD (nixl) explicitly recorded as unvalidated acceptable to mark #153 and this PR ready?
      • Should we add an SM120 capability skip guard for those 27 cases hitting paged FA4 on SM120 (either in #153 or a separate test-hygiene PR)?
      • @bjf-frz, I’d appreciate your guidance on these two questions. I’m happy to follow your recommendation for both this PR and [REFACTOR]M7 3/3: validate sampling and RNG contracts on CUDA #153.
  2. Design decisions for maintainer confirmation:
    • RNG protocol (2K + 1 uniforms per round + float64 inverse-CDF): Decouples RNG consumption from acceptance outcomes so submission needs no host sync; fixed-seed replay holds within one execution path, while outputs are intentionally not token-identical to synchronous speculation or across eager/CUDA Graph modes.
    • Numerical error semantics (atomic fatal round cleanup): Invalid distribution/residual at a live request's stop point fails the round through existing fatal cleanup prior to scheduler EOS truncation, rather than isolating individual requests.
    • Scope: Local async block speculation supports random sampling; pipeline mode and PD decode stage keep rejecting temperature != 0.

Test Plan

vllm-rlt Version / Commit: 740a097244f11a2d32a584ac383fd257fa576e05 (feat/m7-async-spec-sampling, base 444638beca89e95af6efb34a39fd3edf01c4020b, version 0.2.0).

GPU evidence for 740a097 includes the November 15 primary run (m7-gpu-validation-20261115) and the November 7 independent-seed replication (m7-gpu-validation-20261107, recorded before a docstring/documentation-only update to 740a097 with identical executable code and tests). Both GPU environments used one NVIDIA GeForce RTX 5090 (32,607 MiB, SM120), Python 3.12.3, PyTorch 2.8.0+cu128, Triton 3.4.0, --attention-backend triton, and model Ouro-1.4B (574fa66, BF16 unless marked FP32). The 11/15 primary run used driver 595.71.05; the 11/07 replication used drivers 595.71.05 and 580.76.05. The archives record GPU UUIDs, manifests, and exact commands.

Test Result

Functional and regression checks

Check Result
pre-commit run --all-files (Ruff lint, Ruff format, CPU regression hooks) Passed (exit 0)
Full Linux CPU test suite (-m "not gpu") 1004 passed, 7 skipped, 659 deselected (740a097); 976 passed, 35 skipped (11/07 environment)
New & related M7 CUDA cases (eager + CUDA Graph) 61 passed, 0 skipped (both 11/15 and 11/07)
Existing greedy / Graph / lifecycle / input-prep CUDA selection 121 passed, 3 skipped, 27 failed (both 11/15 and 11/07)
Same selection excluding unavailable SM120 FA4 parameterizations 121 passed, 2 skipped, 0 failed (both 11/15 and 11/07)
Same 27 failing node IDs on unmodified base 444638b 27/27 failed with the same failing node IDs and SM120 FA4 rejection message
#153's 43 CUDA cases on integration snapshot (740a097 production + 5110f42 tests) 43 passed, 0 failed (both 11/15 and 11/07)
Three potentially flaky assertions, 15 repetitions each (11/07 archive) 45/45 passed
Real-model BF16 greedy equality, base 444638b vs candidate (10 questions × 128 tokens) 10/10 token-identical at c=1 and c=4 (both 11/15 and 11/07)

The 61 CUDA cases cover two in-flight rounds, real CUDA Graph capture/replay, ragged and K=0 rounds, EOS/abort/ID reuse, numerical-error cleanup, and request-local RNG isolation. A layout regression verifies that when a non-argmax draft is accepted, the CPU committed sequence, next-round GPU input, and serial KV replay remain identical. CUDA skips require dual GPUs or nixl; CPU skips require optional lm_eval or cached tokenizer/weight fixtures.

Real-model quality and distribution

GSM8K smoke: 10 fixed questions × 3 seeds ({1, 2, 3}), 512-token limit, temperature=0.7, top_p=0.95. Both runs produced flexible-extract scores of nospec 22/30, sync 24/30, and async 20/30 (strict-match 12/30, 13/30, 12/30), with 1 length-truncated sample on async (id=150, seed=2). Async vs sync had 3 wins / 7 losses / 20 ties (exact McNemar p = 0.344); async vs nospec had 4 / 6 / 20 (p = 0.754). No paired comparison reached significance; this small smoke only checks against large regressions.

Distribution diagnostic (FP32): One fixed prompt, temperature=1.0, top_k=20, K=3, draft depths 1 and 2. Reference distributions use serial replay with the existing Sampler.probabilities (independent of the new batched_probabilities). nospec uses disjoint seeds. At the pre-frozen α = 0.01, t1 is a prerequisite control and both t2 primary criteria must pass:

Check (Formal 5,000 seeds per group, α = 0.01) 11/15 Primary (offset=2e6) 11/07 Independent Replication (offset=0)
Prerequisite t1 control $p$ (nospec / speculative 4-group) 0.198 / 0.687 0.728 / 0.180
Primary 1: async t2 goodness-of-fit $p$ vs serial reference 0.207 (PASS) 0.500 (PASS)
Primary 2: async vs nospec t2 two-sample homogeneity $p$ 0.747 (PASS) 0.371 (PASS)
async t2 strata / sample coverage (>= 200 samples/stratum) 6 strata / 79.0% (3,952 / 5,000) 6 strata / 78.1% (3,905 / 5,000)
async_d1 (draft_loops=1) t2 GOF $p$ / vs nospec homogeneity $p$ 0.507 / 0.510 0.274 / 0.034

Both formal primary criteria passed in the 11/15 primary run and were independently reproduced on a disjoint seed range in the 11/07 run; no distribution deviation was detected under the tested FP32 configuration and sample sizes.

Expanded 20,000-seed runs, per-group p-values, and distribution notes
Metric 11/15 Primary 5k (offset=2e6) 11/15 Expanded 20k (offset=4e6, disjoint) 11/07 Replication 5k (offset=0) 11/07 Expanded 20k (offset=0, includes 5k)
t1 control $p$ (nospec / speculative 4-group) 0.198 / 0.687 0.125 / 0.346 0.728 / 0.180 0.313 / 0.0355
t1 out-of-support / chi2_excluded 0 / False 0 / False 0 / False 0 / False
Primary 1: async t2 GOF $p$ vs reference 0.207 0.752 0.500 0.324
Primary 2: async vs nospec t2 homogeneity $p$ 0.747 0.980 0.371 0.258
t2 aggregate GOF $p$ (nospec / sync / async / sync_d1 / async_d1) 0.597 / 0.568 / 0.207 / 0.819 / 0.507 0.398 / 0.391 / 0.752 / 0.841 / 0.145 0.524 / 0.326 / 0.500 / 0.379 / 0.274 0.420 / 0.497 / 0.324 / 0.216 / 0.903
async vs sync t2 homogeneity $p$ 0.346 0.589 0.385 0.135
async_d1 vs nospec t2 homogeneity $p$ 0.510 0.323 0.034 0.878
async t2 sample coverage (>= 200 samples/stratum) 6 strata / 3,952 (79.0%) 14 strata / 19,192 (96.0%) 6 strata / 3,905 (78.1%) 14 strata / 19,167 (95.8%)
Exploratory t3 aggregate GOF $p$ (nospec / sync / async / sync_d1 / async_d1) Not run 0.968 / 0.238 / 0.878 / 0.801 / 0.076 Not run 0.625 / 0.797 / 0.391 / 0.279 / 0.458
First-position rejection probability $\mathrm{TV}(p, q)$ (d2 / d1) 0.0851 / 0.3534 0.0856 / 0.3541 0.0863 / 0.3573 0.0856 / 0.3537
Mean accepted / drafted ratio (sync / async; sync_d1 / async_d1) 0.805 / 0.806; 0.497 / 0.524 0.805 / 0.806; 0.503 / 0.519 0.809 / 0.808; 0.502 / 0.512 0.807 / 0.808; 0.504 / 0.521
  • Disjoint seed ranges: The 11/15 5k (offset=2e6, nospec 3e6) and 11/15 20k (offset=4e6, nospec 5e6) seed ranges are disjoint from each other and from the 11/07 runs (offset=0, nospec 1e6, where the 20k run expands the first 5k).
  • Exploratory t3 strata: In the 11/15 20k run, all five group-level t3 aggregates pass (p = 0.076–0.968). The isolated async_d1 t3 stratum (t1, t2) = (26764, 13971) that had p = 0.0068 in the 11/07 20k run did not replicate in the 11/15 independent 20k run (p = 0.1734).
  • Scope note: Passing statistical tests does not prove exact distribution equality or qualify BF16 serving distributions. chi2_excluded_samples=0 counts only out-of-support exclusions; strata below 200 samples are reflected in the coverage row above.

Performance, memory and limits

  • Greedy non-regression (P1, base 444638b vs candidate 740a097, Ouro-1.4B BF16, 64 tokens, ABBA n=4 measured runs across 2 processes per side): Committed throughput mean ± SD [min, max] is 25.63 ± 1.69 [24.16, 27.10] → 26.95 ± 0.25 [26.68, 27.18] tok/s (+5.14%) at c=1 and 100.91 ± 7.26 [94.54, 107.26] → 105.96 ± 0.86 [105.11, 106.82] tok/s (+5.01%) at c=4; TPOT is 38.78 ± 2.56 [36.55, 41.00] → 36.76 ± 0.34 [36.45, 37.12] ms (-5.21%) at c=1 and 39.40 ± 2.85 [36.91, 41.88] → 37.38 ± 0.31 [37.07, 37.69] ms (-5.13%) at c=4. Measure ranges overlap on both sides and satisfy the pre-frozen no_regression criterion. In the 5-step ops-only window (c=1), base and candidate execute the exact same 69,665 kernels and 35 copies (69,700 total GPU events, 184.78 ms kernel CUDA time / 184.80 ms total GPU activity time on both), with identical draft acceptance (184/196 = 93.9% in P1, 46/49 in the profiled request) and identical peak allocation (3,140.27 MiB at c=1, 3,552.86 MiB at c=4).
  • Random-path cost & memory: No speed claim is made for the random path. In the 5-step full profile of random_async (c=1, 188.87 ms kernel CUDA time / 189.11 ms total GPU activity time), speculative_sampling.py:8 batched_probabilities accounts for 2.965 ms self CUDA time (~1.6%); by phase, draft sampling activity takes 2.007 ms (525 GPU events, 1.1%) and accept activity takes 0.311 ms (220 GPU events, 0.2%). Peak allocated memory delta (random_async vs greedy) is +9.74 MiB at c=1 (3,150.00 vs 3,140.27 MiB; 28,398 vs 28,388 MiB under auto-KV) and +39.52 MiB at c=4 (3,592.38 vs 3,552.86 MiB; 28,219 vs 28,179 MiB under auto-KV), within the existing 36.3 / 145.0 MiB speculative scratch reserve (11/07 transient allocation delta was 11.9 / 48.2 MiB). Unprofiled P2–P4 measurements and profiler tables are recorded in the collapsible details below.
  • Not covered: HTTP serving, CUDA Graph performance benchmarks, long contexts, concurrency above 4, flash_attn_4 / non-SM120 GPUs, PD/NIXL/multi-GPU, and BF16 distribution checks.

Single-request (c=1) greedy A/B timeline (base 444638b vs candidate 740a097, 64-token request, total trace span ~3.93 s; zoom panel aligned at step2, the first speculative decode round):

timeline-greedy-c1
Unprofiled A/B measurements (P1–P4), c=1 timeline numbers, and profiler attribution (ops-only & full)

1. Unprofiled BF16 measurements (m7-gpu-validation-20261115/logs/G2-summary.md, 64 tokens, n=4 measured reps across 2 processes/side)

(Recorded as observations without speed claims across sampling modes; short-request random throughput reflects stochastic draft acceptance counts and round counts.)

Comparison c Side A tok/s (mean ± SD [min, max]) Side B tok/s (mean ± SD [min, max]) $\Delta$ tok/s Side A TPOT ms Side B TPOT ms Peak MiB (A / B) Drafted Accepted (A / B)
P1: Base greedy (444638b) vs Cand greedy (740a097) 1 25.63 ± 1.69 [24.16, 27.10] 26.95 ± 0.25 [26.68, 27.18] +5.14% 38.78 ± 2.56 36.76 ± 0.34 3140.27 / 3140.27 184/196 (93.9%) / 184/196 (93.9%)
P1: Base greedy (444638b) vs Cand greedy (740a097) 4 100.91 ± 7.26 [94.54, 107.26] 105.96 ± 0.86 [105.11, 106.82] +5.01% 39.40 ± 2.85 37.38 ± 0.31 3552.86 / 3552.86 736/784 (93.9%) / 736/784 (93.9%)
P2: Cand random async vs random sync 1 21.10 ± 0.44 [20.68, 21.48] 23.89 ± 2.95 [21.05, 26.57] +13.25% 47.21 ± 0.98 42.00 ± 5.21 3150.00 / 3131.35 168/240 (70.0%, 84 rds) / 184/196 (93.9%, 68 rds)
P2: Cand random async vs random sync 4 78.38 ± 1.19 [76.70, 79.48] 84.27 ± 7.17 [77.71, 90.54] +7.52% 50.13 ± 0.78 46.10 ± 3.94 3592.38 / 3529.33 692/908 (76.2%, 316 rds) / 720/832 (86.5%, 288 rds)
P3: Cand random async vs nospec (ordinary async) 1 21.10 ± 0.44 [20.68, 21.48] 16.82 ± 0.05 [16.77, 16.89] -20.29% 47.21 ± 0.98 59.47 ± 0.18 3150.00 / 3140.13 168/240 (70.0%) / n/a
P3: Cand random async vs nospec (ordinary async) 4 78.38 ± 1.19 [76.70, 79.48] 64.55 ± 0.67 [63.57, 65.10] -17.64% 50.13 ± 0.78 61.96 ± 0.66 3592.38 / 3552.17 692/908 (76.2%) / n/a
P4: Cand random async vs Cand greedy 1 21.10 ± 0.44 [20.68, 21.48] 27.03 ± 0.07 [26.96, 27.11] +28.14% 47.21 ± 0.98 36.65 ± 0.10 3150.00 / 3140.27 168/240 (70.0%) / 184/196 (93.9%)
P4: Cand random async vs Cand greedy 4 78.38 ± 1.19 [76.70, 79.48] 106.68 ± 0.12 [106.55, 106.83] +36.12% 50.13 ± 0.78 37.12 ± 0.04 3592.38 / 3552.86 692/908 (76.2%) / 736/784 (93.9%)

2. Single-request (c=1) full-request timeline summary (timeline/render.log, 64 tokens)

Phase Side Wall Span (ms) Kernels GPU Time (ms) GPU Busy
Prefill Base 444638b 97.64 5,535 18.06 18.5%
Prefill Candidate 740a097 95.32 (-2.31) 5,535 (+0) 18.03 (-0.03) 18.9%
Decode Base 444638b 3,830.58 231,301 614.49 16.0%
Decode Candidate 740a097 3,835.79 (+5.21) 231,301 (+0) 614.49 (+0.00) 16.0%

3. Scheduled 5-step (steps 3..7) ops-only phase breakdown (logs/analysis-*-ops-only-c1.md)

(Phase counts from analyze_profile.py include CUDA kernels plus device/pinned memcpy and memset events; for greedy c=1, both base and candidate execute 69,665 kernels + 35 copies = 69,700 GPU events, with 184.78 ms kernel time / 184.80 ms total GPU activity.)

Profile (c=1, 5 active steps) Total GPU Events (Kernels + Copy) Total GPU Activity (ms) draft Events / ms verify Events / ms accept Events / ms collect Events / ms other (recurrent forward) Events / ms
greedy / base (444638b) 69,700 (69,665 + 35) 184.803 (184.78 kernel) 0 / 0.000 5,285 / 33.738 75 / 0.091 15 / 0.007 64,325 / 150.967
greedy / cand (740a097) 69,700 (69,665 + 35) 184.801 (184.78 kernel) 0 / 0.000 5,285 / 33.782 75 / 0.091 15 / 0.007 64,325 / 150.921
random_async / cand (740a097) 71,970 (71,795 + 175) 188.668 525 / 1.997 (1.1%) 5,460 / 33.865 (17.9%) 220 / 0.310 (0.2%) 15 / 0.006 65,750 / 152.490 (80.8%)
random_sync / cand (740a097) 71,204 (70,937 + 267) 187.913 826 / 3.099 (1.6%) 5,549 / 33.781 (18.0%) 66 / 0.166 (0.1%) 72 / 0.024 64,691 / 150.843 (80.3%)

4. Scheduled 5-step (steps 3..7) full mode source-level stack attribution (logs/analysis-greedy-cand-full-c1.md, logs/analysis-random_async-cand-full-c1.md)

vllm_rlt Location (file:line func) greedy / full / cand c=1 Self CUDA ms (Events) random_async / full / cand c=1 Self CUDA ms (Events)
async_speculative.py:772 submit 180.076 (291,467) 184.269 (300,069)
executor.py:165 forward 179.373 (288,607) 179.318 (288,609)
async_speculative.py:449 _execute_recurrent_step 177.351 (289,972) 177.274 (289,974)
ouro.py:283 recurrent_prepared 177.006 (288,362) 176.930 (288,364)
ouro.py:210 forward 176.200 (285,612) 176.123 (285,614)
stage_engine.py:740 _step_async_speculative 143.966 (233,209) 147.292 (240,091)
async_speculative.py:598 _draft 136.191 (231,585) 138.598 (236,232)
async_speculative.py:668 _verify 43.786 (59,152) 44.839 (60,217)
normalization.py:20 forward 44.105 (121,257) 44.176 (121,244)
rotary_embedding.py:14 apply_rotary_pos_emb 16.447 (43,202) 16.488 (43,203)
kv_cache_manager.py:663 _write_prepared 6.308 (25,920) 6.346 (25,909)
speculative_sampling.py:8 batched_probabilities — (0) 2.965 (3,810)
async_speculative.py:588 _probabilities — (0) 2.965 (3,810)
activation.py:20 silu_and_mul_reference 2.657 (4,800) 2.656 (4,800)

Evidence

Archive Size SHA-256 Contents
m7-gpu-validation-20261115-summary.zip 21 MB c9cd4dd80bb755c6be844038388a8b1441117ff4c4fb87f13505dac932ae5d9c 11/15 G0–G8 logs, JSON/CSV, scripts, timeline-greedy-c1.svg, 13 profile checks and summaries (check.json, summary.json), 6 stack files (stacks.txt), 7 analyses (analysis.json), and a self-contained SHA256SUMS (295 files; omits only the 32 raw *.gz trace dumps to fit GitHub's 25 MB limit)
m7-gpu-validation-20261107.zip 1.4 MB a88e24e4de5f388fdc45063b6fc86393169e8b166945ed93de8e84526c4734a4 11/07 independent replication logs, distribution/GSM8K/greedy JSONs, scripts, and SHA256SUMS (152 files)

Add batched_probabilities in vllm_rlt/worker/speculative_sampling.py. It
builds every row's temperature/top-k/top-p distribution from per-row
parameters in one pass, matching Sampler.probabilities row by row, with
no host read. Sampler is unchanged.

Signed-off-by: ycjcx123 <ycjcx123@gmail.com>
Compare batched_probabilities with the row-wise Sampler.probabilities,
covering tied top-k logits, mixed truncating and non-truncating rows and
BF16 logits. Check that inputs stay unmodified and host arguments are
validated.

Signed-off-by: ycjcx123 <ycjcx123@gmail.com>
Add categorical, an inverse-CDF draw with a float64 CDF that always
lands on a positive-weight token, and speculative_accept, which runs
rejection sampling over ragged draft rows without host reads.

The first draft position that is not a valid accept is the stop point: a
valid rejection draws the correction from max(p - q, 0), an invalid
distribution there sets the request's error flag, and later positions
are ignored.

Signed-off-by: ycjcx123 <ycjcx123@gmail.com>
Check categorical's support and float64 CDF edge cases. Check
speculative_accept against analytic target distributions with
single-draft quadrature and ragged multi-draft Monte Carlo, plus
hand-checked stop-point cases.

Signed-off-by: ycjcx123 <ycjcx123@gmail.com>
Plan random rows and parameters on the host, sample drafts from q, build
p at verify and decide committed tokens with speculative_accept. A
random round returns the accepted drafts plus the correction or bonus
token, so the CPU commits what the GPU continues from; greedy rows are
unchanged.

Each random request draws 2K + 1 uniforms per round, so no host sync is
needed. Rounds with random rows append per-request error flags; a live
flag fails the round through the existing fatal cleanup before the
scheduler update. The entry point still rejects random requests.

Signed-off-by: ycjcx123 <ycjcx123@gmail.com>
Pin the host plan for mixed, all-greedy and zero-draft batches, the
float64 bit-cast uploads and explicit-range result decoding, including
stale flags in reused buffers. Check round-level results against a
per-request reference and the engine's handling of stale and live flags.

Signed-off-by: ycjcx123 <ycjcx123@gmail.com>
Lift the greedy-only check for async block speculation. Pipeline mode
and the PD decode stage keep rejecting random sampling, and the existing
rejection test now asserts the pipeline case.

Signed-off-by: ycjcx123 <ycjcx123@gmail.com>
Check that accepted non-argmax drafts keep the CPU sequence, next-round
GPU input and per-depth KV consistent with a serial replay. Cover
in-flight, ragged, zero-draft and CUDA Graph rounds, per-request RNG
ownership and replay, and fatal errors with stale-flag handling.

Signed-off-by: ycjcx123 <ycjcx123@gmail.com>
Through the public API only: async block engines accept random sampling
and pipeline engines reject it, random requests complete and close
cleanly, the same seed replays, and greedy output is unchanged next to
random requests. A new file avoids the test files changed by InferMatrixLab#153.

Signed-off-by: ycjcx123 <ycjcx123@gmail.com>
Document the algorithm, result layout, RNG protocol and reproducibility
scope, numerical-error semantics and support matrix. Eager and CUDA
Graph execution are marked validated on one RTX 5090 (SM120) with
Triton.

Signed-off-by: ycjcx123 <ycjcx123@gmail.com>
@ycjcx123 ycjcx123 changed the title [Feat]Sampler for Async spec [FEAT] Support random sampling in async block speculation Nov 15, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant