flashinfer_paged_kv_cache_batch_decode_attention_quickstart.py

1import torch
2import flashinfer
3
4# Define dimensions
5num_heads = 32
6head_dim = 128
7num_pages = 1000
8page_size = 16
9batch_size = 8
10
11# Initialize data on GPU
12data_type = torch.float16
13device = torch.device("cuda:0")
14
15# Setup KV cache
16kv_cache = torch.randn(num_pages, 2, page_size, num_heads, head_dim, dtype=data_type, device=device)
17kv_page_indices = torch.arange(num_pages, dtype=torch.int32, device=device)
18kv_page_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device=device) * (num_pages // batch_size)
19kv_last_page_len = torch.full((batch_size,), page_size, dtype=torch.int32, device=device)
20
21# Setup Query
22q = torch.randn(batch_size, num_heads, head_dim, dtype=data_type, device=device)
23
24# Initialize the wrapper
25workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=device)
26wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(workspace_buffer, "NHD")
27
28# Plan and execute the attention kernel
29wrapper.plan(
30    kv_page_indptr,
31    kv_page_indices,
32    kv_last_page_len,
33    num_heads,
34    num_heads, # num_kv_heads
35    head_dim,
36    page_size,
37    data_type,
38)
39
40output = wrapper.run(q, kv_cache)
41
42print("Output shape:", output.shape)