[wdykas]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Adds an opt-in fp64_sampling_noise for the torch sampling backend. When it is on, low-probability tokens are drawn at their softmax rate: float64 race noise moves the representable tail from about 20 nats below a row's top token to about 40 nats. Default behavior and speed are unchanged.

Problem

TorchSampling.sample_from_logits draws argmax(p / E) with E = -log(U) (megatron/core/inference/sampling/torch_sampling.py:136-159 at 13a9d46). U is a float32 uniform, from q.exponential_() for unseeded rows or from tl.uint_to_uniform_float in request_seed_noise.py:35-39 for request-seeded rows. A float32 uniform has only 24 bits near 1, so the smallest nonzero E is about 6e-8, and the largest is about 22 (about 87 for the seeded path's clamped zero). A token D nats below the top token wins only when E_top / E_token > e^D. That ratio cannot exceed about e^20, so:

• a token more than about 20 nats below the top token is never drawn;
• tokens about 16-20 nats below it are drawn at the wrong rate, because the noise is coarsely quantized there.

In wide-vocabulary rows those tail tokens add up. The sampled law is then not the softmax that log_probs_kernel reports, and no log-prob check can detect it. The code comment says the race keeps the low-probability tail; the noise precision is what loses it.

Reproduced on main (13a9d46), GB200. Rows are [0, -D, ..., -D] in float32 with V = 131072. Under the exact law, a draw misses token 0 with probability (V-1) r / (1 + (V-1) r), r = e^-D. Each cell gives the number of draws that missed token 0; p is the two-sided Poisson p-value.

r = e^-D	draws	expected	torch, unseeded	torch, request-seeded	FlashInfer sampling_from_logits	float64 noise (this PR)
1e-7	2.0e5	2597	2937 (p 7e-11)	1527 (p 2e-114)	1322 (p 4e-168)	2657 (p 0.24)
1e-8	1.0e6	1313	37	7	1	1270 (p 0.24)
1e-9	1.0e6	132	0 (p 1e-57)	0	0	123 (p 0.49)

The change

• InferenceConfig.fp64_sampling_noise (default False); CLI --inference-dynamic-batching-fp64-sampling-noise. The controller passes it to TorchSampling(..., fp64_noise=...).
• TorchSampling.sample_from_logits(..., fp64_noise=True) casts the float32 softmax to float64 and runs the same race in float64:
  • unseeded rows use the shared generator's float64 exponential_ (53-bit uniforms);
  • request-seeded rows use fill_request_seed_noise on a float64 buffer. The Triton kernel takes two Philox words (it used one before and dropped the other three) and forms a 52-bit midpoint uniform (k + 1/2) / 2^52, exact in float64. The noise then lies in [1.1e-16, 36.7] and needs no endpoint clamp. The variate is still keyed by (request seed, absolute position, vocabulary index), so seeded draws stay batch-position invariant.
• The sampled law is unchanged: it is still the float32 softmax of the filtered logits. Only the noise is wider. The option does change the random stream, so seeded outputs differ from the float32 default.
• The FlashInfer backend draws float32 uniforms inside its kernels, so fp64_sampling_noise=True with sampling_backend='flashinfer' raises in InferenceConfig.__post_init__. A FlashInfer-side fix would belong in FlashInfer.

Tests

tests/unit_tests/inference/sampling/test_fp64_sampling_noise.py (GPU):

• test_fp64_noise_matches_tail_rate: unseeded and request-seeded, at r = 1e-7, 1e-8 and 1e-9, 262k draws of V = 131072 each. The miss count must lie inside the two-sided 1e-6 Poisson interval of the exact rate.
• test_fp32_noise_misses_tail (non-vacuity): the default float32 noise falls below that interval at r = 1e-9 (expected 34, observed 0).
• test_fp64_noise_chi_square: 1M draws over a 16-token row; chi-square below the 1e-6 quantile (56.5, 15 dof). Unseeded and seeded.
• test_fp64_seeded_draws_are_batch_position_invariant: for no filter, top-p 0.9, and top-k 50 at T 0.7, seeded draws are bitwise equal across repeated calls, single-row calls, and a permuted batch with unseeded batch mates.
• test_sample_kernel_uses_configured_noise: TorchSampling(..., fp64_noise=True).sample_kernel equals the float64 sample_from_logits.
• test_fp64_noise_requires_torch_backend.

tests/unit_tests/determinism/kernels/test_request_seed_sampling.py::test_fp64_request_noise_replays_and_preserves_logical_coordinates: the float64 Philox noise replays bit-exactly under side-stream contention, and single-row launches equal batched ones. It also checks the noise range and that the float32 stream is unchanged. The manifest entry inference_request_seed_sampling notes both widths.

Run on GB200, one GPU: the new file plus test_request_seed_sampling.py (both), test_torch_sampling.py, and test_manifest.py: 323 passed. test_inference_config.py, training/config and test_inference_initialization.py: 388 passed. 3 failures in training/config/test_instantiate_utils.py / test_target_allowlist.py also fail on unmodified main in this environment. Black, isort, ruff and pylint (10.00/10) pass on the changed files.

CUDA graphs: the torch sampler is not captured in a graph on main (TorchSampling.sample_kernel ignores eager / cache_key), so the option adds no graph path.

Cost

Sampler time per call, sample_from_logits on bf16 logits, no filter, median of 30, GB200:

V	rows	unseeded fp32 → fp64	seeded fp32 → fp64
131072	1	154 → 153 µs	157 → 172 µs
131072	8	135 → 148 µs	221 → 225 µs
131072	64	139 → 211 µs	227 → 261 µs
131072	256	372 → 670 µs	488 → 737 µs
131072	1024	1313 → 2486 µs	1617 → 2582 µs
262144	64	246 → 398 µs	326 → 453 µs
262144	256	704 → 1300 µs	880 → 1372 µs

Below about 64 rows the cost is within launch overhead. At larger batches it is about +1.1 µs per row at V = 131072, because the float64 noise and ratio are materialized: twice the bytes of the float32 path. A follow-up could avoid that with a fused kernel that generates the noise in registers and reduces each row in place. It would need several blocks per row, since one program per row is latency-bound at small batch.

Issue tracking

Standalone fix; no linked issue. Related: #7182 (float32 exponential race), #7828 (request-local seeds), #8290 (request-local FlashInfer seeds).

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• If this PR adds or changes a GPU kernel (Triton, jit_fuser/torch.compile, CUDA extension, TE or external-library dispatch, or a scatter/index accumulation), I have added or updated its bit-exact determinism test and registered it in tests/unit_tests/determinism/kernels/manifest.py (guide)
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR (ran black, isort, ruff and pylint on the changed files)

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[wdykas]

/ok to test 2b66a7f