Skip to content

vllm.models.step5.nvidia.ops.sparse_prefill

Compact sparse-prefill KV gathering and reusable host planning metadata.

Functions:

  • gather_prefill_kv –

    Gather selected tokens in exactly the boolean-indexed reference order.

  • make_indptr –

    Compute exact compact lengths from the selector's cardinality contract.

  • make_prefill_chunks –

    Share CPU plans and one GPU prefix copy across all sparse layers.

gather_prefill_kv(cache, regions, block_table, owners, positions, cu_kv, total_tokens, max_length, *, output_dtype=None, k_scale=None, v_scale=None)

Gather selected tokens in exactly the boolean-indexed reference order.

Source code in vllm/models/step5/nvidia/ops/sparse_prefill.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
def gather_prefill_kv(
    cache,
    regions,
    block_table,
    owners,
    positions,
    cu_kv,
    total_tokens,
    max_length,
    *,
    output_dtype=None,
    k_scale=None,
    v_scale=None,
):
    """Gather selected tokens in exactly the boolean-indexed reference order."""
    fp8_cache = cache.dtype in (torch.uint8, torch.float8_e4m3fn)
    if fp8_cache:
        if output_dtype not in (torch.bfloat16, torch.float16):
            raise ValueError("FP8 gather requires a model output dtype")
        if k_scale is None or v_scale is None:
            raise ValueError("FP8 gather requires K/V scales")
        cache = cache.view(torch.float8_e4m3fn)
    elif cache.dtype not in (torch.bfloat16, torch.float16):
        raise ValueError("Unsupported sparse KV cache dtype")
    rows, count = regions.shape
    dim = cache.shape[-1] // 2
    key = torch.empty(
        (total_tokens, 1, dim), device=cache.device, dtype=output_dtype or cache.dtype
    )
    value = torch.empty_like(key)
    if rows and total_tokens:
        # Integer views make bit preservation explicit, including NaN payloads.
        _gather_prefill_kv[(triton.cdiv(max_length, 16), rows)](
            cache if fp8_cache else cache.view(torch.int16),
            regions,
            block_table,
            owners,
            positions,
            cu_kv,
            key if fp8_cache else key.view(torch.int16),
            value if fp8_cache else value.view(torch.int16),
            k_scale,
            v_scale,
            fp8_cache,
            cache.stride(0),
            cache.stride(2),
            cache.stride(3),
            *regions.stride(),
            count,
            *block_table.stride(),
            owners.stride(0),
            positions.stride(0),
            cache.shape[2],
            dim,
            16,
            triton.next_power_of_2(dim),
        )
    return key, value

make_indptr(positions, region_count)

Compute exact compact lengths from the selector's cardinality contract.

Source code in vllm/models/step5/nvidia/ops/sparse_prefill.py 
75
76
77
78
79
80
81
82
83
def make_indptr(positions: list[int], region_count: int) -> torch.Tensor:
    """Compute exact compact lengths from the selector's cardinality contract."""
    if region_count < 1 or any(position < 0 for position in positions):
        raise ValueError("Prefill needs nonnegative positions and a current region")
    result = [0]
    for position in positions:
        count = min(position // 8, region_count - 1) * 8 + position % 8 + 1
        result.append(result[-1] + count)
    return torch.tensor(result, dtype=torch.int32, device="cpu")

make_prefill_chunks(positions, topk, device)

Share CPU plans and one GPU prefix copy across all sparse layers.

Source code in vllm/models/step5/nvidia/ops/sparse_prefill.py 
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
def make_prefill_chunks(
    positions: list[int], topk: int, device: torch.device
) -> list[PrefillChunk]:
    """Share CPU plans and one GPU prefix copy across all sparse layers."""
    cu_k = make_indptr(positions, topk + 1)
    device_cu_k = cu_k.pin_memory().to(device=device, non_blocking=True)
    chunks = []
    for start in range(0, len(positions), 32):
        end = min(start + 32, len(positions))
        local_cu_k = cu_k[start : end + 1] - cu_k[start]
        chunks.append(
            PrefillChunk(
                torch.arange(end - start + 1, dtype=torch.int32, device="cpu"),
                local_cu_k,
                device_cu_k[start : end + 1],
                int(local_cu_k[-1]),
                int((local_cu_k[1:] - local_cu_k[:-1]).max()),
            )
        )
    return chunks