47 lines
1.8 KiB
Python
47 lines
1.8 KiB
Python
#!/usr/bin/env python3
|
|
"""Probe ixformer paged attention KV cache shape requirements."""
|
|
import torch
|
|
import ixformer
|
|
|
|
num_heads = 4
|
|
num_kv_heads = 1
|
|
head_dim = 256
|
|
block_size = 16
|
|
num_blocks = 4
|
|
context_len = num_blocks * block_size
|
|
head_mapping = torch.zeros(num_heads, dtype=torch.int32, device="cuda")
|
|
scale = head_dim ** -0.5
|
|
query = torch.randn(1, num_heads, head_dim, device="cuda", dtype=torch.float16)
|
|
context_lens = torch.tensor([context_len], device="cuda", dtype=torch.int32)
|
|
block_tables = torch.arange(num_blocks, device="cuda", dtype=torch.int32).unsqueeze(0)
|
|
|
|
# Read the ixformer vllm source for the correct layout
|
|
import inspect
|
|
src_file = "/usr/local/corex/lib64/python3/dist-packages/ixformer/functions/vllm.py"
|
|
try:
|
|
with open(src_file) as f:
|
|
print(f"=== {src_file} ===")
|
|
print(f.read())
|
|
except:
|
|
print(f"Cannot read {src_file}")
|
|
|
|
# Try different 5D layouts
|
|
print("\n=== Testing 5D KV cache layouts ===")
|
|
for x in [1, 2, 4, 8, 16]:
|
|
if head_dim % x != 0:
|
|
continue
|
|
# Layout: (num_blocks, num_kv_heads, head_dim//x, block_size, x)
|
|
kc = torch.randn(num_blocks, num_kv_heads, head_dim // x, block_size, x,
|
|
device="cuda", dtype=torch.float16)
|
|
vc = torch.randn(num_blocks, num_kv_heads, head_dim // x, block_size, x,
|
|
device="cuda", dtype=torch.float16)
|
|
out = torch.empty(1, num_heads, head_dim, device="cuda", dtype=torch.float16)
|
|
try:
|
|
ixformer.vllm_single_query_cached_kv_attention(
|
|
out, query, kc, vc, head_mapping, scale,
|
|
block_tables, context_lens, block_size, context_len)
|
|
print(f" x={x:2d} shape={kc.shape}: OK nan={out.isnan().any().item()}")
|
|
except Exception as e:
|
|
err = str(e)[:80]
|
|
print(f" x={x:2d} shape={kc.shape}: {err}")
|