| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for the Model Runner V2 Gumbel-max sampling kernel. |
| 4 | |
| 5 | Accuracy: define a target categorical distribution as a non-negative int64 |
| 6 | count tensor summing to N, turn it into logits (= log(count)), sample many |
| 7 | times with `gumbel_sample`, and check the empirical distribution matches. |
| 8 | |
| 9 | The count tensor is deliberately heavy-tailed (one dominant token, the rest |
| 10 | ~18 logits below). That tail is the sensitive part: the fp32 Gumbel noise must |
| 11 | reach ~18 to ever sample it. A flat distribution would keep every token within |
| 12 | a few logits of the top and would not exercise the noise tail at all. |
| 13 | """ |
| 14 | |
| 15 | import math |
| 16 | |
| 17 | import pytest |
| 18 | import torch |
| 19 | |
| 20 | pytest.importorskip("triton") |
| 21 | if not torch.cuda.is_available(): |
| 22 | pytest.skip("CUDA required for Gumbel sampler tests", allow_module_level=True) |
| 23 | |
| 24 | from vllm.triton_utils import tl, triton |
| 25 | from vllm.v1.worker.gpu.sample.gumbel import ( |
| 26 | _uniform64_from_random53, |
| 27 | gumbel_sample, |
| 28 | murmur3_hash32, |
| 29 | ) |
| 30 | |
| 31 | DEVICE = "cuda" |
| 32 | VOCAB_SIZE = 200_000 |
| 33 | NUM_SAMPLES = 500_000 |
| 34 | # Dominant token is exp(HEAD_LOG_GAP)x larger than the unit-count tail, so the |
| 35 | # tail sits ~HEAD_LOG_GAP logits below the top. |
| 36 | HEAD_LOG_GAP = 18.0 |
| 37 | # 10-sigma band: a correct sampler effectively never trips it. |
| 38 | Z_TOLERANCE = 10.0 |
| 39 | |
| 40 | |
| 41 | @triton.jit |
| 42 | def _murmur3_known_answer_kernel( |
| 43 | seeds_ptr, |
| 44 | positions_ptr, |
| 45 | offsets_ptr, |
| 46 | hashes_ptr, |
| 47 | domain_hashes_ptr, |
| 48 | numel: tl.constexpr, |
| 49 | BLOCK_SIZE: tl.constexpr, |
| 50 | ): |
| 51 | idx = tl.arange(0, BLOCK_SIZE) |
| 52 | mask = idx < numel |
| 53 | seed = tl.load(seeds_ptr + idx, mask=mask) |
| 54 | pos = tl.load(positions_ptr + idx, mask=mask) |
| 55 | offset = tl.load(offsets_ptr + idx, mask=mask) |
| 56 | hash32 = murmur3_hash32(seed, pos, offset).to(tl.int64) |
| 57 | domain_hash32 = murmur3_hash32(seed, pos, offset, domain=0x9E3779B9).to(tl.int64) |
| 58 | tl.store(hashes_ptr + idx, hash32, mask=mask) |
| 59 | tl.store(domain_hashes_ptr + idx, domain_hash32, mask=mask) |
| 60 | |
| 61 | |
| 62 | @triton.jit |
| 63 | def _uniform64_endpoint_kernel( |
| 64 | random53_ptr, |
| 65 | uniform_ptr, |
| 66 | numel: tl.constexpr, |
| 67 | BLOCK_SIZE: tl.constexpr, |
| 68 | ): |
| 69 | idx = tl.arange(0, BLOCK_SIZE) |
| 70 | mask = idx < numel |
| 71 | random53 = tl.load(random53_ptr + idx, mask=mask).to(tl.uint64) |
| 72 | tl.store(uniform_ptr + idx, _uniform64_from_random53(random53), mask=mask) |
| 73 | |
| 74 | |
| 75 | def test_murmur3_known_answers(): |
| 76 | seeds = torch.tensor( |
| 77 | [0, 1, -1, 0x123456789ABCDEF, -(1 << 63), (1 << 63) - 1], |
| 78 | dtype=torch.int64, |
| 79 | device=DEVICE, |
| 80 | ) |
| 81 | positions = torch.tensor( |
| 82 | [0, 2, 17, 123456789, 1 << 30, (1 << 30) + 7], |
| 83 | dtype=torch.int64, |
| 84 | device=DEVICE, |
| 85 | ) |
| 86 | offsets = torch.tensor( |
| 87 | [0, 3, 1023, 199999, 4096, 0xFFFFFFFF], |
| 88 | dtype=torch.int64, |
| 89 | device=DEVICE, |
| 90 | ) |
| 91 | hashes = torch.empty_like(seeds) |
| 92 | domain_hashes = torch.empty_like(seeds) |
| 93 | |
| 94 | _murmur3_known_answer_kernel[(1,)]( |
| 95 | seeds, |
| 96 | positions, |
| 97 | offsets, |
| 98 | hashes, |
| 99 | domain_hashes, |
| 100 | numel=seeds.numel(), |
| 101 | BLOCK_SIZE=8, |
| 102 | ) |
| 103 | |
| 104 | assert hashes.tolist() == [ |
| 105 | 0x8134CDF8, |
| 106 | 0x23B53F91, |
| 107 | 0xC8C13C14, |
| 108 | 0x5DDE974F, |
| 109 | 0xF1A9355B, |
| 110 | 0xD952F9AA, |
| 111 | ] |
| 112 | assert domain_hashes.tolist() == [ |
| 113 | 0xD28AD962, |
| 114 | 0xE3AAF818, |
| 115 | 0xD223C51C, |
| 116 | 0x0547D2FE, |
| 117 | 0x3AED160F, |
| 118 | 0x49E8EA7F, |
| 119 | ] |
| 120 | |
| 121 | |
| 122 | def test_uniform64_excludes_endpoints(): |
| 123 | random53 = torch.tensor([0, (1 << 53) - 1], dtype=torch.int64, device=DEVICE) |
| 124 | uniform = torch.empty(2, dtype=torch.float64, device=DEVICE) |
| 125 | |
| 126 | _uniform64_endpoint_kernel[(1,)]( |
| 127 | random53, |
| 128 | uniform, |
| 129 | numel=random53.numel(), |
| 130 | BLOCK_SIZE=2, |
| 131 | ) |
| 132 | |
| 133 | assert torch.isfinite(uniform).all() |
| 134 | assert ((uniform > 0.0) & (uniform < 1.0)).all() |
| 135 | assert uniform.tolist() == [2.0**-54, 1.0 - 2.0**-53] |
| 136 | |
| 137 | |
| 138 | def _make_heavy_tailed_counts(seed: int = 1234) -> torch.Tensor: |
| 139 | """Non-negative int64 counts of shape [VOCAB_SIZE]; target prob = counts/N.""" |
| 140 | gen = torch.Generator(device=DEVICE).manual_seed(seed) |
| 141 | counts = torch.randint( |
| 142 | 1, 4, (VOCAB_SIZE,), generator=gen, dtype=torch.int64, device=DEVICE |
| 143 | ) |
| 144 | counts[0] = round(math.exp(HEAD_LOG_GAP)) # dominant token |
| 145 | return counts |
| 146 | |
| 147 | |
| 148 | def _counts_to_logits(counts: torch.Tensor) -> torch.Tensor: |
| 149 | # softmax(log(count)) == count / sum(count); count 0 -> logit -inf -> prob 0. |
| 150 | return counts.double().log().to(torch.float32) |
| 151 | |
| 152 | |
| 153 | def _sample( |
| 154 | logits_1d: torch.Tensor, |
| 155 | num_samples: int, |
| 156 | *, |
| 157 | use_fp64: bool = False, |
| 158 | temperature: float = 1.0, |
| 159 | is_drafting: bool = False, |
| 160 | ) -> torch.Tensor: |
| 161 | """Sample `num_samples` tokens from one logit vector. |
| 162 | |
| 163 | Fixed seed with a distinct `pos` per sample gives independent draws; the |
| 164 | logits are broadcast with a 0-stride view to avoid materializing |
| 165 | [num_samples, vocab_size]. |
| 166 | """ |
| 167 | vocab_size = logits_1d.shape[0] |
| 168 | logits = logits_1d.unsqueeze(0).expand(num_samples, vocab_size) |
| 169 | idx_mapping = torch.zeros(num_samples, dtype=torch.int32, device=DEVICE) |
| 170 | temp = torch.tensor([temperature], dtype=torch.float32, device=DEVICE) |
| 171 | seed = torch.tensor([0xABCD], dtype=torch.int64, device=DEVICE) |
| 172 | pos = torch.arange(num_samples, dtype=torch.int64, device=DEVICE) |
| 173 | return gumbel_sample( |
| 174 | logits, |
| 175 | idx_mapping, |
| 176 | temp, |
| 177 | seed, |
| 178 | pos, |
| 179 | apply_temperature=True, |
| 180 | is_drafting=is_drafting, |
| 181 | use_fp64=use_fp64, |
| 182 | ) |
| 183 | |
| 184 | |
| 185 | def _z_score(observed: int, expected: float, num_trials: int) -> float: |
| 186 | p = expected / num_trials |
| 187 | return (observed - expected) / math.sqrt(num_trials * p * (1 - p)) |
| 188 | |
| 189 | |
| 190 | def _sample_histogram( |
| 191 | logits_1d: torch.Tensor, |
| 192 | num_samples: int, |
| 193 | *, |
| 194 | chunk: int = 1_000_000, |
| 195 | is_drafting: bool = False, |
| 196 | ) -> torch.Tensor: |
| 197 | """Histogram of `num_samples` draws, accumulated in chunks. |
| 198 | |
| 199 | Chunking keeps the kernel's per-sample scratch ([chunk, num_blocks]) bounded |
| 200 | so a large sample count does not blow up memory. |
| 201 | """ |
| 202 | vocab_size = logits_1d.shape[0] |
| 203 | hist = torch.zeros(vocab_size, dtype=torch.float64, device=DEVICE) |
| 204 | for start in range(0, num_samples, chunk): |
| 205 | size = min(chunk, num_samples - start) |
| 206 | logits = logits_1d.unsqueeze(0).expand(size, vocab_size) |
| 207 | idx_mapping = torch.zeros(size, dtype=torch.int32, device=DEVICE) |
| 208 | temp = torch.tensor([1.0], dtype=torch.float32, device=DEVICE) |
| 209 | seed = torch.tensor([0xABCD], dtype=torch.int64, device=DEVICE) |
| 210 | pos = torch.arange(start, start + size, dtype=torch.int64, device=DEVICE) |
| 211 | out = gumbel_sample( |
| 212 | logits, |
| 213 | idx_mapping, |
| 214 | temp, |
| 215 | seed, |
| 216 | pos, |
| 217 | apply_temperature=True, |
| 218 | is_drafting=is_drafting, |
| 219 | ) |
| 220 | hist += torch.bincount(out, minlength=vocab_size).double() |
| 221 | return hist |
| 222 | |
| 223 | |
| 224 | # ----------------------------- Accuracy ------------------------------------ |
| 225 | |
| 226 | |
| 227 | @pytest.mark.parametrize("use_fp64", [False, True]) |
| 228 | def test_sampling_matches_target_distribution(use_fp64: bool): |
| 229 | counts = _make_heavy_tailed_counts() |
| 230 | total = counts.sum().item() |
| 231 | logits = _counts_to_logits(counts) |
| 232 | |
| 233 | sampled = _sample(logits, NUM_SAMPLES, use_fp64=use_fp64) |
| 234 | assert sampled.min() >= 0 and sampled.max() < VOCAB_SIZE |
| 235 | |
| 236 | # The dominant token (index 0) and the aggregate tail are the two |
| 237 | # statistically resolvable bins (individual tail tokens are far below the |
| 238 | # ~5/N detectability floor). The tail mass is small but well above noise, |
| 239 | # and it lives beyond the fp32 Gumbel cap -- the regime sensitive to noise |
| 240 | # precision -- so matching it is the meaningful check. |
| 241 | tail_prob = (total - counts[0].item()) / total |
| 242 | tail_count = (sampled != 0).sum().item() |
| 243 | z = _z_score(tail_count, NUM_SAMPLES * tail_prob, NUM_SAMPLES) |
| 244 | assert abs(z) < Z_TOLERANCE, ( |
| 245 | f"sampled tail mass {tail_count / NUM_SAMPLES:.3e} != target " |
| 246 | f"{tail_prob:.3e} (z={z:.2f})" |
| 247 | ) |
| 248 | |
| 249 | |
| 250 | def test_full_vocab_distribution_fidelity(): |
| 251 | """The sampled distribution matches the target across the WHOLE vocab. |
| 252 | |
| 253 | A near-flat count tensor makes every one of the 200K bins individually |
| 254 | measurable. With ~20 samples/bin, a goodness-of-fit over all bins checks |
| 255 | that no part of the vocab is over- or under-represented (the heavy-tailed |
| 256 | test above only resolves head vs aggregate tail). Empirically the fp32 |
| 257 | sampler is as faithful here as torch.multinomial; the residual error is the |
| 258 | multinomial sampling-noise floor, not the kernel. |
| 259 | """ |
| 260 | gen = torch.Generator(device=DEVICE).manual_seed(2024) |
| 261 | counts = torch.randint( |
| 262 | 500, 1500, (VOCAB_SIZE,), generator=gen, dtype=torch.int64, device=DEVICE |
| 263 | ) |
| 264 | total = counts.sum().item() |
| 265 | logits = _counts_to_logits(counts) |
| 266 | |
| 267 | num_samples = 4_000_000 |
| 268 | hist = _sample_histogram(logits, num_samples) |
| 269 | |
| 270 | # Diversity: essentially every token must be reachable (no starved region). |
| 271 | coverage = (hist > 0).sum().item() / VOCAB_SIZE |
| 272 | assert coverage > 0.99, f"only {coverage:.4f} of the vocab was ever sampled" |
| 273 | |
| 274 | # Goodness-of-fit across all bins (each has expected count >= ~10). |
| 275 | expected = (counts.double() / total) * num_samples |
| 276 | chi2 = (((hist - expected) ** 2) / expected).sum().item() |
| 277 | df = VOCAB_SIZE - 1 |
| 278 | assert chi2 < df + 10 * math.sqrt(2 * df), f"chi2={chi2:.0f}, df={df}" |
| 279 | |
| 280 | |
| 281 | # --------------------------- Noise streams --------------------------------- |
| 282 | |
| 283 | |
| 284 | def test_drafting_uses_a_separate_noise_stream(): |
| 285 | """is_drafting salts the RNG position: same inputs, different draws. |
| 286 | |
| 287 | The draft proposal and the residual resample after a rejection must be |
| 288 | independent. They key noise by the same (seed, pos), so only the salt keeps |
| 289 | them apart -- without it the resample inherits the very noise vector that |
| 290 | picked the rejected proposal. See test_stochastic_rejection_sample in |
| 291 | tests/v1/spec_decode/test_rejection_sampler_utils.py for the distributional |
| 292 | consequence. |
| 293 | |
| 294 | Salting the position must not distort the draw either, so both streams are |
| 295 | checked against the target's far-tail mass, which sits HEAD_LOG_GAP logits |
| 296 | below the head -- the regime where fp32 Gumbel precision matters. |
| 297 | """ |
| 298 | counts = _make_heavy_tailed_counts() |
| 299 | total = counts.sum().item() |
| 300 | logits = _counts_to_logits(counts) |
| 301 | tail_prob = (total - counts[0].item()) / total |
| 302 | |
| 303 | target = _sample(logits, NUM_SAMPLES, is_drafting=False) |
| 304 | draft = _sample(logits, NUM_SAMPLES, is_drafting=True) |
| 305 | |
| 306 | # The head dominates, so the streams agree on most draws by construction. |
| 307 | # Compare instead which draws leave the head: shared noise makes that |
| 308 | # identical, independent noise makes them differ on ~2p(1-p) of draws. |
| 309 | target_tail = target != 0 |
| 310 | draft_tail = draft != 0 |
| 311 | disagree = (target_tail != draft_tail).double().mean().item() |
| 312 | assert disagree > tail_prob, ( |
| 313 | f"streams leave the head on the same draws ({disagree:.3e} disagreement " |
| 314 | f"vs tail mass {tail_prob:.3e}); the draft salt is not taking effect" |
| 315 | ) |
| 316 | |
| 317 | # Both streams must still reproduce the target's tail mass. |
| 318 | for name, tail in (("target", target_tail), ("draft", draft_tail)): |
| 319 | tail_count = tail.sum().item() |
| 320 | z = _z_score(tail_count, NUM_SAMPLES * tail_prob, NUM_SAMPLES) |
| 321 | assert abs(z) < Z_TOLERANCE, ( |
| 322 | f"{name} tail mass {tail_count / NUM_SAMPLES:.3e} != " |
| 323 | f"{tail_prob:.3e} (z={z:.2f})" |
| 324 | ) |
| 325 | |
| 326 | # The draft stream is reproducible. |
| 327 | assert torch.equal(draft, _sample(logits, NUM_SAMPLES, is_drafting=True)) |
| 328 | |
| 329 | |
| 330 | # ----------------------------- Edge cases ---------------------------------- |
| 331 | |
| 332 | |
| 333 | def test_greedy_temperature_zero_returns_argmax(): |
| 334 | """Temperature == 0 skips Gumbel noise and returns the exact argmax.""" |
| 335 | torch.manual_seed(0) |
| 336 | num_reqs = 128 |
| 337 | logits = torch.randn(num_reqs, VOCAB_SIZE, device=DEVICE, dtype=torch.float32) |
| 338 | idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=DEVICE) |
| 339 | temp = torch.zeros(num_reqs, dtype=torch.float32, device=DEVICE) |
| 340 | seed = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE) |
| 341 | pos = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE) |
| 342 | |
| 343 | sampled = gumbel_sample( |
| 344 | logits, idx_mapping, temp, seed, pos, apply_temperature=True, is_drafting=False |
| 345 | ) |
| 346 | assert torch.equal(sampled, logits.argmax(dim=-1)) |
| 347 | |
| 348 | |
| 349 | def test_zero_count_tokens_are_never_sampled(): |
| 350 | """Count 0 -> -inf logit -> probability 0; must never be selected.""" |
| 351 | counts = _make_heavy_tailed_counts(seed=7) |
| 352 | zeroed = torch.arange(1, VOCAB_SIZE, 2, device=DEVICE) # odd indices (not head) |
| 353 | counts[zeroed] = 0 |
| 354 | logits = _counts_to_logits(counts) |
| 355 | |
| 356 | sampled = _sample(logits, NUM_SAMPLES) |
| 357 | assert sampled.min() >= 0 and sampled.max() < VOCAB_SIZE |
| 358 | assert not torch.isin(sampled, zeroed).any(), "sampled a zero-probability token" |
| 359 | |
| 360 | |
| 361 | def test_single_nonzero_token_is_always_sampled(): |
| 362 | """A lone finite logit must win every draw, regardless of its index.""" |
| 363 | counts = torch.zeros(VOCAB_SIZE, dtype=torch.int64, device=DEVICE) |
| 364 | counts[123_456] = 1000 |
| 365 | logits = _counts_to_logits(counts) |
| 366 | |
| 367 | sampled = _sample(logits, 10_000) |
| 368 | assert (sampled == 123_456).all() |
| 369 | |
| 370 | |
| 371 | @pytest.mark.parametrize("vocab_size", [1, 999, 1024, 4097]) |
| 372 | def test_vocab_size_not_multiple_of_block(vocab_size: int): |
| 373 | """Per-block tail masking for non-block-aligned vocab; all bins measurable.""" |
| 374 | gen = torch.Generator(device=DEVICE).manual_seed(vocab_size) |
| 375 | counts = torch.randint( |
| 376 | 20, 200, (vocab_size,), generator=gen, dtype=torch.int64, device=DEVICE |
| 377 | ) |
| 378 | total = counts.sum().item() |
| 379 | logits = _counts_to_logits(counts) |
| 380 | num_samples = max(40 * vocab_size, 50_000) |
| 381 | |
| 382 | sampled = _sample(logits, num_samples) |
| 383 | assert sampled.min() >= 0 and sampled.max() < vocab_size |
| 384 | |
| 385 | observed = torch.bincount(sampled, minlength=vocab_size).double() |
| 386 | expected = (counts.double() / total) * num_samples |
| 387 | chi2 = (((observed - expected) ** 2) / expected).sum().item() |
| 388 | df = vocab_size - 1 |
| 389 | if df >= 1: |
| 390 | assert chi2 < df + 10 * math.sqrt(2 * df), f"chi2={chi2:.1f}, df={df}" |
| 391 | |
| 392 | |
| 393 | # ----------------------------- Logits cache -------------------------------- |
| 394 | |
| 395 | |
| 396 | def _float_bits(t: torch.Tensor) -> torch.Tensor: |
| 397 | """Reinterpret floats as same-width ints, so equality is truly bitwise.""" |
| 398 | int_dtype = torch.int32 if t.dtype == torch.float32 else torch.int16 |
| 399 | return t.view(int_dtype) |
| 400 | |
| 401 | |
| 402 | @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) |
| 403 | @pytest.mark.parametrize("per_token_col", [False, True]) |
| 404 | def test_logits_cache_stores_input_logits_bitwise( |
| 405 | dtype: torch.dtype, per_token_col: bool |
| 406 | ): |
| 407 | """`logits_cache` must receive the input logits, pre-temperature and bit-exact. |
| 408 | |
| 409 | The cache is kept in the head's dtype, which is only lossless because the |
| 410 | stored value is the input logit itself. Storing `logit / temp` instead would |
| 411 | generally not be representable there, and the rejection sampler -- which |
| 412 | divides by the same temperature on load -- would then verify against a `q` |
| 413 | the draft never sampled from. Temperatures 0.0 and 1.0 are included because |
| 414 | both make the divide a no-op and would mask such a bug. |
| 415 | |
| 416 | Both column-addressing modes are covered: a 0-d column (one draft step per |
| 417 | call, as MTP/EAGLE do) and a [num_tokens] column (as DFlash does, sampling |
| 418 | every step in one call). |
| 419 | """ |
| 420 | torch.manual_seed(0) |
| 421 | num_reqs, vocab_size, num_steps = 8, 4099, 3 |
| 422 | logits = torch.randn(num_reqs, vocab_size, device=DEVICE, dtype=dtype) |
| 423 | idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=DEVICE) |
| 424 | temp = torch.tensor( |
| 425 | [0.0, 0.1, 0.5, 1.0, 1.5, 2.0, 0.7, 1.0], dtype=torch.float32, device=DEVICE |
| 426 | ) |
| 427 | seed = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE) |
| 428 | pos = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE) |
| 429 | |
| 430 | cache = torch.zeros(num_reqs, num_steps, vocab_size, device=DEVICE, dtype=dtype) |
| 431 | if per_token_col: |
| 432 | # Each token lands in its own column, cycling through the steps. |
| 433 | cols = torch.arange(num_reqs, dtype=torch.int32, device=DEVICE) % num_steps |
| 434 | else: |
| 435 | cols = torch.tensor(1, dtype=torch.int32, device=DEVICE) |
| 436 | |
| 437 | gumbel_sample( |
| 438 | logits, |
| 439 | idx_mapping, |
| 440 | temp, |
| 441 | seed, |
| 442 | pos, |
| 443 | apply_temperature=True, |
| 444 | is_drafting=True, |
| 445 | logits_cache=cache, |
| 446 | logits_cache_col=cols, |
| 447 | ) |
| 448 | |
| 449 | if per_token_col: |
| 450 | stored = cache[torch.arange(num_reqs, device=DEVICE), cols.long()] |
| 451 | else: |
| 452 | stored = cache[:, 1] |
| 453 | # Untouched columns must stay untouched. |
| 454 | assert not cache[:, 0].any() and not cache[:, 2].any() |
| 455 | assert torch.equal(_float_bits(stored), _float_bits(logits)), ( |
| 456 | "cached logits differ from the input logits" |
| 457 | ) |
| 458 | |
| 459 | |
| 460 | @pytest.mark.parametrize("extra_cache_cols", [0, 1]) |
| 461 | def test_logits_cache_columns_stay_separate_across_steps(extra_cache_cols: int): |
| 462 | """Each drafting step must land in its own cache column. |
| 463 | |
| 464 | `extra_cache_cols=1` is the shape a draft produces when it adds an |
| 465 | input-only mask/noise row to its embedding table but keeps a full-width |
| 466 | output head: N-wide logits cached into N+1-wide rows. Striding cache |
| 467 | columns by the logits width instead of the cache's own row width leaves |
| 468 | step 0 correct and silently misaligns every step after it. |
| 469 | """ |
| 470 | torch.manual_seed(0) |
| 471 | num_reqs, vocab_size, num_steps = 4, 1031, 3 |
| 472 | idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=DEVICE) |
| 473 | temp = torch.ones(num_reqs, dtype=torch.float32, device=DEVICE) |
| 474 | seed = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE) |
| 475 | pos = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE) |
| 476 | |
| 477 | cache = torch.zeros( |
| 478 | num_reqs, num_steps, vocab_size + extra_cache_cols, device=DEVICE |
| 479 | ) |
| 480 | cols = torch.arange(num_steps, dtype=torch.int32, device=DEVICE) |
| 481 | per_step = [ |
| 482 | torch.randn(num_reqs, vocab_size, device=DEVICE) for _ in range(num_steps) |
| 483 | ] |
| 484 | |
| 485 | for step, logits in enumerate(per_step): |
| 486 | gumbel_sample( |
| 487 | logits, |
| 488 | idx_mapping, |
| 489 | temp, |
| 490 | seed, |
| 491 | pos, |
| 492 | apply_temperature=True, |
| 493 | is_drafting=True, |
| 494 | logits_cache=cache, |
| 495 | logits_cache_col=cols[step], |
| 496 | ) |
| 497 | |
| 498 | for step, logits in enumerate(per_step): |
| 499 | stored = cache[:, step, :vocab_size] |
| 500 | assert torch.equal(_float_bits(stored), _float_bits(logits)), ( |
| 501 | f"step {step} was overwritten by a later step" |
| 502 | ) |
| 503 | # Columns past the sampled width belong to no step and stay untouched. |
| 504 | assert not cache[:, :, vocab_size:].any() |
| 505 | |
| 506 | |
| 507 | def test_logits_cache_narrower_than_logits_is_rejected(): |
| 508 | """A cache too narrow to hold a step would silently drop its tail.""" |
| 509 | num_reqs, vocab_size, num_steps = 2, 64, 3 |
| 510 | logits = torch.randn(num_reqs, vocab_size, device=DEVICE) |
| 511 | idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=DEVICE) |
| 512 | temp = torch.ones(num_reqs, dtype=torch.float32, device=DEVICE) |
| 513 | seed = torch.zeros(num_reqs, dtype=torch.int64, device=DEVICE) |
| 514 | pos = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE) |
| 515 | cache = torch.zeros(num_reqs, num_steps, vocab_size - 1, device=DEVICE) |
| 516 | |
| 517 | with pytest.raises(AssertionError, match="narrower"): |
| 518 | gumbel_sample( |
| 519 | logits, |
| 520 | idx_mapping, |
| 521 | temp, |
| 522 | seed, |
| 523 | pos, |
| 524 | apply_temperature=True, |
| 525 | is_drafting=True, |
| 526 | logits_cache=cache, |
| 527 | logits_cache_col=torch.tensor(0, dtype=torch.int32, device=DEVICE), |
| 528 | ) |