#024

paged_attention_decode

bf16 rtp-llm · decode · aiter.paged_attention_rocm · importance 4.4%

Reference Implementation

reference.py

import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):

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

    def forward(self, query: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, block_tables: torch.Tensor, seq_lens: torch.Tensor) -> torch.Tensor:
        (num_seqs, num_query_heads, head_size) = query.shape
        num_kv_heads = v_cache.shape[1]
        block_size = v_cache.shape[3]
        num_queries_per_kv = num_query_heads // num_kv_heads
        scale = head_size ** (-0.5)
        k_cache_flat = k_cache.permute(0, 3, 1, 2, 4).contiguous().view(-1, num_kv_heads, head_size)
        v_cache_flat = v_cache.permute(0, 3, 1, 2).contiguous().view(-1, num_kv_heads, head_size)
        seq_lens_list = seq_lens.tolist()
        block_tables_list = block_tables.tolist()
        output = torch.zeros_like(query)
        for seq_idx in range(num_seqs):
            ctx_len = int(seq_lens_list[seq_idx])
            if ctx_len <= 0:
                continue
            blocks = block_tables_list[seq_idx]
            blocks_t = torch.as_tensor(blocks[:(ctx_len + block_size - 1) // block_size], dtype=torch.long, device=query.device)
            within = torch.arange(ctx_len, device=query.device, dtype=torch.long)
            idx = blocks_t[within // block_size] * block_size + within % block_size
            keys = k_cache_flat.index_select(0, idx).float()
            values = v_cache_flat.index_select(0, idx).float()
            if num_queries_per_kv > 1:
                keys = keys.repeat_interleave(num_queries_per_kv, dim=1)
                values = values.repeat_interleave(num_queries_per_kv, dim=1)
            q_seq = query[seq_idx].float().unsqueeze(0).unsqueeze(2)
            ks = keys.transpose(0, 1).unsqueeze(0)
            vs = values.transpose(0, 1).unsqueeze(0)
            os = F.scaled_dot_product_attention(q_seq, ks, vs, is_causal=False)
            output[seq_idx] = os.squeeze(0).squeeze(1).to(query.dtype)
        return output

Shapes

T_SOL hardware:

#	num_seqs	ctx_len	num_query_heads	num_kv_heads	head_size	block_size	dtype	T_SOL(XPU-A)	T_Prod	S
0	256	64	16	16	128	16	bf16	26.12 us	152.60 us	17.1%
1	1	128	16	2	256	16	bf16	0.15 us	17.10 us	0.9%
2	128	64	3	3	128	16	bf16	2.49 us	24.70 us	10.1%
3	4	128	32	4	128	16	bf16	0.31 us	15.60 us	2.0%
4	4	128	32	8	128	16	fp16	0.61 us	15.70 us	3.9%
5	1	128	4	1	128	16	bf16	0.04 us	15.20 us	0.3%
6	4	128	40	8	128	16	fp16	0.61 us	15.80 us	3.9%
7	4	128	64	8	128	16	fp16	0.62 us	15.50 us	4.0%

Input Generation

input.py

import torch

def _make_inputs(num_seqs: int, ctx_len: int, num_query_heads: int, num_kv_heads: int, head_size: int, block_size: int=16, dtype: str='bf16') -> dict[str, torch.Tensor]:
    dt = torch.bfloat16 if dtype == 'bf16' else torch.float16
    max_num_blocks_per_seq = (ctx_len + block_size - 1) // block_size
    num_blocks = max_num_blocks_per_seq * num_seqs + 16
    query = torch.randn(num_seqs, num_query_heads, head_size, dtype=dt, device='cuda')
    x = 16 // dt.itemsize
    k_cache = torch.randn(num_blocks, num_kv_heads, head_size // x, block_size, x, dtype=dt, device='cuda')
    v_cache = torch.randn(num_blocks, num_kv_heads, head_size, block_size, dtype=dt, device='cuda')
    block_tables = torch.zeros(num_seqs, max_num_blocks_per_seq, dtype=torch.int32, device='cuda')
    for i in range(num_seqs):
        for j in range(max_num_blocks_per_seq):
            block_tables[i, j] = i * max_num_blocks_per_seq + j
    seq_lens = torch.full((num_seqs,), ctx_len, dtype=torch.int32, device='cuda')
    return {'query': query, 'k_cache': k_cache, 'v_cache': v_cache, 'block_tables': block_tables, 'seq_lens': seq_lens}