| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """CUDA proof for fp32 exponential-race tail truncation. |
| 4 | |
| 5 | This script is intentionally not a unit test. It is a reproducible, GPU-only |
| 6 | statistical proof for the hidden Gumbel-max idiom: |
| 7 | |
| 8 | q.exponential_() |
| 9 | sample = (probs / q).argmax() |
| 10 | |
| 11 | For q ~ Exp(1), this is equivalent to argmax(log(probs) + Gumbel). On CUDA, |
| 12 | fp32 exponential samples inherit a 24-bit uniform lower-tail cutoff, so very |
| 13 | small q values are impossible. The many-tail experiment below chooses a case |
| 14 | where a correct sampler should select a low-probability tail token dozens of |
| 15 | times, while fp32 q cannot select one. |
| 16 | """ |
| 17 | |
| 18 | from __future__ import annotations |
| 19 | |
| 20 | import argparse |
| 21 | import math |
| 22 | import time |
| 23 | |
| 24 | import torch |
| 25 | |
| 26 | |
| 27 | def _seed(seed: int) -> None: |
| 28 | torch.manual_seed(seed) |
| 29 | |
| 30 | |
| 31 | def measure_exponential_lower_tail( |
| 32 | *, |
| 33 | device: torch.device, |
| 34 | samples: int, |
| 35 | chunk_size: int, |
| 36 | seed: int, |
| 37 | ) -> None: |
| 38 | threshold = 2.0**-24 |
| 39 | print(f"lower-tail threshold: {threshold:.18e}") |
| 40 | for dtype in (torch.float32, torch.float64): |
| 41 | _seed(seed) |
| 42 | count_below = 0 |
| 43 | min_q = float("inf") |
| 44 | max_q = 0.0 |
| 45 | start = time.perf_counter() |
| 46 | remaining = samples |
| 47 | while remaining > 0: |
| 48 | n = min(chunk_size, remaining) |
| 49 | q = torch.empty((n,), dtype=dtype, device=device) |
| 50 | q.exponential_() |
| 51 | count_below += int((q < threshold).sum().item()) |
| 52 | min_q = min(min_q, float(q.min().item())) |
| 53 | max_q = max(max_q, float(q.max().item())) |
| 54 | remaining -= n |
| 55 | torch.accelerator.synchronize() |
| 56 | elapsed = time.perf_counter() - start |
| 57 | print( |
| 58 | f"{dtype}: samples={samples} count(q < 2^-24)={count_below} " |
| 59 | f"min={min_q:.18e} max={max_q:.6f} elapsed={elapsed:.2f}s" |
| 60 | ) |
| 61 | |
| 62 | |
| 63 | def run_many_tail_race( |
| 64 | *, |
| 65 | device: torch.device, |
| 66 | trials: int, |
| 67 | num_tail_tokens: int, |
| 68 | gap: float, |
| 69 | chunk_trials: int, |
| 70 | seed: int, |
| 71 | ) -> None: |
| 72 | p_tail = math.exp(-gap) |
| 73 | expected_tail_hits = ( |
| 74 | trials * (num_tail_tokens * p_tail) / (1.0 + num_tail_tokens * p_tail) |
| 75 | ) |
| 76 | print( |
| 77 | "many-tail race: " |
| 78 | f"trials={trials} num_tail_tokens={num_tail_tokens} gap={gap} " |
| 79 | f"expected_tail_hits={expected_tail_hits:.4f}" |
| 80 | ) |
| 81 | |
| 82 | for dtype in (torch.float32, torch.float64): |
| 83 | _seed(seed) |
| 84 | hits = 0 |
| 85 | p0 = torch.tensor(1.0, dtype=dtype, device=device) |
| 86 | pt = torch.tensor(p_tail, dtype=dtype, device=device) |
| 87 | start = time.perf_counter() |
| 88 | remaining = trials |
| 89 | while remaining > 0: |
| 90 | batch = min(chunk_trials, remaining) |
| 91 | q0 = torch.empty((batch,), dtype=dtype, device=device) |
| 92 | q0.exponential_() |
| 93 | qt = torch.empty((batch, num_tail_tokens), dtype=dtype, device=device) |
| 94 | qt.exponential_() |
| 95 | head_score = p0 / q0 |
| 96 | tail_score = (pt / qt).amax(dim=-1) |
| 97 | hits += int((tail_score > head_score).sum().item()) |
| 98 | remaining -= batch |
| 99 | torch.accelerator.synchronize() |
| 100 | elapsed = time.perf_counter() - start |
| 101 | print(f"{dtype}: tail_hits={hits} elapsed={elapsed:.2f}s") |
| 102 | |
| 103 | |
| 104 | def main() -> None: |
| 105 | parser = argparse.ArgumentParser() |
| 106 | parser.add_argument("--lower-tail-samples", type=int, default=200_000_000) |
| 107 | parser.add_argument("--lower-tail-chunk-size", type=int, default=10_000_000) |
| 108 | parser.add_argument("--race-trials", type=int, default=100_000) |
| 109 | parser.add_argument("--race-tail-tokens", type=int, default=262_144) |
| 110 | parser.add_argument("--race-gap", type=float, default=20.5) |
| 111 | parser.add_argument("--race-chunk-trials", type=int, default=64) |
| 112 | parser.add_argument("--seed", type=int, default=2026) |
| 113 | args = parser.parse_args() |
| 114 | |
| 115 | if not torch.accelerator.is_available(): |
| 116 | raise RuntimeError("CUDA is required for this proof.") |
| 117 | |
| 118 | device = torch.accelerator.current_accelerator() |
| 119 | if device.type != "cuda": |
| 120 | raise RuntimeError("CUDA is required for this proof.") |
| 121 | |
| 122 | print(f"torch={torch.__version__} cuda={torch.version.cuda}") |
| 123 | print(f"device={device}") |
| 124 | measure_exponential_lower_tail( |
| 125 | device=device, |
| 126 | samples=args.lower_tail_samples, |
| 127 | chunk_size=args.lower_tail_chunk_size, |
| 128 | seed=args.seed, |
| 129 | ) |
| 130 | run_many_tail_race( |
| 131 | device=device, |
| 132 | trials=args.race_trials, |
| 133 | num_tail_tokens=args.race_tail_tokens, |
| 134 | gap=args.race_gap, |
| 135 | chunk_trials=args.race_chunk_trials, |
| 136 | seed=args.seed, |
| 137 | ) |
| 138 | |
| 139 | |
| 140 | if __name__ == "__main__": |
| 141 | main() |