gpt-o3 / tritonb12b97

import math
import torch
import triton
import triton.language as tl


# ----------------------------------------------------------------------------- #
#                               Triton Kernel                                   #
# ----------------------------------------------------------------------------- #
@triton.jit
def gqa_paged_decode_kernel(
    q_ptr,                       # *bf16  – [num_qo_heads, head_dim]
    k_cache_ptr,                 # *bf16  – [num_pages, num_kv_heads, head_dim]
    v_cache_ptr,                 # *bf16  – [num_pages, num_kv_heads, head_dim]
    kv_indices_ptr,              # *i32   – [seq_len]
    out_ptr,                     # *bf16  – [num_qo_heads, head_dim]
    lse_ptr,                     # *fp32  – [num_qo_heads]

    sm_scale: tl.float32,        # soft-max scale
    seq_len: tl.constexpr,       # number of tokens for this sequence

    stride_q_head: tl.constexpr,
    stride_k_page: tl.constexpr,
    stride_k_head: tl.constexpr,
    stride_v_page: tl.constexpr,
    stride_v_head: tl.constexpr,
    stride_out_head: tl.constexpr,
    stride_lse_head: tl.constexpr,

    gqa_ratio: tl.constexpr,     # 8 for 32/4
    BLOCK_TOKENS: tl.constexpr,  # 128
    BLOCK_D: tl.constexpr,       # 128
):
    # --------------------------------------------------------------------- #
    # Each program instance processes one query/output (qo) head.
    # --------------------------------------------------------------------- #
    pid   = tl.program_id(axis=0)
    h_qo  = pid                          # 0 … 31
    h_kv  = h_qo // gqa_ratio            # corresponding KV head (0 … 3)

    d_offsets = tl.arange(0, BLOCK_D)

    # --------------------------------------------------------------------- #
    # Load query vector (bf16 → fp32)
    # --------------------------------------------------------------------- #
    q_head_ptr   = q_ptr + h_qo * stride_q_head
    q_vec        = tl.load(q_head_ptr + d_offsets).to(tl.float32)      # [D]

    # Initial values for online softmax reduction
    neg_inf = -1.0e30
    m_val   = tl.full((), neg_inf, dtype=tl.float32)                   # running max

    # --------------------------------------------------------------------- #
    # Pass-1 : find maximum logit
    # --------------------------------------------------------------------- #
    for t0 in range(0, seq_len, BLOCK_TOKENS):
        tok_offsets = tl.arange(0, BLOCK_TOKENS) + t0                  # [T]
        tok_mask    = tok_offsets < seq_len

        page_ids = tl.load(kv_indices_ptr + tok_offsets, mask=tok_mask, other=0)

        k_ptrs = (
            k_cache_ptr
            + page_ids[:, None] * stride_k_page
            + h_kv             * stride_k_head
            + d_offsets[None, :]
        )                                                              # [T, D] pointers

        k_block = tl.load(k_ptrs, mask=tok_mask[:, None], other=0).to(tl.float32)

        logits     = tl.sum(k_block * q_vec[None, :], axis=1) * sm_scale
        logits     = tl.where(tok_mask, logits, neg_inf)
        block_max  = tl.max(logits, axis=0)
        m_val      = tl.maximum(m_val, block_max)

    # --------------------------------------------------------------------- #
    # Pass-2 : compute lse, softmax, output
    # --------------------------------------------------------------------- #
    lse_denom = tl.zeros((), tl.float32)
    out_acc   = tl.zeros((BLOCK_D, ), tl.float32)

    for t0 in range(0, seq_len, BLOCK_TOKENS):
        tok_offsets = tl.arange(0, BLOCK_TOKENS) + t0
        tok_mask    = tok_offsets < seq_len

        page_ids = tl.load(kv_indices_ptr + tok_offsets, mask=tok_mask, other=0)

        k_ptrs = (
            k_cache_ptr
            + page_ids[:, None] * stride_k_page
            + h_kv             * stride_k_head
            + d_offsets[None, :]
        )
        v_ptrs = (
            v_cache_ptr
            + page_ids[:, None] * stride_v_page
            + h_kv             * stride_v_head
            + d_offsets[None, :]
        )

        k_block = tl.load(k_ptrs, mask=tok_mask[:, None], other=0).to(tl.float32)
        v_block = tl.load(v_ptrs, mask=tok_mask[:, None], other=0).to(tl.float32)

        logits  = tl.sum(k_block * q_vec[None, :], axis=1) * sm_scale
        logits  = logits - m_val
        logits  = tl.where(tok_mask, logits, neg_inf)

        weights = tl.exp(logits)                                        # [T]
        weights = tl.where(tok_mask, weights, 0.0)

        lse_denom += tl.sum(weights, axis=0)

        out_acc += tl.sum(v_block * weights[:, None], axis=0)           # [D]

    out_final = out_acc / lse_denom                                     # [D]

    # --------------------------------------------------------------------- #
    # Write back output and LSE (convert back to bf16 where needed)
    # --------------------------------------------------------------------- #
    out_head_ptr = out_ptr + h_qo * stride_out_head
    tl.store(out_head_ptr + d_offsets, out_final.to(tl.bfloat16))

    lse_head_ptr = lse_ptr + h_qo * stride_lse_head
    INV_LN2      = 1.4426950408889634                                   # 1 / ln(2)
    l_val        = (tl.log(lse_denom) + m_val) * INV_LN2
    tl.store(lse_head_ptr, l_val)


# ----------------------------------------------------------------------------- #
#                                Python Wrapper                                #
# ----------------------------------------------------------------------------- #
@torch.no_grad()
def run(
    q: torch.Tensor,
    k_cache: torch.Tensor,
    v_cache: torch.Tensor,
    kv_indptr: torch.Tensor,
    kv_indices: torch.Tensor,
    sm_scale: float | None = None,
):
    """
    Parameters follow the specification in the problem statement.
    Returns
    -------
    output : bfloat16[batch, 32, 128]
    lse    :  float32[batch, 32]
    """
    if sm_scale is None:
        sm_scale = 1.0 / math.sqrt(128.0)

    # --------------------------------------------------------------------- #
    # Device handling
    # --------------------------------------------------------------------- #
    orig_device = q.device
    cuda_available = torch.cuda.is_available()

    def to_cuda(t):
        if t.device.type == 'cuda':
            return t
        if not cuda_available:
            raise RuntimeError("CUDA is required but not available.")
        return t.cuda(non_blocking=True)

    q          = to_cuda(q).contiguous()
    k_cache    = to_cuda(k_cache).contiguous()
    v_cache    = to_cuda(v_cache).contiguous()
    kv_indptr  = to_cuda(kv_indptr).contiguous()
    kv_indices = to_cuda(kv_indices).contiguous()

    # --------------------------------------------------------------------- #
    # Assertions & shapes
    # --------------------------------------------------------------------- #
    batch_size, num_qo_heads, head_dim = q.shape
    num_pages, page_sz, num_kv_heads, _ = k_cache.shape

    assert num_qo_heads == 32,  "num_qo_heads must be 32"
    assert num_kv_heads == 4,   "num_kv_heads must be 4"
    assert head_dim      == 128
    assert page_sz       == 1
    assert kv_indptr.numel() == batch_size + 1
    assert kv_indices.numel() == kv_indptr[-1].item()

    # Flatten caches (page_size == 1)
    k_cache_flat = k_cache.squeeze(1).contiguous()   # [num_pages, 4, 128]
    v_cache_flat = v_cache.squeeze(1).contiguous()

    # Allocate output tensors
    out = torch.zeros_like(q, dtype=torch.bfloat16, device=q.device)
    lse = torch.full((batch_size, num_qo_heads),
                     -float("inf"), dtype=torch.float32, device=q.device)

    # Strides (in elements, not bytes)
    stride_q_head   = head_dim
    stride_k_head   = head_dim
    stride_k_page   = num_kv_heads * head_dim      # 4 * 128
    stride_v_head   = head_dim
    stride_v_page   = num_kv_heads * head_dim
    stride_out_head = head_dim
    stride_lse_head = 1

    gqa_ratio     = num_qo_heads // num_kv_heads   # 8
    BLOCK_TOKENS  = 128
    BLOCK_D       = 128

    # --------------------------------------------------------------------- #
    # Launch per-sequence kernel
    # --------------------------------------------------------------------- #
    for b in range(batch_size):
        start   = kv_indptr[b].item()
        end     = kv_indptr[b + 1].item()
        seq_len = end - start

        if seq_len == 0:
            continue

        seq_indices = kv_indices[start:end].contiguous()

        grid = (num_qo_heads,)

        gqa_paged_decode_kernel[grid](
            q[b],
            k_cache_flat,
            v_cache_flat,
            seq_indices,
            out[b],
            lse[b],
            sm_scale,
            seq_len,

            stride_q_head,
            stride_k_page,
            stride_k_head,
            stride_v_page,
            stride_v_head,
            stride_out_head,
            stride_lse_head,

            gqa_ratio,
            BLOCK_TOKENS,
            BLOCK_D,

            num_warps = 8,
            num_stages = 4,
        )

    # --------------------------------------------------------------------- #
    # Move outputs back if the original tensors were on CPU
    # --------------------------------------------------------------------- #
    if orig_device.type == 'cpu':
        out = out.to(orig_device)
        lse = lse.to(orig_device)

    return out, lse