Paged Attention GLM-5.2

"""Custom paged-attention decode kernel for B200 (SM100, HBM3e).

Flash-decoding: one CTA per (batch, kv_head, seq_partition). Each KV element is
read exactly once (GQA query heads in a group share the load). Split-K along the
sequence gives enough parallelism to fill all 148 SMs even at batch=4. A second
small kernel does the flash reduction across partitions.

No SDPA / vLLM / FlashInfer calls -- pure Triton.
"""
import math

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

OP_TYPE = "attention"
SUPPORTED_PRECISIONS = ["bf16"]
HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]

# --- Shape knobs (overridden by check.py / benchmark.py from shapes.py) ----
BATCH = 8
NUM_HEADS = 32
NUM_KV_HEADS = 8
HEAD_DIM = 128
SEQ_LEN = 1024
PAGE_SIZE = 16

_NUM_SMS = 148  # B200


# ===========================================================================
# Decode kernel: produces per-partition partial (o, m, l) or final output.
# Grid: (B, Hkv, num_partitions)
# ===========================================================================
@triton.autotune(
    configs=[
        triton.Config({"BLOCK_N": 32},  num_warps=4, num_stages=4),
        triton.Config({"BLOCK_N": 32},  num_warps=8, num_stages=5),
        triton.Config({"BLOCK_N": 64},  num_warps=4, num_stages=4),
        triton.Config({"BLOCK_N": 64},  num_warps=8, num_stages=4),
        triton.Config({"BLOCK_N": 64},  num_warps=8, num_stages=5),
        triton.Config({"BLOCK_N": 64},  num_warps=8, num_stages=6),
        triton.Config({"BLOCK_N": 128}, num_warps=8, num_stages=3),
        triton.Config({"BLOCK_N": 128}, num_warps=8, num_stages=4),
        triton.Config({"BLOCK_N": 128}, num_warps=8, num_stages=5),
        triton.Config({"BLOCK_N": 128}, num_warps=16, num_stages=3),
        triton.Config({"BLOCK_N": 128}, num_warps=16, num_stages=4),
    ],
    key=["G", "D", "P", "WRITE_FINAL"],
)
@triton.jit
def _decode_kernel(
    Q_ptr, KV_ptr, BT_ptr, SL_ptr,
    OUT_ptr,
    PO_ptr, PM_ptr, PL_ptr,
    q_sb, q_sh,
    kv_sblk, kv_spg, kv_sh,
    bt_sb,
    out_sb, out_sh,
    po_sb, po_sh, po_sp,
    pm_sb, pm_sh, pm_sp,
    pl_sb, pl_sh, pl_sp,
    sm_scale,
    PART_SEQ,
    G: tl.constexpr,
    D: tl.constexpr,
    HKV: tl.constexpr,
    P: tl.constexpr,
    KVD: tl.constexpr,
    MQ: tl.constexpr,
    BLOCK_N: tl.constexpr,
    WRITE_FINAL: tl.constexpr,
):
    pid_b = tl.program_id(0)
    pid_h = tl.program_id(1)
    pid_p = tl.program_id(2)

    Lb = tl.load(SL_ptr + pid_b)
    start = pid_p * PART_SEQ
    # exclusive end of this partition (clamped to seq len); masks the inner loop
    # so interior partitions don't read into their neighbour's token range.
    part_end = tl.minimum((pid_p + 1) * PART_SEQ, Lb)

    # ---- Q tile (MQ, D) bf16; rows >= G are padded with 0 ----
    offs_q = tl.arange(0, MQ)
    g_mask = offs_q < G
    q_rows = pid_h * G + offs_q
    offs_d = tl.arange(0, D)
    q_ptrs = Q_ptr + pid_b * q_sb + q_rows[:, None] * q_sh + offs_d[None, :]
    Q = tl.load(q_ptrs, mask=g_mask[:, None], other=0.0)

    # ---- accumulators ----
    m_i = tl.full([MQ], float("-inf"), dtype=tl.float32)
    l_i = tl.full([MQ], 0.0, dtype=tl.float32)
    o_i = tl.zeros([MQ, D], dtype=tl.float32)

    if start < Lb:
        for start_n in range(0, PART_SEQ, BLOCK_N):
            n = start + start_n + tl.arange(0, BLOCK_N)
            mask_n = n < part_end
            page = n // P
            off_pg = n % P
            bid = tl.load(BT_ptr + pid_b * bt_sb + page, mask=mask_n, other=0)
            phys = bid * P + off_pg
            base = phys * (HKV * KVD) + pid_h * KVD
            kv_ptrs = KV_ptr + base[:, None] + offs_d[None, :]
            K = tl.load(kv_ptrs, mask=mask_n[:, None], other=0.0, eviction_policy="evict_first")
            qk = tl.dot(Q, tl.trans(K)) * sm_scale          # (MQ, BLOCK_N)
            qk = tl.where(mask_n[None, :], qk, float("-inf"))
            m_ij = tl.maximum(m_i, tl.max(qk, axis=1))
            qk = qk - m_ij[:, None]
            p = tl.exp(qk)                                   # (MQ, BLOCK_N)
            alpha = tl.exp(m_i - m_ij)
            l_ij = tl.sum(p, axis=1)
            V = tl.load(kv_ptrs + D, mask=mask_n[:, None], other=0.0, eviction_policy="evict_first")
            o_i = o_i * alpha[:, None] + tl.dot(p.to(tl.bfloat16), V)
            l_i = l_i * alpha + l_ij
            m_i = m_ij

    if WRITE_FINAL:
        o = o_i / l_i[:, None]
        out_ptrs = OUT_ptr + pid_b * out_sb + q_rows[:, None] * out_sh + offs_d[None, :]
        tl.store(out_ptrs, o.to(tl.bfloat16), mask=g_mask[:, None])
    else:
        po_ptrs = PO_ptr + pid_b * po_sb + pid_h * po_sh + pid_p * po_sp + offs_q[:, None] * D + offs_d[None, :]
        tl.store(po_ptrs, o_i, mask=g_mask[:, None])
        pm_ptrs = PM_ptr + pid_b * pm_sb + pid_h * pm_sh + pid_p * pm_sp + offs_q
        tl.store(pm_ptrs, m_i, mask=g_mask)
        pl_ptrs = PL_ptr + pid_b * pl_sb + pid_h * pl_sh + pid_p * pl_sp + offs_q
        tl.store(pl_ptrs, l_i, mask=g_mask)


# ===========================================================================
# Reduce kernel: flash-reduce per-partition partials -> final output.
# Grid: (B, H) -- one CTA per query head for high occupancy on small batches.
# ===========================================================================
@triton.jit
def _reduce_kernel(
    OUT_ptr, PO_ptr, PM_ptr, PL_ptr,
    out_sb, out_sh,
    po_sb, po_sh, po_sp,
    pm_sb, pm_sh, pm_sp,
    pl_sb, pl_sh, pl_sp,
    NP,
    G: tl.constexpr,
    D: tl.constexpr,
    BLOCK_P: tl.constexpr,
):
    pid_b = tl.program_id(0)
    pid_h = tl.program_id(1)          # query head index in [0, H)
    hkv = pid_h // G
    g = pid_h % G

    offs_d = tl.arange(0, D)
    offs_p = tl.arange(0, BLOCK_P)
    p_mask = offs_p < NP

    # partial max for this head: (BLOCK_P,)
    pm_ptrs = PM_ptr + pid_b * pm_sb + hkv * pm_sh + offs_p * pm_sp + g
    pm = tl.load(pm_ptrs, mask=p_mask, other=float("-inf"))
    m_g = tl.max(pm, axis=0)

    factor = tl.exp(pm - m_g)                                 # (BLOCK_P,)

    pl_ptrs = PL_ptr + pid_b * pl_sb + hkv * pl_sh + offs_p * pl_sp + g
    pl = tl.load(pl_ptrs, mask=p_mask, other=0.0)
    l_g = tl.sum(factor * pl, axis=0)

    # partial output for this head: (BLOCK_P, D)
    po_ptrs = (PO_ptr + pid_b * po_sb + hkv * po_sh
               + offs_p[:, None] * po_sp + g * D + offs_d[None, :])
    po = tl.load(po_ptrs, mask=p_mask[:, None], other=0.0)
    o_g = tl.sum(factor[:, None] * po, axis=0)               # (D,)

    out = o_g / l_g
    out_ptrs = OUT_ptr + pid_b * out_sb + pid_h * out_sh + offs_d
    tl.store(out_ptrs, out.to(tl.bfloat16))


def _next_pow2(x):
    p = 1
    while p < x:
        p <<= 1
    return p


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)
        import os as _os
        _mq_override = _os.environ.get("KBH_MQ")
        if _mq_override:
            self.MQ = int(_mq_override)
        else:
            self.MQ = _next_pow2(max(self.group_size, 1))

        # ---- choose split-K parallelism ----
        # Decode prefers a modest number of fairly large CTAs (per-CTA softmax
        # overhead + split-K reduce cost penalize very high partition counts).
        # Empirically ~1 wave (B*Hkv*npart ~= 128) is the sweet spot across the
        # shape sweep. Going higher only helps the pure-read ceiling, not decode.
        import os
        base = batch * num_kv_heads
        target_ctas = 128
        npart = (target_ctas + base - 1) // base
        npart = max(npart, 1)
        min_tokens_per_part = 64
        max_npart = max(1, seq_len // min_tokens_per_part)
        npart = min(npart, max_npart)
        npart = max(npart, 1)
        env_np = os.environ.get("KBH_NPART")
        if env_np:
            npart = int(env_np)
        self.npart = npart
        self.part_seq = (seq_len + npart - 1) // npart

        self.register_buffer("_dummy", torch.zeros(1, dtype=torch.bfloat16), persistent=False)

        # Preallocate output and partial buffers (static addresses for CUDA graphs).
        self._out = torch.empty(batch, num_heads, head_dim, dtype=torch.bfloat16, device="cuda")
        if npart > 1:
            self._po = torch.empty(batch, num_kv_heads, npart, self.MQ, head_dim, dtype=torch.float32, device="cuda")
            self._pm = torch.empty(batch, num_kv_heads, npart, self.MQ, dtype=torch.float32, device="cuda")
            self._pl = torch.empty(batch, num_kv_heads, npart, self.MQ, dtype=torch.float32, device="cuda")
        else:
            self._po = self._pm = self._pl = None

        self._graph = None
        self._graph_key = None

    def _launch(self, out, 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
        KVD = 2 * D
        MQ = self.MQ
        npart = self.npart
        part_seq = self.part_seq

        grid = (B, Hkv, npart)
        if npart == 1:
            _decode_kernel[grid](
                query, kv_cache, block_table, seq_lens,
                out,
                None, None, None,
                query.stride(0), query.stride(1),
                kv_cache.stride(0), kv_cache.stride(1), kv_cache.stride(2),
                block_table.stride(0),
                out.stride(0), out.stride(1),
                0, 0, 0, 0, 0, 0, 0, 0, 0,
                self.scale, part_seq,
                G=G, D=D, HKV=Hkv, P=P, KVD=KVD, MQ=MQ,
                WRITE_FINAL=True,
            )
        else:
            po, pm, pl = self._po, self._pm, self._pl
            _decode_kernel[grid](
                query, kv_cache, block_table, seq_lens,
                out,
                po, pm, pl,
                query.stride(0), query.stride(1),
                kv_cache.stride(0), kv_cache.stride(1), kv_cache.stride(2),
                block_table.stride(0),
                out.stride(0), out.stride(1),
                po.stride(0), po.stride(1), po.stride(2),
                pm.stride(0), pm.stride(1), pm.stride(2),
                pl.stride(0), pl.stride(1), pl.stride(2),
                self.scale, part_seq,
                G=G, D=D, HKV=Hkv, P=P, KVD=KVD, MQ=MQ,
                WRITE_FINAL=False,
            )
            _reduce_kernel[(B, self.num_heads)](
                out, po, pm, pl,
                out.stride(0), out.stride(1),
                po.stride(0), po.stride(1), po.stride(2),
                pm.stride(0), pm.stride(1), pm.stride(2),
                pl.stride(0), pl.stride(1), pl.stride(2),
                npart,
                G=G, D=D, BLOCK_P=_next_pow2(npart),
            )

    def forward(self, query, kv_cache, block_table, seq_lens):
        out = self._out
        key = (query.data_ptr(), kv_cache.data_ptr(),
               block_table.data_ptr(), seq_lens.data_ptr())

        # If we're already inside an external CUDA-graph capture, just launch
        # eagerly so the kernel launches get captured by the outer graph.
        capturing = torch.cuda.is_current_stream_capturing()

        if not capturing and self._graph is not None and key == self._graph_key:
            self._graph.replay()
            return out

        # Eager launch (also resolves Triton autotune before capture).
        self._launch(out, query, kv_cache, block_table, seq_lens)

        # Try to capture a graph for this set of input addresses so subsequent
        # calls replay with near-zero launch overhead. Falls back to eager.
        if not capturing and self._graph is None:
            try:
                torch.cuda.synchronize()
                g = torch.cuda.CUDAGraph()
                with torch.cuda.graph(g):
                    self._launch(out, query, kv_cache, block_table, seq_lens)
                self._graph = g
                self._graph_key = key
            except Exception:
                self._graph = None
        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]