Skip to content

vllm.benchmarks.forward.kv_init

KV blocks of a benchmark case: initialization and the blocks a step writes.

Only the case's blocks are touched, group by group: groups can alias one allocation with different layouts, and the batch builder keeps their block IDs disjoint. Zero KV times within a few percent of random KV, even for sparse attention (DeepSeek-V4.1, ~100K context), so zeros are the only init for now.

Functions:

  • touched_kv_views –

    Views a step may write: attention tails and all recurrent state slots.

_block_dim(tensor, num_blocks)

Return (block dim, kernel blocks per allocator block).

Source code in vllm/benchmarks/forward/kv_init.py 
32
33
34
35
36
37
def _block_dim(tensor: torch.Tensor, num_blocks: int) -> tuple[int, int]:
    """Return (block dim, kernel blocks per allocator block)."""
    for dim in range(min(2, tensor.dim())):
        if tensor.shape[dim] >= num_blocks and tensor.shape[dim] % num_blocks == 0:
            return dim, tensor.shape[dim] // num_blocks
    raise ValueError(f"No block dim of size {num_blocks} in {tuple(tensor.shape)}")

touched_kv_views(kv_caches, attn_groups, batch, num_blocks, dcp_size, lookahead=0)

Views a step may write: attention tails and all recurrent state slots.

Source code in vllm/benchmarks/forward/kv_init.py 
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
def touched_kv_views(
    kv_caches: dict[str, Any],
    attn_groups: list["AttentionGroup"],
    batch: BenchBatch,
    num_blocks: int,
    dcp_size: int,
    lookahead: int = 0,
) -> list[torch.Tensor]:
    """Views a step may write: attention tails and all recurrent state slots."""
    scheduled_tokens = [
        batch.expected[req_id][0] + lookahead for req_id in batch.req_ids
    ]
    views = []
    for group in attn_groups:
        ranges = batch.group_block_ranges[group.kv_cache_group_id]
        block_size = kv_block_size(group.kv_cache_spec, dcp_size)
        for tensor in _group_tensors(kv_caches, group, num_blocks):
            dim, factor = _block_dim(tensor, num_blocks)
            for (begin, end), tokens in zip(ranges, scheduled_tokens):
                first = (
                    begin
                    if isinstance(group.kv_cache_spec, MambaSpec)
                    else max(begin, end - (-(-tokens // block_size) + 2))
                )
                views.append(tensor.narrow(dim, first * factor, (end - first) * factor))
    return views