flashinfer_single_decode_attention_with_kv_cache_quickstart.py

1import torch
2import flashinfer
3
4# Set up dimensions
5batch_size = 128
6num_qo_heads = 32
7num_kv_heads = 32
8head_dim = 128
9seq_len = 1024
10
11# Create sample input tensors on CUDA
12q = torch.randn(batch_size, num_qo_heads, head_dim).half().to(0)
13k = torch.randn(batch_size, seq_len, num_kv_heads, head_dim).half().to(0)
14v = torch.randn(batch_size, seq_len, num_kv_heads, head_dim).half().to(0)
15
16# Run single-request decode attention
17# FlashInfer's single_decode_with_kv_cache handles the attention computation efficiently
18output = flashinfer.single_decode_with_kv_cache(q, k, v)
19
20print(f"Output shape: {output.shape}")