Support exact DFlash rollout distribution capture for score centering

[jvmncs]

Score centering needs the behavior-policy probabilities for every committed token. This dependency branch enables complete DFlash filtered-support capture and lets Miles verify that the loaded worker supports stochastic verification before collecting training rollouts.

The base is pinned sglang-miles commit 2b7313dc21a57251f568553247d36d78e585a614. The fork base branch preserves that snapshot, so this review contains only the explicit dependencies below. stitch-sglang-v0.5.20 is a patch reference, not the integration baseline. Companion implementation: modal-projects/miles#11, stacked on upstream score-centering PR sgl-project#3732.

Changes

• Backport merged upstream #34201 for DFlash/DSpark support-probability capture.
• Preserve the three commits from open upstream #41955 at 4df763c9de4152606b03d70350bc4bd8572506f5. These expose a common filtering order and change its default to top_k_first. Experiments explicitly select the same order in both arms; changing the default is not mathematically required by score centering.
• Report dflash_sampling_verify_available from the loaded worker in internal_states. Miles combines this capability with effective acceptance thresholds and environment checks; ordinary model-logprob requests retain their existing semantics.
• Reject an unexpected greedy fallback when a worker advertised stochastic verification but lacks its sampled proposal distribution.
• Preserve all ties at the top-k cutoff before top-p, matching ordinary filtering. Capture the complete positive support up to sampling_mask_max_tokens, reporting OVERFLOW instead of truncating it.
• Preserve each disabled top-p row (p >= 1) exactly during renormalization. A live mixed batch exposed AIR pruning low-probability tokens when float32 input mass rounded above one. An exact fixture reproduced the discrepancy; ordinary sampling retained those tokens. The shared helper and DFlash sparse path now use the per-row identity guard.

Validation

Final tested revision: 79f36a8a1b49f84d8ade21b8f9a4cae0db6d6b4d.

Check	Result
Focused CPU regressions	44 tests plus 13 subtests passed; CPU compilation disabled because the host lacks Python development headers
Final H200 CUDA suite	29 passed; target-only and DFlash2 rejection kernels, zero-probability exclusions, cutoff ties, top-p identity, and bonus-token CDF checks across the 248,320-token Qwen vocabulary
Qwen3.8-27B + inco DFlash2	12 K/filter/temperature settings; 84 requests plus two longer gradient-coverage supplements; 1,720 generated tokens
Committed-token alignment	All 1,634 retained post-prefill rows aligned, covering accepted drafts, rejection bonuses, and all-accepted bonuses
Metadata versus actual verifier	Maximum logprob error 4.40e-6; independent float64 probability-reference maximum L1 error 1.85e-6; unchanged gates 2e-5
Estimator arithmetic	90 none/TIS/MIS saved-logit gradient checks across all 12 settings; maximum absolute gradient error 9.32e-14
API/production collection	12 direct-engine and 12 actual-router cases passed, including native/OpenAI K32/K128, filtered support, EOS and string stops
Dependency-only regular control	Patched versus unmodified pinned base, both using joint: 36 cold/warm comparisons across 18 prefixes, maximum recorded probability L1 6.04e-8, maximum logprob difference 2.38e-7, identical support
Trace-free capture-enabled throughput	Median DFlash2/regular ratios 2.28x (K32), 2.20x (K128), 1.91x (filtered); three repeats of eight concurrent 256-token responses, with ranges and outliers retained

The real-model test uses pinned target/draft revisions, BF16, FA3 full attention, float32 Mamba state, and explicit top_k_first. Live DFlash2 selects the chain/rejection verifier; synthetic CUDA tests cover that and the distinct target-only path. The tested image's unchanged sglang-router 0.3.2 passes the required metadata through; this does not establish compatibility with a separately rebuilt newer gateway.

Same-prefix model-forward comparisons show material differences across decode, prefill, and cache paths: maximum recorded-probability L1 0.1050 for spec-source replay, with differences also present in regular-source controls. Passing the exact-verifier checks does not establish model numerical equivalence, training-quality parity, or universal model/draft compatibility. All four companion Miles training arms (regular/DFlash2 × unfiltered/filtered) passed three-step integration audits: nonzero gradients, finite score-centering metrics, two successful target handoffs per arm, and valid metadata on 650,816 generated token rows. Both DFlash2 arms preserve private draft weights across updates. Matching prompts and initial language weights were checked from artifacts. These short runs do not establish training-quality parity.

The recorded results include the numerical limitations and full training table. Reproduction commands and interpretation criteria are in the companion experiment README. The complete inference bundle preserves all 526 artifact files, including failed runs, raw traces, timings, controls, and source manifests; SHA256 f17a3c4545c5ffd4e9779956f7f2ed8bcff2f2f8750a147f12ecf4d543f5a3c8.

[jvmncs]

Status update: modal-projects/miles#13 replaces modal-projects/miles#11 and supports unfiltered sampling only. It needs no change from this PR.

• The capability report (dflash_sampling_verify_available) is replaced by Add opt-in error for greedy fallback in speculative verification #73. That PR rejects sampled DFLASH requests that would fall back to greedy verification.
• The [RL, Spec] Introduce top-p mask capture for spec and add DFlash/DSpark impl sgl-project/sglang#34201 backport, the cutoff-tie and overflow fixes, and the top-p identity fix are needed only for filtered sampling with DFLASH. They will move to a separate, smaller PR for that follow-up.
• The [Sampling] Default to top_k_first and expose filtering order sgl-project/sglang#41955 filter-order commits are not needed by Miles.

---

Generated by Claude Code