claude-opus-4-1 / tritonde54a2

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

@triton.jit
def gqa_paged_decode_kernel(
    q_ptr, k_cache_ptr, v_cache_ptr,
    kv_indptr_ptr, kv_indices_ptr,
    output_ptr, lse_ptr,
    sm_scale,
    batch_size, num_pages,
    stride_qb, stride_qh, stride_qd,
    stride_kp, stride_kh, stride_kd,
    stride_vp, stride_vh, stride_vd,
    stride_ob, stride_oh, stride_od,
    stride_lseb, stride_lseh,
    log2_e,
    BLOCK_SIZE: tl.constexpr,
    HEAD_DIM: tl.constexpr,
    NUM_QO_HEADS: tl.constexpr,
    NUM_KV_HEADS: tl.constexpr,
    GQA_RATIO: tl.constexpr,
):
    # Get batch and head indices
    bid = tl.program_id(0)
    hid = tl.program_id(1)
    
    if bid >= batch_size or hid >= NUM_QO_HEADS:
        return
    
    # Get KV head index for GQA
    kv_hid = hid // GQA_RATIO
    
    # Get page range for this batch
    page_start = tl.load(kv_indptr_ptr + bid)
    page_end = tl.load(kv_indptr_ptr + bid + 1)
    num_tokens = page_end - page_start
    
    if num_tokens <= 0:
        # No KV cache for this batch element
        d_range = tl.arange(0, HEAD_DIM)
        output_offset = bid * stride_ob + hid * stride_oh + d_range * stride_od
        tl.store(output_ptr + output_offset, tl.zeros((HEAD_DIM,), dtype=tl.bfloat16))
        lse_offset = bid * stride_lseb + hid * stride_lseh
        tl.store(lse_ptr + lse_offset, float('-inf'))
        return
    
    # Load query vector for this head
    d_range = tl.arange(0, HEAD_DIM)
    q_offset = bid * stride_qb + hid * stride_qh + d_range * stride_qd
    q = tl.load(q_ptr + q_offset).to(tl.float32)
    
    # Initialize accumulators for online softmax
    m_i = float('-inf')
    l_i = 0.0
    acc = tl.zeros((HEAD_DIM,), dtype=tl.float32)
    
    # Process KV cache in blocks
    for token_idx in range(0, num_tokens, BLOCK_SIZE):
        # Create mask for valid tokens in this block
        token_range = tl.arange(0, BLOCK_SIZE)
        token_pos = token_idx + token_range
        mask = token_pos < num_tokens
        
        # Load page indices for this block
        page_indices = tl.load(
            kv_indices_ptr + page_start + token_pos,
            mask=mask,
            other=0
        )
        
        # Compute attention scores for this block
        scores = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
        
        # Load K vectors and compute dot products
        for i in range(BLOCK_SIZE):
            if token_idx + i < num_tokens:
                page_id = tl.load(kv_indices_ptr + page_start + token_idx + i)
                k_offset = page_id * stride_kp + kv_hid * stride_kh + d_range * stride_kd
                k = tl.load(k_cache_ptr + k_offset).to(tl.float32)
                score = tl.sum(q * k) * sm_scale
                scores = tl.where(token_range == i, score, scores)
        
        # Apply mask to scores
        scores = tl.where(mask, scores, float('-inf'))
        
        # Find maximum score in this block
        m_ij = tl.max(scores, axis=0)
        m_new = tl.maximum(m_i, m_ij)
        
        # Compute exponentials with numerical stability
        exp_scores = tl.exp(scores - m_new)
        exp_scores = tl.where(mask, exp_scores, 0.0)
        
        # Update running statistics
        alpha = tl.exp(m_i - m_new)
        l_new = alpha * l_i + tl.sum(exp_scores)
        
        # Scale previous accumulator
        acc = acc * alpha
        
        # Accumulate weighted V vectors
        for i in range(BLOCK_SIZE):
            if token_idx + i < num_tokens:
                page_id = tl.load(kv_indices_ptr + page_start + token_idx + i)
                v_offset = page_id * stride_vp + kv_hid * stride_vh + d_range * stride_vd
                v = tl.load(v_cache_ptr + v_offset).to(tl.float32)
                weight = tl.where(token_range == i, exp_scores, 0.0)
                acc = acc + tl.sum(weight) * v
        
        # Update state
        m_i = m_new
        l_i = l_new
    
    # Normalize output
    output = acc / l_i
    
    # Store output
    output_offset = bid * stride_ob + hid * stride_oh + d_range * stride_od
    tl.store(output_ptr + output_offset, output.to(tl.bfloat16))
    
    # Compute and store LSE (2-based)
    # lse = (m_i + log(l_i)) / log(2) = (m_i + log(l_i)) * log2(e)
    lse = (m_i + tl.log(l_i)) * log2_e
    lse_offset = bid * stride_lseb + hid * stride_lseh
    tl.store(lse_ptr + lse_offset, lse)

def run(q, k_cache, v_cache, kv_indptr, kv_indices, sm_scale=None):
    # Handle device management
    device = None
    if q.is_cuda:
        device = q.device
    elif torch.cuda.is_available():
        device = torch.device('cuda')
        q = q.cuda()
        k_cache = k_cache.cuda() if not k_cache.is_cuda else k_cache
        v_cache = v_cache.cuda() if not v_cache.is_cuda else v_cache
        kv_indptr = kv_indptr.cuda() if not kv_indptr.is_cuda else kv_indptr
        kv_indices = kv_indices.cuda() if not kv_indices.is_cuda else kv_indices
    else:
        raise RuntimeError("CUDA is not available for GPU tensors")
    
    # Get dimensions
    batch_size, num_qo_heads, head_dim = q.shape
    num_pages, page_size, num_kv_heads, _ = k_cache.shape
    
    # Verify constants
    assert num_qo_heads == 32, f"num_qo_heads must be 32, got {num_qo_heads}"
    assert num_kv_heads == 4, f"num_kv_heads must be 4, got {num_kv_heads}"
    assert head_dim == 128, f"head_dim must be 128, got {head_dim}"
    assert page_size == 1, f"page_size must be 1, got {page_size}"
    
    # Compute GQA ratio
    gqa_ratio = num_qo_heads // num_kv_heads
    
    # Set default sm_scale if not provided
    if sm_scale is None:
        sm_scale = 1.0 / math.sqrt(head_dim)
    
    # Allocate output tensors
    output = torch.zeros((batch_size, num_qo_heads, head_dim), dtype=torch.bfloat16, device=device)
    lse = torch.full((batch_size, num_qo_heads), -float('inf'), dtype=torch.float32, device=device)
    
    # Squeeze page_size dimension since it's 1
    k_cache_flat = k_cache.squeeze(1)
    v_cache_flat = v_cache.squeeze(1)
    
    # Define block size for token processing
    BLOCK_SIZE = 64
    
    # Precompute log2(e) for LSE calculation
    log2_e = 1.0 / math.log(2.0)
    
    # Launch kernel
    grid = (batch_size, num_qo_heads)
    
    gqa_paged_decode_kernel[grid](
        q, k_cache_flat, v_cache_flat,
        kv_indptr, kv_indices,
        output, lse,
        sm_scale,
        batch_size, num_pages,
        q.stride(0), q.stride(1), q.stride(2),
        k_cache_flat.stride(0), k_cache_flat.stride(1), k_cache_flat.stride(2),
        v_cache_flat.stride(0), v_cache_flat.stride(1), v_cache_flat.stride(2),
        output.stride(0), output.stride(1), output.stride(2),
        lse.stride(0), lse.stride(1),
        log2_e,
        BLOCK_SIZE=BLOCK_SIZE,
        HEAD_DIM=head_dim,
        NUM_QO_HEADS=num_qo_heads,
        NUM_KV_HEADS=num_kv_heads,
        GQA_RATIO=gqa_ratio,
    )
    
    # Move outputs back to original device if needed
    if not q.is_cuda and torch.cuda.is_available():
        output = output.cpu()
        lse = lse.cpu()
    
    return output, lse