vllm.models.kimi_k3.amd.ops.fp4_indexer_prefill

FP4 prefill: gather compact index keys once, then score global causal rows.

Functions:

• prefill_chunks –

  Reuse each request set's keys and exchange all query chunks together.

• prefill_workspace_bytes –

  Peak prefill scratch of one chunk holding at most max_keys keys.

• prepare_prefill_metadata –

  Build read-only geometry once for the layers sharing this attention group.

• stage_keys –

  Preserve quantized bytes and swizzled scale rows, including merged strides.

prefill_chunks(q, weights, cache, out, metadata, world, interleave, weight_scale)

Reuse each request set's keys and exchange all query chunks together.

Source code in vllm/models/kimi_k3/amd/ops/fp4_indexer_prefill.py

def prefill_chunks(q, weights, cache, out, metadata, world, interleave, weight_scale):
    """Reuse each request set's keys and exchange all query chunks together."""
    gathered_keys = None
    pending: list[tuple[int, int, torch.Tensor]] = []
    for chunk, plan in zip(metadata.chunks, metadata.plans, strict=True):
        cache_key = (chunk.block_table.data_ptr(), chunk.num_reqs, chunk.total_seq_lens)
        if gathered_keys is None or gathered_keys[0] != cache_key:
            # Release the previous request set's pages before staging the next.
            gathered_keys = None
            gathered_keys = (
                cache_key,
                stage_keys(cache, chunk, world, interleave, plan.all_cu),
            )
        prefill_indexer(
            q,
            weights,
            cache,
            out,
            chunk,
            world,
            interleave,
            weight_scale,
            staged_keys=gathered_keys[1],
            pending=pending,
            plan=plan,
        )
    if not pending:
        return
    mine = (
        torch.cat([rows for _, _, rows in pending])
        if len(pending) > 1
        else pending[0][2]
    )
    selected = get_dcp_group().all_gather(mine, dim=0) if world > 1 else mine
    selected = selected.view(world, mine.shape[0], -1)
    offset = 0
    for start, end, rows in pending:
        count = rows.shape[0]
        out[start:end].copy_(
            selected[:, offset : offset + count].reshape(-1, out.shape[1])[
                : end - start
            ]
        )
        offset += count

prefill_workspace_bytes(max_keys, max_reqs, world, interleave)

Peak prefill scratch of one chunk holding at most max_keys keys.

The indexer builder caps each chunk's keys at its prefill buffer size and _prepare_chunk keeps each logits tile under 1 GiB. Staging holds the local packed keys, their DCP all-gather and the restored pages at once; a rank's share exceeds an even split by at most one interleave run per request.

Source code in vllm/models/kimi_k3/amd/ops/fp4_indexer_prefill.py

def prefill_workspace_bytes(max_keys, max_reqs, world, interleave):
    """Peak prefill scratch of one chunk holding at most `max_keys` keys.

    The indexer builder caps each chunk's keys at its prefill buffer size and
    _prepare_chunk keeps each logits tile under 1 GiB. Staging holds the local
    packed keys, their DCP all-gather and the restored pages at once; a rank's
    share exceeds an even split by at most one interleave run per request.
    """
    pages = triton.cdiv(max_keys, 64) * 64
    if world == 1:
        return 1024**3 + 68 * (max_keys + pages)
    local = triton.cdiv(max_keys, world) + max_reqs * interleave
    return 1024**3 + 68 * ((world + 1) * local + pages)

prepare_prefill_metadata(metadata, world, rank, interleave)

Build read-only geometry once for the layers sharing this attention group.

Source code in vllm/models/kimi_k3/amd/ops/fp4_indexer_prefill.py

def prepare_prefill_metadata(metadata, world, rank, interleave):
    """Build read-only geometry once for the layers sharing this attention group."""
    return FP4PrefillMetadata(
        chunks=metadata.chunks,
        max_prefill_seq_len=metadata.max_prefill_seq_len,
        shardq_gather_indices=metadata.shardq_gather_indices,
        plans=[
            _prepare_chunk(chunk, world, rank, interleave) for chunk in metadata.chunks
        ],
    )

stage_keys(cache, chunk, world, interleave, all_cu=None)

Preserve quantized bytes and swizzled scale rows, including merged strides.

Source code in vllm/models/kimi_k3/amd/ops/fp4_indexer_prefill.py

def stage_keys(cache, chunk, world, interleave, all_cu=None):
    """Preserve quantized bytes and swizzled scale rows, including merged strides."""
    device = cache.device
    cu = chunk.cu_seq_lens
    if all_cu is None:
        all_cu = _local_key_cu(chunk, world, interleave)
    max_local = chunk.max_local_total_seq_lens if world > 1 else chunk.total_seq_lens
    packed = torch.empty((max_local, 68), device=device, dtype=torch.uint8)
    local_width = (
        (chunk.max_seqlen_k + world * interleave - 1)
        // (world * interleave)
        * interleave
    )
    _pack_keys[(chunk.num_reqs, triton.cdiv(local_width, 32))](
        cache,
        chunk.block_table,
        chunk.local_cu_seq_lens if world > 1 else chunk.cu_seq_lens,
        packed,
        cache.stride(0),
        chunk.block_table.stride(0),
        32,
    )
    gathered = get_dcp_group().all_gather(packed, dim=0) if world > 1 else packed
    pages = triton.cdiv(chunk.total_seq_lens, 64)
    payload = torch.zeros((pages, 1, 4, 64, 16), device=device, dtype=torch.uint8)
    scales = torch.zeros((pages, 1, 4, 64), device=device, dtype=torch.uint8)
    _stage_keys[(triton.cdiv(chunk.total_seq_lens, 32),)](
        gathered,
        chunk.token_to_seq,
        cu,
        all_cu,
        payload,
        scales,
        chunk.total_seq_lens,
        max_local,
        world,
        interleave,
        32,
    )
    return payload, scales