Skip to content

vLLM

summary_update

`vllm.models.step5.nvidia.ops.summary_update` ¶

CSA region summaries, compressed once when a region's eighth row arrives.

A region summary is an online softmax over the region's eight index-key rows, accumulated in position order. Only closed regions are ever scored, so the raw [key | z] rows of each request's open region wait in a circular buffer (Step5CSATailRing) and the summary is written once, by the row that closes the region. A rejected draft that closed a region leaves a summary nobody reads; it is rewritten when the region closes again.

Functions:

build_prefill_groups –

Summary update metadata on device, without host slots or a sync.
summary_layout –

(bytes between blocks, tokens per block) of a [B, R, 256] cache.
update_prefill –

Summarize the regions this step closes, then store open-region rows.
update_verify –

Uniform MTP verify update; graph capturable, no host metadata.

`_replay_ring(Ring, base, first, last, dim, RING_CAPACITY)` ¶

Accumulate the ring rows at positions [first, last) in order.

Source code in vllm/models/step5/nvidia/ops/summary_update.py

@triton.jit
def _replay_ring(Ring, base, first, last, dim, RING_CAPACITY: tl.constexpr):
    """Accumulate the ring rows at positions ``[first, last)`` in order."""
    numerator = tl.zeros((256,), tl.float32)
    denominator = tl.zeros((256,), tl.float32)
    maximum = tl.max(tl.full((256,), -float("inf"), tl.float32), 0)
    for position in range(first, last):
        row = Ring + (base + position % RING_CAPACITY) * 512
        key = tl.load(row + dim).to(tl.float32)
        z = tl.load(row + 256 + dim).to(tl.float32)
        numerator, denominator, maximum = _accumulate(
            numerator, denominator, maximum, key, z
        )
    return numerator, denominator, maximum

`_store_summary(Mean, slot, dim, numerator, denominator, MEAN_BLOCK, BLOCK_SIZE)` ¶

Write the summary of the region holding KV slot; blocks of Mean are MEAN_BLOCK bytes apart.

Source code in vllm/models/step5/nvidia/ops/summary_update.py

@triton.jit
def _store_summary(
    Mean,
    slot,
    dim,
    numerator,
    denominator,
    MEAN_BLOCK: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
):
    """Write the summary of the region holding KV ``slot``; blocks of ``Mean``
    are ``MEAN_BLOCK`` bytes apart."""
    summary = tl.div_rn(numerator, _maximum(denominator, 1e-20)).to(tl.bfloat16)
    summary_fp8 = summary.to(tl.float8e4nv).to(tl.uint8, bitcast=True)
    address = slot // BLOCK_SIZE * MEAN_BLOCK + slot % BLOCK_SIZE // 8 * 256
    tl.store(Mean + address + dim, summary_fp8)

`_update_decode_values(slot, position, base, dim, key, z, Mean, Ring, RING_CAPACITY, MEAN_BLOCK, BLOCK_SIZE)` ¶

Store an open-region row, or summarize the region this row closes.

Source code in vllm/models/step5/nvidia/ops/summary_update.py

@triton.jit
def _update_decode_values(
    slot,
    position,
    base,
    dim,
    key,
    z,
    Mean,
    Ring,
    RING_CAPACITY: tl.constexpr,
    MEAN_BLOCK: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
):
    """Store an open-region row, or summarize the region this row closes."""
    if (slot >= 0) & (base >= 0):
        if position % 8 == 7:
            numerator, denominator, maximum = _replay_ring(
                Ring, base, position - 7, position, dim, RING_CAPACITY
            )
            numerator, denominator, maximum = _accumulate(
                numerator, denominator, maximum, key, z
            )
            _store_summary(
                Mean, slot, dim, numerator, denominator, MEAN_BLOCK, BLOCK_SIZE
            )
        else:
            row = Ring + (base + position % RING_CAPACITY) * 512
            tl.store(row + dim, key.to(tl.bfloat16))
            tl.store(row + 256 + dim, z.to(tl.bfloat16))

`build_prefill_groups(slots, positions, owners, query_start_loc)` ¶

Summary update metadata on device, without host slots or a sync.

Returns [rows, 10] metadata [region, first offset, rows[8]]; non-head rows have region -1.

Source code in vllm/models/step5/nvidia/ops/summary_update.py

def build_prefill_groups(slots, positions, owners, query_start_loc) -> torch.Tensor:
    """Summary update metadata on device, without host slots or a sync.

    Returns ``[rows, 10]`` metadata ``[region, first offset, rows[8]]``;
    non-head rows have region -1.
    """
    rows = slots.shape[0]
    metadata = torch.empty((rows, 10), dtype=torch.int64, device=slots.device)
    if rows:
        _build_prefill_groups[(rows,)](
            slots,
            positions,
            owners,
            query_start_loc,
            metadata,
            positions.stride(0),
        )
    return metadata

`summary_layout(mean)` ¶

(bytes between blocks, tokens per block) of a [B, R, 256] cache.

Source code in vllm/models/step5/nvidia/ops/summary_update.py

def summary_layout(mean: torch.Tensor) -> tuple[int, int]:
    """(bytes between blocks, tokens per block) of a ``[B, R, 256]`` cache."""
    assert mean.stride(1) == 256 and mean.stride(2) == 1
    return mean.stride(0), mean.shape[1] * 8

`update_prefill(key, z, metadata, positions, ring_base, ring_slots, mean, ring, capacity)` ¶

Summarize the regions this step closes, then store open-region rows.

Every ring read of the step precedes its writes, so a long chunk cannot overwrite open-region rows it still has to read.

Source code in vllm/models/step5/nvidia/ops/summary_update.py

def update_prefill(
    key, z, metadata, positions, ring_base, ring_slots, mean, ring, capacity
):
    """Summarize the regions this step closes, then store open-region rows.

    Every ring read of the step precedes its writes, so a long chunk cannot
    overwrite open-region rows it still has to read.
    """
    rows = metadata.shape[0]
    if not rows:
        return
    strides = (key.stride(0), key.stride(2), z.stride(0), z.stride(2))
    _update_prefill[(rows,)](
        key,
        z,
        metadata,
        positions,
        ring_base,
        mean,
        ring,
        positions.stride(0),
        *strides,
        capacity,
        *summary_layout(mean),
        num_warps=4,
        enable_fp_fusion=False,
    )
    _write_ring[(ring_slots.shape[0],)](key, z, ring_slots, ring, *strides, num_warps=4)

`update_verify(key, z, slots, query_start_loc, seq_lens, ring_base, mean, ring, capacity)` ¶

Uniform MTP verify update; graph capturable, no host metadata.

Source code in vllm/models/step5/nvidia/ops/summary_update.py

def update_verify(
    key, z, slots, query_start_loc, seq_lens, ring_base, mean, ring, capacity
):
    """Uniform MTP verify update; graph capturable, no host metadata."""
    requests = seq_lens.shape[0]
    if requests:
        _update_verify[(requests,)](
            key,
            z,
            slots,
            query_start_loc,
            seq_lens,
            ring_base,
            mean,
            ring,
            key.stride(0),
            key.stride(2),
            z.stride(0),
            z.stride(2),
            capacity,
            *summary_layout(mean),
            num_warps=4,
            enable_fp_fusion=False,
        )