#004

chunk_delta_rule_output

bf16 sglang · prefill · sglang.srt.layers.attention.fla.chunk_o · importance 0.8%

Reference Implementation

reference.py

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

class Model(nn.Module):

    def __init__(self, chunk_size: int=64, scale: float | None=None) -> None:
        super().__init__()
        self.chunk_size = int(chunk_size)
        self.scale = scale

    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, h: torch.Tensor, g: torch.Tensor) -> torch.Tensor:
        (batch, tokens, num_q_heads, key_dim) = q.shape
        (num_heads, value_dim) = (v.shape[2], v.shape[3])
        heads_per_q = max(num_heads // num_q_heads, 1)
        cs = self.chunk_size
        scale = self.scale if self.scale is not None else key_dim ** (-0.5)
        num_chunks = math.ceil(tokens / cs)
        tokens_padded = num_chunks * cs
        pad = tokens_padded - tokens
        head_to_qhead = torch.tensor([min(hd // heads_per_q, num_q_heads - 1) for hd in range(num_heads)], dtype=torch.long, device=q.device)
        q_h = q.index_select(2, head_to_qhead)
        k_h = k.index_select(2, head_to_qhead)
        if pad > 0:
            q_h = F.pad(q_h, (0, 0, 0, 0, 0, pad))
            k_h = F.pad(k_h, (0, 0, 0, 0, 0, pad))
            v = F.pad(v, (0, 0, 0, 0, 0, pad))
            g = F.pad(g, (0, 0, 0, pad))

        def chunkify(x: torch.Tensor) -> torch.Tensor:
            (B, T, H, D) = x.shape
            return x.reshape(B, num_chunks, cs, H, D).permute(0, 3, 1, 2, 4).reshape(B * H * num_chunks, cs, D)
        q_blk = chunkify(q_h)
        k_blk = chunkify(k_h)
        v_blk = chunkify(v)
        g_blk = chunkify(g.unsqueeze(-1)).float().squeeze(-1)
        h_blk = h.permute(0, 2, 1, 3, 4).reshape(batch * num_heads * num_chunks, key_dim, value_dim)
        o = torch.bmm(q_blk.float(), h_blk.float())
        attn = torch.bmm(q_blk.float(), k_blk.float().transpose(-2, -1))
        exp_g = torch.exp(g_blk)
        o = o * exp_g.unsqueeze(-1)
        g_diff = g_blk.unsqueeze(-1) - g_blk.unsqueeze(-2)
        attn = attn * torch.exp(torch.where(g_diff <= 0, g_diff, torch.tensor(float('-inf'), device=g_diff.device)))
        attn = torch.tril(attn)
        o_full = (o + torch.bmm(attn.to(v.dtype).float(), v_blk.float())) * scale
        out = o_full.reshape(batch, num_heads, num_chunks, cs, value_dim).permute(0, 2, 3, 1, 4).reshape(batch, tokens_padded, num_heads, value_dim).to(v.dtype)
        return out[:, :tokens]

Shapes

T_SOL hardware:

#	token_count	num_v_heads	value_dim	num_q_heads	chunk_size	key_dim	batch_size	T_SOL(XPU-A)	T_Prod	S
0	64	32	128	16	64	128	1	0.58 us	20.20 us	2.9%
1	128	32	128	16	64	128	1	1.16 us	23.90 us	4.9%
2	256	32	128	16	64	128	1	2.33 us	37.10 us	6.3%
3	512	32	128	16	64	128	1	4.65 us	57.90 us	8.0%
4	1024	32	128	16	64	128	1	9.30 us	97.20 us	9.6%
5	2048	32	128	16	64	128	1	18.61 us	190.50 us	9.8%
6	4096	32	128	16	64	128	1	37.21 us	372.90 us	10.0%
7	8192	32	128	16	64	128	1	74.42 us	725.50 us	10.3%
8	704	16	128	8	64	128	1	3.20 us	46.70 us	6.9%
9	1024	16	128	8	64	128	1	4.65 us	59.90 us	7.8%
10	2112	16	128	8	64	128	1	9.59 us	103.10 us	9.3%
11	4096	16	128	8	64	128	1	18.61 us	195.80 us	9.5%
12	8256	16	128	8	64	128	1	37.50 us	385.00 us	9.7%
13	4195	32	128	16	64	128	1	38.37 us	403.60 us	9.5%
14	11027	32	128	16	64	128	1	100.58 us	1.06 ms	9.5%
15	14807	32	128	16	64	128	1	134.89 us	1.29 ms	10.4%

Input Generation

input.py

import math
import torch

def _make_inputs(batch_size: int, token_count: int, num_q_heads: int, num_v_heads: int, key_dim: int, value_dim: int, chunk_size: int=64) -> dict[str, torch.Tensor]:
    q = torch.randn(batch_size, token_count, num_q_heads, key_dim, dtype=torch.bfloat16, device='cuda') * 0.02
    k = torch.randn_like(q)
    v = torch.randn(batch_size, token_count, num_v_heads, value_dim, dtype=torch.bfloat16, device='cuda') * 0.02
    num_chunks = math.ceil(token_count / chunk_size)
    h = torch.randn(batch_size, num_chunks, num_v_heads, key_dim, value_dim, dtype=torch.bfloat16, device='cuda') * 0.02
    g = -torch.rand(batch_size, token_count, num_v_heads, dtype=torch.float32, device='cuda')
    return {'q': q, 'k': k, 'v': v, 'h': h, 'g': g}