Skip to content

vllm.model_executor.models.diffusion_gemma_sampler

One-pass row statistics for the DiffusionGemma denoise sampler.

The step needs, for every canvas position, the argmax of the temperature-scaled logits, a Gumbel-max sample from them, the entropy of their softmax, and the softmax itself in the model dtype for the self-conditioning matmul. As PyTorch ops these are several passes over a [rows, vocab] fp32 tensor plus a same-sized noise tensor. The Triton kernel reads each row once with an online max and sum, draws the noise inline, and writes the probabilities once.

Functions:

sample_row_stats(logits, temps, canvas_len, seed, probs_dtype)

Argmax, Gumbel-max sample, entropy and (optionally) softmax per row.

logits is [rows, vocab] and row i uses temperature temps[i // canvas_len]. A zero temperature is greedy: the sample is the argmax, and the entropy is the reference's, which clamps the temperature at 1e-10.

Source code in vllm/model_executor/models/diffusion_gemma_sampler.py 
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
def sample_row_stats(
    logits: torch.Tensor,
    temps: torch.Tensor,
    canvas_len: int,
    seed: int,
    probs_dtype: torch.dtype | None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
    """Argmax, Gumbel-max sample, entropy and (optionally) softmax per row.

    ``logits`` is ``[rows, vocab]`` and row ``i`` uses temperature
    ``temps[i // canvas_len]``. A zero temperature is greedy: the sample is
    the argmax, and the entropy is the reference's, which clamps the
    temperature at 1e-10.
    """
    rows, vocab = logits.shape
    device = logits.device
    argmax = torch.empty(rows, dtype=torch.int64, device=device)
    sample = torch.empty(rows, dtype=torch.int64, device=device)
    entropy = torch.empty(rows, dtype=torch.float32, device=device)
    if probs_dtype is not None:
        probs = torch.empty(rows, vocab, dtype=probs_dtype, device=device)
        probs_arg, probs_stride = probs, probs.stride(0)
    else:
        probs = None
        probs_arg, probs_stride = entropy, 0
    if rows == 0:
        return argmax, sample, entropy, probs
    _row_stats_kernel[(rows,)](
        logits,
        logits.stride(0),
        temps,
        argmax,
        sample,
        entropy,
        probs_arg,
        probs_stride,
        seed,
        vocab,
        canvas_len,
        WRITE_PROBS=probs is not None,
        BLOCK=4096,
        num_warps=8,
    )
    return argmax, sample, entropy, probs

sample_row_stats_reference(logits, temps, canvas_len, probs_dtype)

The PyTorch form of the same statistics, for tests and non-CUDA runs.

Source code in vllm/model_executor/models/diffusion_gemma_sampler.py 
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
def sample_row_stats_reference(
    logits: torch.Tensor,
    temps: torch.Tensor,
    canvas_len: int,
    probs_dtype: torch.dtype | None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
    """The PyTorch form of the same statistics, for tests and non-CUDA runs."""
    temp = temps.repeat_interleave(canvas_len).float()[:, None]
    scaled = logits.float() / temp.clamp(min=1e-10)
    u = torch.rand_like(scaled).clamp(min=1e-20)
    noisy = scaled + (-torch.log(-torch.log(u))) * (temp > 0).float()
    log_probs = scaled.log_softmax(dim=-1)
    probs = log_probs.exp()
    # Masked (-inf) columns: 0 * -inf is NaN, so they are dropped from the sum.
    entropy = -torch.where(probs > 0, probs * log_probs, 0.0).sum(dim=-1)
    out_probs = probs.to(probs_dtype) if probs_dtype is not None else None
    return scaled.argmax(dim=-1), noisy.argmax(dim=-1), entropy, out_probs