gemini-2.5-pro / tritonpr9imz

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


@triton.autotune(
    configs=[
        triton.Config({'BLOCK_KV_LEN': 16}, num_warps=4),
        triton.Config({'BLOCK_KV_LEN': 32}, num_warps=4),
        triton.Config({'BLOCK_KV_LEN': 64}, num_warps=4),
        triton.Config({'BLOCK_KV_LEN': 128}, num_warps=4),
        triton.Config({'BLOCK_KV_LEN': 256}, num_warps=4),
        triton.Config({'BLOCK_KV_LEN': 16}, num_warps=8),
        triton.Config({'BLOCK_KV_LEN': 32}, num_warps=8),
        triton.Config({'BLOCK_KV_LEN': 64}, num_warps=8),
        triton.Config({'BLOCK_KV_LEN': 128}, num_warps=8),
        triton.Config({'BLOCK_KV_LEN': 256}, num_warps=8),
    ],
    key=['HEAD_DIM'],
)
@triton.jit
def gqa_paged_decode_h32_kv8_d128_ps1_kernel(
    # Pointers to Tensors
    Q_ptr, K_cache_ptr, V_cache_ptr,
    kv_indptr_ptr, kv_indices_ptr,
    sm_scale,
    Output_ptr, LSE_ptr,
    # Stride Info
    stride_q_bs, stride_q_h,
    stride_k_num_pages, stride_k_ps, stride_k_h,
    stride_v_num_pages, stride_v_ps, stride_v_h,
    stride_o_bs, stride_o_h,
    stride_lse_bs,
    # Compile-time Constants
    NUM_QO_HEADS: tl.constexpr,
    NUM_KV_HEADS: tl.constexpr,
    HEAD_DIM: tl.constexpr,
    GQA_RATIO: tl.constexpr,
    PAGE_SIZE: tl.constexpr,
    BLOCK_D: tl.constexpr,
    BLOCK_KV_LEN: tl.constexpr,
):
    """
    Triton kernel for GQA paged decode attention.

    This kernel computes attention for a batch of query vectors against their
    corresponding key/value history, which is stored in a paged cache. Each
    program instance handles one query head for one sequence in the batch.

    Grid: (batch_size, num_qo_heads)
    - pid_b (program_id 0): batch index
    - pid_h (program_id 1): query head index

    Key Optimizations for B200:
    - Processes the variable-length KV sequence in fixed-size blocks (BLOCK_KV_LEN)
      to increase arithmetic intensity and hide memory latency from gathers.
    - Online softmax algorithm is used to compute attention scores and output
      in a single pass over the KV sequence, avoiding materialization of the
      full attention matrix.
    - For this decode-style kernel (one query vector per program), the Q@K.T and P@V
      operations are GEMV-like. Since tl.dot has minimum shape requirements (e.g., M>=16)
      for Tensor Core usage that are not met by a single vector, these operations
      are implemented using efficient element-wise operations and reductions.
    - Autotuning is enabled for BLOCK_KV_LEN and num_warps to find the optimal
      configuration for the target hardware.
    """
    # 1. Get program IDs for batch and query head
    pid_b = tl.program_id(0)
    pid_h = tl.program_id(1)

    # 2. Determine KV sequence length and handle empty sequences
    page_start = tl.load(kv_indptr_ptr + pid_b)
    page_end = tl.load(kv_indptr_ptr + pid_b + 1)
    seq_len = page_end - page_start

    # Early exit for sequences with no KV history
    if seq_len == 0:
        offs_d = tl.arange(0, BLOCK_D)
        out_ptr = Output_ptr + pid_b * stride_o_bs + pid_h * stride_o_h + offs_d
        lse_ptr = LSE_ptr + pid_b * stride_lse_bs + pid_h

        tl.store(out_ptr, tl.zeros([BLOCK_D], dtype=tl.bfloat16), mask=offs_d < BLOCK_D)
        tl.store(lse_ptr, -float('inf'))
        return

    # 3. Load query vector
    offs_d = tl.arange(0, BLOCK_D)
    q_ptr = Q_ptr + pid_b * stride_q_bs + pid_h * stride_q_h + offs_d
    q = tl.load(q_ptr, mask=offs_d < BLOCK_D).to(tl.float32)[None, :]  # Shape: [1, BLOCK_D]

    # 4. Initialize accumulators for online softmax
    # FIX: Use scalar accumulators for m_i and l_i to avoid shape errors
    # during the final tl.store operation for the LSE scalar.
    acc_o = tl.zeros([BLOCK_D], dtype=tl.float32)
    m_i = -float('inf')
    l_i = 0.0

    # 5. Determine the corresponding KV head for GQA
    kv_head_idx = pid_h // GQA_RATIO

    # 6. Loop over the KV sequence in blocks
    for offset in range(0, seq_len, BLOCK_KV_LEN):
        # a. Create masks and pointers for the current block
        offs_kv_block = offset + tl.arange(0, BLOCK_KV_LEN)
        mask_kv_block = offs_kv_block < seq_len
        indices_ptr = kv_indices_ptr + page_start + offs_kv_block

        # b. Gather page indices for K and V caches
        page_indices = tl.load(indices_ptr, mask=mask_kv_block, other=0)

        # c. Gather K vectors for the block
        offs_k_h = kv_head_idx * stride_k_h
        k_ptrs = K_cache_ptr + page_indices[:, None] * stride_k_num_pages + offs_k_h + offs_d[None, :]
        k = tl.load(k_ptrs, mask=mask_kv_block[:, None] & (offs_d[None, :] < BLOCK_D), other=0.0)

        # d. Compute scores S = Q @ K.T
        s_block = tl.sum(q * k.to(tl.float32), axis=1) * sm_scale
        # FIX: Keep scores as a 1D tensor [BLOCK_KV_LEN] for scalar reduction.
        s = tl.where(mask_kv_block, s_block, -float('inf'))

        # e. Online softmax update (with scalar state)
        m_block_max = tl.max(s, axis=0)
        m_curr = tl.maximum(m_i, m_block_max)
        p = tl.exp(s - m_curr)
        l_i_exp = tl.exp(m_i - m_curr)
        l_curr = l_i_exp * l_i + tl.sum(p, axis=0)

        # f. Gather V vectors for the block
        offs_v_h = kv_head_idx * stride_v_h
        v_ptrs = V_cache_ptr + page_indices[:, None] * stride_v_num_pages + offs_v_h + offs_d[None, :]
        v = tl.load(v_ptrs, mask=mask_kv_block[:, None] & (offs_d[None, :] < BLOCK_D), other=0.0)

        # g. Update output accumulator
        acc_o = acc_o * l_i_exp
        # FIX: Reshape 1D p to [BLOCK_KV_LEN, 1] for broadcasted matmul-like update.
        update_o = tl.sum(p[:, None] * v.to(tl.float32), axis=0)
        acc_o += update_o

        # h. Update state for the next iteration
        m_i = m_curr
        l_i = l_curr

    # 7. Finalize and store results
    l_i_safe = tl.where(l_i == 0, 1.0, l_i)
    acc_o = acc_o / l_i_safe

    # Calculate 2-based log-sum-exp
    lse = m_i + tl.log(l_i_safe)
    lse = lse / 0.6931471805599453

    # Store output vector and LSE value
    out_ptr = Output_ptr + pid_b * stride_o_bs + pid_h * stride_o_h + offs_d
    tl.store(out_ptr, acc_o.to(tl.bfloat16), mask=offs_d < BLOCK_D)
    # FIX: Storing the scalar `lse` value to a scalar pointer is now valid.
    tl.store(LSE_ptr + pid_b * stride_lse_bs + pid_h, lse)


def _gqa_paged_decode_h32_kv8_d128_ps1_launcher(q, k_cache, v_cache, kv_indptr, kv_indices, sm_scale):
    """
    Host-side wrapper for the Triton kernel.

    This function handles device management, output tensor allocation, grid
    computation, and kernel invocation. It ensures all tensors are on the
    same CUDA device before launching the kernel and moves the results back
    to the original device of the `q` tensor.
    """
    # 1. Validate inputs and extract dimensions
    assert q.dtype == torch.bfloat16
    assert k_cache.dtype == torch.bfloat16
    assert v_cache.dtype == torch.bfloat16
    assert kv_indptr.dtype == torch.int32
    assert kv_indices.dtype == torch.int32

    batch_size, num_qo_heads, head_dim = q.shape
    num_pages, page_size, num_kv_heads, _ = k_cache.shape

    # Check problem-specific constants
    assert num_qo_heads == 32, f"Expected num_qo_heads=32, got {num_qo_heads}"
    assert num_kv_heads == 8, f"Expected num_kv_heads=8, got {num_kv_heads}"
    assert head_dim == 128, f"Expected head_dim=128, got {head_dim}"
    assert page_size == 1, f"Expected page_size=1, got {page_size}"

    # 2. Set default sm_scale if not provided
    if sm_scale is None:
        sm_scale = 1.0 / math.sqrt(head_dim)

    # 3. Complete device management
    original_q_device = q.device

    # Find a common CUDA device for computation
    target_device = None
    for t in [q, k_cache, v_cache, kv_indptr, kv_indices]:
        if t.is_cuda:
            target_device = t.device
            break
    if target_device is None:
        if torch.cuda.is_available():
            target_device = torch.device("cuda")
        else:
            raise RuntimeError("This Triton kernel requires a CUDA-enabled GPU, but none was found.")

    # Move all tensors to the target CUDA device
    q = q.to(target_device)
    k_cache = k_cache.to(target_device)
    v_cache = v_cache.to(target_device)
    kv_indptr = kv_indptr.to(target_device)
    kv_indices = kv_indices.to(target_device)

    # 4. Allocate output tensors on the target device
    output = torch.empty_like(q)
    lse = torch.empty((batch_size, num_qo_heads), dtype=torch.float32, device=target_device)

    # 5. Set up grid and launch kernel
    grid = (batch_size, num_qo_heads)

    gqa_paged_decode_h32_kv8_d128_ps1_kernel[grid](
        q, k_cache, v_cache,
        kv_indptr, kv_indices,
        float(sm_scale),
        output, lse,
        # Strides
        q.stride(0), q.stride(1),
        k_cache.stride(0), k_cache.stride(1), k_cache.stride(2),
        v_cache.stride(0), v_cache.stride(1), v_cache.stride(2),
        output.stride(0), output.stride(1),
        lse.stride(0),
        # Constants
        NUM_QO_HEADS=num_qo_heads,
        NUM_KV_HEADS=num_kv_heads,
        HEAD_DIM=head_dim,
        GQA_RATIO=num_qo_heads // num_kv_heads,
        PAGE_SIZE=page_size,
        BLOCK_D=head_dim,
    )

    # 6. Move results back to the original device of `q`
    output = output.to(original_q_device)
    lse = lse.to(original_q_device)

    return output, lse

def run(*args, **kwargs):
    """
    Public entry point for the GQA paged decode attention kernel.

    This function acts as a flexible interface, accepting both positional and
    keyword arguments and forwarding them to the core launcher function.

    Args:
        q (torch.Tensor): Query tensor of shape [batch_size, 32, 128] and dtype bfloat16.
        k_cache (torch.Tensor): Key cache of shape [num_pages, 1, 8, 128] and dtype bfloat16.
        v_cache (torch.Tensor): Value cache of shape [num_pages, 1, 8, 128] and dtype bfloat16.
        kv_indptr (torch.Tensor): KV page offsets of shape [batch_size + 1] and dtype int32.
        kv_indices (torch.Tensor): Page IDs of shape [num_kv_indices] and dtype int32.
        sm_scale (float, optional): Softmax scale. Defaults to 1/sqrt(head_dim).

    Returns:
        Tuple[torch.Tensor, torch.Tensor]:
            - output: The attention output tensor of shape [batch_size, 32, 128] and dtype bfloat16.
            - lse: The log-sum-exp of attention logits (base 2) of shape [batch_size, 32] and dtype float32.
    """
    return _gqa_paged_decode_h32_kv8_d128_ps1_launcher(*args, **kwargs)