kernel.py

"""
Custom Triton flash-attention forward kernel optimized for H100.
Manual config selection per (head_dim, seq_len). Falls back to SDPA otherwise.
"""

KERNEL_TYPE = "flash_attention"

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


@triton.jit
def _attn_fwd_kernel(
    Q, K, V, sm_scale, Out,
    stride_qz, stride_qh, stride_qm, stride_qk,
    stride_kz, stride_kh, stride_kn, stride_kk,
    stride_vz, stride_vh, stride_vn, stride_vk,
    stride_oz, stride_oh, stride_om, stride_ok,
    Z, H, N_CTX,
    HEAD_DIM: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    start_m = tl.program_id(0)
    off_hz = tl.program_id(1)
    off_z = off_hz // H
    off_h = off_hz % H

    q_off = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh
    k_off = off_z.to(tl.int64) * stride_kz + off_h.to(tl.int64) * stride_kh
    v_off = off_z.to(tl.int64) * stride_vz + off_h.to(tl.int64) * stride_vh
    o_off = off_z.to(tl.int64) * stride_oz + off_h.to(tl.int64) * stride_oh

    Q_block_ptr = tl.make_block_ptr(
        base=Q + q_off, shape=(N_CTX, HEAD_DIM),
        strides=(stride_qm, stride_qk),
        offsets=(start_m * BLOCK_M, 0),
        block_shape=(BLOCK_M, HEAD_DIM), order=(1, 0),
    )
    K_block_ptr = tl.make_block_ptr(
        base=K + k_off, shape=(HEAD_DIM, N_CTX),
        strides=(stride_kk, stride_kn),
        offsets=(0, 0),
        block_shape=(HEAD_DIM, BLOCK_N), order=(0, 1),
    )
    V_block_ptr = tl.make_block_ptr(
        base=V + v_off, shape=(N_CTX, HEAD_DIM),
        strides=(stride_vn, stride_vk),
        offsets=(0, 0),
        block_shape=(BLOCK_N, HEAD_DIM), order=(1, 0),
    )
    O_block_ptr = tl.make_block_ptr(
        base=Out + o_off, shape=(N_CTX, HEAD_DIM),
        strides=(stride_om, stride_ok),
        offsets=(start_m * BLOCK_M, 0),
        block_shape=(BLOCK_M, HEAD_DIM), order=(1, 0),
    )

    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = tl.arange(0, BLOCK_N)

    m_i = tl.full([BLOCK_M], -float("inf"), dtype=tl.float32)
    l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
    acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)

    qk_scale = sm_scale * 1.44269504  # 1/log(2)
    q = tl.load(Q_block_ptr)

    hi = (start_m + 1) * BLOCK_M

    # Off-diagonal: no causal mask needed
    for start_n in range(0, start_m * BLOCK_M, BLOCK_N):
        start_n = tl.multiple_of(start_n, BLOCK_N)
        k = tl.load(K_block_ptr)
        qk = tl.dot(q, k) * qk_scale
        m_ij = tl.maximum(m_i, tl.max(qk, 1))
        qk = qk - m_ij[:, None]
        p = tl.math.exp2(qk)
        l_ij = tl.sum(p, 1)
        alpha = tl.math.exp2(m_i - m_ij)
        l_i = l_i * alpha + l_ij
        acc = acc * alpha[:, None]
        v = tl.load(V_block_ptr)
        p = p.to(v.dtype)
        acc = tl.dot(p, v, acc)
        m_i = m_ij
        K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
        V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))

    # Diagonal blocks: apply causal mask
    for start_n in range(start_m * BLOCK_M, hi, BLOCK_N):
        start_n = tl.multiple_of(start_n, BLOCK_N)
        k = tl.load(K_block_ptr)
        qk = tl.dot(q, k) * qk_scale
        mask = offs_m[:, None] >= (start_n + offs_n[None, :])
        qk = tl.where(mask, qk, -1.0e6)
        m_ij = tl.maximum(m_i, tl.max(qk, 1))
        qk = qk - m_ij[:, None]
        p = tl.math.exp2(qk)
        l_ij = tl.sum(p, 1)
        alpha = tl.math.exp2(m_i - m_ij)
        l_i = l_i * alpha + l_ij
        acc = acc * alpha[:, None]
        v = tl.load(V_block_ptr)
        p = p.to(v.dtype)
        acc = tl.dot(p, v, acc)
        m_i = m_ij
        K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
        V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))

    acc = acc / l_i[:, None]
    tl.store(O_block_ptr, acc.to(Out.type.element_ty))


def _select_config(D: int, S: int):
    """Return (BLOCK_M, BLOCK_N, num_warps, num_stages)."""
    if D == 64:
        if S <= 64:
            return 64, 64, 4, 3
        if S <= 128:
            return 128, 64, 4, 3
        # large
        return 128, 128, 4, 3
    else:  # D == 128
        if S <= 128:
            return 64, 32, 4, 3
        return 128, 64, 8, 3


def _triton_attention(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor) -> torch.Tensor:
    B, H, S, D = Q.shape
    sm_scale = D ** -0.5
    o = torch.empty_like(Q)

    BM, BN, nw, ns = _select_config(D, S)

    # Need S divisible by BM and BN
    if S % BM != 0 or S % BN != 0:
        return F.scaled_dot_product_attention(Q, K, V, is_causal=True)

    grid = (triton.cdiv(S, BM), B * H, 1)
    _attn_fwd_kernel[grid](
        Q, K, V, sm_scale, o,
        Q.stride(0), Q.stride(1), Q.stride(2), Q.stride(3),
        K.stride(0), K.stride(1), K.stride(2), K.stride(3),
        V.stride(0), V.stride(1), V.stride(2), V.stride(3),
        o.stride(0), o.stride(1), o.stride(2), o.stride(3),
        B, H, S,
        HEAD_DIM=D,
        BLOCK_M=BM,
        BLOCK_N=BN,
        num_warps=nw,
        num_stages=ns,
    )
    return o


def kernel_fn(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor) -> torch.Tensor:
    if Q.dim() != 4 or Q.dtype not in (torch.float16, torch.bfloat16):
        return F.scaled_dot_product_attention(Q, K, V, is_causal=True)
    D = Q.shape[-1]
    if D not in (64, 128):
        return F.scaled_dot_product_attention(Q, K, V, is_causal=True)
    return _triton_attention(Q, K, V)