| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """One-pass row statistics for the DiffusionGemma denoise sampler. |
| 4 | |
| 5 | The step needs, for every canvas position, the argmax of the |
| 6 | temperature-scaled logits, a Gumbel-max sample from them, the entropy of |
| 7 | their softmax, and the softmax itself in the model dtype for the |
| 8 | self-conditioning matmul. As PyTorch ops these are several passes over a |
| 9 | ``[rows, vocab]`` fp32 tensor plus a same-sized noise tensor. The Triton |
| 10 | kernel reads each row once with an online max and sum, draws the noise |
| 11 | inline, and writes the probabilities once. |
| 12 | """ |
| 13 | |
| 14 | import torch |
| 15 | |
| 16 | from vllm.triton_utils import tl, triton |
| 17 | |
| 18 | |
| 19 | @triton.jit |
| 20 | def _row_stats_kernel( |
| 21 | logits_ptr, |
| 22 | logits_stride, |
| 23 | temps_ptr, |
| 24 | argmax_ptr, |
| 25 | sample_ptr, |
| 26 | entropy_ptr, |
| 27 | probs_ptr, |
| 28 | probs_stride, |
| 29 | seed, |
| 30 | V, |
| 31 | CL, |
| 32 | WRITE_PROBS: tl.constexpr, |
| 33 | BLOCK: tl.constexpr, |
| 34 | ): |
| 35 | row = tl.program_id(0) |
| 36 | temp = tl.load(temps_ptr + row // CL).to(tl.float32) |
| 37 | inv = 1.0 / tl.maximum(temp, 1e-10) |
| 38 | noise = (temp > 0).to(tl.float32) |
| 39 | base = logits_ptr + row.to(tl.int64) * logits_stride |
| 40 | rng_base = row.to(tl.int64) * V |
| 41 | |
| 42 | # The running max stays in raw logit units and the temperature is applied |
| 43 | # to differences from it. Scaling the row first breaks at a zero |
| 44 | # temperature: the row is multiplied by 1e10, fp32 spacing there is about |
| 45 | # a thousand, and exp(x - max) becomes 0 or inf for the max element itself |
| 46 | # whenever the two sides round differently. |
| 47 | m = float("-inf") # running max of the raw row |
| 48 | z = 0.0 # sum exp(d), d = (x - m) * inv |
| 49 | s = 0.0 # sum d exp(d) |
| 50 | best = float("-inf") |
| 51 | best_idx = 0 |
| 52 | best_noisy = float("-inf") |
| 53 | best_noisy_idx = 0 |
| 54 | for start in range(0, V, BLOCK): |
| 55 | offs = start + tl.arange(0, BLOCK) |
| 56 | mask = offs < V |
| 57 | x = tl.load(base + offs, mask=mask, other=float("-inf")).to(tl.float32) |
| 58 | |
| 59 | block_max = tl.max(x, 0) |
| 60 | new_m = tl.maximum(m, block_max) |
| 61 | # Until a block holds a finite logit the running max stays -inf, and |
| 62 | # top_k/top_p leave whole blocks masked. x - new_m would then be |
| 63 | # -inf - -inf = NaN and poison z for the rest of the row, so subtract |
| 64 | # 0 instead: every d is still -inf and every e 0. |
| 65 | safe_m = tl.where(new_m > float("-inf"), new_m, 0.0) |
| 66 | # Rescale the running sums to the new max. While m is -inf, z and s |
| 67 | # are 0, so alpha and the shift are zeroed rather than computed from |
| 68 | # -inf - -inf. |
| 69 | alpha = tl.where(m > float("-inf"), tl.exp((m - safe_m) * inv), 0.0) |
| 70 | shift = tl.where(z > 0, (m - safe_m) * inv * z, 0.0) |
| 71 | d = (x - safe_m) * inv |
| 72 | e = tl.exp(d) |
| 73 | z = z * alpha + tl.sum(e, 0) |
| 74 | # A -inf logit (top_k/top_p, or the row's padding) has d = -inf and |
| 75 | # e = 0, and -inf * 0 is NaN. The where drops it from the sum. |
| 76 | s = (s + shift) * alpha + tl.sum(tl.where(e > 0, d * e, 0.0), 0) |
| 77 | m = new_m |
| 78 | |
| 79 | block_idx = tl.argmax(x, 0) |
| 80 | take = block_max > best |
| 81 | best_idx = tl.where(take, start + block_idx, best_idx) |
| 82 | best = tl.where(take, block_max, best) |
| 83 | |
| 84 | u = tl.rand(seed, rng_base + offs) |
| 85 | u = tl.maximum(u, 1e-20) |
| 86 | gumbel = -tl.log(-tl.log(u)) |
| 87 | noisy = tl.where(mask, x * inv + gumbel * noise, float("-inf")) |
| 88 | noisy_max = tl.max(noisy, 0) |
| 89 | noisy_idx = tl.argmax(noisy, 0) |
| 90 | take_n = noisy_max > best_noisy |
| 91 | best_noisy_idx = tl.where(take_n, start + noisy_idx, best_noisy_idx) |
| 92 | best_noisy = tl.where(take_n, noisy_max, best_noisy) |
| 93 | |
| 94 | # With d = (x - m) * inv: p = e^d / z, log Z = log z in d units, and |
| 95 | # H = -sum p (d - log z) = log z - s / z. |
| 96 | # tl.store casts to the pointee dtype. mypy type-checks the kernel body |
| 97 | # as Python and rejects .to() on the int accumulators. |
| 98 | tl.store(argmax_ptr + row, best_idx) |
| 99 | tl.store(sample_ptr + row, best_noisy_idx) |
| 100 | tl.store(entropy_ptr + row, tl.log(z) - s / z) |
| 101 | |
| 102 | if WRITE_PROBS: |
| 103 | out = probs_ptr + row.to(tl.int64) * probs_stride |
| 104 | for start in range(0, V, BLOCK): |
| 105 | offs = start + tl.arange(0, BLOCK) |
| 106 | mask = offs < V |
| 107 | x = tl.load(base + offs, mask=mask, other=float("-inf")).to(tl.float32) |
| 108 | p = tl.exp((x - m) * inv) / z |
| 109 | tl.store(out + offs, p.to(probs_ptr.dtype.element_ty), mask=mask) |
| 110 | |
| 111 | |
| 112 | def sample_row_stats( |
| 113 | logits: torch.Tensor, |
| 114 | temps: torch.Tensor, |
| 115 | canvas_len: int, |
| 116 | seed: int, |
| 117 | probs_dtype: torch.dtype | None, |
| 118 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: |
| 119 | """Argmax, Gumbel-max sample, entropy and (optionally) softmax per row. |
| 120 | |
| 121 | ``logits`` is ``[rows, vocab]`` and row ``i`` uses temperature |
| 122 | ``temps[i // canvas_len]``. A zero temperature is greedy: the sample is |
| 123 | the argmax, and the entropy is the reference's, which clamps the |
| 124 | temperature at 1e-10. |
| 125 | """ |
| 126 | rows, vocab = logits.shape |
| 127 | device = logits.device |
| 128 | argmax = torch.empty(rows, dtype=torch.int64, device=device) |
| 129 | sample = torch.empty(rows, dtype=torch.int64, device=device) |
| 130 | entropy = torch.empty(rows, dtype=torch.float32, device=device) |
| 131 | if probs_dtype is not None: |
| 132 | probs = torch.empty(rows, vocab, dtype=probs_dtype, device=device) |
| 133 | probs_arg, probs_stride = probs, probs.stride(0) |
| 134 | else: |
| 135 | probs = None |
| 136 | probs_arg, probs_stride = entropy, 0 |
| 137 | if rows == 0: |
| 138 | return argmax, sample, entropy, probs |
| 139 | _row_stats_kernel[(rows,)]( |
| 140 | logits, |
| 141 | logits.stride(0), |
| 142 | temps, |
| 143 | argmax, |
| 144 | sample, |
| 145 | entropy, |
| 146 | probs_arg, |
| 147 | probs_stride, |
| 148 | seed, |
| 149 | vocab, |
| 150 | canvas_len, |
| 151 | WRITE_PROBS=probs is not None, |
| 152 | BLOCK=4096, |
| 153 | num_warps=8, |
| 154 | ) |
| 155 | return argmax, sample, entropy, probs |
| 156 | |
| 157 | |
| 158 | def sample_row_stats_reference( |
| 159 | logits: torch.Tensor, |
| 160 | temps: torch.Tensor, |
| 161 | canvas_len: int, |
| 162 | probs_dtype: torch.dtype | None, |
| 163 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: |
| 164 | """The PyTorch form of the same statistics, for tests and non-CUDA runs.""" |
| 165 | temp = temps.repeat_interleave(canvas_len).float()[:, None] |
| 166 | scaled = logits.float() / temp.clamp(min=1e-10) |
| 167 | u = torch.rand_like(scaled).clamp(min=1e-20) |
| 168 | noisy = scaled + (-torch.log(-torch.log(u))) * (temp > 0).float() |
| 169 | log_probs = scaled.log_softmax(dim=-1) |
| 170 | probs = log_probs.exp() |
| 171 | # Masked (-inf) columns: 0 * -inf is NaN, so they are dropped from the sum. |
| 172 | entropy = -torch.where(probs > 0, probs * log_probs, 0.0).sum(dim=-1) |
| 173 | out_probs = probs.to(probs_dtype) if probs_dtype is not None else None |
| 174 | return scaled.argmax(dim=-1), noisy.argmax(dim=-1), entropy, out_probs |