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.

The logprobs of a converging slot are written to a per-slot stash by stash_converged_logprobs, which skips the rows of slots that did not converge.

Functions:

• sample_row_stats –

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

• sample_row_stats_reference –

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

• stash_converged_logprobs –

  Write the logprobs of the converged requests' rows to the stash.

_lp_topk_kernel(logits_ptr, logits_stride, src_ptr, keep_ptr, valid_ptr, temps_ptr, slots_ptr, num_ids_ptr, block_max_ptr, cand_ptr, K, W, V, NUM_CHUNKS, NUM_BLOCKS, HAS_IDS, BLOCK, BLOCKS_PER_CHUNK, BLOCKS_POW2, KP)

The KP largest keys of one chunk of a row, unordered.

A key orders by value, then by lower index: the fp32 bits mapped to an order-preserving int32 in the high word, the inverted index in the low. Keys are distinct, so the top-k under them is unique.

Source code in vllm/model_executor/models/diffusion_gemma_sampler.py

@triton.jit
def _lp_topk_kernel(
    logits_ptr,
    logits_stride,
    src_ptr,
    keep_ptr,
    valid_ptr,
    temps_ptr,
    slots_ptr,
    num_ids_ptr,
    block_max_ptr,
    cand_ptr,
    K,
    W,
    V,
    NUM_CHUNKS,
    NUM_BLOCKS,
    HAS_IDS: tl.constexpr,
    BLOCK: tl.constexpr,
    BLOCKS_PER_CHUNK: tl.constexpr,
    BLOCKS_POW2: tl.constexpr,
    KP: tl.constexpr,
):
    """The KP largest keys of one chunk of a row, unordered.

    A key orders by value, then by lower index: the fp32 bits mapped to an
    order-preserving int32 in the high word, the inverted index in the low.
    Keys are distinct, so the top-k under them is unique.
    """
    j = tl.program_id(0)
    pos = tl.program_id(1)
    chunk = tl.program_id(2)
    if _lp_skip_topk(keep_ptr, valid_ptr, slots_ptr, num_ids_ptr, j, pos, HAS_IDS):
        return
    r = (j * W + pos).to(tl.int64)
    # K block maxima are K distinct elements at or above the K-th largest
    # block max, so every top-k element is too: the rest of the row is
    # skipped without a reduction over it.
    b = tl.arange(0, BLOCKS_POW2)
    block_max = tl.load(
        block_max_ptr + r * NUM_BLOCKS + b, mask=b < NUM_BLOCKS, other=float("-inf")
    )
    sorted_max = tl.sort(_lp_key32(block_max), descending=True)
    kth = tl.max(tl.where(b == K - 1, sorted_max, -0x7FFFFFFF - 1), 0)
    threshold = kth.to(tl.int64) << 32
    # Distinct placeholders below every real key.
    best = _NO_KEY + 1 + tl.arange(0, KP).to(tl.int64)
    best_min = tl.min(best, 0)
    for blk in range(BLOCKS_PER_CHUNK):
        offs = (chunk * BLOCKS_PER_CHUNK + blk) * BLOCK + tl.arange(0, BLOCK)
        x = _lp_load(
            logits_ptr,
            logits_stride,
            src_ptr,
            temps_ptr,
            j,
            pos,
            W,
            offs,
            V,
            float("-inf"),
        )
        key = (_lp_key32(x).to(tl.int64) << 32) | (0x7FFFFFFF - offs).to(tl.int64)
        key = tl.where((offs < V) & (key >= threshold), key, _NO_KEY)
        key_max = tl.max(key, 0)
        while key_max > best_min:
            best = tl.where(best == best_min, key_max, best)
            best_min = tl.min(best, 0)
            key = tl.where(key == key_max, _NO_KEY, key)
            key_max = tl.max(key, 0)
    out = (r * NUM_CHUNKS + chunk) * KP + tl.arange(0, KP)
    tl.store(cand_ptr + out, best)

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

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

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

stash_converged_logprobs(logits, canvas_len, src, slots, keep, valid, temps, num_topk, num_token_ids, token_ids, out_ids, out_logprobs, out_ranks, logits_mode)

Write the logprobs of the converged requests' rows to the stash.

Request j owns logits rows src[j] * canvas_len + p for p < valid[j], divided by temps[j]. Where keep[j], they go to out_*[slots[j], p], laid out as compute_topk_scores lays out a row: the argmax, then the request's token_ids if it has any, else the num_topk largest (tied ids by lower id), padded with token 0 at -inf. Rows not kept exit at once; the others are split into chunks so a few kept rows still fill the GPU.

Source code in vllm/model_executor/models/diffusion_gemma_sampler.py

def stash_converged_logprobs(
    logits: torch.Tensor,
    canvas_len: int,
    src: torch.Tensor,
    slots: torch.Tensor,
    keep: torch.Tensor,
    valid: torch.Tensor,
    temps: torch.Tensor,
    num_topk: int,
    num_token_ids: torch.Tensor,
    token_ids: torch.Tensor | None,
    out_ids: torch.Tensor,
    out_logprobs: torch.Tensor,
    out_ranks: torch.Tensor,
    logits_mode: bool,
) -> None:
    """Write the logprobs of the converged requests' rows to the stash.

    Request `j` owns logits rows `src[j] * canvas_len + p` for
    `p < valid[j]`, divided by `temps[j]`. Where `keep[j]`, they go to
    `out_*[slots[j], p]`, laid out as compute_topk_scores lays out a row:
    the argmax, then the request's `token_ids` if it has any, else the
    `num_topk` largest (tied ids by lower id), padded with token 0 at -inf.
    Rows not kept exit at once; the others are split into chunks so a few
    kept rows still fill the GPU.
    """
    m = slots.shape[0]
    vocab = logits.shape[-1]
    assert logits.stride(-1) == 1
    assert temps.dtype == torch.float32
    assert num_topk + 1 <= out_logprobs.shape[-1]
    block = _LP_BLOCK
    num_blocks = triton.cdiv(vocab, block)
    blocks_per_chunk = min(_LP_BLOCKS_PER_CHUNK, num_blocks)
    num_chunks = triton.cdiv(num_blocks, blocks_per_chunk)
    rows = m * canvas_len
    device = logits.device
    chunk_max = torch.empty(rows, num_chunks, dtype=torch.float32, device=device)
    chunk_arg = torch.empty(rows, num_chunks, dtype=torch.int64, device=device)
    block_sum = torch.empty(rows, num_blocks, dtype=torch.float32, device=device)
    block_cnt = torch.empty(rows, num_blocks, dtype=torch.int32, device=device)
    block_max = (
        torch.empty(rows, num_blocks, dtype=torch.float32, device=device)
        if num_topk > 0
        else chunk_max
    )
    has_ids = token_ids is not None
    ids_arg = token_ids if token_ids is not None else num_token_ids
    common = (logits, logits.stride(0), src, keep, valid, temps)
    grid = (m, canvas_len, num_chunks)
    _lp_max_kernel[grid](
        *common,
        chunk_max,
        chunk_arg,
        block_max,
        canvas_len,
        vocab,
        num_chunks,
        num_blocks,
        BLOCK=block,
        BLOCKS_PER_CHUNK=blocks_per_chunk,
        STORE_BLOCK_MAX=num_topk > 0,
        num_warps=4,
    )
    chunks_pow2 = triton.next_power_of_2(num_chunks)
    _lp_sum_kernel[grid](
        *common,
        chunk_max,
        block_sum,
        block_cnt,
        canvas_len,
        vocab,
        num_chunks,
        num_blocks,
        BLOCK=block,
        BLOCKS_PER_CHUNK=blocks_per_chunk,
        CHUNKS_POW2=chunks_pow2,
        num_warps=4,
    )
    topk = chunk_arg
    if num_topk > 0:
        kp = triton.next_power_of_2(num_topk)
        cand = torch.empty(rows, num_chunks, kp, dtype=torch.int64, device=device)
        _lp_topk_kernel[grid](
            *common,
            slots,
            num_token_ids,
            block_max,
            cand,
            num_topk,
            canvas_len,
            vocab,
            num_chunks,
            num_blocks,
            HAS_IDS=has_ids,
            BLOCK=block,
            BLOCKS_PER_CHUNK=blocks_per_chunk,
            BLOCKS_POW2=triton.next_power_of_2(max(num_blocks, num_topk)),
            KP=kp,
            num_warps=4,
        )
        topk = torch.empty(rows, num_topk, dtype=torch.int32, device=device)
        _lp_topk_merge_kernel[(m, canvas_len)](
            keep,
            valid,
            slots,
            num_token_ids,
            cand,
            topk,
            num_topk,
            canvas_len,
            num_chunks * kp,
            HAS_IDS=has_ids,
            CAND_POW2=triton.next_power_of_2(num_chunks * kp),
        )
    stash_width = out_logprobs.shape[-1]
    _lp_write_kernel[(m, canvas_len)](
        logits,
        logits.stride(0),
        src,
        slots,
        keep,
        valid,
        temps,
        chunk_max,
        chunk_arg,
        block_sum,
        block_cnt,
        topk,
        num_token_ids,
        ids_arg,
        ids_arg.stride(0) if has_ids else 0,
        out_ids,
        out_logprobs,
        out_ranks,
        num_topk,
        canvas_len,
        out_logprobs.shape[1],
        vocab,
        stash_width,
        num_chunks,
        num_blocks,
        HAS_IDS=has_ids,
        LOGITS_MODE=logits_mode,
        CHUNKS_POW2=chunks_pow2,
        BLOCKS_POW2=triton.next_power_of_2(num_blocks),
        SUM_GROUP=16,
        COL_BLOCK=min(triton.next_power_of_2(stash_width), 1024),
    )