#005

chunk_gated_delta_rule_state

bf16 sglang · prefill · sglang.srt.layers.attention.fla.chunk_delta_h · importance 1.6%

Reference Implementation

reference.py

import math
import torch
import torch.nn as nn

class Model(nn.Module):

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

    def forward(self, k: torch.Tensor, w: torch.Tensor, u: torch.Tensor, g: torch.Tensor, initial_state: torch.Tensor, initial_state_indices: torch.Tensor) -> dict[str, torch.Tensor]:
        (batch, tokens, num_k_heads, key_dim) = k.shape
        (num_heads, value_dim) = (u.shape[2], u.shape[3])
        heads_per_k = max(num_heads // num_k_heads, 1)
        cs = self.chunk_size
        num_chunks = math.ceil(tokens / cs)
        head_to_khead = torch.tensor([min(hd // heads_per_k, num_k_heads - 1) for hd in range(num_heads)], dtype=torch.long, device=k.device)
        k_h = k.index_select(2, head_to_khead)
        final_state = initial_state.float().clone()
        idx_list = initial_state_indices.tolist()
        state_indices = torch.tensor(idx_list, dtype=torch.long, device=k.device)
        h = final_state.index_select(0, state_indices).clone()
        h_chunks = torch.empty(batch, num_chunks, num_heads, key_dim, value_dim, dtype=k.dtype, device=k.device)
        v_new = torch.empty_like(u)
        BH = batch * num_heads
        h_flat = h.reshape(BH, key_dim, value_dim)
        for chunk_idx in range(num_chunks):
            start = chunk_idx * cs
            end = min(start + cs, tokens)
            cur_cs = end - start
            h_chunks[:, chunk_idx] = h_flat.reshape(batch, num_heads, key_dim, value_dim).to(k.dtype)
            w_blk = w[:, start:end].float()
            u_blk = u[:, start:end].float()
            k_blk = k_h[:, start:end].float()
            g_blk = g[:, start:end].float()
            w_bh = w_blk.permute(0, 2, 1, 3).reshape(BH, cur_cs, key_dim)
            u_bh = u_blk.permute(0, 2, 1, 3).reshape(BH, cur_cs, value_dim)
            k_bh = k_blk.permute(0, 2, 1, 3).reshape(BH, cur_cs, key_dim)
            g_bh = g_blk.permute(0, 2, 1).reshape(BH, cur_cs)
            v_blk = u_bh - torch.bmm(w_bh.to(k.dtype).float(), h_flat.to(k.dtype).float())
            v_blk_out = v_blk.reshape(batch, num_heads, cur_cs, value_dim).permute(0, 2, 1, 3)
            v_new[:, start:end] = v_blk_out.to(u.dtype)
            g_last = g_bh[:, -1:]
            g_diff = g_last - g_bh
            v_blk_scaled = v_blk * torch.exp(torch.where(g_diff <= 0, g_diff, torch.tensor(float('-inf'), device=g_diff.device))).unsqueeze(-1)
            v_blk_scaled = v_blk_scaled.to(k.dtype)
            h_flat = h_flat * torch.exp(g_last).unsqueeze(-1)
            h_flat = h_flat + torch.bmm(k_bh.to(k.dtype).transpose(-2, -1).float(), v_blk_scaled.float())
        h_back = h_flat.reshape(batch, num_heads, key_dim, value_dim)
        final_state.index_copy_(0, state_indices, h_back)
        return {'h': h_chunks, 'v_new': v_new, 'final_state': final_state}

Shapes

T_SOL hardware:

#	token_count	num_v_heads	state_count	value_dim	key_dim	num_k_heads	batch_size	T_SOL(XPU-A)	T_Prod	S
0	2048	2	474	128	128	1	1	1.17 us	93.20 us	1.3%
1	2048	20	474	128	128	10	1	11.67 us	184.80 us	6.3%
2	2048	22	474	128	128	11	1	12.84 us	269.80 us	4.8%
3	2048	24	474	128	128	12	1	14.01 us	290.50 us	4.8%
4	2048	26	474	128	128	13	1	15.18 us	284.60 us	5.3%
5	2048	28	474	128	128	14	1	16.34 us	293.60 us	5.6%
6	2048	30	474	128	128	15	1	17.51 us	293.20 us	6.0%
7	2048	32	474	128	128	16	1	18.68 us	416.00 us	4.5%
8	2048	4	474	128	128	2	1	2.33 us	95.00 us	2.5%
9	4096	4	474	128	128	2	1	4.67 us	179.00 us	2.6%
10	8192	4	474	128	128	2	1	9.34 us	348.40 us	2.7%
11	16384	4	474	128	128	2	1	18.68 us	767.60 us	2.4%
12	24576	4	474	128	128	2	1	28.02 us	1.20 ms	2.3%
13	2048	6	474	128	128	3	1	3.50 us	94.20 us	3.7%
14	512	8	474	128	128	4	1	1.29 us	32.20 us	4.0%
15	1024	8	474	128	128	4	1	2.38 us	54.10 us	4.4%
16	2048	8	474	128	128	4	1	4.67 us	97.10 us	4.8%
17	2048	10	474	128	128	5	1	5.84 us	95.70 us	6.1%
18	2048	12	474	128	128	6	1	7.00 us	184.30 us	3.8%
19	2048	14	474	128	128	7	1	8.17 us	181.90 us	4.5%
20	1024	16	474	128	128	8	1	4.76 us	101.60 us	4.7%
21	2048	16	474	128	128	8	1	9.34 us	188.40 us	5.0%
22	2048	18	474	128	128	9	1	10.51 us	182.60 us	5.8%
23	4195	32	1116	128	128	16	1	38.26 us	899.50 us	4.3%
24	11027	32	1116	128	128	16	1	100.57 us	2.35 ms	4.3%

Input Generation

input.py

import torch

def _make_inputs(batch_size: int, token_count: int, num_k_heads: int, num_v_heads: int, key_dim: int, value_dim: int, state_count: int) -> dict[str, torch.Tensor]:
    k = torch.randn(batch_size, token_count, num_k_heads, key_dim, dtype=torch.bfloat16, device='cuda') * 0.02
    w = torch.randn(batch_size, token_count, num_v_heads, key_dim, dtype=torch.bfloat16, device='cuda') * 0.02
    u = torch.randn(batch_size, token_count, num_v_heads, value_dim, dtype=torch.bfloat16, device='cuda') * 0.02
    g = -torch.rand(batch_size, token_count, num_v_heads, dtype=torch.float32, device='cuda')
    initial_state = torch.randn(state_count, num_v_heads, key_dim, value_dim, dtype=torch.float32, device='cuda') * 0.02
    initial_state_indices = torch.arange(batch_size, dtype=torch.int32, device='cuda') % state_count
    return {'k': k, 'w': w, 'u': u, 'g': g, 'initial_state': initial_state, 'initial_state_indices': initial_state_indices}