#026

reshape_and_cache

bf16 vllm · both · vllm._C_cache_ops.reshape_and_cache_flash · importance 4.4%

Reference Implementation

reference.py

import torch
import torch.nn as nn

class Model(nn.Module):

    def __init__(self) -> None:
        super().__init__()

    def forward(self, key: torch.Tensor, value: torch.Tensor, key_cache: torch.Tensor, value_cache: torch.Tensor, slot_mapping: torch.Tensor) -> dict[str, torch.Tensor]:
        key_out = key_cache.clone()
        value_out = value_cache.clone()
        valid = slot_mapping >= 0
        slots = slot_mapping[valid].long()
        key_out.view(-1, key.shape[1], key.shape[2])[slots] = key[valid].to(key_out.dtype)
        value_out.view(-1, value.shape[1], value.shape[2])[slots] = value[valid].to(value_out.dtype)
        return {'key_cache': key_out, 'value_cache': value_out}

Shapes

T_SOL hardware:

#	shape	block_size	T_SOL(XPU-A)	T_Prod	S
0	[1, 4, 128]	16	0.01 us	7.50 us	0.1%
1	[144, 2, 256]	16	0.17 us	7.70 us	2.2%
2	[248, 16, 128]	16	1.18 us	9.80 us	12.0%
3	[520, 16, 128]	16	2.44 us	14.20 us	17.2%
4	[1015, 16, 128]	16	4.74 us	21.20 us	22.4%
5	[2044, 16, 128]	16	9.49 us	29.20 us	32.5%
6	[4108, 4, 128]	16	4.77 us	33.70 us	14.2%
7	[8192, 4, 128]	16	9.51 us	60.30 us	15.8%

Input Generation

input.py

import torch

def _make_inputs(shape: list[int], block_size: int=16) -> dict[str, torch.Tensor]:
    (num_tokens, num_heads, head_size) = (int(dim) for dim in shape)
    block_size = int(block_size)
    num_blocks = (num_tokens + block_size - 1) // block_size
    key = torch.randn((num_tokens, num_heads, head_size), dtype=torch.bfloat16, device='cuda')
    value = torch.randn((num_tokens, num_heads, head_size), dtype=torch.bfloat16, device='cuda')
    key_cache = torch.zeros((num_blocks, block_size, num_heads, head_size), dtype=torch.bfloat16, device='cuda')
    value_cache = torch.zeros_like(key_cache)
    slot_mapping = torch.arange(num_tokens, dtype=torch.long, device='cuda')
    return {'key': key, 'value': value, 'key_cache': key_cache, 'value_cache': value_cache, 'slot_mapping': slot_mapping}