#001

attention_forward

bf16 vllm · prefill · ck_tile FmhaFwdKernel / attn_fwd raw trace · importance 2.2%

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, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cu_seqlens_q: torch.Tensor, cu_seqlens_k: torch.Tensor) -> torch.Tensor:
        out = torch.empty_like(q)
        q_bounds = cu_seqlens_q.tolist()
        k_bounds = cu_seqlens_k.tolist()
        for seq_idx in range(len(q_bounds) - 1):
            qs_start, qs_end = q_bounds[seq_idx], q_bounds[seq_idx + 1]
            ks_start, ks_end = k_bounds[seq_idx], k_bounds[seq_idx + 1]
            if qs_end <= qs_start:
                continue
            qs = q[qs_start:qs_end].transpose(0, 1).unsqueeze(0).float()
            ks = k[ks_start:ks_end].transpose(0, 1).unsqueeze(0).float()
            vs = v[ks_start:ks_end].transpose(0, 1).unsqueeze(0).float()
            os = F.scaled_dot_product_attention(qs, ks, vs, is_causal=False)
            out[qs_start:qs_end] = os.squeeze(0).transpose(0, 1).to(q.dtype)
        return out

Shapes

T_SOL hardware:

#	num_heads	head_dim	seq_lens	T_SOL(XPU-A)	T_Prod	S
0	16	72	[720] × 1	10.27 us	68.40 us	15.0%
1	16	72	[1200] × 1	28.52 us	108.10 us	26.4%
2	16	72	[2116] × 1	88.68 us	332.70 us	26.7%
3	16	72	[3844] × 1	292.67 us	1.02 ms	28.7%
4	16	72	[8136] × 1	1.31 ms	3.79 ms	34.6%
5	16	72	[17296] × 1	5.93 ms	17.07 ms	34.7%
6	16	72	[24368] × 1	11.76 ms	33.27 ms	35.3%
7	16	72	[49596] × 1	48.72 ms	134.51 ms	36.2%
8	16	72	[65556] × 1	85.12 ms	237.71 ms	35.8%
9	32	128	[16742] × 1	19.74 ms	37.32 ms	52.9%
10	32	128	[30793] × 1	66.78 ms	124.19 ms	53.8%
11	4	72	[276] × 1	0.38 us	43.20 us	0.9%
12	4	72	[600] × 1	1.78 us	47.40 us	3.8%
13	4	72	[1012] × 1	5.07 us	57.00 us	8.9%
14	4	72	[2024] × 1	20.28 us	105.50 us	19.2%
15	4	72	[4100] × 1	83.24 us	308.70 us	27.0%
16	4	72	[8184] × 1	331.65 us	1.17 ms	28.3%
17	4	72	[15476] × 1	1.19 ms	3.75 ms	31.7%
18	4	72	[24952] × 1	3.08 ms	8.55 ms	36.1%
19	4	72	[40560] × 1	8.15 ms	21.90 ms	37.2%
20	8	128	[1] × 1	0.00 us	14.40 us	0.0%

Input Generation

input.py

import torch

def _make_inputs(seq_lens: list[int], num_heads: int, head_dim: int) -> dict[str, torch.Tensor]:
    total_tokens = sum(seq_lens)
    q = torch.randn(total_tokens, num_heads, head_dim, dtype=torch.bfloat16, device='cuda')
    k = torch.randn(total_tokens, num_heads, head_dim, dtype=torch.bfloat16, device='cuda')
    v = torch.randn(total_tokens, num_heads, head_dim, dtype=torch.bfloat16, device='cuda')
    cu_seqlens_q = torch.zeros(len(seq_lens) + 1, dtype=torch.int32, device='cuda')
    cu_seqlens_k = torch.zeros(len(seq_lens) + 1, dtype=torch.int32, device='cuda')
    for i, sl in enumerate(seq_lens):
        cu_seqlens_q[i + 1] = cu_seqlens_q[i] + sl
        cu_seqlens_k[i + 1] = cu_seqlens_k[i] + sl
    return {'q': q, 'k': k, 'v': v, 'cu_seqlens_q': cu_seqlens_q, 'cu_seqlens_k': cu_seqlens_k}