iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Perf] DiffusionGemma: one-pass sampler statistics kernel (#58226)

Browse Source 
Signed-off-by: Matt Mastracci <matthew@mastracci.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
Matt Mastracci
2026-09-25 14:14:55 -04:00
committed by GitHub
co-authored by Claude Fable 5.1
parent e3dd5b5a75
commit 353e491efd
4 changed files with 386 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/sample/test_diffusion_gemma_reads.py
+13 -4
View File
@@ -12,7 +12,9 @@ from vllm.model_executor.models.diffusion_gemma import (
DiffusionGemmaRequestStates,
_compiled_sample_step,
_concat_logprob_stashes,
_denoise_temperature,
_mask_rows_to_allowed,
sample_row_stats_reference,
)
from vllm.platforms import current_platform
from vllm.v1.outputs import LogprobsTensors
@@ -132,8 +134,18 @@ def _denoise_once(
decode_idx = torch.arange(n, dtype=torch.int64, device=device)
sampled = torch.zeros(n, CL, dtype=torch.int32, device=device)[:, :width]
num_sampled = torch.zeros(n, dtype=torch.int32, device=device)
_compiled_sample_step(
temp = _denoise_temperature(states.step, decode_slots, float(MAX_STEPS), 0.5, 1.0)
argmax, sample, entropy, probs = sample_row_stats_reference(
torch.zeros(n * width, VOCAB, device=device),
temp,
width,
embed_dtype if compute_sc else None,
)
_compiled_sample_step(
sample.view(n, width),
argmax.view(n, width),
entropy.view(n, width),
probs.view(n, width, -1) if probs is not None else None,
decode_slots,
decode_idx,
decode_slots,
@@ -159,9 +171,6 @@ def _denoise_once(
sampled,
num_sampled,
torch.zeros(MAX_REQS, CL, dtype=torch.int64, device=device),
max_denoising_steps=float(MAX_STEPS),
t_min=0.5,
t_max=1.0,
confidence_threshold=0.1,
vocab_size=VOCAB,
CL=width,
tests/v1/sample/test_diffusion_sampler_stats.py
+129
View File
@@ -0,0 +1,129 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""The fused row statistics of the DiffusionGemma sampler against the PyTorch
reference: argmax, entropy and softmax must match. A zero temperature must
sample the argmax and a positive one must sample from the distribution."""
import math
import pytest
import torch
from vllm.model_executor.models.diffusion_gemma_sampler import (
sample_row_stats,
sample_row_stats_reference,
)
pytestmark = pytest.mark.skipif(
not torch.cuda.is_available(), reason="the fused kernel needs a GPU"
)
@pytest.mark.parametrize("vocab", [5000, 9001])
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
def test_stats_match_reference(vocab: int, dtype: torch.dtype):
torch.manual_seed(0)
num_decode, cl = 9, 4
logits = (torch.randn(num_decode * cl, vocab, device="cuda") * 3).to(dtype)
temps = torch.tensor([0.0, 0.7, 2.0, 1.0, 0.0, 0.3, 1.5, 1.0, 0.0], device="cuda")
argmax, sample, entropy, probs = sample_row_stats(
logits, temps, cl, seed=1234, probs_dtype=torch.bfloat16
)
ref_argmax, _, ref_entropy, ref_probs = sample_row_stats_reference(
logits, temps, cl, probs_dtype=torch.bfloat16
)
assert torch.equal(argmax, ref_argmax)
torch.testing.assert_close(entropy, ref_entropy, atol=2e-3, rtol=1e-3)
torch.testing.assert_close(probs.float(), ref_probs.float(), atol=5e-3, rtol=2e-2)
assert probs.dtype == torch.bfloat16
assert probs.shape == logits.shape
# A padded (all-zero) row is uniform: max entropy, argmax 0.
zero = torch.zeros(1, vocab, device="cuda")
a0, _, e0, _ = sample_row_stats(zero, temps[:1] + 1.0, 1, 1, None)
assert a0.item() == 0
assert abs(e0.item() - torch.log(torch.tensor(float(vocab))).item()) < 1e-3
def test_zero_temperature_is_greedy_and_probs_optional():
torch.manual_seed(0)
logits = torch.randn(12, 4097, device="cuda")
temps = torch.zeros(12, device="cuda")
argmax, sample, entropy, probs = sample_row_stats(logits, temps, 1, 7, None)
assert torch.equal(sample, argmax)
assert probs is None
assert torch.equal(argmax, logits.argmax(dim=-1))
# Greedy entropy is that of the reference's clamped temperature.
_, _, ref_entropy, _ = sample_row_stats_reference(logits, temps, 1, None)
torch.testing.assert_close(entropy, ref_entropy, atol=2e-3, rtol=1e-3)
def test_positive_temperature_samples_from_the_distribution():
torch.manual_seed(0)
rows, vocab = 2000, 4099
# A peaked row samples its argmax almost always and a flat row rarely.
peaked = torch.zeros(rows, vocab, device="cuda")
peaked[:, 17] = 16.0 # p(17) = e^16 / (e^16 + 4098) > 0.999
flat = torch.zeros(rows, vocab, device="cuda")
temps = torch.ones(rows, device="cuda")
_, s_peaked, _, _ = sample_row_stats(peaked, temps, 1, 99, None)
_, s_flat, _, _ = sample_row_stats(flat, temps, 1, 99, None)
assert (s_peaked == 17).float().mean().item() > 0.99
assert s_flat.unique().numel() > rows // 2
# The same seed repeats the draws and another seed changes them.
_, s_again, _, _ = sample_row_stats(flat, temps, 1, 99, None)
_, s_other, _, _ = sample_row_stats(flat, temps, 1, 100, None)
assert torch.equal(s_flat, s_again)
assert not torch.equal(s_flat, s_other)
def test_empty_batch():
logits = torch.empty(0, 128, device="cuda")
temps = torch.empty(0, device="cuda")
argmax, sample, entropy, probs = sample_row_stats(
logits, temps, 1, 1, torch.bfloat16
)
assert argmax.numel() == sample.numel() == entropy.numel() == 0
assert probs is not None and probs.shape == (0, 128)
def test_masked_logits_keep_a_finite_entropy():
"""top_k/top_p leave -inf logits. Masked columns gave NaN entropy, and a
row with one live column has zero entropy."""
logits = torch.full((4, 5000), float("-inf"), device="cuda")
logits[:, 3] = 0.0
logits[2, 9] = 0.0 # two live columns: entropy log 2
temps = torch.tensor([1.0, 0.0, 1.0, 0.7], device="cuda")
argmax, sample, entropy, probs = sample_row_stats(
logits, temps, 1, 5, torch.float32
)
assert torch.isfinite(entropy).all()
assert entropy[0].item() < 1e-6 and entropy[1].item() < 1e-6
assert abs(entropy[2].item() - math.log(2.0)) < 1e-4
assert torch.equal(argmax[[0, 1, 3]], torch.tensor([3, 3, 3], device="cuda"))
assert probs[0, 3].item() == 1.0 and probs[0].sum().item() == 1.0
_, _, ref_entropy, _ = sample_row_stats_reference(logits, temps, 1, None)
assert torch.isfinite(ref_entropy).all()
def test_leading_masked_block_keeps_a_finite_entropy():
"""top_k leaves most of the vocabulary -inf, so whole leading blocks can be
masked before the first live logit. The running max was -inf through
them, and -inf - -inf poisoned the row with NaN."""
logits = torch.full((3, 9001), float("-inf"), device="cuda")
logits[:, 8500] = 0.0 # the only live column, past the first block
logits[1, 8600] = 0.0 # two live columns: entropy log 2
logits[2, 4] = -1.0 # a live column in the first block as well
temps = torch.tensor([1.0, 1.0, 0.7], device="cuda")
argmax, sample, entropy, probs = sample_row_stats(
logits, temps, 1, 5, torch.float32
)
_, _, ref_entropy, ref_probs = sample_row_stats_reference(
logits, temps, 1, torch.float32
)
assert torch.isfinite(entropy).all() and torch.isfinite(probs).all()
assert entropy[0].item() < 1e-6
assert abs(entropy[1].item() - math.log(2.0)) < 1e-4
torch.testing.assert_close(entropy, ref_entropy, atol=1e-4, rtol=1e-4)
torch.testing.assert_close(probs, ref_probs, atol=1e-5, rtol=1e-4)
assert argmax.tolist() == [8500, 8500, 8500]
assert sample[0].item() == 8500
vllm/model_executor/models/diffusion_gemma.py
+70 -53
View File
@@ -37,6 +37,10 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.model_executor.layers.vocab_parallel_embedding import (
ParallelLMHead,
)
from vllm.model_executor.models.diffusion_gemma_sampler import (
sample_row_stats,
sample_row_stats_reference,
)
from vllm.model_executor.models.gemma4 import Gemma4Model
from vllm.model_executor.models.gemma4_mm import (
Gemma4DummyInputsBuilder,
@@ -447,9 +451,27 @@ def _mask_rows_to_allowed(
# eager: about twice as slow for a self-conditioned step.
@torch._dynamo.config.patch(recompile_limit=64)
@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
def _denoise_temperature(
step_tensor: torch.Tensor,
slots: torch.Tensor,
max_denoising_steps: float,
t_min: float,
t_max: float,
) -> torch.Tensor:
"""The schedule's temperature for each slot at its current step."""
steps_f = step_tensor[slots].float()
remaining = (max_denoising_steps - steps_f).clamp(min=1.0)
return t_min + (t_max - t_min) * (remaining / max_denoising_steps)
def _compiled_sample_step(
# Logits from the model [num_decode * CL, vocab]
logits: torch.Tensor,
# Per-position statistics of the temperature-scaled logits, from
# sample_row_stats: [num_decode, CL] each, and the softmax
# [num_decode, CL, vocab] in the embedding dtype when compute_sc.
new_tokens: torch.Tensor,
argmax_tokens: torch.Tensor,
token_entropy: torch.Tensor,
probs: torch.Tensor | None,
# Request mapping
decode_slots: torch.Tensor, # [num_decode] int64 → slot indices
decode_idx: torch.Tensor, # [num_decode] int64 → position in num_reqs
@@ -475,9 +497,6 @@ def _compiled_sample_step(
num_sampled: torch.Tensor, # [num_reqs]
draft_tokens: torch.Tensor, # [max_num_reqs, >=CL]
# Scalar config
max_denoising_steps: float,
t_min: float,
t_max: float,
confidence_threshold: float,
vocab_size: int,
CL: int,
@@ -492,46 +511,15 @@ def _compiled_sample_step(
tp_size: int,
tp_group_name: str,
compute_sc: bool = True,
# Constrained reads: ``logits`` has one column per allowed id and
# ``embed_weight`` holds those ids' rows, so every argmax is an index
# into ``allowed``. ``vocab_size`` stays the real vocabulary for renoise.
allowed: torch.Tensor | None = None,
) -> torch.Tensor:
"""Compiled decode step: temperature → Gumbel sample → probs/confidence →
accept/renoise → convergence, all as vectorized PyTorch ops.
Returns the temperature-scaled logits ``[num_decode, CL, vocab]`` so the
caller can compute logprobs outside the compiled region."""
) -> None:
"""Compiled decode step: confidence → accept/renoise → convergence, as
vectorized PyTorch ops over [num_decode, CL] tensors. The per-position
statistics (argmax, Gumbel-max sample, entropy, softmax) come from one
pass over the logits in ``sample_row_stats``."""
num_decode = decode_slots.shape[0]
device = decode_slots.device
# ---- Phase 1: Temperature schedule ----
steps_f = step_tensor[decode_slots].float()
remaining = (max_denoising_steps - steps_f).clamp(min=1.0)
temp = t_min + (t_max - t_min) * (remaining / max_denoising_steps)
# ---- Phase 2: Temperature scaling + Gumbel-max sampling ----
logits_3d = logits.reshape(num_decode, CL, -1).float()
scaled = logits_3d / temp[:, None, None].clamp(min=1e-10)
# Gumbel-max trick: argmax(logits/T + Gumbel) ~ sample from softmax(logits/T)
u = torch.rand_like(scaled).clamp(min=1e-20)
gumbel = -torch.log(-torch.log(u))
# Zero noise when temp==0 (greedy)
noisy = scaled + gumbel * (temp[:, None, None] > 0).float()
new_tokens = noisy.view(-1, noisy.shape[-1]).argmax(dim=-1).view(num_decode, CL)
argmax_tokens = (
scaled.view(-1, scaled.shape[-1]).argmax(dim=-1).view(num_decode, CL)
)
if allowed is not None:
new_tokens = allowed[new_tokens]
argmax_tokens = allowed[argmax_tokens]
# ---- Phase 3: Probs, self-conditioning, confidence ----
log_probs = scaled.log_softmax(dim=-1)
probs = log_probs.exp()
token_entropy = -(probs * log_probs).sum(dim=-1) # [num_decode, CL]
# ---- Phase 3: Confidence ----
# A canvas truncated near max_model_len is zero-padded up to CL by the
# caller; those padded rows are uniform (max entropy, argmax 0), so they
# never trigger early convergence and are stable, and only the real
@@ -626,6 +614,7 @@ def _compiled_sample_step(
# [.., vocab] probs avoids a giant persistent buffer.
sc_keep = (is_denoise & ~is_encoder_phase[decode_slots])[:, None, None]
if compute_sc:
assert probs is not None
# Self-conditioning soft embed = probs @ embed_tokens.weight. Under
# tensor parallelism the embedding is vocab-sharded ([vocab/tp,
# hidden]) while probs spans the full vocab, so each rank multiplies
@@ -664,8 +653,6 @@ def _compiled_sample_step(
# ---- Phase 7: Copy canvas → draft_tokens for all slots ----
draft_tokens[all_slots, :CL] = canvas[all_slots]
return scaled
class DiffusionGemmaRequestStates:
"""Pre-allocated GPU tensors for DiffusionGemma per-request state.
@@ -1513,9 +1500,10 @@ class DiffusionSampler:
run_end = run_start
while run_end < num_decode and widths_np[order[run_end]] == W:
run_end += 1
# ~10 transient fp32 copies of [tile * W, vocab] inside the step
# (eager peaks at ~8; pad for allocator overhead and small tensors).
budget = max(1, int(free * 0.5) // max(W * self.vocab_size * 4 * 10, 1))
# Transient [tile * W, vocab] tensors: the softmax in the embedding
# dtype for self-conditioning, plus the fp32 scaled logits when
# logprobs are wanted (pad for allocator overhead).
budget = max(1, int(free * 0.5) // max(W * self.vocab_size * 4 * 3, 1))
for t0 in range(run_start, run_end, budget):
sel_np = order[t0 : min(t0 + budget, run_end)]
n = len(sel_np)
@@ -1547,8 +1535,40 @@ class DiffusionSampler:
)
)
scaled = _compiled_sample_step(
tile_logits,
temp = _denoise_temperature(
states.step,
tile_slots,
float(states.max_denoising_steps),
self.t_min,
self.t_max,
)
probs_dtype = self.embed_weight.dtype if compute_sc else None
if tile_logits.is_cuda:
seed = int(torch.randint(0, 2**31 - 1, (1,)).item())
stats = sample_row_stats(tile_logits, temp, W, seed, probs_dtype)
else:
stats = sample_row_stats_reference(
tile_logits, temp, W, probs_dtype
)
argmax_rows, sample_rows, entropy_rows, probs_rows = stats
if allowed is not None:
# K-space picks back to token ids. The softmax stays K wide
# for the K-row self-conditioning matmul.
argmax_rows = allowed[argmax_rows]
sample_rows = allowed[sample_rows]
probs = probs_rows.view(n, W, -1) if probs_rows is not None else None
# Only the logprob stash reads the scaled logits.
scaled = None
if want_logprobs:
scaled = tile_logits.float().view(n, W, -1) / temp[
:, None, None
].clamp(min=1e-10)
_compiled_sample_step(
sample_rows.view(n, W),
argmax_rows.view(n, W),
entropy_rows.view(n, W),
probs,
tile_slots,
decode_idx[sel],
all_slots,
@@ -1573,9 +1593,6 @@ class DiffusionSampler:
num_sampled,
self.req_states.draft_tokens,
# Config
max_denoising_steps=float(states.max_denoising_steps),
t_min=self.t_min,
t_max=self.t_max,
confidence_threshold=self.confidence_threshold,
vocab_size=vocab_size,
CL=W,
@@ -1586,7 +1603,6 @@ class DiffusionSampler:
tp_size=self.tp_size,
tp_group_name=self.tp_group_name,
compute_sc=compute_sc,
allowed=allowed,
)
# Stash newly converged logprobs, including reads that emit now.
@@ -1596,6 +1612,7 @@ class DiffusionSampler:
)
just_converged = converged_mask & ~is_committing[sel]
if just_converged.any():
assert scaled is not None
flat_logits = scaled.reshape(-1, scaled.shape[-1])
argmax_tokens = scaled.argmax(dim=-1)
raw_flat: torch.Tensor | None = None
vllm/model_executor/models/diffusion_gemma_sampler.py
+174
View File
@@ -0,0 +1,174 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""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.
"""
import torch
from vllm.triton_utils import tl, triton
@triton.jit
def _row_stats_kernel(
logits_ptr,
logits_stride,
temps_ptr,
argmax_ptr,
sample_ptr,
entropy_ptr,
probs_ptr,
probs_stride,
seed,
V,
CL,
WRITE_PROBS: tl.constexpr,
BLOCK: tl.constexpr,
):
row = tl.program_id(0)
temp = tl.load(temps_ptr + row // CL).to(tl.float32)
inv = 1.0 / tl.maximum(temp, 1e-10)
noise = (temp > 0).to(tl.float32)
base = logits_ptr + row.to(tl.int64) * logits_stride
rng_base = row.to(tl.int64) * V
# The running max stays in raw logit units and the temperature is applied
# to differences from it. Scaling the row first breaks at a zero
# temperature: the row is multiplied by 1e10, fp32 spacing there is about
# a thousand, and exp(x - max) becomes 0 or inf for the max element itself
# whenever the two sides round differently.
m = float("-inf") # running max of the raw row
z = 0.0 # sum exp(d), d = (x - m) * inv
s = 0.0 # sum d exp(d)
best = float("-inf")
best_idx = 0
best_noisy = float("-inf")
best_noisy_idx = 0
for start in range(0, V, BLOCK):
offs = start + tl.arange(0, BLOCK)
mask = offs < V
x = tl.load(base + offs, mask=mask, other=float("-inf")).to(tl.float32)
block_max = tl.max(x, 0)
new_m = tl.maximum(m, block_max)
# Until a block holds a finite logit the running max stays -inf, and
# top_k/top_p leave whole blocks masked. x - new_m would then be
# -inf - -inf = NaN and poison z for the rest of the row, so subtract
# 0 instead: every d is still -inf and every e 0.
safe_m = tl.where(new_m > float("-inf"), new_m, 0.0)
# Rescale the running sums to the new max. While m is -inf, z and s
# are 0, so alpha and the shift are zeroed rather than computed from
# -inf - -inf.
alpha = tl.where(m > float("-inf"), tl.exp((m - safe_m) * inv), 0.0)
shift = tl.where(z > 0, (m - safe_m) * inv * z, 0.0)
d = (x - safe_m) * inv
e = tl.exp(d)
z = z * alpha + tl.sum(e, 0)
# A -inf logit (top_k/top_p, or the row's padding) has d = -inf and
# e = 0, and -inf * 0 is NaN. The where drops it from the sum.
s = (s + shift) * alpha + tl.sum(tl.where(e > 0, d * e, 0.0), 0)
m = new_m
block_idx = tl.argmax(x, 0)
take = block_max > best
best_idx = tl.where(take, start + block_idx, best_idx)
best = tl.where(take, block_max, best)
u = tl.rand(seed, rng_base + offs)
u = tl.maximum(u, 1e-20)
gumbel = -tl.log(-tl.log(u))
noisy = tl.where(mask, x * inv + gumbel * noise, float("-inf"))
noisy_max = tl.max(noisy, 0)
noisy_idx = tl.argmax(noisy, 0)
take_n = noisy_max > best_noisy
best_noisy_idx = tl.where(take_n, start + noisy_idx, best_noisy_idx)
best_noisy = tl.where(take_n, noisy_max, best_noisy)
# With d = (x - m) * inv: p = e^d / z, log Z = log z in d units, and
# H = -sum p (d - log z) = log z - s / z.
# tl.store casts to the pointee dtype. mypy type-checks the kernel body
# as Python and rejects .to() on the int accumulators.
tl.store(argmax_ptr + row, best_idx)
tl.store(sample_ptr + row, best_noisy_idx)
tl.store(entropy_ptr + row, tl.log(z) - s / z)
if WRITE_PROBS:
out = probs_ptr + row.to(tl.int64) * probs_stride
for start in range(0, V, BLOCK):
offs = start + tl.arange(0, BLOCK)
mask = offs < V
x = tl.load(base + offs, mask=mask, other=float("-inf")).to(tl.float32)
p = tl.exp((x - m) * inv) / z
tl.store(out + offs, p.to(probs_ptr.dtype.element_ty), mask=mask)
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
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