vllm.models.step5.nvidia.ops.prefill_indexer_metadata

Step-local indexer metadata shared across sparse attention layers.

Functions:

• prefill_indexer_chunk_size –

  Bound FP32 logits to 128 MiB for the model's 1M-token context.

• prepare_prefill_indexer –

  Prepare bounded FI query chunks once; reuse them for every DSA layer.

prefill_indexer_chunk_size(regions)

Bound FP32 logits to 128 MiB for the model's 1M-token context.

Each chunk costs a few host launches per DSA layer, so chunks are as large as the bound allows.

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

def prefill_indexer_chunk_size(regions: int) -> int:
    """Bound FP32 logits to 128 MiB for the model's 1M-token context.

    Each chunk costs a few host launches per DSA layer, so chunks are as large
    as the bound allows.
    """
    padded_regions = triton.cdiv(max(1, regions), 256) * 256
    rows = max(1, (128 << 20) // (4 * padded_regions))
    return max(256, min(4096, 1 << (rows.bit_length() - 1)))

prepare_prefill_indexer(owners, positions, regions, topk, num_sms, *, paged_prefill=None)

Prepare bounded FI query chunks once; reuse them for every DSA layer.

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

def prepare_prefill_indexer(
    owners: torch.Tensor,
    positions: torch.Tensor,
    regions: int,
    topk: int,
    num_sms: int,
    *,
    paged_prefill: PagedPrefillMetadata | None = None,
) -> list[PrefillIndexerMetadata]:
    """Prepare bounded FI query chunks once; reuse them for every DSA layer."""
    rows = positions.numel()
    chunk_size = prefill_indexer_chunk_size(regions)
    need_scores = regions > topk
    paged_kwargs = {}
    if paged_prefill is not None:
        if need_scores or paged_prefill.block_size not in (16, 32, 64, 128):
            raise ValueError("Shared pages require unscored prefill")
        paged_kwargs = dict(
            PREPARE_PAGES=True,
            PageOwners=paged_prefill.owners,
            PageTable=paged_prefill.block_table,
            PageIndices=paged_prefill.indices,
            PageIndptr=paged_prefill.indptr,
            PAGE_OWNER_STRIDE=paged_prefill.owners.stride(0),
            PAGE_TABLE_ROW=paged_prefill.block_table.stride(0),
            PAGE_TABLE_COL=paged_prefill.block_table.stride(1),
            PAGE_INDPTR_STRIDE=paged_prefill.indptr.stride(0),
            BLOCK_SIZE=paged_prefill.block_size,
            PAGE_STRIDE=paged_prefill.page_stride,
        )
    pages = triton.cdiv(regions, 128) * 4 if need_scores else 0
    current = torch.empty(rows, device=positions.device, dtype=torch.int64)
    lengths = torch.empty(rows, device=positions.device, dtype=torch.int32)
    topk_lengths = torch.empty_like(lengths)
    table = torch.empty((rows, pages), device=positions.device, dtype=torch.int32)
    schedules = (
        torch.empty(
            (triton.cdiv(rows, chunk_size), num_sms + 1, 2),
            device=positions.device,
            dtype=torch.int32,
        )
        if need_scores
        else None
    )
    selected = (
        torch.empty((rows, regions + 1), device=positions.device, dtype=torch.int64)
        if not need_scores
        else None
    )
    if rows:
        _prepare_prefill_indexer[(max(1, triton.cdiv(pages, 256)), rows)](
            owners,
            positions,
            current,
            lengths,
            topk_lengths,
            table,
            schedules,
            selected,
            rows,
            pages,
            owners.stride(0),
            positions.stride(0),
            num_sms,
            triton.next_power_of_2(num_sms + 1),
            need_scores,
            regions if not need_scores else 0,
            # Sized by topk, not the prompt: every short prompt reuses one variant.
            triton.next_power_of_2(topk + 1) if not need_scores else 1,
            chunk_size,
            chunk_size.bit_length(),
            num_warps=4,
            **paged_kwargs,
        )
    return [
        PrefillIndexerMetadata(
            current[start : start + chunk_size],
            lengths[start : start + chunk_size],
            topk_lengths[start : start + chunk_size],
            table[start : start + chunk_size],
            schedules[start // chunk_size] if schedules is not None else None,
            selected[start : start + chunk_size] if selected is not None else None,
        )
        for start in range(0, rows, chunk_size)
    ]