vllm.models.step5.nvidia.ops.sparse_indexer_scores

FP8 indexer scores directly from paged summary keys.

Functions:

• decode_summary_scores –

  Read raw FP8 summaries without repacking keys or allocating unit scales.

• paged_mqa_logits_sm90 –

  Score Step5's packed FP8 summary pages on Hopper.

decode_summary_scores(query, weights, mean, block_table, slots, seq_lens, capacity, *, num_sms, row_groups=1, query_len=1)

Read raw FP8 summaries without repacking keys or allocating unit scales.

Only the 128-region rounded live prefix is written. Consumers must use lengths/topk_lengths to ignore the unspecified score tail. Each group of row_groups query rows is one token of slots; each query_len tokens are one request of block_table/seq_lens.

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

def decode_summary_scores(
    query: torch.Tensor,
    weights: torch.Tensor,
    mean: torch.Tensor,
    block_table: torch.Tensor,
    slots: torch.Tensor,
    seq_lens: torch.Tensor,
    capacity: int,
    *,
    num_sms: int,
    row_groups: int = 1,
    query_len: int = 1,
) -> tuple[torch.Tensor, DecodeIndexerMetadata]:
    """Read raw FP8 summaries without repacking keys or allocating unit scales.

    Only the 128-region rounded live prefix is written. Consumers must use
    lengths/topk_lengths to ignore the unspecified score tail. Each group of
    ``row_groups`` query rows is one token of ``slots``; each ``query_len``
    tokens are one request of ``block_table``/``seq_lens``.
    """
    # query: [tokens, groups, 8, 256] and weights: [tokens, groups, 8], each
    # group's heads contiguous; groups may be the outer storage dimension.
    tokens, groups = query.shape[:2]
    assert groups == row_groups and query.stride(3) == 1 and query.stride(2) == 256
    assert weights.shape == query.shape[:3] and weights.stride(2) == 1
    assert mean.stride(1) == 256 and mean.stride(2) == 1
    batch = tokens * groups
    scores = torch.empty(
        (batch, triton.cdiv(capacity, 256) * 256),
        device=query.device,
        dtype=torch.float32,
    )
    active = torch.empty(batch, device=query.device, dtype=torch.bool)
    positions = torch.empty(batch, device=query.device, dtype=torch.int64)
    current = torch.empty_like(positions)
    lengths = torch.empty(batch, device=query.device, dtype=torch.int32)
    topk_lengths = torch.empty_like(lengths)
    # Bound idle CTAs for short contexts; each worker walks its live tiles.
    workers = min(triton.cdiv(capacity, 128), triton.cdiv(2 * num_sms, batch))
    _direct_scores[(workers, batch)](
        query,
        weights,
        mean,
        block_table,
        slots,
        seq_lens,
        scores,
        active,
        positions,
        current,
        lengths,
        topk_lengths,
        query.stride(0),
        query.stride(1),
        weights.stride(0),
        weights.stride(1),
        block_table.shape[1],
        block_table.stride(0),
        slots.stride(0),
        seq_lens.stride(0),
        mean.shape[0],
        mean.shape[1],
        mean.stride(0),
        capacity,
        scores.stride(0),
        8,
        128,
        workers,
        row_groups,
        query_len,
        num_warps=4,
        enable_fp_fusion=False,
    )
    metadata = DecodeIndexerMetadata(
        active, positions, current, lengths, block_table, topk_lengths
    )
    return scores[:, :capacity], metadata

paged_mqa_logits_sm90(q, kv_fused, weights, block_tables, seq_lens, max_seq_len, *, schedule_meta=None)

Score Step5's packed FP8 summary pages on Hopper.

The packed page holds all FP8 keys followed by one FP32 scale per key. Query pairs share a page table and length, as in FlashInfer's interface.

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

def paged_mqa_logits_sm90(
    q: torch.Tensor,
    kv_fused: torch.Tensor,
    weights: torch.Tensor,
    block_tables: torch.Tensor,
    seq_lens: torch.Tensor,
    max_seq_len: int,
    *,
    schedule_meta: torch.Tensor | None = None,
) -> torch.Tensor:
    """Score Step5's packed FP8 summary pages on Hopper.

    The packed page holds all FP8 keys followed by one FP32 scale per key.
    Query pairs share a page table and length, as in FlashInfer's interface.
    """
    batch, next_n, heads, dim = q.shape
    assert q.dtype == torch.float8_e4m3fn and dim == 256
    assert heads in (4, 8, 16) and q.stride(-1) == 1
    assert kv_fused.dtype == torch.uint8 and kv_fused.is_contiguous()
    assert kv_fused.shape[2:] == (1, 260)
    assert weights.shape == (batch * next_n, heads)
    output = torch.empty(
        (batch * next_n, max_seq_len), dtype=torch.float32, device=q.device
    )
    if batch and max_seq_len:
        _paged_scores_sm90[(batch * next_n, triton.cdiv(max_seq_len, 32))](
            q,
            kv_fused,
            weights,
            seq_lens,
            block_tables,
            output,
            *q.stride()[:3],
            *weights.stride(),
            kv_fused.stride(0),
            *block_tables.stride(),
            seq_lens.stride(0),
            max_seq_len,
            next_n,
            heads,
            kv_fused.shape[1],
            32,
            num_warps=4,
        )
    return output