vllm.models.step5.nvidia.ops.sparse_indexer

Packing and stable region selection for the sparse indexer.

Functions:

• pack_decode_summary_pages –

  Pack summaries and emit decode metadata in the same launch.

• select_all_regions –

  Select all history and current without computing unused indexer scores.

• select_regions –

  Select stable top-k historical regions and append current.

pack_decode_summary_pages(mean, block_table, regions, slots, seq_lens, *, num_sms=0, live_only=False)

Pack summaries and emit decode metadata in the same launch.

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

def pack_decode_summary_pages(
    mean: torch.Tensor,
    block_table: torch.Tensor,
    regions: int,
    slots: torch.Tensor,
    seq_lens: torch.Tensor,
    *,
    num_sms: int = 0,
    live_only: bool = False,
) -> tuple[torch.Tensor, DecodeIndexerMetadata]:
    """Pack summaries and emit decode metadata in the same launch."""
    packed, _, metadata = _pack_summaries(
        mean, block_table, regions, slots, seq_lens, num_sms, live_only=live_only
    )
    assert metadata is not None
    return packed, metadata

select_all_regions(current, capacity)

Select all history and current without computing unused indexer scores.

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

def select_all_regions(current: torch.Tensor, capacity: int) -> torch.Tensor:
    """Select all history and current without computing unused indexer scores."""
    rows = current.shape[0]
    output = torch.empty((rows, capacity + 1), device=current.device, dtype=torch.long)
    if rows:
        _all_regions[(rows,)](
            current, output, capacity, triton.next_power_of_2(capacity + 1)
        )
    return output

select_regions(scores, current, topk, backend='auto', *, topk_lengths=None, paged_prefill=None, sort=True)

Select stable top-k historical regions and append current.

With sort=False selected regions keep top-k order, left-padded with -1.

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

def select_regions(
    scores: torch.Tensor,
    current: torch.Tensor,
    topk: int,
    backend: str = "auto",
    *,
    topk_lengths: torch.Tensor | None = None,
    paged_prefill: PagedPrefillMetadata | None = None,
    sort: bool = True,
) -> torch.Tensor:
    """Select stable top-k historical regions and append current.

    With ``sort=False`` selected regions keep top-k order, left-padded with -1.
    """
    rows, regions = scores.shape
    k = min(topk, regions)
    if paged_prefill is not None:
        if regions <= topk:
            raise ValueError("Fused prefill pages require scored prefill only")
        if paged_prefill.block_size not in (16, 32, 64, 128):
            raise ValueError("Unsupported prefill cache block size")
    if regions <= topk:
        return select_all_regions(current, k)
    output = torch.empty((rows, k + 1), device=scores.device, dtype=torch.long)
    if not rows:
        return output
    # Step5 requires exact FP32 ordering. The cooperative/persistent backends
    # can select the wrong cutoff for scores separated by one FP32 epsilon.
    if backend == "auto":
        backend = (
            "deep_select"
            if current_platform.is_device_capability_family(100)
            and is_deep_select_supported(scores, k)
            else "per_row"
        )
    if backend in ("cooperative", "persistent"):
        raise ValueError(
            "Step5 requires exact FP32 top-k; use auto, deep_select, or per_row."
        )
    # Only DeepSelect guarantees a deterministic unsorted selection order.
    sort = sort or backend != "deep_select"
    # DeepSelect needs 32-byte row alignment; per_row writes contiguous rows.
    stride = triton.cdiv(k, 8) * 8 if backend == "deep_select" else k
    indices = torch.empty((rows, stride), device=scores.device, dtype=torch.int32)[
        :, :k
    ]
    get_indexer_topk(backend)(
        scores,
        current.to(torch.int32).contiguous() if topk_lengths is None else topk_lengths,
        1,
        indices,
        k,
        regions,
    )
    scratch = torch.empty((rows, k + 1), device=scores.device, dtype=torch.int32)
    paged_kwargs = {}
    if paged_prefill is not None:
        paged_kwargs = dict(
            PREPARE_PREFILL=True,
            Owners=paged_prefill.owners,
            Table=paged_prefill.block_table,
            Pages=paged_prefill.indices,
            Indptr=paged_prefill.indptr,
            OWNER_STRIDE=paged_prefill.owners.stride(0),
            TABLE_ROW=paged_prefill.block_table.stride(0),
            TABLE_COL=paged_prefill.block_table.stride(1),
            INDPTR_STRIDE=paged_prefill.indptr.stride(0),
            BLOCK_SIZE=paged_prefill.block_size,
            PAGE_STRIDE=paged_prefill.page_stride,
        )
    _stable_topk_regions[(rows,)](
        scores,
        current,
        indices,
        scratch,
        output,
        scores.stride(0),
        indices.stride(0),
        k,
        triton.next_power_of_2(k),
        1024,
        triton.next_power_of_2(k + 1),
        **paged_kwargs,
        SORT=sort,
    )
    return output