Function kernel_fn

kernels/flash_attention.py:130–180

Entry point called by bench.py. Must match reference.flash_attention_ref signature. Args: Q: [batch, heads, seq_len, head_dim] K: [batch, heads, seq_len, head_dim] V: [batch, heads, seq_len, head_dim] causal: whether to apply causal masking sm_scale:

(
    Q: torch.Tensor,
    K: torch.Tensor,
    V: torch.Tensor,
    causal: bool = True,
    sm_scale: float = None,
)

def kernel_fn(
    Q: torch.Tensor,
    K: torch.Tensor,
    V: torch.Tensor,
    causal: bool = True,
    sm_scale: float = None,
) -> torch.Tensor:
    """
    Entry point called by bench.py. Must match reference.flash_attention_ref signature.

    Args:
        Q: [batch, heads, seq_len, head_dim]
        K: [batch, heads, seq_len, head_dim]
        V: [batch, heads, seq_len, head_dim]
        causal: whether to apply causal masking
        sm_scale: softmax scale factor, default 1/sqrt(head_dim)
    """
    assert Q.is_cuda and K.is_cuda and V.is_cuda

    Z, H, M_size, D = Q.shape
    _, _, N_size, _ = K.shape

    if sm_scale is None:
        sm_scale = 1.0 / math.sqrt(D)

    O = torch.empty_like(Q)

    # Block sizes -- must be powers of 2
    # D (head_dim) must be a constexpr and power of 2 for tl.trans to work
    assert D in (16, 32, 64, 128, 256), f"Head dim {D} not supported, must be power of 2 in [16..256]"

    BLOCK_M = 64
    BLOCK_N = 64

    grid = (triton.cdiv(M_size, BLOCK_M), H, Z)

    flash_attention_kernel[grid](
        Q, K, V, 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),
        Z, H, M_size, N_size,
        D=D,
        sm_scale=sm_scale,
        IS_CAUSAL=causal,
        BLOCK_M=BLOCK_M,
        BLOCK_N=BLOCK_N,
    )

    return O