gpt-o3 / tritonc3c0cc

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


@triton.jit
def gqa_paged_decode_kernel(
    q_ptr,               # *bf16  [B, 32, 128]
    k_ptr,               # *bf16  [N_pages, 8, 128]  (page_size squeezed)
    v_ptr,               # *bf16  [N_pages, 8, 128]  (page_size squeezed)
    kv_indptr_ptr,       # *int32 [B + 1]
    kv_indices_ptr,      # *int32 [num_kv_indices]
    sm_scale,            # fp32 scalar
    out_ptr,             # *bf16  [B, 32, 128]
    lse_ptr,             # *fp32  [B, 32]
    BLOCK_T: tl.constexpr,
    HEAD_DIM: tl.constexpr,
    NUM_QO_HEADS: tl.constexpr,
    NUM_KV_HEADS: tl.constexpr,
):
    pid = tl.program_id(0)

    batch_idx = pid // NUM_QO_HEADS
    qo_head   = pid % NUM_QO_HEADS
    gqa_ratio = NUM_QO_HEADS // NUM_KV_HEADS
    kv_head   = qo_head // gqa_ratio

    # ---- strides (in elements, not bytes) ----
    stride_q_batch    = NUM_QO_HEADS * HEAD_DIM
    stride_q_head     = HEAD_DIM

    stride_k_page     = NUM_KV_HEADS * HEAD_DIM         # page_size = 1
    stride_k_kv_head  = HEAD_DIM

    stride_v_page     = stride_k_page
    stride_v_kv_head  = HEAD_DIM

    # ---- load query vector ----
    d_offs = tl.arange(0, HEAD_DIM)
    q_ptr_head = q_ptr + batch_idx * stride_q_batch + qo_head * stride_q_head + d_offs
    q_vec = tl.cast(tl.load(q_ptr_head), tl.float32)

    # ---- sequence token range ----
    start = tl.load(kv_indptr_ptr + batch_idx)
    end   = tl.load(kv_indptr_ptr + batch_idx + 1)
    num_tokens = end - start

    # ---- streaming softmax vars ----
    m_val   = tl.full([], -1e30, tl.float32)          # running max
    d_val   = tl.zeros([], tl.float32)                # running sum exp
    o_vec   = tl.zeros([HEAD_DIM], tl.float32)        # running output vector

    offset = tl.zeros([], tl.int32)

    while offset < num_tokens:
        t_offs      = tl.arange(0, BLOCK_T)
        remain      = num_tokens - offset
        tok_mask    = t_offs < remain

        # ---- load page indices ----
        pages = tl.load(kv_indices_ptr + start + offset + t_offs,
                        mask=tok_mask, other=0)

        # ---- gather K / V ----
        k_ptrs = k_ptr + pages[:, None] * stride_k_page + kv_head * stride_k_kv_head + d_offs[None, :]
        v_ptrs = v_ptr + pages[:, None] * stride_v_page + kv_head * stride_v_kv_head + d_offs[None, :]

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

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

        # ---- block softmax ----
        m_block        = tl.max(logits, axis=0)
        exp_logits     = tl.exp(logits - m_block)
        sum_exp_block  = tl.sum(exp_logits, axis=0)
        weighted_v     = tl.sum(exp_logits[:, None] * v_block, axis=0)

        # ---- merge with running values ----
        new_m      = tl.maximum(m_val, m_block)
        alpha_prev = tl.exp(m_val - new_m)
        alpha_blk  = tl.exp(m_block - new_m)

        o_vec = o_vec * alpha_prev + weighted_v * alpha_blk
        d_val = d_val * alpha_prev + sum_exp_block * alpha_blk
        m_val = new_m

        offset += BLOCK_T

    inv_d   = tl.where(d_val == 0, 0.0, 1.0 / d_val)
    out_vec = o_vec * inv_d
    log2e   = 1.4426950408889634
    lse_val = tl.where(d_val == 0,
                       -1e30,
                       (tl.log(d_val) + m_val) * log2e)

    # ---- store ----
    out_ptr_head = out_ptr + batch_idx * stride_q_batch + qo_head * stride_q_head + d_offs
    tl.store(out_ptr_head, tl.cast(out_vec, tl.bfloat16))

    lse_ptr_head = lse_ptr + batch_idx * NUM_QO_HEADS + qo_head
    tl.store(lse_ptr_head, lse_val)


def run(q,
        k_cache,
        v_cache,
        kv_indptr,
        kv_indices,
        sm_scale: float | None = None):
    """
    Entry point for gqa_paged_decode_h32_kv8_d128_ps1.
    Returns (output, lse).
    """
    if sm_scale is None:
        sm_scale = 1.0 / math.sqrt(128.0)

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA device is required to run Triton kernels.")

    # move tensors to GPU if necessary
    tensors = [q, k_cache, v_cache, kv_indptr, kv_indices]
    device_tensors = [t.cuda() if not t.is_cuda else t for t in tensors]
    q_dev, k_dev, v_dev, iptr_dev, idx_dev = [t.contiguous() for t in device_tensors]

    batch_size = q_dev.shape[0]
    num_qo_heads = 32
    head_dim = 128

    # squeeze page dimension (=1)
    k_dev_flat = k_dev.squeeze(1).contiguous()
    v_dev_flat = v_dev.squeeze(1).contiguous()

    out_dev = torch.empty((batch_size, num_qo_heads, head_dim),
                          dtype=torch.bfloat16,
                          device=q_dev.device)
    lse_dev = torch.empty((batch_size, num_qo_heads),
                          dtype=torch.float32,
                          device=q_dev.device)

    # launch kernel
    BLOCK_T = 128
    grid = (batch_size * num_qo_heads,)

    gqa_paged_decode_kernel[grid](
        q_dev, k_dev_flat, v_dev_flat,
        iptr_dev, idx_dev,
        sm_scale,
        out_dev, lse_dev,
        BLOCK_T=BLOCK_T,
        HEAD_DIM=128,
        NUM_QO_HEADS=32,
        NUM_KV_HEADS=8,
        num_warps=4,
        num_stages=4,
    )

    # move back to original device if needed
    if not q.is_cuda:
        return out_dev.cpu(), lse_dev.cpu()
    return out_dev, lse_dev