Paged Attention GLM-5.2

"""Custom paged-attention decode kernel for H100 PCIe (SM90 Hopper).

Memory-bound single-query decode. Strategy: grouped flash-decoding.

  * One CTA per (seq_split, batch, kv_head) -> processes ALL G query heads in
    that GQA group so the KV cache for that kv_head is streamed exactly once
    (the dominant cost). G = num_heads // num_kv_heads.
  * Adaptive sequence splitting so the grid fills the 114 SMs even at small
    batch*kv_head (e.g. B=4,Hkv=8 => 32 pairs => split seq 4096 into chunks).
  * Two-kernel flash-decode: split kernel writes fp32 partial (O, lse); a tiny
    reduce kernel merges the splits via logsumexp. When num_splits==1 the split
    kernel normalises and writes the final bf16 output directly (no reduce).
  * KV cache layout (num_blocks, page_size, num_kv_heads, head_dim*2) packs
    [K|V] on the last dim, so one gather per token pulls both. We iterate in
    page-sized tiles so each page is a single contiguous load.
"""
import math

import torch
import torch.nn as nn
import triton
import triton.language as tl

OP_TYPE = "attention"
SUPPORTED_PRECISIONS = ["bf16"]
HARDWARE_REQUIRED = ["H100"]

BATCH = 8
NUM_HEADS = 32
NUM_KV_HEADS = 8
HEAD_DIM = 128
SEQ_LEN = 1024
PAGE_SIZE = 16


@triton.jit
def _attn_split_kernel(
    Q_ptr, KV_ptr, BT_ptr, SL_ptr,
    OPART_ptr, LSE_ptr, O_ptr,
    # strides
    qsb, qsh,
    kvsb, kvsp, kvsh,
    btsb,
    opsb, opbb, ophb,
    lssb, lsbb,
    osb, ohb,
    # runtime ints
    pages_per_split,
    pages_per_seq,
    sm_scale,
    # constexprs
    G_C: tl.constexpr,
    D_C: tl.constexpr,
    P_C: tl.constexpr,
    N_C: tl.constexpr,          # pages per inner-loop tile
    BLOCK_SEQ: tl.constexpr,    # N_C * P_C
    NUM_SPLITS: tl.constexpr,
):
    split_idx = tl.program_id(0)
    b = tl.program_id(1)
    hkv = tl.program_id(2)

    seq_len = tl.load(SL_ptr + b)

    g_ar = tl.arange(0, G_C)
    d_ar = tl.arange(0, D_C)
    n_ar = tl.arange(0, N_C)
    p_ar = tl.arange(0, P_C)

    # Load Q for all G heads of this kv_head group: [G_C, D_C]
    q_ptrs = Q_ptr + b * qsb + (hkv * G_C + g_ar[:, None]) * qsh + d_ar[None, :]
    q = tl.load(q_ptrs)

    first_page = split_idx * pages_per_split
    last_page = first_page + pages_per_split
    last_page = tl.minimum(last_page, pages_per_seq)

    m_i = tl.full([G_C], -float("inf"), tl.float32)
    l_i = tl.zeros([G_C], tl.float32)
    acc = tl.zeros([G_C, D_C], tl.float32)

    for p0 in range(first_page, last_page, N_C):
        page_idx = p0 + n_ar                                   # [N_C]
        page_valid = page_idx < pages_per_seq
        phys = tl.load(BT_ptr + b * btsb + page_idx,
                       mask=page_valid, other=0)              # [N_C] int32

        # token index for each (page, in-page-offset): [N_C, P_C]
        tidx = page_idx[:, None] * P_C + p_ar[None, :]
        tok_valid = tidx < seq_len

        base = (KV_ptr
                + phys[:, None, None] * kvsb
                + p_ar[None, :, None] * kvsp
                + hkv * kvsh)
        K = tl.reshape(tl.load(base + d_ar[None, None, :],
                               mask=tok_valid[:, :, None], other=0.0),
                       (BLOCK_SEQ, D_C))
        V = tl.reshape(tl.load(base + D_C + d_ar[None, None, :],
                               mask=tok_valid[:, :, None], other=0.0),
                       (BLOCK_SEQ, D_C))

        s = tl.dot(q, tl.trans(K)) * sm_scale                  # [G_C, BLOCK_SEQ]
        tok_valid_1d = tl.reshape(tok_valid, (BLOCK_SEQ,))
        s = tl.where(tok_valid_1d[None, :], s, -float("inf"))

        m_n = tl.maximum(m_i, tl.max(s, 1))
        alpha = tl.exp(m_i - m_n)
        p = tl.exp(s - m_n[:, None])
        acc = acc * alpha[:, None]
        acc = tl.dot(p.to(tl.bfloat16), V, acc=acc)
        l_i = l_i * alpha + tl.sum(p, 1)
        m_i = m_n

    if NUM_SPLITS == 1:
        o = acc / l_i[:, None]
        o_ptrs = O_ptr + b * osb + (hkv * G_C + g_ar[:, None]) * ohb + d_ar[None, :]
        tl.store(o_ptrs, o.to(tl.bfloat16))
    else:
        lse = m_i + tl.log(l_i)
        opart = acc / l_i[:, None]                              # normalised partial
        op_ptrs = (OPART_ptr + split_idx * opsb + b * opbb
                   + (hkv * G_C + g_ar[:, None]) * ophb + d_ar[None, :])
        tl.store(op_ptrs, opart)                                # fp32
        lse_ptrs = LSE_ptr + split_idx * lssb + b * lsbb + (hkv * G_C + g_ar)
        tl.store(lse_ptrs, lse)


@triton.jit
def _reduce_kernel(
    OPART_ptr, LSE_ptr, O_ptr,
    opsb, opbb, ophb,
    lssb, lsbb,
    osb, ohb,
    NUM_SPLITS,
    G_C: tl.constexpr,
    D_C: tl.constexpr,
    BLOCK_S: tl.constexpr,
):
    b = tl.program_id(0)
    hkv = tl.program_id(1)

    g_ar = tl.arange(0, G_C)
    d_ar = tl.arange(0, D_C)
    s_ar = tl.arange(0, BLOCK_S)
    s_valid = s_ar < NUM_SPLITS

    lse_ptrs = (LSE_ptr + s_ar[:, None] * lssb + b * lsbb
                + (hkv * G_C + g_ar[None, :]))
    lse = tl.load(lse_ptrs, mask=s_valid[:, None], other=-float("inf"))   # [BLOCK_S, G_C]
    m_g = tl.max(lse, 0)
    alpha = tl.exp(lse - m_g[None, :])
    l_g = tl.sum(alpha, 0)

    op_ptrs = (OPART_ptr + s_ar[:, None, None] * opsb + b * opbb
               + (hkv * G_C + g_ar[None, :, None]) * ophb + d_ar[None, None, :])
    opart = tl.load(op_ptrs, mask=s_valid[:, None, None], other=0.0)      # [BLOCK_S,G_C,D_C]
    o = tl.sum(alpha[:, :, None] * opart, 0)                              # [G_C, D_C]
    o = o / l_g[:, None]
    o_ptrs = O_ptr + b * osb + (hkv * G_C + g_ar[:, None]) * ohb + d_ar[None, :]
    tl.store(o_ptrs, o.to(tl.bfloat16))


class Model(nn.Module):
    def __init__(self, batch, num_heads, num_kv_heads, head_dim, seq_len, page_size):
        super().__init__()
        assert num_heads % num_kv_heads == 0
        self.batch = batch
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.head_dim = head_dim
        self.seq_len = seq_len
        self.page_size = page_size
        self.group_size = num_heads // num_kv_heads
        self.scale = 1.0 / math.sqrt(head_dim)
        self.register_buffer(
            "_dummy", torch.zeros(1, dtype=torch.bfloat16), persistent=False
        )

    def forward(self, query, kv_cache, block_table, seq_lens):
        B, H, D = query.shape
        Hkv = self.num_kv_heads
        G = self.group_size
        P = self.page_size

        out = torch.empty(B, H, D, dtype=query.dtype, device=query.device)

        total_pages = block_table.shape[1]
        TARGET_BLOCKS = 512
        bhkv = B * Hkv
        ns = max(1, (TARGET_BLOCKS + bhkv - 1) // bhkv)
        ns = min(ns, total_pages)
        pages_per_split = (total_pages + ns - 1) // ns
        num_splits = (total_pages + pages_per_split - 1) // pages_per_split

        qsb, qsh = query.stride(0), query.stride(1)
        kvsb, kvsp, kvsh = kv_cache.stride(0), kv_cache.stride(1), kv_cache.stride(2)
        btsb = block_table.stride(0)
        osb, ohb = out.stride(0), out.stride(1)

        N_C = 4
        BLOCK_SEQ = N_C * P
        grid = (num_splits, B, Hkv)

        if num_splits == 1:
            _attn_split_kernel[grid](
                query, kv_cache, block_table, seq_lens,
                None, None, out,
                qsb, qsh, kvsb, kvsp, kvsh, btsb,
                0, 0, 0, 0, 0, osb, ohb,
                pages_per_split, total_pages, self.scale,
                G_C=G, D_C=D, P_C=P, N_C=N_C, BLOCK_SEQ=BLOCK_SEQ,
                NUM_SPLITS=1, num_warps=4, num_stages=2,
            )
        else:
            opart = torch.empty(num_splits, B, H, D, dtype=torch.float32,
                                device=query.device)
            lse = torch.empty(num_splits, B, H, dtype=torch.float32,
                              device=query.device)
            opsb, opbb, ophb = opart.stride(0), opart.stride(1), opart.stride(2)
            lssb, lsbb = lse.stride(0), lse.stride(1)
            _attn_split_kernel[grid](
                query, kv_cache, block_table, seq_lens,
                opart, lse, out,
                qsb, qsh, kvsb, kvsp, kvsh, btsb,
                opsb, opbb, ophb, lssb, lsbb, osb, ohb,
                pages_per_split, total_pages, self.scale,
                G_C=G, D_C=D, P_C=P, N_C=N_C, BLOCK_SEQ=BLOCK_SEQ,
                NUM_SPLITS=num_splits, num_warps=4, num_stages=2,
            )
            BLOCK_S = triton.next_power_of_2(num_splits)
            _reduce_kernel[(B, Hkv)](
                opart, lse, out,
                opsb, opbb, ophb, lssb, lsbb, osb, ohb,
                num_splits, G_C=G, D_C=D, BLOCK_S=BLOCK_S, num_warps=4,
            )
        return out


def get_inputs():
    B = BATCH
    H = NUM_HEADS
    Hkv = NUM_KV_HEADS
    D = HEAD_DIM
    L = SEQ_LEN
    P = PAGE_SIZE

    pages_per_seq = (L + P - 1) // P
    total_pages = max(B * pages_per_seq + 8, 64)

    query = torch.randn(B, H, D, dtype=torch.bfloat16) * 0.1
    kv_cache = torch.randn(total_pages, P, Hkv, 2 * D, dtype=torch.bfloat16) * 0.1

    perm = torch.randperm(total_pages)[: B * pages_per_seq].reshape(B, pages_per_seq).int()
    block_table = perm.contiguous()
    seq_lens = torch.full((B,), L, dtype=torch.int32)

    return [query, kv_cache, block_table, seq_lens]


def get_init_inputs():
    return [BATCH, NUM_HEADS, NUM_KV_HEADS, HEAD_DIM, SEQ_LEN, PAGE_SIZE]