| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """The fused row statistics of the DiffusionGemma sampler against the PyTorch |
| 4 | reference: argmax, entropy and softmax must match. A zero temperature must |
| 5 | sample the argmax and a positive one must sample from the distribution.""" |
| 6 | |
| 7 | import math |
| 8 | |
| 9 | import pytest |
| 10 | import torch |
| 11 | |
| 12 | from vllm.model_executor.models.diffusion_gemma_sampler import ( |
| 13 | sample_row_stats, |
| 14 | sample_row_stats_reference, |
| 15 | ) |
| 16 | |
| 17 | pytestmark = pytest.mark.skipif( |
| 18 | not torch.cuda.is_available(), reason="the fused kernel needs a GPU" |
| 19 | ) |
| 20 | |
| 21 | |
| 22 | @pytest.mark.parametrize("vocab", [5000, 9001]) |
| 23 | @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) |
| 24 | def test_stats_match_reference(vocab: int, dtype: torch.dtype): |
| 25 | torch.manual_seed(0) |
| 26 | num_decode, cl = 9, 4 |
| 27 | logits = (torch.randn(num_decode * cl, vocab, device="cuda") * 3).to(dtype) |
| 28 | temps = torch.tensor([0.0, 0.7, 2.0, 1.0, 0.0, 0.3, 1.5, 1.0, 0.0], device="cuda") |
| 29 | argmax, sample, entropy, probs = sample_row_stats( |
| 30 | logits, temps, cl, seed=1234, probs_dtype=torch.bfloat16 |
| 31 | ) |
| 32 | ref_argmax, _, ref_entropy, ref_probs = sample_row_stats_reference( |
| 33 | logits, temps, cl, probs_dtype=torch.bfloat16 |
| 34 | ) |
| 35 | assert torch.equal(argmax, ref_argmax) |
| 36 | torch.testing.assert_close(entropy, ref_entropy, atol=2e-3, rtol=1e-3) |
| 37 | assert probs is not None |
| 38 | assert ref_probs is not None |
| 39 | torch.testing.assert_close(probs.float(), ref_probs.float(), atol=5e-3, rtol=2e-2) |
| 40 | assert probs.dtype == torch.bfloat16 |
| 41 | assert probs.shape == logits.shape |
| 42 | # A padded (all-zero) row is uniform: max entropy, argmax 0. |
| 43 | zero = torch.zeros(1, vocab, device="cuda") |
| 44 | a0, _, e0, _ = sample_row_stats(zero, temps[:1] + 1.0, 1, 1, None) |
| 45 | assert a0.item() == 0 |
| 46 | assert abs(e0.item() - torch.log(torch.tensor(float(vocab))).item()) < 1e-3 |
| 47 | |
| 48 | |
| 49 | def test_zero_temperature_is_greedy_and_probs_optional(): |
| 50 | torch.manual_seed(0) |
| 51 | logits = torch.randn(12, 4097, device="cuda") |
| 52 | temps = torch.zeros(12, device="cuda") |
| 53 | argmax, sample, entropy, probs = sample_row_stats(logits, temps, 1, 7, None) |
| 54 | assert torch.equal(sample, argmax) |
| 55 | assert probs is None |
| 56 | assert torch.equal(argmax, logits.argmax(dim=-1)) |
| 57 | # Greedy entropy is that of the reference's clamped temperature. |
| 58 | _, _, ref_entropy, _ = sample_row_stats_reference(logits, temps, 1, None) |
| 59 | torch.testing.assert_close(entropy, ref_entropy, atol=2e-3, rtol=1e-3) |
| 60 | |
| 61 | |
| 62 | def test_positive_temperature_samples_from_the_distribution(): |
| 63 | torch.manual_seed(0) |
| 64 | rows, vocab = 2000, 4099 |
| 65 | # A peaked row samples its argmax almost always and a flat row rarely. |
| 66 | peaked = torch.zeros(rows, vocab, device="cuda") |
| 67 | peaked[:, 17] = 16.0 # p(17) = e^16 / (e^16 + 4098) > 0.999 |
| 68 | flat = torch.zeros(rows, vocab, device="cuda") |
| 69 | temps = torch.ones(rows, device="cuda") |
| 70 | _, s_peaked, _, _ = sample_row_stats(peaked, temps, 1, 99, None) |
| 71 | _, s_flat, _, _ = sample_row_stats(flat, temps, 1, 99, None) |
| 72 | assert (s_peaked == 17).float().mean().item() > 0.99 |
| 73 | assert s_flat.unique().numel() > rows // 2 |
| 74 | # The same seed repeats the draws and another seed changes them. |
| 75 | _, s_again, _, _ = sample_row_stats(flat, temps, 1, 99, None) |
| 76 | _, s_other, _, _ = sample_row_stats(flat, temps, 1, 100, None) |
| 77 | assert torch.equal(s_flat, s_again) |
| 78 | assert not torch.equal(s_flat, s_other) |
| 79 | |
| 80 | |
| 81 | def test_empty_batch(): |
| 82 | logits = torch.empty(0, 128, device="cuda") |
| 83 | temps = torch.empty(0, device="cuda") |
| 84 | argmax, sample, entropy, probs = sample_row_stats( |
| 85 | logits, temps, 1, 1, torch.bfloat16 |
| 86 | ) |
| 87 | assert argmax.numel() == sample.numel() == entropy.numel() == 0 |
| 88 | assert probs is not None and probs.shape == (0, 128) |
| 89 | |
| 90 | |
| 91 | def test_masked_logits_keep_a_finite_entropy(): |
| 92 | """top_k/top_p leave -inf logits. Masked columns gave NaN entropy, and a |
| 93 | row with one live column has zero entropy.""" |
| 94 | logits = torch.full((4, 5000), float("-inf"), device="cuda") |
| 95 | logits[:, 3] = 0.0 |
| 96 | logits[2, 9] = 0.0 # two live columns: entropy log 2 |
| 97 | temps = torch.tensor([1.0, 0.0, 1.0, 0.7], device="cuda") |
| 98 | argmax, sample, entropy, probs = sample_row_stats( |
| 99 | logits, temps, 1, 5, torch.float32 |
| 100 | ) |
| 101 | assert torch.isfinite(entropy).all() |
| 102 | assert entropy[0].item() < 1e-6 and entropy[1].item() < 1e-6 |
| 103 | assert abs(entropy[2].item() - math.log(2.0)) < 1e-4 |
| 104 | assert torch.equal(argmax[[0, 1, 3]], torch.tensor([3, 3, 3], device="cuda")) |
| 105 | assert probs is not None |
| 106 | assert probs[0, 3].item() == 1.0 and probs[0].sum().item() == 1.0 |
| 107 | _, _, ref_entropy, _ = sample_row_stats_reference(logits, temps, 1, None) |
| 108 | assert torch.isfinite(ref_entropy).all() |
| 109 | |
| 110 | |
| 111 | def test_leading_masked_block_keeps_a_finite_entropy(): |
| 112 | """top_k leaves most of the vocabulary -inf, so whole leading blocks can be |
| 113 | masked before the first live logit. The running max was -inf through |
| 114 | them, and -inf - -inf poisoned the row with NaN.""" |
| 115 | logits = torch.full((3, 9001), float("-inf"), device="cuda") |
| 116 | logits[:, 8500] = 0.0 # the only live column, past the first block |
| 117 | logits[1, 8600] = 0.0 # two live columns: entropy log 2 |
| 118 | logits[2, 4] = -1.0 # a live column in the first block as well |
| 119 | temps = torch.tensor([1.0, 1.0, 0.7], device="cuda") |
| 120 | argmax, sample, entropy, probs = sample_row_stats( |
| 121 | logits, temps, 1, 5, torch.float32 |
| 122 | ) |
| 123 | _, _, ref_entropy, ref_probs = sample_row_stats_reference( |
| 124 | logits, temps, 1, torch.float32 |
| 125 | ) |
| 126 | assert torch.isfinite(entropy).all() and torch.isfinite(probs).all() |
| 127 | assert entropy[0].item() < 1e-6 |
| 128 | assert abs(entropy[1].item() - math.log(2.0)) < 1e-4 |
| 129 | torch.testing.assert_close(entropy, ref_entropy, atol=1e-4, rtol=1e-4) |
| 130 | torch.testing.assert_close(probs, ref_probs, atol=1e-5, rtol=1e-4) |
| 131 | assert argmax.tolist() == [8500, 8500, 8500] |
| 132 | assert sample[0].item() == 8500 |