Skip to content

vllm.v1.attention.ops.prefix_prefill

_paged_kv_cache_offsets(B_Loc, cur_batch, token_indices, token_valid, offs_d, cur_kv_head, x, stride_b_loc_b, stride_b_loc_s, stride_k_cache_bs, stride_k_cache_h, stride_k_cache_d, stride_k_cache_bl, stride_k_cache_x, stride_v_cache_bs, stride_v_cache_h, stride_v_cache_d, stride_v_cache_bl, PHYSICAL_BLOCK_SIZE, BLOCK_DMODEL, INTERLEAVED_V_PACK_FACTOR=0, MASK_BLOCK_TABLE=False)

Compute paged K/V cache element offsets for a tile of token positions.

token_indices holds the absolute sequence positions of the tile. The logical block for each token is looked up in B_Loc (handling cross-block tiles), then the physical element offsets are built with the same layout the kernel uses everywhere: K cache: [num_blocks, num_kv_heads, head_size // x, block_size, x] V cache: [num_blocks, num_kv_heads, head_size, block_size], or the interleaved layout [num_blocks, num_kv_heads, block_size // x, head_size, x] when INTERLEAVED_V_PACK_FACTOR (= x) > 0. Returns (off_k, off_v) with shapes [D, N] and [N, D] respectively.

When MASK_BLOCK_TABLE is set, the block-table load is masked with token_valid so padded lanes (a tile that extends past the end of the sequence) cannot read a block-table entry past this batch's row. Invalid lanes resolve to block 0 (always in-bounds) and are dropped by the caller's later K/V load mask. When it is not set, token_valid is unused and the load is emitted exactly as before.

Source code in vllm/v1/attention/ops/prefix_prefill.py 
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 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
 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
@triton.jit
def _paged_kv_cache_offsets(
    B_Loc,
    cur_batch,
    token_indices,
    token_valid,
    offs_d,
    cur_kv_head,
    x,
    stride_b_loc_b,
    stride_b_loc_s,
    stride_k_cache_bs,
    stride_k_cache_h,
    stride_k_cache_d,
    stride_k_cache_bl,
    stride_k_cache_x,
    stride_v_cache_bs,
    stride_v_cache_h,
    stride_v_cache_d,
    stride_v_cache_bl,
    PHYSICAL_BLOCK_SIZE: tl.constexpr,
    BLOCK_DMODEL: tl.constexpr,
    INTERLEAVED_V_PACK_FACTOR: tl.constexpr = 0,
    MASK_BLOCK_TABLE: tl.constexpr = False,
):
    """Compute paged K/V cache element offsets for a tile of token positions.

    `token_indices` holds the absolute sequence positions of the tile. The
    logical block for each token is looked up in `B_Loc` (handling cross-block
    tiles), then the physical element offsets are built with the same layout the
    kernel uses everywhere:
      K cache: [num_blocks, num_kv_heads, head_size // x, block_size, x]
      V cache: [num_blocks, num_kv_heads, head_size, block_size], or the
        interleaved layout [num_blocks, num_kv_heads, block_size // x,
        head_size, x] when `INTERLEAVED_V_PACK_FACTOR` (= x) > 0.
    Returns `(off_k, off_v)` with shapes [D, N] and [N, D] respectively.

    When `MASK_BLOCK_TABLE` is set, the block-table load is masked with
    `token_valid` so padded lanes (a tile that extends past the end of the
    sequence) cannot read a block-table entry past this batch's row. Invalid
    lanes resolve to block 0 (always in-bounds) and are dropped by the caller's
    later K/V load mask. When it is not set, `token_valid` is unused and the
    load is emitted exactly as before.
    """
    bn_logical = token_indices // PHYSICAL_BLOCK_SIZE
    if MASK_BLOCK_TABLE:
        bn = tl.load(
            B_Loc + cur_batch * stride_b_loc_b + bn_logical * stride_b_loc_s,
            mask=token_valid,
            other=0,
        ).to(tl.int64)
    else:
        bn = tl.load(
            B_Loc + cur_batch * stride_b_loc_b + bn_logical * stride_b_loc_s
        ).to(tl.int64)
    internal = token_indices % PHYSICAL_BLOCK_SIZE
    off_k = (
        bn[None, :] * stride_k_cache_bs
        + cur_kv_head * stride_k_cache_h
        + (offs_d[:, None] // x) * stride_k_cache_d
        + internal[None, :] * stride_k_cache_bl
        + (offs_d[:, None] % x) * stride_k_cache_x
    )
    if INTERLEAVED_V_PACK_FACTOR > 0:
        # Interleaved V-cache: [num_blocks, num_kv_heads, block_size/x,
        # head_size, x] flattened onto the 4D stride layout.
        off_v = (
            bn[:, None] * stride_v_cache_bs
            + cur_kv_head * stride_v_cache_h
            + (internal[:, None] // INTERLEAVED_V_PACK_FACTOR)
            * BLOCK_DMODEL
            * INTERLEAVED_V_PACK_FACTOR
            + offs_d[None, :] * INTERLEAVED_V_PACK_FACTOR
            + (internal[:, None] % INTERLEAVED_V_PACK_FACTOR)
        )
    else:
        off_v = (
            bn[:, None] * stride_v_cache_bs
            + cur_kv_head * stride_v_cache_h
            + offs_d[None, :] * stride_v_cache_d
            + internal[:, None] * stride_v_cache_bl
        )
    return off_k, off_v