Skip to content

vllm.distributed.kv_transfer.kv_checksum.kernels

Triton kernel for the KV block checksums (see checksum.py) of several layers of one layout, in one launch.

Functions:

  • layers_checksum –

    Add each layer's weighted piece sum of each block row into its column

layers_checksum(pieces, layers, block_ids, first_tokens, valid_tokens, out, aligned)

Add each layer's weighted piece sum of each block row into its column of out.

Parameters:

  • pieces

    (Tensor) –

    The first layer's cache as integer pieces [num_blocks, kernel_blocks, heads, tokens, pieces], read as unsigned. Every layer's cache has its shape, strides and dtype.

  • layers

    (Tensor) –

    [num_layers, 3] int64 on the device: each layer's cache address, layer * H + head of its first head, and its column of out.

  • block_ids

    (Tensor) –

    [n] int64 logical block ids.

  • first_tokens

    (Tensor) –

    [n] int64 position in the request of each block's first token.

  • valid_tokens

    (Tensor) –

    [n] int64 tokens to include from each block's start.

  • out

    (Tensor) –

    [n, num_columns] int64 accumulator with unit column stride.

  • aligned

    (bool) –

    Every address in layers is a multiple of 16 bytes.

Source code in vllm/distributed/kv_transfer/kv_checksum/kernels.py 
 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
144
145
146
def layers_checksum(
    pieces: torch.Tensor,
    layers: torch.Tensor,
    block_ids: torch.Tensor,
    first_tokens: torch.Tensor,
    valid_tokens: torch.Tensor,
    out: torch.Tensor,
    aligned: bool,
) -> None:
    """Add each layer's weighted piece sum of each block row into its column
    of ``out``.

    Args:
        pieces: The first layer's cache as integer pieces
            ``[num_blocks, kernel_blocks, heads, tokens, pieces]``, read as
            unsigned. Every layer's cache has its shape, strides and dtype.
        layers: ``[num_layers, 3]`` int64 on the device: each layer's cache
            address, ``layer * H + head`` of its first head, and its column
            of ``out``.
        block_ids: ``[n]`` int64 logical block ids.
        first_tokens: ``[n]`` int64 position in the request of each block's
            first token.
        valid_tokens: ``[n]`` int64 tokens to include from each block's start.
        out: ``[n, num_columns]`` int64 accumulator with unit column stride.
        aligned: Every address in ``layers`` is a multiple of 16 bytes.

    """
    _, kernel_blocks, heads, tokens, num_pieces = pieces.shape
    num_pieces_in_block = kernel_blocks * heads * tokens * num_pieces
    grid = (block_ids.shape[0], triton.cdiv(num_pieces_in_block, _CHUNK), len(layers))
    _layers_checksum_kernel[grid](
        pieces,
        *pieces.stride(),
        layers,
        block_ids,
        first_tokens,
        valid_tokens,
        out,
        out.stride(0),
        kernel_blocks,
        heads,
        tokens,
        num_pieces,
        PIECE_BITS=8 * pieces.element_size(),
        CHUNK=_CHUNK,
        BLOCK=min(triton.next_power_of_2(num_pieces_in_block), 1024),
        SINGLE_LAYER=len(layers) == 1,
        ALIGNED=aligned,
    )