Uploaded using `kernel-builder`.

build/torch-rocm/__init__.py

@@ -0,0 +1,32 @@
+"""AITER Flash Attention kernels for AMD ROCm (Triton MHA).
+
+Repackaged from the `aiter.ops.triton.attention.mha` module of the
+`ROCm/aiter <https://github.com/ROCm/aiter>`_ project as a self-contained
+Hugging Face Hub kernel. Exposes the two FlashAttention entry points used by
+the `transformers` ROCm fallback.
+"""
+
+from .mha import (
+    flash_attn_func,
+    flash_attn_varlen_func,
+    flash_attn_with_kvcache,
+    mha_set_impl,
+    mha_set_use_fused_bwd_kernel,
+    mha_set_use_int64_strides,
+)
+
+
+__kernel_metadata__ = {
+    "license": "mit",
+}
+
+
+__all__ = [
+    "__kernel_metadata__",
+    "flash_attn_func",
+    "flash_attn_varlen_func",
+    "flash_attn_with_kvcache",
+    "mha_set_impl",
+    "mha_set_use_fused_bwd_kernel",
+    "mha_set_use_int64_strides",
+]

build/torch-rocm/_kernels/flash_attn_triton_amd/__init__.py

@@ -0,0 +1,4 @@
+from . import interface_v2 as flash_attn_2
+from . import interface_v3 as flash_attn_3
+
+__all__ = ["flash_attn_2", "flash_attn_3"]

build/torch-rocm/_kernels/flash_attn_triton_amd/common.py

@@ -0,0 +1,573 @@
+"""
+Triton kernel helper functions shared across flash attention modules.
+
+This module contains Triton JIT-compiled helper functions that are used within
+the main attention kernels (fwd_prefill, fwd_decode, bwd). These are kept
+separate from utils.py to allow stricter type checking on pure Python utilities.
+"""
+
+from typing import Literal, Optional, Tuple, Union
+
+import torch
+import triton
+import triton.language as tl
+
+from .utils import DEBUG, get_shape_from_layout, get_stride_from_layout, is_fp8
+
+
+@triton.jit
+def compute_alibi_block(
+    alibi_slope, seqlen_q, seqlen_k, offs_m, offs_n, transpose=False
+):
+    """
+    Compute ALiBi (Attention with Linear Biases) block.
+
+    When seqlen_k and seqlen_q are different, the diagonal sticks to the
+    bottom right of the attention matrix.
+    """
+    # e.g. alibi_slope = 1, seqlen_q = 2, seqlen_k = 5
+    # offs_m = [0, 1], offs_n = [0, 1, 2, 3, 4]
+    # Result: [[-3, -2, -1, 0, -1], [-4, -3, -2, -1, 0]]
+    relative_pos_block = offs_m[:, None] + seqlen_k - seqlen_q - offs_n[None, :]
+    alibi_block = -1 * alibi_slope * tl.abs(relative_pos_block)
+    if transpose:
+        return alibi_block.T
+    else:
+        return alibi_block
+
+
+@triton.jit
+def compute_fp8_scaling_factors(x, fp8_max: tl.constexpr):
+    """Compute FP8 scaling and descaling factors for a block."""
+    x_amax = tl.max(tl.abs(x))
+    x_amax = tl.where(x_amax <= 1e-9, 1e-9, x_amax)
+    scale_x = fp8_max / x_amax
+    descale_x = x_amax / fp8_max
+    return scale_x, descale_x
+
+
+@triton.jit
+def _cast_varlen_to_fp8_kernel_2d(
+    X,
+    X_fp8,
+    Descale,
+    cu_seqlens,
+    H,
+    MAX_SEQLEN,
+    stride_batch,
+    stride_seq,
+    stride_head,
+    stride_dim,
+    stride_out_batch,
+    stride_out_seq,
+    stride_out_head,
+    stride_out_dim,
+    stride_desc_batch,
+    stride_desc_head,
+    FP8_CLAMP_VAL,
+    FP8_MAX,
+    BLOCK_SIZE: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    ACTUAL_HEAD_DIM: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    """Cast tensor to FP8 with per-(batch, head) scaling."""
+    b_id = tl.program_id(0)
+    h_id = tl.program_id(1)
+
+    # Get sequence bounds for this batch
+    if IS_VARLEN:
+        seq_start = tl.load(cu_seqlens + b_id)
+        seq_end = tl.load(cu_seqlens + b_id + 1)
+        seqlen = seq_end - seq_start
+    else:
+        seq_start = 0
+        seqlen = MAX_SEQLEN
+
+    # initialize max value tracker
+    x_max_val = 0.0
+
+    # STEP 1: Find max absolute value across the entire sequence
+    num_of_blocks = tl.cdiv(seqlen, BLOCK_SIZE)
+    for blk_idx in range(0, num_of_blocks):
+        offs_seq = blk_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+        offs_dim = tl.arange(0, HEAD_DIM)
+
+        mask_seq = offs_seq[:, None] < seqlen
+        if ACTUAL_HEAD_DIM != HEAD_DIM:
+            mask_dim = offs_dim[None, :] < ACTUAL_HEAD_DIM
+            mask_seq = mask_seq & mask_dim
+
+        adj_x = (
+            b_id * stride_batch
+            + h_id * stride_head
+            + seq_start * stride_seq
+            + offs_seq[:, None] * stride_seq
+            + offs_dim[None, :] * stride_dim
+        )
+        x_block = tl.load(X + adj_x, mask=mask_seq, other=0.0)
+        block_max = tl.max(tl.abs(x_block))
+        x_max_val = tl.maximum(x_max_val, block_max)
+
+    # clamp to avoid division by zero
+    x_max_val = tl.maximum(x_max_val, FP8_CLAMP_VAL)
+
+    # compute scale and descale factors
+    scale = FP8_MAX / x_max_val
+    descale = x_max_val / FP8_MAX
+
+    # store descale factor
+    desc_ptr = Descale + b_id * stride_desc_batch + h_id
+    tl.store(desc_ptr, descale)
+
+    # STEP 2: Apply scaling and convert to FP8
+    for blk_idx in range(0, num_of_blocks):
+        offs_seq = blk_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+        offs_dim = tl.arange(0, HEAD_DIM)
+
+        mask_seq = offs_seq[:, None] < seqlen
+        if ACTUAL_HEAD_DIM != HEAD_DIM:
+            mask_dim = offs_dim[None, :] < ACTUAL_HEAD_DIM
+            mask_seq = mask_seq & mask_dim
+
+        addr = (
+            b_id * stride_batch
+            + h_id * stride_head
+            + seq_start * stride_seq
+            + offs_seq[:, None] * stride_seq
+            + offs_dim[None, :] * stride_dim
+        )
+        x_block = tl.load(X + addr, mask=mask_seq, other=0.0)
+        x_fp8_block = (x_block * scale).to(X_fp8.type.element_ty)
+
+        addr_out = (
+            b_id * stride_out_batch
+            + h_id * stride_out_head
+            + seq_start * stride_out_seq
+            + offs_seq[:, None] * stride_out_seq
+            + offs_dim[None, :] * stride_out_dim
+        )
+        tl.store(X_fp8 + addr_out, x_fp8_block, mask=mask_seq)
+
+
+@triton.jit
+def _rotary_kernel(
+    OUT,
+    X,
+    COS,
+    SIN,
+    CU_SEQLENS,
+    SEQLEN_OFFSETS,
+    seqlen,
+    nheads,
+    seqlen_ro,
+    stride_out_batch,
+    stride_out_seqlen,
+    stride_out_nheads,
+    stride_out_headdim,
+    stride_x_batch,
+    stride_x_seqlen,
+    stride_x_nheads,
+    stride_x_headdim,
+    ROTARY_DIM: tl.constexpr,
+    IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    INTERLEAVED: tl.constexpr,
+    CONJUGATE: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+):
+    """Apply rotary positional embeddings."""
+    BLOCK_K: tl.constexpr = triton.next_power_of_2(ROTARY_DIM)
+    ROTARY_DIM_HALF = ROTARY_DIM // 2
+    pid_head = tl.program_id(axis=0)
+    pid_m = tl.program_id(axis=1)
+    pid_batch = tl.program_id(axis=2)
+
+    if not IS_VARLEN:
+        X = X + pid_batch * stride_x_batch
+        OUT = OUT + pid_batch * stride_out_batch
+    else:
+        start_idx = tl.load(CU_SEQLENS + pid_batch)
+        seqlen = tl.load(CU_SEQLENS + pid_batch + 1) - start_idx
+        X = X + start_idx * stride_x_seqlen
+        OUT = OUT + start_idx * stride_out_seqlen
+
+    if pid_m * BLOCK_M >= seqlen:
+        return
+
+    rh = pid_head * BLOCK_H + tl.arange(0, BLOCK_H)
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    if not IS_SEQLEN_OFFSETS_TENSOR:
+        rm_cs = rm + SEQLEN_OFFSETS
+    else:
+        rm_cs = rm + tl.load(SEQLEN_OFFSETS + pid_batch)
+
+    rk_half = tl.arange(0, BLOCK_K // 2)
+    COS = COS + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :])
+    SIN = SIN + (rm_cs[:, None] * ROTARY_DIM_HALF + rk_half[None, :])
+    mask_cs = (rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < ROTARY_DIM_HALF)
+    cos = tl.load(COS, mask=mask_cs, other=1.0).to(tl.float32)
+    sin = tl.load(SIN, mask=mask_cs, other=0.0).to(tl.float32)
+    if CONJUGATE:
+        sin = -sin
+
+    if not INTERLEAVED:
+        X = X + (
+            rh[:, None, None] * stride_x_nheads
+            + rm[None, :, None] * stride_x_seqlen
+            + rk_half[None, None, :] * stride_x_headdim
+        )
+        OUT = OUT + (
+            rh[:, None, None] * stride_out_nheads
+            + rm[None, :, None] * stride_out_seqlen
+            + rk_half[None, None, :] * stride_out_headdim
+        )
+        mask = (
+            (rh[:, None, None] < nheads)
+            & (rm[None, :, None] < seqlen)
+            & (rk_half[None, None, :] < ROTARY_DIM_HALF)
+        )
+        x0 = tl.load(X, mask=mask, other=0.0).to(tl.float32)
+        x1 = tl.load(X + ROTARY_DIM_HALF * stride_x_headdim, mask=mask, other=0.0).to(
+            tl.float32
+        )
+        o0 = x0 * cos - x1 * sin
+        o1 = x0 * sin + x1 * cos
+        tl.store(OUT, o0, mask=mask)
+        tl.store(OUT + ROTARY_DIM_HALF * stride_out_headdim, o1, mask=mask)
+    else:
+        rk = tl.arange(0, BLOCK_K)
+        X = X + (
+            rh[:, None, None] * stride_x_nheads
+            + rm[None, :, None] * stride_x_seqlen
+            + rk[None, None, :] * stride_x_headdim
+        )
+        OUT = OUT + (
+            rh[:, None, None] * stride_out_nheads
+            + rm[None, :, None] * stride_out_seqlen
+            + rk[None, None, :] * stride_out_headdim
+        )
+        mask = (
+            (rh[:, None, None] < nheads)
+            & (rm[None, :, None] < seqlen)
+            & (rk[None, None, :] < ROTARY_DIM)
+        )
+        x = tl.load(X, mask=mask, other=0.0).to(tl.float32)
+        x0, x1 = tl.split(tl.reshape(x, [BLOCK_H, BLOCK_M, BLOCK_K // 2, 2]))
+        o0 = x0 * cos - x1 * sin
+        o1 = x0 * sin + x1 * cos
+        o = tl.reshape(tl.join(o0, o1), [BLOCK_H, BLOCK_M, BLOCK_K])
+        tl.store(OUT, o, mask=mask)
+
+
+# -------------------------------
+# Python wrappers for Triton kernels
+# -------------------------------
+
+
+def cast_to_fp8(
+    x: torch.Tensor,
+    fp8_dtype: torch.dtype,
+    layout: Literal["bshd", "thd"],
+    clamp_val: float = 1e-9,
+    cu_seqlens: Optional[torch.Tensor] = None,
+    max_seqlen: Optional[int] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Cast tensor to FP8 with per-(batch, head) scaling factors."""
+    if DEBUG > 0:
+        print()
+        print("cast_to_fp8")
+        print("x:", x, x.shape)
+        print("fp8_dtype:", fp8_dtype)
+        print("cu_seqlens:", cu_seqlens)
+        print("max_seqlen:", max_seqlen)
+        print("clamp_val:", clamp_val)
+
+    assert x.dtype in {
+        torch.float16,
+        torch.float32,
+        torch.float64,
+        torch.bfloat16,
+    } and is_fp8(fp8_dtype), f"Cannot cast {x.dtype} to {fp8_dtype}"
+
+    batch, max_seqlen_final, num_heads, head_dim = get_shape_from_layout(
+        x, layout, cu_seqlens, max_seqlen
+    )
+    is_varlen = layout == "thd"
+    fp8_max = torch.finfo(fp8_dtype).max
+
+    padded_head_dim = 1 << (head_dim - 1).bit_length()
+    padded_head_dim = max(padded_head_dim, 32)
+
+    x_fp8 = torch.zeros_like(x, dtype=fp8_dtype)
+    descale_factors = torch.zeros(
+        (batch, num_heads), device=x.device, dtype=torch.float32
+    )
+    BLOCK_SIZE = 128
+
+    stride_batch, stride_head, stride_seq, stride_dim = get_stride_from_layout(
+        x, layout
+    )
+    stride_out_batch, stride_out_head, stride_out_seq, stride_out_dim = (
+        get_stride_from_layout(x_fp8, layout)
+    )
+    stride_desc_batch, stride_desc_head = descale_factors.stride()
+
+    grid = (batch, num_heads)
+    _cast_varlen_to_fp8_kernel_2d[grid](
+        x,
+        x_fp8,
+        descale_factors,
+        cu_seqlens,
+        num_heads,
+        max_seqlen_final,
+        stride_batch,
+        stride_seq,
+        stride_head,
+        stride_dim,
+        stride_out_batch,
+        stride_out_seq,
+        stride_out_head,
+        stride_out_dim,
+        stride_desc_batch,
+        stride_desc_head,
+        clamp_val,
+        fp8_max,
+        BLOCK_SIZE=BLOCK_SIZE,
+        HEAD_DIM=padded_head_dim,
+        ACTUAL_HEAD_DIM=head_dim,
+        IS_VARLEN=is_varlen,
+    )
+
+    return x_fp8, descale_factors
+
+
+def _apply_rotary_kernel(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    seqlen_offsets: Union[int, torch.Tensor] = 0,
+    cu_seqlens: Optional[torch.Tensor] = None,
+    max_seqlen: Optional[int] = None,
+    interleaved: bool = False,
+    inplace: bool = False,
+    conjugate: bool = False,
+) -> torch.Tensor:
+    """Apply rotary positional embeddings using Triton kernel."""
+    is_varlen = cu_seqlens is not None
+    if not is_varlen:
+        batch, seqlen, nheads, headdim = x.shape
+    else:
+        assert (
+            max_seqlen is not None
+        ), "If cu_seqlens is passed, max_seqlen must also be provided"
+        total_seqlen, nheads, headdim = x.shape
+        assert cu_seqlens is not None
+        batch_p_1 = cu_seqlens.shape[0]
+        batch = batch_p_1 - 1
+        seqlen = max_seqlen
+    seqlen_ro, rotary_dim_half = cos.shape
+    assert sin.shape == cos.shape
+    rotary_dim = 2 * rotary_dim_half
+    assert rotary_dim <= headdim
+    assert headdim <= 256
+    assert seqlen_ro >= seqlen
+
+    cos, sin = cos.contiguous(), sin.contiguous()
+    if isinstance(seqlen_offsets, torch.Tensor):
+        assert seqlen_offsets.shape == (batch,)
+        assert seqlen_offsets.dtype in (torch.int32, torch.int64)
+        seqlen_offsets = seqlen_offsets.contiguous()
+    else:
+        assert seqlen_offsets + seqlen <= seqlen_ro
+
+    out = torch.empty_like(x) if not inplace else x
+    if rotary_dim < headdim and not inplace:
+        out[..., rotary_dim:].copy_(x[..., rotary_dim:])
+
+    BLOCK_M = 8 if rotary_dim <= 128 else 4
+    grid = (
+        triton.cdiv(nheads, 2),
+        triton.cdiv(seqlen, BLOCK_M),
+        batch,
+    )
+
+    with torch.cuda.device(x.device.index):
+        torch.library.wrap_triton(_rotary_kernel)[grid](
+            out,
+            x,
+            cos,
+            sin,
+            cu_seqlens,
+            seqlen_offsets,
+            seqlen,
+            nheads,
+            seqlen_ro,
+            out.stride(0) if not is_varlen else 0,
+            out.stride(-3),
+            out.stride(-2),
+            out.stride(-1),
+            x.stride(0) if not is_varlen else 0,
+            x.stride(-3),
+            x.stride(-2),
+            x.stride(-1),
+            rotary_dim,
+            isinstance(seqlen_offsets, torch.Tensor),
+            is_varlen,
+            interleaved,
+            conjugate,
+            BLOCK_M=BLOCK_M,
+            BLOCK_H=2,
+        )
+    return out
+
+
+class _ApplyRotary(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        interleaved: bool,
+        inplace: bool,
+        seqlen_offsets: Union[int, torch.Tensor],
+        cu_seqlens: Optional[torch.Tensor],
+        max_seqlen: Optional[int],
+    ) -> torch.Tensor:
+        out = _apply_rotary_kernel(
+            x,
+            cos,
+            sin,
+            seqlen_offsets=seqlen_offsets,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+            interleaved=interleaved,
+            inplace=inplace,
+            conjugate=False,
+        )
+        if isinstance(seqlen_offsets, int):
+            ctx.save_for_backward(cos, sin, cu_seqlens)
+            ctx.seqlen_offsets = seqlen_offsets
+        else:
+            ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets)
+            ctx.seqlen_offsets = None
+        ctx.interleaved = interleaved
+        ctx.inplace = inplace
+        ctx.max_seqlen = max_seqlen
+        return out if not inplace else x
+
+    @staticmethod
+    def backward(
+        ctx, do: torch.Tensor
+    ) -> tuple[torch.Tensor, None, None, None, None, None, None, None]:
+        seqlen_offsets = ctx.seqlen_offsets
+        if seqlen_offsets is None:
+            cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors
+        else:
+            cos, sin, cu_seqlens = ctx.saved_tensors
+        dx = _apply_rotary_kernel(
+            do,
+            cos,
+            sin,
+            seqlen_offsets=seqlen_offsets,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=ctx.max_seqlen,
+            interleaved=ctx.interleaved,
+            inplace=ctx.inplace,
+            conjugate=True,
+        )
+        return dx, None, None, None, None, None, None, None
+
+
+def apply_rotary_emb(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    interleaved: bool = False,
+    inplace: bool = False,
+    seqlen_offsets: Union[int, torch.Tensor] = 0,
+    cu_seqlens: Optional[torch.Tensor] = None,
+    max_seqlen: Optional[int] = None,
+) -> torch.Tensor:
+    """Apply rotary embeddings to tensor x.
+
+    Args:
+        x: (B, S, H, D) if `cu_seqlens` is None else (total_S, H, D).
+        cos, sin: (S_rotary, rotary_dim/2)
+        interleaved: GPT-J style if True.
+        inplace: modify x in place.
+        seqlen_offsets: int or (B,) tensor of starting offsets per sequence.
+        cu_seqlens: (B+1,) tensor enabling varlen mode.
+        max_seqlen: required when `cu_seqlens` is provided.
+    """
+    original_dtype = x.dtype
+    is_fp8_input = is_fp8(original_dtype)
+    if is_fp8_input:
+        target_dtype = (
+            torch.bfloat16
+            if cos.dtype == torch.bfloat16 or torch.cuda.is_bf16_supported()
+            else torch.float16
+        )
+        x_up = x.to(target_dtype)
+        cos_up = cos.to(target_dtype) if cos.dtype != target_dtype else cos
+        sin_up = sin.to(target_dtype) if sin.dtype != target_dtype else sin
+        out_up = _ApplyRotary.apply(
+            x_up,
+            cos_up,
+            sin_up,
+            interleaved,
+            False,
+            seqlen_offsets,
+            cu_seqlens,
+            max_seqlen,
+        )
+        if inplace:
+            x.copy_(out_up.to(original_dtype))
+            return x
+        return out_up.to(original_dtype)
+    else:
+        return _ApplyRotary.apply(
+            x, cos, sin, interleaved, inplace, seqlen_offsets, cu_seqlens, max_seqlen
+        )
+
+
+def apply_rotary(
+    q: torch.Tensor,
+    k_new: Optional[torch.Tensor],
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    *,
+    causal: bool,
+    local: bool,
+    interleaved: bool = False,
+    seqlen_offsets: Union[int, torch.Tensor] = 0,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    """Apply rotary embeddings to q and optionally k_new.
+
+    Policy:
+      - If causal OR local attention: apply rotary directly on (B, S, H, D).
+      - Else (non-causal global): flatten heads into sequence, apply, unflatten.
+      - k_new is always rotated directly when provided.
+    """
+    assert q.ndim == 4, f"Expected q shape (B,S,H,D), got {q.shape}"
+    B, S, H, D = q.shape
+    use_flatten = (not causal) and (not local)
+
+    if use_flatten:
+        q_flat = q.reshape(B, S * H, D).unsqueeze(1)
+        q_flat = apply_rotary_emb(
+            q_flat, cos, sin, interleaved=interleaved, seqlen_offsets=seqlen_offsets
+        )
+        q = q_flat.view(B, 1, S * H, D).reshape(B, S, H, D)
+    else:
+        q = apply_rotary_emb(
+            q, cos, sin, interleaved=interleaved, seqlen_offsets=seqlen_offsets
+        )
+
+    if k_new is not None:
+        k_new = apply_rotary_emb(
+            k_new, cos, sin, interleaved=interleaved, seqlen_offsets=seqlen_offsets
+        )
+    return q, k_new

build/torch-rocm/_kernels/flash_attn_triton_amd/fwd_decode.py

@@ -0,0 +1,1475 @@
+import warnings
+import torch
+import triton
+import triton.language as tl
+from typing import Literal, Optional
+from .common import apply_rotary
+from .utils import (
+    DEBUG,
+    AUTOTUNE,
+    AutotuneMode,
+    get_arch,
+    get_padded_headsize,
+    get_shape_from_layout,
+    get_stride_from_layout,
+    is_fp8,
+)
+
+FWD_DECODE_AUTOTUNE_KEYS = [
+    "N_CTX_Q",
+    "N_CTX_K",
+    "ACTUAL_BLOCK_DMODEL",
+    "H_q",
+    "H_kv",
+    "IS_CAUSAL",
+    "IS_GQA",
+]
+
+# Maximum BLOCK_M across all configs (for intermediate tensor allocation)
+MAX_BLOCK_M = 64
+
+
+def get_fwd_decode_configs(mode: AutotuneMode):
+    """
+    Returns configs for both the splitK kernel and reduce kernel.
+
+    Returns:
+        (splitk_configs, reduce_config): Tuple of config lists for each kernel
+    """
+
+    if mode == "off":
+        arch = get_arch()
+        if arch.is_rdna:
+            return (
+                [
+                    triton.Config(
+                        {"BLOCK_M": 32, "BLOCK_N": 32},
+                        num_stages=1,
+                        num_warps=4,
+                    ),
+                ],
+                [triton.Config({}, num_stages=1, num_warps=4)],
+            )
+        elif arch.is_cdna:
+            return (
+                [
+                    triton.Config(
+                        {"BLOCK_M": 64, "BLOCK_N": 64, "waves_per_eu": 1},
+                        num_stages=1,
+                        num_warps=4,
+                    ),
+                ],
+                [triton.Config({}, num_stages=1, num_warps=4)],
+            )
+        else:
+            return (
+                [
+                    triton.Config(
+                        {"BLOCK_M": 64, "BLOCK_N": 64, "waves_per_eu": 1},
+                        num_stages=1,
+                        num_warps=4,
+                    ),
+                ],
+                [triton.Config({}, num_stages=1, num_warps=4)],
+            )
+
+    elif mode == "on":
+        arch = get_arch()
+        if arch.is_rdna:
+            return (
+                [
+                    triton.Config(
+                        {"BLOCK_M": 32, "BLOCK_N": 32},
+                        num_stages=1,
+                        num_warps=4,
+                    ),
+                ],
+                [triton.Config({}, num_stages=1, num_warps=4)],
+            )
+        elif arch.is_cdna:
+            return (
+                [
+                    triton.Config(
+                        {"BLOCK_M": 64, "BLOCK_N": 64, "waves_per_eu": 1},
+                        num_stages=1,
+                        num_warps=4,
+                    ),
+                ],
+                [triton.Config({}, num_stages=1, num_warps=4)],
+            )
+        else:
+            return (
+                [
+                    triton.Config(
+                        {"BLOCK_M": 64, "BLOCK_N": 64, "waves_per_eu": 1},
+                        num_stages=1,
+                        num_warps=4,
+                    ),
+                ],
+                [triton.Config({}, num_stages=1, num_warps=4)],
+            )
+
+    else:  # sweep
+        arch = get_arch()
+        splitk_configs = []
+
+        BLOCK_M_OPTIONS = [64, 32, 16]
+        BLOCK_N_OPTIONS = [128, 64, 32, 16]
+        NUM_WARPS_OPTIONS = [2, 4]
+        NUM_STAGES_OPTIONS = [1]
+        WAVES_PER_EU_OPTIONS = [4, 2, 1]
+
+        # Ensure BLOCK_M options don't exceed MAX_BLOCK_M
+        assert all(
+            bm <= MAX_BLOCK_M for bm in BLOCK_M_OPTIONS
+        ), f"BLOCK_M_OPTIONS {BLOCK_M_OPTIONS} exceeds MAX_BLOCK_M {MAX_BLOCK_M}"
+
+        for bm in BLOCK_M_OPTIONS:
+            for bn in BLOCK_N_OPTIONS:
+                for waves in WAVES_PER_EU_OPTIONS:
+                    for nw in NUM_WARPS_OPTIONS:
+                        for ns in NUM_STAGES_OPTIONS:
+                            splitk_configs.append(
+                                triton.Config(
+                                    {
+                                        "BLOCK_M": bm,
+                                        "BLOCK_N": bn,
+                                        "waves_per_eu": waves,
+                                    },
+                                    num_stages=ns,
+                                    num_warps=nw,
+                                )
+                            )
+
+        # Reduce kernel configs - sweep num_warps
+        NUM_WARPS_REDUCE_OPTIONS = [2, 4]
+        reduce_configs = [
+            triton.Config({}, num_stages=1, num_warps=nw)
+            for nw in NUM_WARPS_REDUCE_OPTIONS
+        ]
+
+        return splitk_configs, reduce_configs
+
+
+fwd_decode_splitk_configs, fwd_decode_reduce_configs = get_fwd_decode_configs(AUTOTUNE)
+
+
+@triton.jit
+def _attn_fwd_inner(
+    q,
+    kT,
+    v,
+    pos,
+    col_mask,
+    m_i,
+    l_i,
+    acc,
+    pid_m,
+    q_descale,
+    k_descale,
+    v_descale,  # FP8 scaling factors
+    alibi_slope,
+    apply_col_mask,
+    IS_FP8: tl.constexpr,  # FP8 flag
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    N_CTX_Q: tl.constexpr,
+    N_CTX_K_FINAL: tl.constexpr,
+    USE_ALIBI: tl.constexpr,
+    USE_SLIDING_WINDOW: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+    WINDOW_SIZE_LEFT: tl.constexpr,
+    WINDOW_SIZE_RIGHT: tl.constexpr,
+):
+    # -- compute qk ---
+    qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+    if IS_FP8:
+        qk += tl.dot(q, kT) * q_descale * k_descale  # Apply FP8 scaling
+    else:
+        qk = tl.dot(q, kT, acc=qk)  # noqa: F821
+
+    if USE_ALIBI:
+        row_idx = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+        col_idx = pos + tl.arange(0, BLOCK_N)
+
+        # Compute relative positions
+        relative_pos = row_idx[:, None] + N_CTX_K_FINAL - (N_CTX_Q + col_idx[None, :])
+        relative_pos = tl.abs(relative_pos)
+
+        # Compute ALiBi bias
+        alibi_bias = -1 * alibi_slope * relative_pos
+        qk += alibi_bias * 1.44269504
+
+    # ------------------------------------------------------------------
+    # masking
+    # ------------------------------------------------------------------
+    if USE_SLIDING_WINDOW:
+        row_idx = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)  # q positions
+        col_idx = pos + tl.arange(0, BLOCK_N)  # k positions
+        row = row_idx[:, None]  # [M,1]
+        col = col_idx[None, :]  # [1,N]
+
+        if IS_CAUSAL:
+            # -------- causal + window --------
+            diag = N_CTX_K_FINAL - N_CTX_Q  # sk-sq
+            causal_ok = col <= row + diag
+            if WINDOW_SIZE_LEFT < 0:  # only right window
+                win_ok = col <= row + diag + WINDOW_SIZE_RIGHT
+            else:  # both sides
+                win_ok = (col >= row + diag - WINDOW_SIZE_LEFT) & (
+                    col <= row + diag + WINDOW_SIZE_RIGHT
+                )
+            mask = ~(causal_ok & win_ok)  # True => -inf
+        else:
+            # -------- non-causal window --------
+            sk, sq = N_CTX_K_FINAL, N_CTX_Q
+            if WINDOW_SIZE_LEFT < 0:
+                mask = col > row + (sk - sq) + WINDOW_SIZE_RIGHT
+            else:
+                right = tl.minimum(row + (sk - sq) + WINDOW_SIZE_RIGHT, sk)
+                left = row + (sk - sq) - WINDOW_SIZE_LEFT
+                mask = (col > right) | (col < left)
+        qk = tl.where(mask, float("-inf"), qk)
+    else:
+        if IS_CAUSAL:
+            row_idx = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+            col_idx = pos + tl.arange(0, BLOCK_N)
+
+            # create a N_CTX_Q x kv_len causal mask
+            col_offset = N_CTX_K_FINAL - N_CTX_Q
+            causal_mask = row_idx[:, None] >= (col_idx[None, :] - col_offset)
+
+            # Apply the mask
+            qk = tl.where(causal_mask, qk, float("-inf"))
+
+    # Column mask (tail / variable-length). Instead of recomputing an arange each time,
+    # we accept a precomputed mask from the caller (col_valid_mask).
+    if apply_col_mask:
+        # Expect col_mask shape: [BLOCK_N]. True where column is within sequence.
+        qk = tl.where(col_mask[None, :], qk, float("-inf"))
+
+    m_i_new = tl.maximum(m_i, tl.max(qk, 1))  # per-row max so far
+
+    # rows that are *all* -inf after masking
+    valid = m_i_new > float("-inf")
+
+    # scale previous partial sums safely
+    alpha = tl.where(valid, tl.math.exp2(m_i - m_i_new), 0.0)
+
+    # subtract the row max only on valid rows
+    qk = tl.where(valid[:, None], qk - m_i_new[:, None], float("-inf"))
+    p = tl.math.exp2(qk)
+
+    # -- update m_i and l_i --
+    l_i = l_i * alpha + tl.sum(p, 1)
+    m_i = m_i_new
+    p = p.to(q.dtype)
+
+    # -- scale and update acc --
+    acc *= alpha[:, None]
+    if IS_FP8:
+        acc += tl.dot(p.to(v.dtype), v) * v_descale  # Apply FP8 scaling for V
+    else:
+        acc = tl.dot(p.to(v.dtype), v, acc=acc)
+
+    return m_i, l_i, acc
+
+
+@triton.autotune(
+    configs=fwd_decode_splitk_configs,
+    key=FWD_DECODE_AUTOTUNE_KEYS,
+    use_cuda_graph=True,
+)
+@triton.jit
+def _fwd_kernel_splitK(
+    Q,
+    K,
+    V,
+    Q_Descale,  # FP8 descale factors for Q
+    K_Descale,  # FP8 descale factors for K
+    V_Descale,  # FP8 descale factors for V
+    sm_scale,
+    Out_splitK,  # [B*H*G, split_k, Mq, K]
+    Metadata,  # [B*H*G, 2, split_k, M_ceil] contains [mi, li]
+    K_new,
+    V_new,
+    Cache_seqlens,
+    Cache_batch_idx,
+    Block_table,
+    Alibi_slopes,
+    stride_qz,
+    stride_qm,
+    stride_qg,
+    stride_qh,
+    stride_qd,
+    stride_kz,
+    stride_kn,
+    stride_kg,
+    stride_kh,
+    stride_kd,
+    stride_vz,
+    stride_vn,
+    stride_vg,
+    stride_vh,
+    stride_vd,
+    stride_osk_zhg,
+    stride_osk_s,
+    stride_osk_m,
+    stride_osk_d,
+    stride_mzhg,
+    stride_m2,
+    stride_ms,
+    stride_mm,
+    stride_kn_z,
+    stride_kn_n,
+    stride_kn_g,
+    stride_kn_h,
+    stride_kn_d,
+    stride_vn_z,
+    stride_vn_n,
+    stride_vn_g,
+    stride_vn_h,
+    stride_vn_d,
+    stride_bt_b,
+    stride_bt_s,
+    stride_az,
+    stride_ah,
+    stride_q_descale_z,  # FP8 descale strides
+    stride_q_descale_h,
+    stride_k_descale_z,
+    stride_k_descale_h,
+    stride_v_descale_z,
+    stride_v_descale_h,
+    Z,
+    N_CTX_Q,
+    N_CTX_K,
+    N_CTX_NEW,
+    BLOCK_N_PER_SPLIT: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    H_q: tl.constexpr,
+    H_kv: tl.constexpr,
+    G_q: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    ACTUAL_BLOCK_DMODEL: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    USE_CACHE_SEQLENs: tl.constexpr,
+    USE_CACHE_BATCH_IDX: tl.constexpr,
+    NEW_KV: tl.constexpr,
+    IS_GQA: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+    USE_ALIBI: tl.constexpr,
+    PADDED_HEAD: tl.constexpr,
+    GROUP_SIZE: tl.constexpr,
+    USE_SLIDING_WINDOW: tl.constexpr,
+    WINDOW_SIZE_LEFT: tl.constexpr,
+    WINDOW_SIZE_RIGHT: tl.constexpr,
+    USE_BLOCK_TABLE: tl.constexpr,
+    IS_FP8: tl.constexpr,  # FP8 flag
+):
+    # Cast strides to int64 to prevent overflow in pointer arithmetic.
+    # stride * index can exceed int32 range for large tensors (e.g. batch=64, sk=8192, d=56).
+    # See mha.py USE_INT64_STRIDES for the same pattern used in the prefill kernel.
+    stride_qz_i64 = tl.cast(stride_qz, tl.int64)
+    stride_qm_i64 = tl.cast(stride_qm, tl.int64)
+    stride_qg_i64 = tl.cast(stride_qg, tl.int64)
+    stride_qh_i64 = tl.cast(stride_qh, tl.int64)
+    stride_qd_i64 = tl.cast(stride_qd, tl.int64)
+    stride_kz_i64 = tl.cast(stride_kz, tl.int64)
+    stride_kn_i64 = tl.cast(stride_kn, tl.int64)
+    stride_kg_i64 = tl.cast(stride_kg, tl.int64)
+    stride_kh_i64 = tl.cast(stride_kh, tl.int64)
+    stride_kd_i64 = tl.cast(stride_kd, tl.int64)
+    stride_vz_i64 = tl.cast(stride_vz, tl.int64)
+    stride_vn_i64 = tl.cast(stride_vn, tl.int64)
+    stride_vg_i64 = tl.cast(stride_vg, tl.int64)
+    stride_vh_i64 = tl.cast(stride_vh, tl.int64)
+    stride_vd_i64 = tl.cast(stride_vd, tl.int64)
+    stride_osk_zhg_i64 = tl.cast(stride_osk_zhg, tl.int64)
+    stride_osk_s_i64 = tl.cast(stride_osk_s, tl.int64)
+    stride_osk_m_i64 = tl.cast(stride_osk_m, tl.int64)
+    stride_osk_d_i64 = tl.cast(stride_osk_d, tl.int64)
+    stride_mzhg_i64 = tl.cast(stride_mzhg, tl.int64)
+    stride_ms_i64 = tl.cast(stride_ms, tl.int64)
+    stride_m2_i64 = tl.cast(stride_m2, tl.int64)
+
+    # get program ids
+    pid_m = tl.program_id(0)
+    pid_zhg = tl.program_id(1)
+    pid_splitk = tl.program_id(2)
+
+    # compute z, h and g ids
+    z_id = pid_zhg // (H_q * G_q)
+    hq_id = (pid_zhg // G_q) % H_q
+    g_id = pid_zhg % G_q
+
+    # is gqa
+    if IS_GQA:
+        hk_id = hq_id // GROUP_SIZE
+        hv_id = hk_id
+    else:
+        hk_id = hq_id
+        hv_id = hq_id
+
+    # Load FP8 descale factors if needed
+    if IS_FP8:
+        if IS_GQA:
+            # For MQA/GQA, q_descale uses the same indexing as k/v (hk_id)
+            q_descale = tl.load(
+                Q_Descale + z_id * stride_q_descale_z + hk_id * stride_q_descale_h
+            )
+        else:
+            # For MHA, q_descale uses hq_id
+            q_descale = tl.load(
+                Q_Descale + z_id * stride_q_descale_z + hq_id * stride_q_descale_h
+            )
+        k_descale = tl.load(
+            K_Descale + z_id * stride_k_descale_z + hk_id * stride_k_descale_h
+        )
+        v_descale = tl.load(
+            V_Descale + z_id * stride_v_descale_z + hv_id * stride_v_descale_h
+        )
+    else:
+        q_descale, k_descale, v_descale = 1.0, 1.0, 1.0
+
+    # figure out seqlens
+    lo = pid_splitk * BLOCK_N_PER_SPLIT
+    if USE_CACHE_SEQLENs:
+        cache_seqlen_last_idx = tl.load(Cache_seqlens + z_id)
+        N_CTX_K_FINAL = cache_seqlen_last_idx
+    else:
+        N_CTX_K_FINAL = N_CTX_K
+    hi = tl.minimum((pid_splitk + 1) * BLOCK_N_PER_SPLIT, N_CTX_K_FINAL)
+
+    # pick batch index
+    if USE_CACHE_BATCH_IDX:
+        cache_batch_idx = tl.load(Cache_batch_idx + z_id)
+    else:
+        cache_batch_idx = z_id
+
+    # compute offsets
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = tl.arange(0, BLOCK_N)
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+
+    # compute ptrs
+    q_offset = Q + hq_id * stride_qh_i64 + z_id * stride_qz_i64 + g_id * stride_qg_i64
+    q_ptrs = (
+        q_offset + offs_m[:, None] * stride_qm_i64 + offs_d[None, :] * stride_qd_i64
+    )
+
+    # Handle block table for paged attention
+    if USE_BLOCK_TABLE:
+        # K and V now point to paged cache
+        # Each batch has its own block table row
+        block_table_ptr = Block_table + z_id * stride_bt_b
+    else:
+        k_offset = (
+            K
+            + hk_id * stride_kh_i64
+            + cache_batch_idx * stride_kz_i64
+            + g_id * stride_kg_i64
+        )
+        v_offset = (
+            V
+            + hv_id * stride_vh_i64
+            + cache_batch_idx * stride_vz_i64
+            + g_id * stride_vg_i64
+        )
+
+    # compute masks
+    if PADDED_HEAD:
+        q_mask = (offs_m < N_CTX_Q)[:, None] & (offs_d < ACTUAL_BLOCK_DMODEL)[None, :]
+        kT_mask = (offs_d < ACTUAL_BLOCK_DMODEL)[:, None] & (offs_n < N_CTX_K_FINAL)[
+            None, :
+        ]
+        v_mask = (offs_n < N_CTX_K_FINAL)[:, None] & (offs_d < ACTUAL_BLOCK_DMODEL)[
+            None, :
+        ]
+        osk_mask = (offs_m < N_CTX_Q)[:, None] & (offs_d < ACTUAL_BLOCK_DMODEL)[None, :]
+    else:
+        q_mask = (offs_m < N_CTX_Q)[:, None]
+        kT_mask = (offs_n < N_CTX_K_FINAL)[None, :]
+        v_mask = (offs_n < N_CTX_K_FINAL)[:, None]
+        osk_mask = (offs_m < N_CTX_Q)[:, None]
+
+    # scale sm_scale by log_2(e) and use
+    # 2^x instead of exp in the loop because CSE and LICM
+    # don't work as expected with `exp` in the loop
+    qk_scale = sm_scale * 1.44269504
+
+    # load q: it will stay in SRAM throughout
+    q = tl.load(q_ptrs, mask=q_mask, other=0.0)
+    q = (q * qk_scale).to(q.dtype)
+
+    # load ALiBi slope if enabled
+    if USE_ALIBI:
+        a_offset = z_id * stride_az + hq_id * stride_ah
+        alibi_slope = tl.load(Alibi_slopes + a_offset)
+    else:
+        alibi_slope = None
+
+    # initialize pointer to m and l
+    m_i = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
+    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)  # noqa: F821
+
+    # loop over k, v and update accumulator
+    if USE_BLOCK_TABLE:
+        # Paged attention: process all KV blocks from cache
+        # Note: Cache should be updated externally before calling this kernel
+        num_kv_blocks = (N_CTX_K_FINAL + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K
+
+        for block_idx in range(num_kv_blocks):
+            # Calculate sequence range for this block
+            block_start = block_idx * BLOCK_SIZE_K
+            block_end = tl.minimum(block_start + BLOCK_SIZE_K, N_CTX_K_FINAL)
+
+            # Check if block overlaps with our split-k range [lo, hi)
+            if block_end > lo and block_start < hi:
+                # Load physical block number
+                physical_block = tl.load(block_table_ptr + block_idx * stride_bt_s)
+
+                # Calculate the range within this block that overlaps with [lo, hi)
+                process_start = tl.maximum(lo - block_start, 0)
+                process_end = tl.minimum(hi - block_start, BLOCK_SIZE_K)
+                process_end = tl.minimum(process_end, block_end - block_start)
+
+                # Instead of forcing a floor alignment to BLOCK_N (which can still skip
+                # part of the intended range if start falls mid-tile for small splits),
+                # start from the raw (possibly unaligned) process_start rounded *down* but
+                # allow the loop to begin earlier (at most BLOCK_N before) so that any
+                # partial tile overlapping [lo, hi) is covered. Masking below will remove
+                # columns < lo or >= hi ensuring numerically identical coverage without
+                # duplication.
+                aligned_start = (process_start // BLOCK_N) * BLOCK_N
+                if aligned_start > 0 and aligned_start + BLOCK_N > process_start:
+                    # ensure we include the tile that contains process_start
+                    process_start = aligned_start
+                else:
+                    process_start = aligned_start
+
+                for offset in range(process_start, process_end, BLOCK_N):
+                    # Current position (may begin slightly before logical split range; masking fixes it)
+                    pos = block_start + offset
+                    # Proceed unconditionally; masking below enforces [lo, hi)
+                    # Calculate base addresses for K and V in this physical block
+                    k_base = (
+                        K
+                        + physical_block * BLOCK_SIZE_K * stride_kn_i64
+                        + hk_id * stride_kh_i64
+                        + g_id * stride_kg_i64
+                    )
+                    v_base = (
+                        V
+                        + physical_block * BLOCK_SIZE_K * stride_vn_i64
+                        + hv_id * stride_vh_i64
+                        + g_id * stride_vg_i64
+                    )
+
+                    # Offsets within the current block
+                    block_offs = offset + offs_n
+
+                    # Masks for valid data respecting:
+                    #   (1) global key length (seq_mask)
+                    #   (2) block bounds (block_mask)
+                    #   (3) current split range [lo, hi)
+                    seq_mask = (pos + offs_n) < N_CTX_K_FINAL
+                    block_mask = block_offs < BLOCK_SIZE_K
+                    end_mask = block_offs < process_end
+                    split_mask = ((pos + offs_n) >= lo) & ((pos + offs_n) < hi)
+                    col_mask = seq_mask & block_mask & end_mask & split_mask
+
+                    # Apply masks
+                    kT_mask_final = kT_mask & col_mask[None, :]
+                    v_mask_final = v_mask & col_mask[:, None]
+
+                    # Load K and V
+                    kT_ptrs = (
+                        k_base
+                        + offs_d[:, None] * stride_kd_i64
+                        + block_offs[None, :] * stride_kn_i64
+                    )
+                    v_ptrs = (
+                        v_base
+                        + block_offs[:, None] * stride_vn_i64
+                        + offs_d[None, :] * stride_vd_i64
+                    )
+
+                    kT = tl.load(kT_ptrs, mask=kT_mask_final, other=0.0)
+                    v = tl.load(v_ptrs, mask=v_mask_final, other=0.0)
+
+                    # Unified inner function handles both paged and contiguous
+                    m_i, l_i, acc = _attn_fwd_inner(
+                        q,
+                        kT,
+                        v,
+                        pos,
+                        col_mask,
+                        m_i,
+                        l_i,
+                        acc,
+                        pid_m,
+                        q_descale,
+                        k_descale,
+                        v_descale,
+                        alibi_slope,
+                        True,
+                        IS_FP8,
+                        BLOCK_M,
+                        BLOCK_N,
+                        N_CTX_Q,
+                        N_CTX_K_FINAL,
+                        USE_ALIBI,
+                        USE_SLIDING_WINDOW,
+                        IS_CAUSAL,
+                        WINDOW_SIZE_LEFT,
+                        WINDOW_SIZE_RIGHT,
+                    )
+    else:
+        # Non-paged attention: process KV from cache
+        # Note: Cache should be updated externally before calling this kernel
+        # Compute bounds check flag once: needed if split size not aligned to BLOCK_N or variable seqlens
+        bounds_checks_n = ((BLOCK_N_PER_SPLIT % BLOCK_N) > 0) | USE_CACHE_SEQLENs
+        # loop over k, v and update accumulator
+        for start_n in range(lo, hi, BLOCK_N):
+            kT_ptrs = (
+                k_offset
+                + offs_d[:, None] * stride_kd_i64
+                + (start_n + offs_n)[None, :] * stride_kn_i64
+            )
+            V_ptrs = (
+                v_offset
+                + (start_n + offs_n)[:, None] * stride_vn_i64
+                + offs_d[None, :] * stride_vd_i64
+            )
+
+            # load k
+            kT = tl.load(kT_ptrs, mask=kT_mask, other=0.0)
+            v = tl.load(V_ptrs, mask=v_mask, other=0.0)
+
+            # Use the same inner loop logic
+            # Precompute column validity mask for this tile (all True for full tiles).
+            # hi is the upper bound of the overall split range; start_n marks this tile's base.
+            col_valid_mask = offs_n < (hi - start_n)
+
+            m_i, l_i, acc = _attn_fwd_inner(
+                q,
+                kT,
+                v,
+                start_n,
+                col_valid_mask,
+                m_i,
+                l_i,
+                acc,
+                pid_m,
+                q_descale,
+                k_descale,
+                v_descale,
+                alibi_slope,
+                bounds_checks_n,
+                IS_FP8,
+                BLOCK_M,
+                BLOCK_N,
+                N_CTX_Q,
+                N_CTX_K_FINAL,
+                USE_ALIBI,
+                USE_SLIDING_WINDOW,
+                IS_CAUSAL,
+                WINDOW_SIZE_LEFT,
+                WINDOW_SIZE_RIGHT,
+            )
+
+    # write back O
+    osk_offset = (
+        Out_splitK + pid_zhg * stride_osk_zhg_i64 + pid_splitk * stride_osk_s_i64
+    )
+    osk_ptrs = (
+        osk_offset
+        + offs_m[:, None] * stride_osk_m_i64
+        + offs_d[None, :] * stride_osk_d_i64
+    )
+    tl.store(
+        osk_ptrs,
+        acc,
+        mask=osk_mask,
+    )
+
+    # write metadata for split-K reduction
+    metadata_offset = Metadata + pid_zhg * stride_mzhg_i64 + pid_splitk * stride_ms_i64
+    metadata_ptr = metadata_offset + offs_m
+    tl.store(metadata_ptr, m_i)
+    tl.store(metadata_ptr + stride_m2_i64, l_i)
+
+
+FWD_DECODE_REDUCE_AUTOTUNE_KEYS = [
+    "BLOCK_DMODEL",
+    "split_k",
+]
+
+
+@triton.autotune(
+    configs=fwd_decode_reduce_configs,
+    key=FWD_DECODE_REDUCE_AUTOTUNE_KEYS,
+    use_cuda_graph=True,
+)
+@triton.jit
+def _splitK_reduce(
+    Out_splitK,  # [B*H*G, split_k, Mq, K]
+    Metadata,  # [B*H*G, 2, split_k, M_ceil] contains [mi, li]
+    Out,  # [B, H, G, M, K]
+    LSE,  # [B*H*G, M]
+    stride_osk_zhg,
+    stride_osk_s,
+    stride_osk_m,
+    stride_osk_k,
+    stride_mzhg,
+    stride_m2,
+    stride_ms,
+    stride_mm,
+    stride_oz,
+    stride_oh,
+    stride_og,
+    stride_om,
+    stride_ok,
+    stride_lse_zhg,
+    stride_lse_m,
+    K_BLOCK_SIZE: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    ACTUAL_BLOCK_DMODEL: tl.constexpr,
+    H: tl.constexpr,
+    G: tl.constexpr,
+    split_k: tl.constexpr,
+    splitK_pow2: tl.constexpr,
+    MASK_SPLITK: tl.constexpr,
+    PADDED_HEAD: tl.constexpr,
+):
+    # Cast strides to int64 to prevent overflow in pointer arithmetic.
+    stride_osk_zhg_i64 = tl.cast(stride_osk_zhg, tl.int64)
+    stride_osk_s_i64 = tl.cast(stride_osk_s, tl.int64)
+    stride_osk_m_i64 = tl.cast(stride_osk_m, tl.int64)
+    stride_osk_k_i64 = tl.cast(stride_osk_k, tl.int64)
+    stride_mzhg_i64 = tl.cast(stride_mzhg, tl.int64)
+    stride_m2_i64 = tl.cast(stride_m2, tl.int64)
+    stride_ms_i64 = tl.cast(stride_ms, tl.int64)
+    stride_mm_i64 = tl.cast(stride_mm, tl.int64)
+    stride_oz_i64 = tl.cast(stride_oz, tl.int64)
+    stride_oh_i64 = tl.cast(stride_oh, tl.int64)
+    stride_og_i64 = tl.cast(stride_og, tl.int64)
+    stride_om_i64 = tl.cast(stride_om, tl.int64)
+    stride_lse_zhg_i64 = tl.cast(stride_lse_zhg, tl.int64)
+
+    # get pids
+    pid_zhg = tl.program_id(0)
+    pid_m = tl.program_id(1)
+    pid_k = tl.program_id(2)
+
+    # compute offsets
+    offs_splitK = tl.arange(0, splitK_pow2)
+    offs_k = pid_k * K_BLOCK_SIZE + tl.arange(0, K_BLOCK_SIZE)
+
+    # compute masks
+    if PADDED_HEAD:
+        o_mask = offs_k < ACTUAL_BLOCK_DMODEL
+    else:
+        o_mask = None
+
+    # compute ptrs
+    metadata_offset = Metadata + pid_zhg * stride_mzhg_i64
+    metadata_ptr = metadata_offset + offs_splitK * stride_ms_i64 + pid_m * stride_mm_i64
+
+    osk_offset = Out_splitK + pid_zhg * stride_osk_zhg_i64 + pid_m * stride_osk_m_i64
+    osk_ptr = (
+        osk_offset
+        + offs_splitK[:, None] * stride_osk_s_i64
+        + offs_k[None, :] * stride_osk_k_i64
+    )
+
+    # read max values of each splitK
+    if MASK_SPLITK:
+        splitK_mask = offs_splitK < split_k
+        l_m = tl.load(metadata_ptr, mask=splitK_mask, other=float("-inf"))
+        l_sum = tl.load(metadata_ptr + stride_m2_i64, mask=splitK_mask, other=0.0)
+        acc = tl.load(osk_ptr, mask=splitK_mask[:, None], other=0.0)
+    else:
+        l_m = tl.load(metadata_ptr)
+        l_sum = tl.load(metadata_ptr + stride_m2_i64)
+        acc = tl.load(osk_ptr)
+
+    g_m = tl.max(l_m, axis=0)
+
+    alpha = tl.where(l_m > float("-inf"), tl.math.exp2(l_m - g_m), 0.0)
+
+    # read sum
+    l_sum *= alpha
+    g_sum = tl.sum(l_sum, axis=0)
+    acc = acc * alpha[:, None]
+
+    g_sum_safe = tl.where(g_sum > 0, g_sum, 1.0)
+    acc_out = tl.sum(acc, axis=0) / g_sum_safe
+
+    # Store output
+    z_id = pid_zhg // (H * G)
+    h_id = (pid_zhg // G) % H
+    g_id = pid_zhg % G
+    out_offset = (
+        Out + z_id * stride_oz_i64 + h_id * stride_oh_i64 + g_id * stride_og_i64
+    )
+    out_ptr = out_offset + pid_m * stride_om_i64 + offs_k
+    tl.store(out_ptr, acc_out, mask=o_mask)
+
+    # Store lse
+    l_ptrs = LSE + pid_zhg * stride_lse_zhg_i64 + pid_m
+    lse_val = tl.where(g_sum > 0, (g_m + tl.math.log2(g_sum)) / 1.44269504, g_m)
+    tl.store(l_ptrs, lse_val)
+
+
+@triton.jit
+def cast_uint32_to_half2(scale_shift):
+    # Extract two float16 packed into one int32
+    scale = scale_shift & 0xFFFF
+    shift = scale_shift >> 16
+    scale = scale.to(tl.uint16).to(tl.float16, bitcast=True)
+    shift = shift.to(tl.uint16).to(tl.float16, bitcast=True)
+    return scale, shift
+
+
+@triton.jit
+def dequantize(
+    x_,
+    scale,
+    shift,
+    PACKED_PER_VAL: tl.constexpr = 8,
+):
+    # PACKED_PER_VAL is the number of values packed into
+    # each element x_. For example, for int4 quantization
+    # and x_ of type int32, PACKED_PER_VAL is 8.
+
+    BLOCK_N: tl.constexpr = x_.shape[0]
+    BLOCK_DMODEL_PACKED: tl.constexpr = x_.shape[1]
+    offsets = tl.arange(0, PACKED_PER_VAL) * 4
+    quant_offset = (
+        x_[:, None, :] >> offsets[None, :, None]
+    )  # (BLOCK_N, PACKED_PER_VAL, D // PACKED_PER_VAL)
+
+    quant_offset = tl.view(
+        quant_offset, (BLOCK_N, BLOCK_DMODEL_PACKED * PACKED_PER_VAL)
+    )
+    # Trick - instead of converting int4 to float16 we view it as float16
+    # and then multiply by 32768 * 512 == 2**24
+    quant_offset = (quant_offset & 0xF).to(tl.uint16).to(tl.float16, bitcast=True)
+    quant_offset = (quant_offset * 32768.0).to(tl.float16)
+    scale_512 = scale * 512
+
+    dequant = quant_offset * scale_512 + shift
+    return dequant
+
+
+def quantize_kv_int4(k: torch.Tensor, num_groups: int = 1) -> torch.Tensor:
+    # Scale and shift are such that quantization linearly maps
+    # int4 values range [0..15] to input values range min(k)..max(k)
+    # individually for every row
+    k = k.reshape(*k.shape[:-1], num_groups, k.shape[-1] // num_groups)
+    max_vals = torch.max(k, dim=-1, keepdim=True).values
+    min_vals = torch.min(k, dim=-1, keepdim=True).values
+    scale_k: torch.Tensor = (max_vals - min_vals) / 15
+
+    shift_k = torch.min(k, dim=-1, keepdim=True).values
+    scale_k = scale_k.to(torch.float16)
+    shift_k = shift_k.to(torch.float16)
+
+    in_bytes = ((k - shift_k.expand(k.shape)) / scale_k.expand(k.shape)) + 0.5
+    in_bytes = in_bytes.to(torch.uint8)
+    in_int4 = in_bytes & 0xF
+    in_int4_packed = in_int4[..., ::2] + (in_int4[..., 1::2] << 4)
+    scale_shift = torch.concat(
+        [scale_k.view(torch.uint8), shift_k.view(torch.uint8)], dim=-1
+    )
+    k_quant = torch.concat(
+        [
+            scale_shift.flatten(start_dim=-2),
+            in_int4_packed.flatten(start_dim=-2),
+        ],
+        dim=-1,
+    ).view(torch.int16)
+    return k_quant
+
+
+def dequantize_kv_fp16(quant_k: torch.Tensor, num_groups: int = 1) -> torch.Tensor:
+    k_i16 = quant_k.view(torch.int16)
+    k_ui8 = k_i16.view(torch.uint8)
+
+    ss_size = num_groups * 4
+    scale_shift_ui8 = k_ui8[..., 0:ss_size]
+    scale_shift_ui8 = scale_shift_ui8.reshape(
+        *scale_shift_ui8.shape[:-1], num_groups, 4
+    )
+    scale = scale_shift_ui8[..., 0:2].view(torch.float16)
+    shift = scale_shift_ui8[..., 2:4].view(torch.float16)
+
+    kv_ui8 = k_ui8[..., ss_size:]
+    k_ui8 = kv_ui8.reshape(*kv_ui8.shape[:-1], num_groups, -1)
+    k1_i4 = k_ui8 & 0xF
+    k2_i4 = (k_ui8 & 0xF0) >> 4
+    k_shape = k1_i4.shape
+    k1_f16 = k1_i4.to(torch.float16) * scale.expand(k_shape) + shift.expand(k_shape)
+    k2_f16 = k2_i4.to(torch.float16) * scale.expand(k_shape) + shift.expand(k_shape)
+
+    out = torch.empty(
+        (*k1_f16.shape[:-1], k1_f16.shape[-1] * 2),
+        dtype=torch.float16,
+        device=quant_k.device,
+    )
+    out[..., ::2] = k1_f16
+    out[..., 1::2] = k2_f16
+    out = out.reshape(*k_shape[:-2], -1)
+
+    return out
+
+
+def get_split_k(B: int, G: int, H: int, Mk: int) -> int:
+    """Heuristic for the number of splits"""
+    bh = max(B * H, 1)  # NOTE: Handle B*h=0 case
+    split_k = max(Mk, 1024) // bh
+    max_chunk_size = 64
+    while split_k > 0 and Mk / split_k < max_chunk_size:
+        split_k = split_k // 2
+    while B * H * G * split_k >= 1024:
+        split_k = split_k // 2
+    split_k = min(split_k, 512)
+    split_k = max(split_k, 1)
+    return split_k
+
+
+def attention_forward_decode_triton_impl(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    k_new: Optional[torch.Tensor],
+    v_new: Optional[torch.Tensor],
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    sm_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    alibi_slopes: Optional[torch.Tensor],
+    layout: Literal["bshd"],
+    cache_seqlens: Optional[torch.Tensor],
+    cache_batch_idx: Optional[torch.Tensor],
+    block_table: Optional[torch.Tensor] = None,
+    q_descale: Optional[torch.Tensor] = None,
+    k_descale: Optional[torch.Tensor] = None,
+    v_descale: Optional[torch.Tensor] = None,
+    # rotary (optional)
+    rotary_cos: Optional[torch.Tensor] = None,
+    rotary_sin: Optional[torch.Tensor] = None,
+    rotary_interleaved: bool = False,
+    seqlens_rotary: Optional[torch.Tensor] = None,
+):
+    # Validate layout at entry
+    if layout != "bshd":
+        raise ValueError(f"{layout} layout is not supported, only 'bshd' is supported")
+
+    # apply rotary embedding
+    if rotary_cos is not None and rotary_sin is not None:
+        # Prefer explicitly provided rotary sequence start offsets if given; fall back to cache_seqlens.
+        seqlen_offsets = (
+            seqlens_rotary
+            if seqlens_rotary is not None
+            else (cache_seqlens if cache_seqlens is not None else 0)
+        )
+        local = (window_size_left != -1) or (window_size_right != -1)
+        q, k_new = apply_rotary(
+            q,
+            k_new,
+            rotary_cos,
+            rotary_sin,
+            causal=causal,
+            local=local,
+            interleaved=rotary_interleaved,
+            seqlen_offsets=seqlen_offsets,
+        )
+
+    # handle cache updates
+    if k_new is not None and v_new is not None:
+        batch_size = k_new.shape[0]
+        seqlen_new = k_new.shape[1]
+        # shared offset within the new block for every batch element
+        token_offsets = torch.arange(seqlen_new, device=k_new.device, dtype=torch.long)
+
+        if block_table is None:
+            # Non-paged attention
+            if cache_seqlens is not None:
+                insert_positions = cache_seqlens.long()[:, None] + token_offsets
+                batch_idx = torch.arange(
+                    batch_size, device=k_new.device, dtype=torch.long
+                )[:, None]
+                k_cache[batch_idx, insert_positions] = k_new
+                v_cache[batch_idx, insert_positions] = v_new
+                cache_seqlens.add_(seqlen_new)
+            else:
+                # Append at the end of existing cache
+                seqlen_cache = k_cache.shape[1]
+                k_cache[:, seqlen_cache - seqlen_new :] = k_new
+                v_cache[:, seqlen_cache - seqlen_new :] = v_new
+        else:
+            # Paged attention
+            block_size = k_cache.shape[1]
+
+            if cache_seqlens is not None:
+                start_idx = cache_seqlens.long()
+            else:
+                # Find the first unoccupied position per batch from the block table.
+                invalid_mask = block_table < 0
+                has_invalid = invalid_mask.any(dim=1)
+                first_invalid_idx = invalid_mask.to(torch.int32).argmax(dim=1).long()
+                start_idx = torch.where(
+                    has_invalid,
+                    first_invalid_idx * block_size,
+                    torch.full_like(
+                        first_invalid_idx, block_table.shape[1] * block_size
+                    ),
+                )
+
+            insert_positions = start_idx[:, None] + token_offsets
+            block_indices = insert_positions // block_size
+            within_block_indices = insert_positions % block_size
+            batch_idx = torch.arange(batch_size, device=k_new.device, dtype=torch.long)[
+                :, None
+            ]
+            physical_blocks = block_table[batch_idx, block_indices]
+
+            flat_phys = physical_blocks.reshape(-1)
+            flat_within = within_block_indices.reshape(-1)
+            k_cache[flat_phys, flat_within] = k_new.reshape(-1, *k_new.shape[2:])
+            v_cache[flat_phys, flat_within] = v_new.reshape(-1, *v_new.shape[2:])
+
+            if cache_seqlens is not None:
+                cache_seqlens.add_(seqlen_new)
+
+    # kernel_configs
+    is_new_kv = False  # Cache has been updated, so no new KV in kernel
+    use_alibi, (stride_az, stride_ah) = True if alibi_slopes is not None else False, (
+        alibi_slopes.stride() if alibi_slopes is not None else (None, None)
+    )
+    use_cache_seqlens = cache_seqlens is not None
+    use_sliding_window = window_size_left != -1 or window_size_right != -1
+    use_block_table = block_table is not None
+
+    # get shapes and strides
+    batch_size, seqlen_q, nheads_q, dim_q = get_shape_from_layout(q, layout)
+    stride_qz, stride_qh, stride_qm, stride_qd = get_stride_from_layout(q, layout)
+
+    # Handle paged KV cache layout
+    if use_block_table:
+        # For paged attention, k_cache and v_cache have shape [num_blocks, block_size, nheads, head_dim]
+        num_blocks_kc, block_size_k, nheads_kc, dim_kc = k_cache.shape
+        num_blocks_vc, block_size_v, nheads_vc, dim_vc = v_cache.shape
+        # Get the actual sequence length from cache_seqlens or block_table
+        if cache_seqlens is not None:
+            seqlen_kc = int(cache_seqlens.max().item())
+        else:
+            # Infer from block_table shape [batch_size, num_blocks_per_seq]
+            assert block_table is not None
+            num_blocks_per_seq = block_table.shape[1]
+            seqlen_kc = num_blocks_per_seq * block_size_k
+        seqlen_vc = seqlen_kc
+
+        # Strides for paged layout
+        stride_kc_z = 0  # No batch dimension in paged cache
+        stride_kc_n = k_cache.stride(1)  # Sequence stride
+        stride_kc_h = k_cache.stride(2)  # Head stride
+        stride_kc_d = k_cache.stride(3)  # Dim stride
+
+        stride_vc_z = 0
+        stride_vc_n = v_cache.stride(1)
+        stride_vc_h = v_cache.stride(2)
+        stride_vc_d = v_cache.stride(3)
+    else:
+        _, seqlen_kc, nheads_kc, dim_kc = get_shape_from_layout(k_cache, layout)
+        stride_kc_z, stride_kc_h, stride_kc_n, stride_kc_d = get_stride_from_layout(
+            k_cache, layout
+        )
+        _, seqlen_vc, nheads_vc, dim_vc = get_shape_from_layout(v_cache, layout)
+        stride_vc_z, stride_vc_h, stride_vc_n, stride_vc_d = get_stride_from_layout(
+            v_cache, layout
+        )
+        block_size_k = 0  # Not used
+    if is_new_kv:
+        _, seqlen_kn, nheads_kn, dim_kn = get_shape_from_layout(k_new, layout)
+        stride_kn_z, stride_kn_h, stride_kn_n, stride_kn_d = get_stride_from_layout(
+            k_new, layout
+        )
+        _, seqlen_vn, nheads_vn, dim_vn = get_shape_from_layout(v_new, layout)
+        stride_vn_z, stride_vn_h, stride_vn_n, stride_vn_d = get_stride_from_layout(
+            v_new, layout
+        )
+    else:
+        _, _seqlen_kn, nheads_kn, _dim_kn = None, None, None, None
+        stride_kn_z, stride_kn_h, stride_kn_n, stride_kn_d = None, None, None, None
+        _, _seqlen_vn, nheads_vn, _dim_vn = None, None, None, None
+        stride_vn_z, stride_vn_h, stride_vn_n, stride_vn_d = None, None, None, None
+    _, seqlen_o, nheads_o, dim_o = get_shape_from_layout(out, layout)
+    stride_oz, stride_oh, stride_om, stride_od = get_stride_from_layout(out, layout)
+    assert (
+        dim_q == dim_kc == dim_vc
+    ), f"Dimensions must match: {dim_q}, {dim_kc}, {dim_vc}"
+
+    # add extra information needed by the kernels
+    (n_group_q, heads_per_group_q), stride_qg = (1, nheads_q), stride_qm
+    (_n_group_k, heads_per_group_k), stride_kc_g = (1, nheads_kc), stride_kc_n
+    (_n_group_v, _heads_per_group_v), stride_vc_g = (1, nheads_vc), stride_vc_n
+    if is_new_kv:
+        (_n_group_kn, _heads_per_group_kn), stride_kn_g = (1, nheads_kn), stride_kn_n
+        (_n_group_vn, _heads_per_group_vn), stride_vn_g = (1, nheads_vn), stride_vn_n
+    else:
+        (_n_group_kn, _heads_per_group_kn), stride_kn_g = (None, None), None
+        (_n_group_vn, _heads_per_group_vn), stride_vn_g = (None, None), None
+    (_n_group_o, _heads_per_group_o), stride_og = (1, nheads_o), stride_om
+
+    # get padded size
+    dim_padded = get_padded_headsize(dim_kc)
+    is_padded_head = dim_padded != dim_kc
+
+    # Handle MQA/GQA case
+    group_size = nheads_q // nheads_kc
+    if group_size > 1:
+        is_gqa = True
+    else:
+        is_gqa = False
+
+    # Use heuristics for split_k
+    if use_block_table:
+        # For paged attention, use the actual sequence length from cache_seqlens
+        max_seqlen = (
+            int(cache_seqlens.max().item())
+            if cache_seqlens is not None
+            else block_size_k
+        )
+        split_k = get_split_k(batch_size, n_group_q, heads_per_group_q, max_seqlen)
+    else:
+        split_k = get_split_k(batch_size, n_group_q, heads_per_group_q, seqlen_kc)
+    split_size = (seqlen_kc + split_k - 1) // split_k
+
+    # setup grid - use lambda to get BLOCK_M from autotune
+    # Use MAX_BLOCK_M for intermediate tensor allocation to ensure enough space
+    seqlen_q_ceil = (seqlen_q + MAX_BLOCK_M - 1) // MAX_BLOCK_M * MAX_BLOCK_M
+
+    def grid(META):
+        return (
+            triton.cdiv(seqlen_q, META["BLOCK_M"]),
+            batch_size * n_group_q * heads_per_group_q,
+            split_k,
+        )
+
+    # create intermediate tensors
+    out_splitk = torch.empty(
+        [
+            batch_size * n_group_q * heads_per_group_q,
+            split_k,
+            seqlen_q_ceil,
+            dim_padded,
+        ],
+        dtype=torch.float32,
+        device=q.device,
+    )
+    metadata = torch.empty(
+        [batch_size * n_group_q * heads_per_group_q, 2, split_k, seqlen_q_ceil],
+        dtype=torch.float32,
+        device=q.device,
+    )
+
+    # Validate pre-allocated softmax_lse tensor
+    # Expected shape after view: (batch_size, n_group_q * heads_per_group_q, seqlen_q)
+    # Internal shape: (batch_size * n_group_q * heads_per_group_q, seqlen_q)
+    expected_h_total = batch_size * n_group_q * heads_per_group_q
+    assert (
+        softmax_lse.shape[0] == batch_size
+    ), f"softmax_lse.shape[0] ({softmax_lse.shape[0]}) must equal batch_size ({batch_size})"
+    assert (
+        softmax_lse.shape[1] == n_group_q * heads_per_group_q
+    ), f"softmax_lse.shape[1] ({softmax_lse.shape[1]}) must equal n_group_q * heads_per_group_q ({n_group_q * heads_per_group_q})"
+    assert (
+        softmax_lse.shape[2] >= seqlen_q
+    ), f"softmax_lse.shape[2] ({softmax_lse.shape[2]}) must be >= seqlen_q ({seqlen_q})"
+    assert (
+        softmax_lse.dtype == torch.float32
+    ), f"softmax_lse must be float32, got {softmax_lse.dtype}"
+    assert softmax_lse.device == q.device, "softmax_lse must be on same device as q"
+
+    # Create internal lse view for kernel use
+    lse = softmax_lse.view(expected_h_total, -1)[:, :seqlen_q].contiguous()
+
+    # get intermediate tensor strides
+    stride_osk_zhg, stride_osk_s, stride_osk_m, stride_osk_d = out_splitk.stride()
+    stride_mzhg, stride_m2, stride_ms, stride_mm = metadata.stride()
+    stride_lse_zhg, stride_lse_m = lse.stride()
+
+    # Block table strides
+    if use_block_table:
+        assert block_table is not None
+        stride_bt_b, stride_bt_s = block_table.stride()
+    else:
+        stride_bt_b, stride_bt_s = 0, 0
+
+    # FP8 support
+    IS_FP8 = is_fp8([q, k_cache, v_cache])
+    if IS_FP8:
+        arch = get_arch()
+        if not arch.supports_fp8:
+            raise RuntimeError(f"{arch.name} does not support FP8")
+        rec_dtype = arch.recommended_fp8_dtype(q.dtype)
+        if (
+            q.dtype != rec_dtype
+            or k_cache.dtype != rec_dtype
+            or v_cache.dtype != rec_dtype
+        ):
+            warnings.warn(
+                f"Use {rec_dtype} data type on {arch}. Got q: {q.dtype}, k: {k_cache.dtype}, v: {v_cache.dtype}",
+                UserWarning,
+            )
+        if (q_descale is None) or (k_descale is None) or (v_descale is None):
+            warnings.warn(
+                "FP8 tensors detected but descale factors not provided. Using default scale of 1.0",
+                UserWarning,
+            )
+            # Create default descale tensors if not provided
+            if q_descale is None:
+                q_descale = torch.ones(
+                    batch_size, nheads_q, dtype=torch.float32, device=q.device
+                )
+            if k_descale is None:
+                k_descale = torch.ones(
+                    batch_size, nheads_kc, dtype=torch.float32, device=q.device
+                )
+            if v_descale is None:
+                v_descale = torch.ones(
+                    batch_size, nheads_vc, dtype=torch.float32, device=q.device
+                )
+        else:
+            # Enforce exact expected shapes; no reshaping or normalization.
+            assert (
+                q_descale.dim() == 2
+                and q_descale.shape[0] == batch_size
+                and q_descale.shape[1] == nheads_kc
+            ), f"q_descale expected shape ({batch_size}, {nheads_kc}) got {tuple(q_descale.shape)}"
+            assert (
+                k_descale.dim() == 2
+                and k_descale.shape[0] == batch_size
+                and k_descale.shape[1] == nheads_kc
+            ), f"k_descale expected shape ({batch_size}, {nheads_kc}) got {tuple(k_descale.shape)}"
+            assert (
+                v_descale.dim() == 2
+                and v_descale.shape[0] == batch_size
+                and v_descale.shape[1] == nheads_kc
+            ), f"v_descale expected shape ({batch_size}, {nheads_kc}) got {tuple(v_descale.shape)}"
+        stride_q_descale_z, stride_q_descale_h = q_descale.stride()
+        stride_k_descale_z, stride_k_descale_h = k_descale.stride()
+        stride_v_descale_z, stride_v_descale_h = v_descale.stride()
+    else:
+        q_descale = None
+        k_descale = None
+        v_descale = None
+        stride_q_descale_z = 0
+        stride_q_descale_h = 0
+        stride_k_descale_z = 0
+        stride_k_descale_h = 0
+        stride_v_descale_z = 0
+        stride_v_descale_h = 0
+
+    if DEBUG:
+        print(
+            "batch_size, seqlen_q, nheads_q, dim_q",
+            (batch_size, seqlen_q, nheads_q, dim_q),
+        )
+        print("_, seqlen_kc, nheads_kc, dim_kc", (_, seqlen_kc, nheads_kc, dim_kc))
+        print("dim_padded:", dim_padded)
+        print(
+            "stride_qz, stride_qm, stride_qg, stride_qh, stride_qd",
+            (stride_qz, stride_qm, stride_qg, stride_qh, stride_qd),
+        )
+        print(
+            "stride_kc_z, stride_kc_n, stride_kc_g, stride_kc_h, stride_kc_d",
+            (stride_kc_z, stride_kc_n, stride_kc_g, stride_kc_h, stride_kc_d),
+        )
+        print(
+            "stride_vc_z, stride_vc_n, stride_vc_g, stride_vc_h, stride_vc_d",
+            (stride_vc_z, stride_vc_n, stride_vc_g, stride_vc_h, stride_vc_d),
+        )
+        if is_new_kv:
+            print(
+                "stride_kn_z, stride_kn_n, stride_kn_g, stride_kn_h, stride_kn_d",
+                (stride_kn_z, stride_kn_n, stride_kn_g, stride_kn_h, stride_kn_d),
+            )
+            print(
+                "stride_vn_z, stride_vn_n, stride_vn_g, stride_vn_h, stride_vn_d",
+                (stride_vn_z, stride_vn_n, stride_vn_g, stride_vn_h, stride_vn_d),
+            )
+        print(
+            "stride_oz, stride_om, stride_og, stride_oh, stride_od",
+            (stride_oz, stride_om, stride_og, stride_oh, stride_od),
+        )
+        print(
+            "stride_osk_zhg, stride_osk_s, stride_osk_m, stride_osk_d",
+            (stride_osk_zhg, stride_osk_s, stride_osk_m, stride_osk_d),
+        )
+        print(
+            "stride_mzhg, stride_m2, stride_ms, stride_mm",
+            (stride_mzhg, stride_m2, stride_ms, stride_mm),
+        )
+        print("stride_lse_zhg, stride_lse_m", (stride_lse_zhg, stride_lse_m))
+
+    _fwd_kernel_splitK[grid](
+        Q=q,
+        K=k_cache,
+        V=v_cache,
+        Q_Descale=q_descale,
+        K_Descale=k_descale,
+        V_Descale=v_descale,
+        sm_scale=sm_scale,
+        Out_splitK=out_splitk,
+        Metadata=metadata,
+        K_new=None,
+        V_new=None,
+        Cache_seqlens=cache_seqlens,
+        Cache_batch_idx=cache_batch_idx,
+        Block_table=block_table,
+        Alibi_slopes=alibi_slopes,
+        # q strides
+        stride_qz=stride_qz,
+        stride_qm=stride_qm,
+        stride_qg=stride_qg,
+        stride_qh=stride_qh,
+        stride_qd=stride_qd,
+        # k strides
+        stride_kz=stride_kc_z,
+        stride_kn=stride_kc_n,
+        stride_kg=stride_kc_g,
+        stride_kh=stride_kc_h,
+        stride_kd=stride_kc_d,
+        # v strides
+        stride_vz=stride_vc_z,
+        stride_vn=stride_vc_n,
+        stride_vg=stride_vc_g,
+        stride_vh=stride_vc_h,
+        stride_vd=stride_vc_d,
+        # out_splitk strides
+        stride_osk_zhg=stride_osk_zhg,
+        stride_osk_s=stride_osk_s,
+        stride_osk_m=stride_osk_m,
+        stride_osk_d=stride_osk_d,
+        # metadata strides
+        stride_mzhg=stride_mzhg,
+        stride_m2=stride_m2,
+        stride_ms=stride_ms,
+        stride_mm=stride_mm,
+        # k_new strides
+        stride_kn_z=stride_kn_z,
+        stride_kn_n=stride_kn_n,
+        stride_kn_g=stride_kn_g,
+        stride_kn_h=stride_kn_h,
+        stride_kn_d=stride_kn_d,
+        # v_new strides
+        stride_vn_z=stride_vn_z,
+        stride_vn_n=stride_vn_n,
+        stride_vn_g=stride_vn_g,
+        stride_vn_h=stride_vn_h,
+        stride_vn_d=stride_vn_d,
+        # block table strides
+        stride_bt_b=stride_bt_b,
+        stride_bt_s=stride_bt_s,
+        # alibi strides
+        stride_az=stride_az,
+        stride_ah=stride_ah,
+        # FP8 descale strides
+        stride_q_descale_z=stride_q_descale_z,
+        stride_q_descale_h=stride_q_descale_h,
+        stride_k_descale_z=stride_k_descale_z,
+        stride_k_descale_h=stride_k_descale_h,
+        stride_v_descale_z=stride_v_descale_z,
+        stride_v_descale_h=stride_v_descale_h,
+        Z=batch_size,
+        H_q=heads_per_group_q,
+        H_kv=heads_per_group_k,
+        G_q=n_group_q,
+        N_CTX_Q=seqlen_q,
+        N_CTX_K=seqlen_kc,
+        N_CTX_NEW=0,  # No new KV, cache already updated
+        BLOCK_N_PER_SPLIT=split_size,
+        BLOCK_SIZE_K=block_size_k if use_block_table else 256,
+        BLOCK_DMODEL=dim_padded,
+        ACTUAL_BLOCK_DMODEL=dim_kc,
+        USE_CACHE_SEQLENs=use_cache_seqlens,
+        USE_CACHE_BATCH_IDX=cache_batch_idx is not None,
+        NEW_KV=False,  # Cache already updated
+        IS_GQA=is_gqa,
+        IS_CAUSAL=causal,
+        USE_ALIBI=use_alibi,
+        PADDED_HEAD=is_padded_head,
+        GROUP_SIZE=group_size,
+        USE_SLIDING_WINDOW=use_sliding_window,
+        WINDOW_SIZE_LEFT=window_size_left,
+        WINDOW_SIZE_RIGHT=window_size_right,
+        USE_BLOCK_TABLE=use_block_table,
+        IS_FP8=IS_FP8,
+    )
+
+    if DEBUG:
+        print("Out_splitK:", out_splitk, out_splitk.shape)
+        print("metadata:", metadata, metadata.shape)
+        print("lse:", lse, lse.shape)
+        print("Out:", out, out.shape)
+
+    # Merge together
+    splitK_pow2 = triton.next_power_of_2(split_k)
+    mask_split_k = splitK_pow2 > split_k
+    if batch_size * n_group_q * heads_per_group_q * seqlen_q >= 512:
+        k_block_num = 1
+    else:
+        k_block_num = 2
+    assert dim_padded % k_block_num == 0
+    k_block_size = dim_padded // k_block_num
+    reduce_grid = (batch_size * n_group_q * heads_per_group_q, seqlen_q, k_block_num)
+
+    if DEBUG:
+        print("splitK_pow2:", splitK_pow2)
+        print("k_block_num:", k_block_num)
+        print("k_block_size:", k_block_size)
+        print("grid:", reduce_grid)
+
+    _splitK_reduce[reduce_grid](
+        out_splitk,
+        metadata,
+        out,
+        lse,
+        # Split-K output strides
+        stride_osk_zhg=stride_osk_zhg,
+        stride_osk_s=stride_osk_s,
+        stride_osk_m=stride_osk_m,
+        stride_osk_k=stride_osk_d,
+        # Metadata strides
+        stride_mzhg=stride_mzhg,
+        stride_m2=stride_m2,
+        stride_ms=stride_ms,
+        stride_mm=stride_mm,
+        # Output tensor strides
+        stride_oz=stride_oz,
+        stride_oh=stride_oh,
+        stride_og=stride_og,
+        stride_om=stride_om,
+        stride_ok=stride_od,
+        # LSE strides
+        stride_lse_zhg=stride_lse_zhg,
+        stride_lse_m=stride_lse_m,
+        K_BLOCK_SIZE=k_block_size,
+        BLOCK_DMODEL=dim_padded,
+        ACTUAL_BLOCK_DMODEL=dim_kc,
+        G=n_group_q,
+        H=heads_per_group_q,
+        # TODO: Tune num_warps
+        split_k=split_k,
+        splitK_pow2=splitK_pow2,
+        MASK_SPLITK=mask_split_k,
+        PADDED_HEAD=is_padded_head,
+    )

build/torch-rocm/_kernels/flash_attn_triton_amd/fwd_prefill.py

@@ -0,0 +1,1875 @@
+import warnings
+import torch
+import triton
+import triton.language as tl
+from typing import Literal, Optional
+from .common import compute_alibi_block, compute_fp8_scaling_factors, apply_rotary
+from .utils import (
+    AUTOTUNE,
+    AutotuneMode,
+    DEBUG,
+    FWD_CONF_OVERRIDE,
+    get_arch,
+    is_fp8,
+    remap_xcd,
+)
+
+FWD_PREFILL_AUTOTUNE_KEYS = [
+    "IS_CAUSAL",
+    "dropout_p",
+    "MAX_SEQLENS_Q",
+    "MAX_SEQLENS_K",
+    "ACTUAL_BLOCK_DMODEL_QK",
+    "ACTUAL_BLOCK_DMODEL_V",
+    "IS_VARLEN",
+    "HQ",
+    "HK",
+]
+
+
+def get_fwd_prefill_configs(mode: AutotuneMode):
+    # NOTE: Tests expect specific BLOCK_N sizes for attention score renormalization:
+    #   - CDNA: BLOCK_N=64
+    #   - RDNA: BLOCK_N=32
+    # See _get_block_size_n_triton() in test_flash_attn_triton_amd.py
+
+    if mode == "off":
+        if FWD_CONF_OVERRIDE:
+            return [FWD_CONF_OVERRIDE]
+        arch = get_arch()
+        if arch.name == "gfx950":
+            return [
+                triton.Config(
+                    {
+                        "BLOCK_M": 128,
+                        "BLOCK_N": 64,
+                        "waves_per_eu": 2,
+                        "PRE_LOAD_V": False,
+                    },
+                    num_stages=1,
+                    num_warps=4,
+                ),
+            ]
+        elif arch.name == "gfx942":
+            return [
+                triton.Config(
+                    {
+                        "BLOCK_M": 128,
+                        "BLOCK_N": 64,
+                        "waves_per_eu": 2,
+                        "PRE_LOAD_V": False,
+                    },
+                    num_stages=1,
+                    num_warps=4,
+                ),
+            ]
+        elif arch.is_rdna:
+            BLOCK_N = 64 if arch.name == "gfx1100" else 32
+            return [
+                triton.Config(
+                    {
+                        "BLOCK_M": 128,
+                        "BLOCK_N": BLOCK_N,
+                        "PRE_LOAD_V": False,
+                        "waves_per_eu": 6,
+                    },
+                    num_stages=1,
+                    num_warps=8,
+                ),
+            ]
+        else:
+            return [
+                triton.Config(
+                    {
+                        "BLOCK_M": 64,
+                        "BLOCK_N": 64,
+                        "waves_per_eu": 2,
+                        "PRE_LOAD_V": False,
+                    },
+                    num_stages=1,
+                    num_warps=4,
+                )
+            ]
+
+    elif mode == "on":
+        if FWD_CONF_OVERRIDE:
+            return [FWD_CONF_OVERRIDE]
+        arch = get_arch()
+        if arch.name == "gfx950":
+            return [
+                triton.Config(
+                    {
+                        "BLOCK_M": 128,
+                        "BLOCK_N": 64,
+                        "waves_per_eu": 2,
+                        "PRE_LOAD_V": False,
+                    },
+                    num_stages=1,
+                    num_warps=4,
+                ),
+                triton.Config(
+                    {
+                        "BLOCK_M": 128,
+                        "BLOCK_N": 64,
+                        "waves_per_eu": 2,
+                        "PRE_LOAD_V": False,
+                    },
+                    num_stages=1,
+                    num_warps=2,
+                ),
+                triton.Config(
+                    {
+                        "BLOCK_M": 128,
+                        "BLOCK_N": 64,
+                        "waves_per_eu": 2,
+                        "PRE_LOAD_V": False,
+                    },
+                    num_stages=2,
+                    num_warps=4,
+                ),
+                triton.Config(
+                    {
+                        "BLOCK_M": 128,
+                        "BLOCK_N": 128,
+                        "waves_per_eu": 2,
+                        "PRE_LOAD_V": False,
+                    },
+                    num_stages=2,
+                    num_warps=4,
+                ),
+            ]
+        elif arch.name == "gfx942":
+            if arch.cu_count < 304:
+                return [
+                    triton.Config(
+                        {
+                            "BLOCK_M": 128,
+                            "BLOCK_N": 64,
+                            "waves_per_eu": 2,
+                            "PRE_LOAD_V": False,
+                        },
+                        num_stages=1,
+                        num_warps=4,
+                    ),
+                ]
+            else:
+                return [
+                    triton.Config(
+                        {
+                            "BLOCK_M": 128,
+                            "BLOCK_N": 64,
+                            "waves_per_eu": 2,
+                            "PRE_LOAD_V": False,
+                        },
+                        num_stages=1,
+                        num_warps=4,
+                    )
+                ]
+        elif arch.is_rdna:
+            BLOCK_N = 64 if arch.name == "gfx1100" else 32
+            return [
+                triton.Config(
+                    {
+                        "BLOCK_M": 128,
+                        "BLOCK_N": BLOCK_N,
+                        "PRE_LOAD_V": False,
+                        "waves_per_eu": 6,
+                    },
+                    num_stages=1,
+                    num_warps=8,
+                ),
+            ]
+        else:
+            return [
+                triton.Config(
+                    {
+                        "BLOCK_M": 64,
+                        "BLOCK_N": 64,
+                        "waves_per_eu": 2,
+                        "PRE_LOAD_V": False,
+                    },
+                    num_stages=1,
+                    num_warps=4,
+                )
+            ]
+
+    else:  # sweep
+        configs = []
+        BLOCK_M_OPTIONS = [128, 64, 32, 16]
+        BLOCK_N_OPTIONS = [128, 64, 32, 16]
+        NUM_WARPS_OPTIONS = [2, 4, 8]
+        NUM_STAGES_OPTIONS = [1, 2]
+        WAVES_PER_EU_OPTIONS = [4, 2, 1]
+        PRE_LOAD_V_OPTIONS = [False]
+        for bm in BLOCK_M_OPTIONS:
+            for bn in BLOCK_N_OPTIONS:
+                for waves in WAVES_PER_EU_OPTIONS:
+                    for nw in NUM_WARPS_OPTIONS:
+                        for ns in NUM_STAGES_OPTIONS:
+                            for preload_v in PRE_LOAD_V_OPTIONS:
+                                configs.append(
+                                    triton.Config(
+                                        {
+                                            "BLOCK_M": bm,
+                                            "BLOCK_N": bn,
+                                            "waves_per_eu": waves,
+                                            "PRE_LOAD_V": preload_v,
+                                        },
+                                        num_stages=ns,
+                                        num_warps=nw,
+                                    )
+                                )
+        return configs
+
+
+fwd_prefill_autotune_configs = get_fwd_prefill_configs(AUTOTUNE)
+
+
+@triton.jit
+def _attn_fwd_inner(
+    acc,
+    l_i,
+    m_i,
+    q,
+    k_base_ptrs,
+    v_base_ptrs,
+    bias_base_ptrs,
+    stride_kn,
+    stride_vk,
+    stride_bn,
+    stride_sn,
+    stride_sm,
+    start_m,
+    seqlen_k,
+    seqlen_q,
+    dropout_p,
+    philox_seed,
+    philox_offset_base,
+    sd_mask,
+    stride_sz,
+    stride_sh,
+    off_z,
+    off_h_q,
+    offs_m,
+    offs_n,
+    offs_d_qk,
+    offs_d_v,
+    block_min,
+    block_max,
+    n_extra_tokens,
+    alibi_slope,
+    q_descale,
+    k_descale,
+    v_descale,
+    IS_FP8: tl.constexpr,
+    FP8_MAX: tl.constexpr,
+    FP8_P_DESCALE: tl.constexpr,
+    APPLY_MASK: tl.constexpr,  # True for masked blocks, False for full blocks
+    IS_CAUSAL: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_DMODEL_QK: tl.constexpr,
+    BLOCK_DMODEL_V: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    PRE_LOAD_V: tl.constexpr,
+    ENABLE_DROPOUT: tl.constexpr,
+    PADDED_HEAD_QK: tl.constexpr,
+    PADDED_HEAD_V: tl.constexpr,
+    ACTUAL_BLOCK_DMODEL_QK: tl.constexpr,
+    ACTUAL_BLOCK_DMODEL_V: tl.constexpr,
+    SM_SCALE: tl.constexpr,
+    USE_ALIBI: tl.constexpr,
+    USE_EXP2: tl.constexpr,
+    RETURN_SCORES: tl.constexpr,
+    USE_SLIDING_WINDOW: tl.constexpr,
+    WINDOW_SIZE_LEFT: tl.constexpr,
+    WINDOW_SIZE_RIGHT: tl.constexpr,
+    ACCUMULATOR_TYPE,
+):
+    """
+    Unified attention forward inner loop.
+
+    APPLY_MASK controls whether causal/window masking is applied:
+    - False: Fast path for full blocks (no masking overhead)
+    - True: Masked path with causal/window masking support
+    """
+    if USE_EXP2:
+        RCP_LN2: tl.constexpr = 1.4426950408889634
+
+    # seqlen diff (only used when APPLY_MASK=True)
+    seqlen_delta_qk = seqlen_k - seqlen_q
+
+    # loop over k, v, and update accumulator
+    for start_n in range(block_min, block_max, BLOCK_N):
+        # get ptrs
+        k_ptrs = k_base_ptrs + start_n * stride_kn
+        v_ptrs = v_base_ptrs + start_n * stride_vk
+
+        kv_offs_n = start_n + tl.arange(0, BLOCK_N)
+
+        # Load K - different masking for APPLY_MASK vs non-masked
+        if APPLY_MASK:
+            # For masked blocks, check seqlen bounds
+            k_mask = kv_offs_n[None, :] < seqlen_k
+            v_mask = kv_offs_n[:, None] < seqlen_k
+            if PADDED_HEAD_QK:
+                k_mask = k_mask & (offs_d_qk[:, None] < ACTUAL_BLOCK_DMODEL_QK)
+            if PADDED_HEAD_V:
+                v_mask = v_mask & (offs_d_v[None, :] < ACTUAL_BLOCK_DMODEL_V)
+            k = tl.load(k_ptrs, mask=k_mask, other=0.0)
+            if PRE_LOAD_V:
+                v = tl.load(v_ptrs, mask=v_mask, other=0.0)
+        else:
+            # For full blocks, only check head dimension padding
+            if PADDED_HEAD_QK:
+                k_mask = offs_d_qk[:, None] < ACTUAL_BLOCK_DMODEL_QK
+                k = tl.load(k_ptrs, mask=k_mask, other=0.0)
+            else:
+                k = tl.load(k_ptrs)
+            if PRE_LOAD_V:
+                if PADDED_HEAD_V:
+                    v_mask = offs_d_v[None, :] < ACTUAL_BLOCK_DMODEL_V
+                    v = tl.load(v_ptrs, mask=v_mask, other=0.0)
+                else:
+                    v = tl.load(v_ptrs)
+
+        # setup qk accumulator
+        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=ACCUMULATOR_TYPE)
+
+        # Apply extra token masking for partial blocks (only when APPLY_MASK=True)
+        if APPLY_MASK:
+            if (n_extra_tokens != 0) and (start_n + BLOCK_N == block_max):
+                boundary_m = tl.full([BLOCK_M], seqlen_k, dtype=tl.int32)
+                size_n = start_n + offs_n[None, :]
+                mask = size_n < boundary_m[:, None]
+                qk = tl.where(mask, qk, float("-inf"))
+
+        # -- compute qk ----
+        if IS_FP8:
+            qk += tl.dot(q, k) * q_descale * k_descale
+        else:
+            qk = tl.dot(q, k, acc=qk)
+        qk_scaled = qk * SM_SCALE
+
+        if USE_ALIBI:
+            # compute the global position of each token within the sequence
+            q_offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+            alibi_block = compute_alibi_block(
+                alibi_slope, seqlen_q, seqlen_k, q_offs_m, kv_offs_n
+            )
+            qk_scaled += alibi_block
+
+        # Apply causal/sliding window masking (only when APPLY_MASK=True)
+        if APPLY_MASK:
+            if USE_SLIDING_WINDOW:
+                if IS_CAUSAL:
+                    # ========== CAUSAL SLIDING WINDOW MASKING ==========
+                    row_idx = offs_m
+                    col_idx = kv_offs_n
+                    row_idx_expanded = row_idx[:, None]
+                    col_idx_expanded = col_idx[None, :]
+
+                    causal_offset = seqlen_k - seqlen_q
+                    causal_mask = col_idx_expanded > (row_idx_expanded + causal_offset)
+
+                    if WINDOW_SIZE_LEFT < 0:
+                        window_mask = col_idx_expanded > (
+                            row_idx_expanded + causal_offset + WINDOW_SIZE_RIGHT
+                        )
+                    else:
+                        left_bound = row_idx_expanded + causal_offset - WINDOW_SIZE_LEFT
+                        right_bound = (
+                            row_idx_expanded + causal_offset + WINDOW_SIZE_RIGHT
+                        )
+                        window_mask = (col_idx_expanded < left_bound) | (
+                            col_idx_expanded > right_bound
+                        )
+
+                    mask = causal_mask | window_mask
+                    qk_scaled = tl.where(mask, float("-inf"), qk_scaled)
+                else:
+                    # ========== NON-CAUSAL SLIDING WINDOW MASKING ==========
+                    row_idx = offs_m
+                    col_idx = kv_offs_n
+                    sk = seqlen_k
+                    sq = seqlen_q
+                    row_idx_expanded = row_idx[:, None]
+                    col_idx_expanded = col_idx[None, :]
+
+                    if WINDOW_SIZE_LEFT < 0:
+                        mask = col_idx_expanded > (
+                            row_idx_expanded + sk - sq + WINDOW_SIZE_RIGHT
+                        )
+                    else:
+                        sk_full = tl.full((1, BLOCK_N), sk, dtype=tl.int32)
+                        right_bound_val = row_idx_expanded + sk - sq + WINDOW_SIZE_RIGHT
+                        right_bound = tl.minimum(right_bound_val, sk_full)
+                        left_bound = row_idx_expanded + sk - sq - WINDOW_SIZE_LEFT
+                        mask = (col_idx_expanded > right_bound) | (
+                            col_idx_expanded < left_bound
+                        )
+
+                    qk_scaled = tl.where(mask, float("-inf"), qk_scaled)
+            else:
+                if IS_CAUSAL:
+                    causal_boundary = start_n + offs_n - seqlen_delta_qk
+                    causal_mask = offs_m[:, None] >= causal_boundary[None, :]
+                    qk_scaled = tl.where(causal_mask, qk_scaled, float("-inf"))
+
+        # compute qk mask for bounds checking
+        qk_mask = (offs_m[:, None] < seqlen_q) & (kv_offs_n[None, :] < seqlen_k)
+
+        # compute bias
+        if bias_base_ptrs is not None:
+            bias_ptrs = bias_base_ptrs + start_n * stride_bn
+            bias = tl.load(bias_ptrs, mask=qk_mask, other=0.0)
+            qk_scaled += bias
+
+        # get max scores so far
+        m_ij = tl.maximum(m_i, tl.max(qk_scaled, 1))
+
+        # scale and subtract max
+        # Handle the case where all values are -inf
+        q_shifted = tl.where(
+            m_ij[:, None] == float("-inf"), float("-inf"), qk_scaled - m_ij[:, None]
+        )
+
+        # Compute scaled QK and softmax probabilities
+        if USE_EXP2:
+            p = tl.math.exp2(q_shifted * RCP_LN2)
+        else:
+            p = tl.math.exp(q_shifted)
+
+        # CAVEAT: Must update l_ij before applying dropout
+        l_ij = tl.sum(p, 1)
+        if ENABLE_DROPOUT:
+            # Compute pointers for this block
+            philox_base = philox_offset_base + off_z * stride_sz + off_h_q * stride_sh
+            philox_ptrs = (
+                philox_base
+                + offs_m[:, None] * stride_sm
+                + kv_offs_n[None, :] * stride_sn
+            )
+
+            # compute dropout mask
+            rng_output = tl.rand(philox_seed, philox_ptrs)
+            dropout_mask = rng_output > dropout_p
+
+            # return scores with negative values for dropped vals (only if RETURN_SCORES is True)
+            if RETURN_SCORES:
+                sd_mask_value = tl.where(dropout_mask, p, -p)
+                sd_mask_base = sd_mask + off_z * stride_sz + off_h_q * stride_sh
+                sd_mask_ptrs = (
+                    sd_mask_base
+                    + offs_m[:, None] * stride_sm
+                    + kv_offs_n[None, :] * stride_sn
+                )
+
+                sd_store_mask = (offs_m[:, None] < seqlen_q) & (
+                    kv_offs_n[None, :] < seqlen_k
+                )
+
+                if APPLY_MASK and IS_CAUSAL:
+                    causal_constraint = kv_offs_n[None, :] <= (
+                        offs_m[:, None] + seqlen_delta_qk
+                    )
+                    sd_store_mask = sd_store_mask & causal_constraint
+
+                if APPLY_MASK and USE_SLIDING_WINDOW:
+                    if WINDOW_SIZE_LEFT < 0:
+                        window_constraint = kv_offs_n[None, :] <= (
+                            offs_m[:, None] + seqlen_delta_qk + WINDOW_SIZE_RIGHT
+                        )
+                    else:
+                        left_bound = (
+                            offs_m[:, None] + seqlen_delta_qk - WINDOW_SIZE_LEFT
+                        )
+                        right_bound = (
+                            offs_m[:, None] + seqlen_delta_qk + WINDOW_SIZE_RIGHT
+                        )
+                        window_constraint = (kv_offs_n[None, :] >= left_bound) & (
+                            kv_offs_n[None, :] <= right_bound
+                        )
+                    sd_store_mask = sd_store_mask & window_constraint
+
+                tl.store(sd_mask_ptrs, sd_mask_value, mask=sd_store_mask)
+
+            # apply dropout mask in place
+            p = tl.where(dropout_mask, p, 0.0)
+        elif RETURN_SCORES:
+            sd_mask_base = sd_mask + off_z * stride_sz + off_h_q * stride_sh
+            sd_mask_ptrs = (
+                sd_mask_base
+                + offs_m[:, None] * stride_sm
+                + kv_offs_n[None, :] * stride_sn
+            )
+
+            sd_store_mask = (offs_m[:, None] < seqlen_q) & (
+                kv_offs_n[None, :] < seqlen_k
+            )
+
+            if APPLY_MASK and IS_CAUSAL:
+                causal_constraint = kv_offs_n[None, :] <= (
+                    offs_m[:, None] + seqlen_delta_qk
+                )
+                sd_store_mask = sd_store_mask & causal_constraint
+
+            if APPLY_MASK and USE_SLIDING_WINDOW:
+                if WINDOW_SIZE_LEFT < 0:
+                    window_constraint = kv_offs_n[None, :] <= (
+                        offs_m[:, None] + seqlen_delta_qk + WINDOW_SIZE_RIGHT
+                    )
+                else:
+                    left_bound = offs_m[:, None] + seqlen_delta_qk - WINDOW_SIZE_LEFT
+                    right_bound = offs_m[:, None] + seqlen_delta_qk + WINDOW_SIZE_RIGHT
+                    window_constraint = (kv_offs_n[None, :] >= left_bound) & (
+                        kv_offs_n[None, :] <= right_bound
+                    )
+                sd_store_mask = sd_store_mask & window_constraint
+
+            tl.store(sd_mask_ptrs, p, mask=sd_store_mask)
+
+        # -- update output accumulator --
+        m_diff = tl.where(m_ij == float("-inf"), float("-inf"), m_i - m_ij)
+        if USE_EXP2:
+            alpha = tl.math.exp2(m_diff * RCP_LN2)
+        else:
+            alpha = tl.math.exp(m_diff)
+        acc = acc * alpha[:, None]
+
+        # Load V if not preloaded
+        if not PRE_LOAD_V:
+            if APPLY_MASK:
+                v_mask = kv_offs_n[:, None] < seqlen_k
+                if PADDED_HEAD_V:
+                    v_mask = v_mask & (offs_d_v[None, :] < ACTUAL_BLOCK_DMODEL_V)
+                v = tl.load(v_ptrs, mask=v_mask, other=0.0)
+            else:
+                if PADDED_HEAD_V:
+                    v_mask = offs_d_v[None, :] < ACTUAL_BLOCK_DMODEL_V
+                    v = tl.load(v_ptrs, mask=v_mask, other=0.0)
+                else:
+                    v = tl.load(v_ptrs)
+
+        # -- update m_i and l_i
+        l_i = l_i * alpha + l_ij
+        m_i = m_ij
+
+        if IS_FP8:
+            if FP8_P_DESCALE:
+                scale_p, descale_p = compute_fp8_scaling_factors(p, FP8_MAX)
+                acc += (
+                    tl.dot((p * scale_p).to(v.type.element_ty), v)
+                    * descale_p
+                    * v_descale
+                )
+            else:
+                acc += tl.dot(p.to(v.type.element_ty), v) * v_descale
+        else:
+            acc = tl.dot(p.to(v.type.element_ty), v, acc=acc)
+
+    return acc, l_i, m_i
+
+
+@triton.jit
+def compute_window_bounds(
+    q_start,
+    q_end,
+    diag,
+    seqlen_k,
+    WINDOW_SIZE_LEFT: tl.constexpr,
+    WINDOW_SIZE_RIGHT: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+):
+    """Calculate the window boundaries for a query block."""
+    # Left boundary
+    if WINDOW_SIZE_LEFT < 0:
+        left_min = 0
+        left_max = 0
+    else:
+        left_min = tl.maximum(0, q_start + diag - WINDOW_SIZE_LEFT)
+        left_max = tl.maximum(0, q_end + diag - WINDOW_SIZE_LEFT)
+
+    # Right boundary
+    if IS_CAUSAL:
+        # Causal cap: col <= row + diag
+        right_min = tl.minimum(seqlen_k - 1, q_start + diag)
+        right_max = tl.minimum(seqlen_k - 1, q_end + diag)
+    else:
+        if WINDOW_SIZE_RIGHT < 0:
+            right_min = tl.minimum(seqlen_k - 1, q_start + diag + WINDOW_SIZE_RIGHT)
+            right_max = tl.minimum(seqlen_k - 1, q_end + diag + WINDOW_SIZE_RIGHT)
+        else:
+            # Non-causal doesn't have the diagonal constraint
+            right_min = tl.minimum(seqlen_k - 1, q_start + diag + WINDOW_SIZE_RIGHT)
+            right_max = tl.minimum(seqlen_k - 1, q_end + diag + WINDOW_SIZE_RIGHT)
+
+    return left_min, left_max, right_min, right_max
+
+
+@triton.jit
+def classify_window_blocks(
+    left_min, left_max, right_min, right_max, BLOCK_N: tl.constexpr
+):
+    """Classify blocks based on window boundaries."""
+    # First and last blocks that have ANY overlap with window
+    first_block = left_min // BLOCK_N
+    last_block = right_max // BLOCK_N
+
+    # First block that is FULLY visible for all rows in Q block
+    full_left_block = left_max // BLOCK_N + (left_max % BLOCK_N != 0)
+    clipped_left = tl.minimum(full_left_block, last_block + 1)
+
+    # Last block that is FULLY visible for all rows in Q block
+    last_full_block_candidate = right_min // BLOCK_N
+    if (last_full_block_candidate + 1) * BLOCK_N - 1 > right_min:
+        last_full_block_candidate -= 1
+    full_right_block = tl.maximum(last_full_block_candidate, clipped_left - 1)
+
+    # Calculate counts
+    n_front_skip_blocks = first_block
+    n_front_masked_blocks = tl.maximum(0, clipped_left - first_block)
+    n_full_blocks = tl.maximum(0, full_right_block - clipped_left + 1)
+    n_back_masked_blocks = tl.maximum(0, last_block - full_right_block)
+
+    return (
+        n_front_skip_blocks,
+        n_front_masked_blocks,
+        n_full_blocks,
+        n_back_masked_blocks,
+        clipped_left,
+    )  # Return clipped_left for padded block handling
+
+
+@triton.jit
+def handle_padded_last_block(
+    n_extra_tokens,
+    last_block,
+    total_k_blocks,
+    clipped_left,
+    n_front_masked_blocks,
+    n_full_blocks,
+    n_back_masked_blocks,
+):
+    """Ensure a padded last K-block is never classified as 'full'.
+
+    We move the padded last block (if visible) into the back-masked bucket.
+    If it's already back-masked, we do nothing.  If it was counted in the
+    front-masked range, we decrement front-masked; if it was counted as full,
+    we decrement full.  Then we increment back-masked.
+    """
+    padded_last_k = (n_extra_tokens != 0) & (last_block == total_k_blocks - 1)
+
+    if padded_last_k:
+        # current 'full' range right edge
+        full_right_block = clipped_left + n_full_blocks - 1
+
+        # If last_block is already beyond full_right_block, it's already in back-masked -> nothing to do
+        last_already_back_masked = last_block > full_right_block
+        if not last_already_back_masked:
+            # If the window starts past last_block, it was counted in front-masked
+            if clipped_left > last_block:
+                n_front_masked_blocks = tl.maximum(0, n_front_masked_blocks - 1)
+            else:
+                # Otherwise it was counted 'full' -> move it out of full
+                n_full_blocks = tl.maximum(0, n_full_blocks - 1)
+            # In both cases we need one more back-masked block
+            n_back_masked_blocks = n_back_masked_blocks + 1
+
+    return n_front_masked_blocks, n_full_blocks, n_back_masked_blocks
+
+
+@triton.jit
+def compute_padding_info(seqlen_k, BLOCK_N: tl.constexpr):
+    """Calculate padding information for the last K block."""
+    # check if we will need to do masking due either BLOCK_N being bigger than seqlen_k or seqlen_k not being a factor of BLOCK_N
+    # n_extra_tokens = 10 % 4 = 2
+    # This means the last K block has 2 valid tokens and 2 padding positions
+    # K blocks visualization:
+    #         Block 0         Block 1         Block 2 (last)
+    #         K0 K1 K2 K3    K4 K5 K6 K7     K8 K9 ?? ??
+    #         ?---------?    ?---------?     ?---? ?---?
+    #         full block     full block      valid  pad
+    if seqlen_k < BLOCK_N:
+        n_extra_tokens = BLOCK_N - seqlen_k
+    elif seqlen_k % BLOCK_N:
+        n_extra_tokens = seqlen_k % BLOCK_N
+    else:
+        n_extra_tokens = 0
+    return n_extra_tokens
+
+
+@triton.jit
+def compute_block_masking(
+    seqlen_k,
+    seqlen_q,
+    start_m,
+    IS_CAUSAL: tl.constexpr,
+    USE_SLIDING_WINDOW: tl.constexpr,
+    WINDOW_SIZE_LEFT: tl.constexpr,
+    WINDOW_SIZE_RIGHT: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """
+    Classify K blocks for attention computation with sliding window support.
+
+    Returns:
+        - n_front_skip_blocks: Blocks completely before the window
+        - n_front_masked_blocks: Blocks partially overlapping window front
+        - n_full_blocks: Blocks completely inside the window
+        - n_back_masked_blocks: Blocks partially overlapping window back
+        - n_extra_tokens: Padding tokens in last K block
+    """
+
+    # common
+    q_start = start_m * BLOCK_M
+    q_end = tl.minimum((start_m + 1) * BLOCK_M - 1, seqlen_q - 1)
+    diag = seqlen_k - seqlen_q
+    total_k_blocks = tl.cdiv(seqlen_k, BLOCK_N)
+    n_extra_tokens = compute_padding_info(seqlen_k, BLOCK_N)
+
+    if USE_SLIDING_WINDOW:
+        # get window bounds
+        left_min, left_max, right_min, right_max = compute_window_bounds(
+            q_start,
+            q_end,
+            diag,
+            seqlen_k,
+            WINDOW_SIZE_LEFT,
+            WINDOW_SIZE_RIGHT,
+            IS_CAUSAL,
+        )
+
+        # window vanishes -> early exit
+        if right_max < left_min:
+            return 0, 0, 0, 0, n_extra_tokens
+
+        # classify blocks
+        (
+            n_front_skip_blocks,
+            n_front_masked_blocks,
+            n_full_blocks,
+            n_back_masked_blocks,
+            clipped_left,
+        ) = classify_window_blocks(left_min, left_max, right_min, right_max, BLOCK_N)
+
+        # handle padded last block if needed
+        if n_extra_tokens != 0:
+            last_block = right_max // BLOCK_N
+            n_front_masked_blocks, n_full_blocks, n_back_masked_blocks = (
+                handle_padded_last_block(
+                    n_extra_tokens,
+                    last_block,
+                    total_k_blocks,
+                    clipped_left,
+                    n_front_masked_blocks,
+                    n_full_blocks,
+                    n_back_masked_blocks,
+                )
+            )
+        return (
+            n_front_skip_blocks,
+            n_front_masked_blocks,
+            n_full_blocks,
+            n_back_masked_blocks,
+            n_extra_tokens,
+        )
+    else:
+        if IS_CAUSAL:
+            # ========== CAUSAL MODE: Classify K Blocks ==========
+            # Calculate causal boundary for this Q block
+            #          [K0 K1 K2 K3] [K4 K5 K6 K7] [K8 K9 ?? ??]
+            # Q0-Q3:   [ 1  0  0  0] [ 0  0  0  0] [ 0  0 -- --]  <- Q0
+            #          [ 1  1  0  0] [ 0  0  0  0] [ 0  0 -- --]  <- Q1
+            #          [ 1  1  1  0] [ 0  0  0  0] [ 0  0 -- --]  <- Q2
+            #          [ 1  1  1  1] [ 1  1  0  0] [ 0  0 -- --]  <- Q3
+            #                           /\ can see up to K5
+            #
+            # Q4-Q7:   [ 1  1  1  1] [ 1  1  1  0] [ 0  0 -- --]  <- Q4
+            #          [ 1  1  1  1] [ 1  1  1  1] [ 0  0 -- --]  <- Q5
+            #          [ 1  1  1  1] [ 1  1  1  1] [ 1  0 -- --]  <- Q6
+            #          [ 1  1  1  1] [ 1  1  1  1] [ 1  1 -- --]  <- Q7
+
+            # ------------------------------------------------------------
+            # 1. figure out, in tokens, the right-most K position
+            #    this Q-block may attend to
+            # ------------------------------------------------------------
+            k_max_token = q_end + diag  # last visible K index
+
+            # this Q-block is entirely above the diagonal => nothing to do
+            if k_max_token < 0:
+                return 0, 0, 0, 0, n_extra_tokens
+
+            k_max_token = tl.minimum(k_max_token, seqlen_k - 1)
+
+            # ------------------------------------------------------------
+            # 2. translate token indices into K-block indices
+            # ------------------------------------------------------------
+            last_visible_k_block = k_max_token // BLOCK_N
+            n_visible_k_blocks = tl.minimum(last_visible_k_block + 1, total_k_blocks)
+
+            # ------------------------------------------------------------
+            # 3. classify those visible blocks
+            #    - we *never* skip or mask blocks in front, because causal
+            #      attention always starts at K0
+            #    - the back side can require several masked blocks:
+            #         o intersection of the causal diagonal with K-grid
+            #           (at most  ?BLOCK_M / BLOCK_N? blocks)
+            #         o plus one for partial K blocks at the causal boundary
+            # ------------------------------------------------------------
+            n_back_masked_blocks = BLOCK_M // BLOCK_N + 1
+            n_back_masked_blocks = tl.minimum(n_back_masked_blocks, n_visible_k_blocks)
+
+            n_front_skip_blocks = 0  # causal never skips the left side
+            n_front_masked_blocks = 0  # ditto
+            n_full_blocks = n_visible_k_blocks - n_back_masked_blocks
+        else:
+            # ========== NON-CAUSAL MODE ==========
+            # Without causal mask, all positions can attend to all positions
+            # Only need to handle the padding in the last block
+            #          [K0 K1 K2 K3] [K4 K5 K6 K7] [K8 K9  ??  ??]
+            # Q0-Q3:   [ 1  1  1  1] [ 1  1  1  1] [ 1  1 -inf -inf]
+            #          [ 1  1  1  1] [ 1  1  1  1] [ 1  1 -inf -inf]
+            #          [ 1  1  1  1] [ 1  1  1  1] [ 1  1 -inf -inf]
+            #          [ 1  1  1  1] [ 1  1  1  1] [ 1  1 -inf -inf]
+            #
+            # Q4-Q7:   [ 1  1  1  1] [ 1  1  1  1] [ 1  1 -inf -inf]
+            #          [ 1  1  1  1] [ 1  1  1  1] [ 1  1 -inf -inf]
+            #          [ 1  1  1  1] [ 1  1  1  1] [ 1  1 -inf -inf]
+            #          [ 1  1  1  1] [ 1  1  1  1] [ 1  1 -inf -inf]
+
+            n_front_skip_blocks = 0  # never skips the left side
+            n_front_masked_blocks = 0  # ditto
+            if n_extra_tokens != 0:
+                n_back_masked_blocks = 1  # Last block needs padding mask
+                n_full_blocks = total_k_blocks - 1
+            else:
+                n_back_masked_blocks = 0  # All blocks are aligned
+                n_full_blocks = total_k_blocks
+
+        return (
+            n_front_skip_blocks,
+            n_front_masked_blocks,
+            n_full_blocks,
+            n_back_masked_blocks,
+            n_extra_tokens,
+        )
+
+
+@triton.autotune(
+    configs=fwd_prefill_autotune_configs,
+    key=FWD_PREFILL_AUTOTUNE_KEYS,
+    use_cuda_graph=True,
+)
+@triton.jit
+def attn_fwd(
+    Q,
+    K,
+    V,
+    bias,
+    Q_Descale,
+    K_Descale,
+    V_Descale,
+    stride_q_descale_z,
+    stride_k_descale_z,
+    stride_v_descale_z,
+    LSE,
+    Out,
+    SD_MASK,
+    ALIBI_SLOPES,
+    stride_qz,
+    stride_qh,
+    stride_qm,
+    stride_qk,
+    stride_kz,
+    stride_kh,
+    stride_kn,
+    stride_kk,
+    stride_vz,
+    stride_vh,
+    stride_vk,
+    stride_vn,
+    stride_oz,
+    stride_oh,
+    stride_om,
+    stride_on,
+    stride_bz,
+    stride_bh,
+    stride_bm,
+    stride_bn,
+    stride_az,
+    stride_ah,
+    stride_sz,
+    stride_sh,
+    stride_sm,
+    stride_sn,
+    stride_lse_z,
+    stride_lse_h,
+    stride_lse_m,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    seqused_q,
+    seqused_k,  # Add seqused parameters
+    dropout_p,
+    philox_seed,
+    philox_offset_base,
+    HQ: tl.constexpr,
+    HK: tl.constexpr,
+    ACTUAL_BLOCK_DMODEL_QK: tl.constexpr,
+    ACTUAL_BLOCK_DMODEL_V: tl.constexpr,
+    MAX_SEQLENS_Q: tl.constexpr,
+    MAX_SEQLENS_K: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    SM_SCALE: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+    USE_SLIDING_WINDOW: tl.constexpr,
+    WINDOW_SIZE_LEFT: tl.constexpr,
+    WINDOW_SIZE_RIGHT: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_DMODEL_QK: tl.constexpr,
+    BLOCK_DMODEL_V: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    PRE_LOAD_V: tl.constexpr,
+    USE_BIAS: tl.constexpr,
+    ENABLE_DROPOUT: tl.constexpr,
+    RETURN_SCORES: tl.constexpr,
+    USE_ALIBI: tl.constexpr,
+    USE_EXP2: tl.constexpr,
+    IS_FP8: tl.constexpr,
+    FP8_MAX: tl.constexpr,
+    FP8_P_DESCALE: tl.constexpr,
+    USE_SEQUSED: tl.constexpr,
+    FORCE_MASKING: tl.constexpr,
+    NUM_XCD: tl.constexpr = 1,
+):
+    # set params
+    ACCUMULATOR_TYPE = tl.float32
+
+    # compute offsets
+    off_h_q = tl.program_id(0)
+    # apply the xcd remapping for the hq dim
+    off_h_q = remap_xcd(off_h_q, HQ, NUM_XCD)
+
+    start_m = tl.program_id(1)
+    off_z = tl.program_id(2)
+    # If MQA / GQA, set the K and V head offsets appropriately.
+    GROUP_SIZE: tl.constexpr = HQ // HK
+    if GROUP_SIZE != 1:
+        off_h_k = off_h_q // GROUP_SIZE
+    else:
+        off_h_k = off_h_q
+    # Determine if we need to mask the heads
+    PADDED_HEAD_QK: tl.constexpr = ACTUAL_BLOCK_DMODEL_QK != BLOCK_DMODEL_QK
+    PADDED_HEAD_V: tl.constexpr = ACTUAL_BLOCK_DMODEL_V != BLOCK_DMODEL_V
+
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = tl.arange(0, BLOCK_N)
+    offs_d_qk = tl.arange(0, BLOCK_DMODEL_QK)
+    offs_d_v = tl.arange(0, BLOCK_DMODEL_V)
+
+    # handle seqlen
+    if IS_VARLEN:
+        cu_seqlens_q_start = tl.load(cu_seqlens_q + off_z)
+        cu_seqlens_q_end = tl.load(cu_seqlens_q + off_z + 1)
+
+        # If seqused is provided, use it to limit the actual sequence length
+        if USE_SEQUSED:
+            actual_seqlen_q = (
+                tl.load(seqused_q + off_z)
+                if seqused_q is not None
+                else cu_seqlens_q_end - cu_seqlens_q_start
+            )
+            seqlen_q = tl.minimum(
+                actual_seqlen_q, cu_seqlens_q_end - cu_seqlens_q_start
+            )
+        else:
+            seqlen_q = cu_seqlens_q_end - cu_seqlens_q_start
+
+        # we have a one-size-fits-all grid in id(0). Some seqlens might be too small for all start_m so for those we return early.
+        if start_m * BLOCK_M > seqlen_q:
+            return
+        cu_seqlens_k_start = tl.load(cu_seqlens_k + off_z)
+        cu_seqlens_k_end = tl.load(cu_seqlens_k + off_z + 1)
+
+        # If seqused is provided, use it to limit the actual sequence length for keys
+        if USE_SEQUSED:
+            actual_seqlen_k = (
+                tl.load(seqused_k + off_z)
+                if seqused_k is not None
+                else cu_seqlens_k_end - cu_seqlens_k_start
+            )
+            seqlen_k = tl.minimum(
+                actual_seqlen_k, cu_seqlens_k_end - cu_seqlens_k_start
+            )
+        else:
+            seqlen_k = cu_seqlens_k_end - cu_seqlens_k_start
+    else:
+        cu_seqlens_q_start = 0
+        cu_seqlens_k_start = 0
+        seqlen_q = MAX_SEQLENS_Q
+        seqlen_k = MAX_SEQLENS_K
+
+    # Load scale factors if IS_FP8.
+    if IS_FP8:
+        # For MQA/GQA (GROUP_SIZE != 1), q_descale uses the same indexing as k/v (off_h_k)
+        # For MHA (GROUP_SIZE == 1), q_descale uses off_h_q (same as off_h_k)
+        if GROUP_SIZE != 1:
+            q_descale = tl.load(
+                Q_Descale + off_z * stride_q_descale_z + off_h_k
+            )  # MQA/GQA: broadcast using k/v head index
+        else:
+            q_descale = tl.load(
+                Q_Descale + off_z * stride_q_descale_z + off_h_q
+            )  # MHA: use q head index
+        k_descale = tl.load(K_Descale + off_z * stride_k_descale_z + off_h_k)
+        v_descale = tl.load(V_Descale + off_z * stride_v_descale_z + off_h_k)
+    else:
+        q_descale, k_descale, v_descale = 1.0, 1.0, 1.0
+
+    # figure out masking pattern
+    (
+        n_front_skip_blocks,
+        n_front_masked_blocks,
+        n_full_blocks,
+        n_back_masked_blocks,
+        n_extra_tokens,
+    ) = compute_block_masking(
+        seqlen_k,
+        seqlen_q,
+        start_m,
+        IS_CAUSAL,
+        USE_SLIDING_WINDOW,
+        WINDOW_SIZE_LEFT,
+        WINDOW_SIZE_RIGHT,
+        BLOCK_M,
+        BLOCK_N,
+    )
+
+    # ============================================================
+    #          PROGRAM EARLY EXIT (All K Blocks Skipped)
+    # ============================================================
+    total_visible_blocks = n_front_masked_blocks + n_full_blocks + n_back_masked_blocks
+    if total_visible_blocks == 0:
+        """
+        No K blocks visible - write zeros and exit.
+        """
+        # Write zeros to output
+        o_offset = (
+            Out
+            + off_z * stride_oz
+            + off_h_q * stride_oh
+            + cu_seqlens_q_start * stride_om
+        )
+        o_ptrs = o_offset + offs_m[:, None] * stride_om + offs_d_v[None, :] * stride_on
+        o_mask = offs_m[:, None] < seqlen_q
+        if PADDED_HEAD_V:
+            o_mask = o_mask & (offs_d_v[None, :] < ACTUAL_BLOCK_DMODEL_V)
+        tl.store(
+            o_ptrs,
+            tl.zeros([BLOCK_M, BLOCK_DMODEL_V], dtype=Out.type.element_ty),
+            mask=o_mask,
+        )
+
+        # Write zeros to LSE
+        l_ptrs = (
+            LSE
+            + off_z * stride_lse_z
+            + off_h_q * stride_lse_h
+            + cu_seqlens_q_start * stride_lse_m
+            + offs_m * stride_lse_m
+        )
+        tl.store(l_ptrs, tl.zeros([BLOCK_M], dtype=tl.float32), mask=offs_m < seqlen_q)
+        return
+
+    # ============================================================
+    #         NORMAL PROCESSING (Some K Blocks Visible)
+    # ============================================================
+    """
+    This program has visible K blocks to process.
+    We'll use two calls to handle different block types efficiently.
+    """
+
+    # Initialize for processing
+    # Compute pointers for all the tensors used in this kernel.
+    q_offset = (
+        Q + off_z * stride_qz + off_h_q * stride_qh + cu_seqlens_q_start * stride_qm
+    )
+    q_ptrs = q_offset + offs_m[:, None] * stride_qm + offs_d_qk[None, :] * stride_qk
+    k_offset = (
+        K + off_z * stride_kz + off_h_k * stride_kh + cu_seqlens_k_start * stride_kn
+    )
+    k_ptrs = k_offset + offs_d_qk[:, None] * stride_kk + offs_n[None, :] * stride_kn
+    v_offset = (
+        V + off_z * stride_vz + off_h_k * stride_vh + cu_seqlens_k_start * stride_vk
+    )
+    v_ptrs = v_offset + offs_n[:, None] * stride_vk + offs_d_v[None, :] * stride_vn
+    if USE_BIAS:
+        # Note: this might get large enough to overflow on some configs
+        bias_offset = off_h_q * stride_bh
+        bias_ptrs = (
+            bias
+            + bias_offset
+            + offs_m[:, None] * stride_bm
+            + offs_n[None, :] * stride_bn
+        )
+    else:
+        bias_ptrs = None
+
+    if USE_ALIBI:
+        a_offset = off_z * stride_az + off_h_q * stride_ah
+        alibi_slope = tl.load(ALIBI_SLOPES + a_offset)
+    else:
+        alibi_slope = None
+
+    # initialize pointer to m and l
+    m_i = tl.full([BLOCK_M], float("-inf"), dtype=ACCUMULATOR_TYPE)
+    l_i = tl.full([BLOCK_M], 1.0, dtype=ACCUMULATOR_TYPE)
+    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL_V], dtype=ACCUMULATOR_TYPE)
+
+    # Q is loaded once at the beginning and shared by all N blocks.
+    q_ptrs_mask = offs_m[:, None] < seqlen_q
+    if PADDED_HEAD_QK:
+        q_ptrs_mask = q_ptrs_mask & (offs_d_qk[None, :] < ACTUAL_BLOCK_DMODEL_QK)
+    q = tl.load(q_ptrs, mask=q_ptrs_mask, other=0.0)
+
+    # ========== Process MASKED K Blocks in the front ==========
+    # NOTE: we use USE_SLIDING_WINDOW as guard because the compiler will crash other wise. front masking is only for sliding window so that is fine.
+    if n_front_masked_blocks > 0 and USE_SLIDING_WINDOW:
+        block_min = n_front_skip_blocks * BLOCK_N
+        block_max = (n_front_skip_blocks + n_front_masked_blocks) * BLOCK_N
+
+        acc, l_i, m_i = _attn_fwd_inner(
+            acc,
+            l_i,
+            m_i,
+            q,
+            k_ptrs,
+            v_ptrs,
+            bias_ptrs,
+            stride_kn,
+            stride_vk,
+            stride_bn,
+            stride_sn,
+            stride_sm,
+            start_m,
+            seqlen_k,
+            seqlen_q,
+            dropout_p,
+            philox_seed,
+            philox_offset_base,
+            SD_MASK,
+            stride_sz,
+            stride_sh,
+            off_z,
+            off_h_q,
+            offs_m,
+            offs_n,
+            offs_d_qk,
+            offs_d_v,
+            block_min,  # Start of front masked blocks
+            block_max,  # End of front masked blocks
+            0,  # n_extra_tokens (0 for front blocks, only relevant for last block)
+            alibi_slope,
+            q_descale,
+            k_descale,
+            v_descale,
+            IS_FP8,
+            FP8_MAX,
+            FP8_P_DESCALE,
+            APPLY_MASK=True,  # Masked blocks
+            IS_CAUSAL=IS_CAUSAL,
+            BLOCK_M=BLOCK_M,
+            BLOCK_DMODEL_QK=BLOCK_DMODEL_QK,
+            BLOCK_DMODEL_V=BLOCK_DMODEL_V,
+            BLOCK_N=BLOCK_N,
+            PRE_LOAD_V=PRE_LOAD_V,
+            ENABLE_DROPOUT=ENABLE_DROPOUT,
+            PADDED_HEAD_QK=PADDED_HEAD_QK,
+            PADDED_HEAD_V=PADDED_HEAD_V,
+            ACTUAL_BLOCK_DMODEL_QK=ACTUAL_BLOCK_DMODEL_QK,
+            ACTUAL_BLOCK_DMODEL_V=ACTUAL_BLOCK_DMODEL_V,
+            SM_SCALE=SM_SCALE,
+            USE_ALIBI=USE_ALIBI,
+            USE_EXP2=USE_EXP2,
+            RETURN_SCORES=RETURN_SCORES,
+            USE_SLIDING_WINDOW=USE_SLIDING_WINDOW,
+            WINDOW_SIZE_LEFT=WINDOW_SIZE_LEFT,
+            WINDOW_SIZE_RIGHT=WINDOW_SIZE_RIGHT,
+            ACCUMULATOR_TYPE=ACCUMULATOR_TYPE,
+        )
+
+    # ========== Process FULL K Blocks (Fast Path) ==========
+    if n_full_blocks > 0:
+        block_min = (n_front_skip_blocks + n_front_masked_blocks) * BLOCK_N
+        block_max = (
+            n_front_skip_blocks + n_front_masked_blocks + n_full_blocks
+        ) * BLOCK_N
+
+        acc, l_i, m_i = _attn_fwd_inner(
+            acc,
+            l_i,
+            m_i,
+            q,
+            k_ptrs,
+            v_ptrs,
+            bias_ptrs,
+            stride_kn,
+            stride_vk,
+            stride_bn,
+            stride_sn,
+            stride_sm,
+            start_m,
+            seqlen_k,
+            seqlen_q,
+            dropout_p,
+            philox_seed,
+            philox_offset_base,
+            SD_MASK,
+            stride_sz,
+            stride_sh,
+            off_z,
+            off_h_q,
+            offs_m,
+            offs_n,
+            offs_d_qk,
+            offs_d_v,
+            block_min,  # Start of range: 0
+            block_max,  # End of range: n_full_blocks * BLOCK_N
+            0,  # n_extra_tokens (not used for full blocks)
+            alibi_slope,
+            q_descale,
+            k_descale,
+            v_descale,
+            IS_FP8,
+            FP8_MAX,
+            FP8_P_DESCALE,
+            APPLY_MASK=FORCE_MASKING,
+            IS_CAUSAL=IS_CAUSAL,
+            BLOCK_M=BLOCK_M,
+            BLOCK_DMODEL_QK=BLOCK_DMODEL_QK,
+            BLOCK_DMODEL_V=BLOCK_DMODEL_V,
+            BLOCK_N=BLOCK_N,
+            PRE_LOAD_V=PRE_LOAD_V,
+            ENABLE_DROPOUT=ENABLE_DROPOUT,
+            PADDED_HEAD_QK=PADDED_HEAD_QK,
+            PADDED_HEAD_V=PADDED_HEAD_V,
+            ACTUAL_BLOCK_DMODEL_QK=ACTUAL_BLOCK_DMODEL_QK,
+            ACTUAL_BLOCK_DMODEL_V=ACTUAL_BLOCK_DMODEL_V,
+            SM_SCALE=SM_SCALE,
+            USE_ALIBI=USE_ALIBI,
+            USE_EXP2=USE_EXP2,
+            RETURN_SCORES=RETURN_SCORES,
+            USE_SLIDING_WINDOW=USE_SLIDING_WINDOW,
+            WINDOW_SIZE_LEFT=WINDOW_SIZE_LEFT,
+            WINDOW_SIZE_RIGHT=WINDOW_SIZE_RIGHT,
+            ACCUMULATOR_TYPE=ACCUMULATOR_TYPE,
+        )
+
+    # ========== Process MASKED K Blocks in the back ==========
+    if n_back_masked_blocks > 0:
+        block_min = (
+            n_front_skip_blocks + n_front_masked_blocks + n_full_blocks
+        ) * BLOCK_N
+        block_max = (
+            n_front_skip_blocks
+            + n_front_masked_blocks
+            + n_full_blocks
+            + n_back_masked_blocks
+        ) * BLOCK_N
+
+        acc, l_i, m_i = _attn_fwd_inner(
+            acc,
+            l_i,
+            m_i,
+            q,
+            k_ptrs,
+            v_ptrs,
+            bias_ptrs,
+            stride_kn,
+            stride_vk,
+            stride_bn,
+            stride_sn,
+            stride_sm,
+            start_m,
+            seqlen_k,
+            seqlen_q,
+            dropout_p,
+            philox_seed,
+            philox_offset_base,
+            SD_MASK,
+            stride_sz,
+            stride_sh,
+            off_z,
+            off_h_q,
+            offs_m,
+            offs_n,
+            offs_d_qk,
+            offs_d_v,
+            block_min,  # Start of range: n_full_blocks * BLOCK_N
+            block_max,  # End of range: n_visible_k_blocks * BLOCK_N
+            n_extra_tokens,  # Padding tokens in last block
+            alibi_slope,
+            q_descale,
+            k_descale,
+            v_descale,
+            IS_FP8,
+            FP8_MAX,
+            FP8_P_DESCALE,
+            APPLY_MASK=True,  # Masked blocks
+            IS_CAUSAL=IS_CAUSAL,  # Use actual causal flag
+            BLOCK_M=BLOCK_M,
+            BLOCK_DMODEL_QK=BLOCK_DMODEL_QK,
+            BLOCK_DMODEL_V=BLOCK_DMODEL_V,
+            BLOCK_N=BLOCK_N,
+            PRE_LOAD_V=PRE_LOAD_V,
+            ENABLE_DROPOUT=ENABLE_DROPOUT,
+            PADDED_HEAD_QK=PADDED_HEAD_QK,
+            PADDED_HEAD_V=PADDED_HEAD_V,
+            ACTUAL_BLOCK_DMODEL_QK=ACTUAL_BLOCK_DMODEL_QK,
+            ACTUAL_BLOCK_DMODEL_V=ACTUAL_BLOCK_DMODEL_V,
+            SM_SCALE=SM_SCALE,
+            USE_ALIBI=USE_ALIBI,
+            USE_EXP2=USE_EXP2,
+            RETURN_SCORES=RETURN_SCORES,
+            USE_SLIDING_WINDOW=USE_SLIDING_WINDOW,
+            WINDOW_SIZE_LEFT=WINDOW_SIZE_LEFT,
+            WINDOW_SIZE_RIGHT=WINDOW_SIZE_RIGHT,
+            ACCUMULATOR_TYPE=ACCUMULATOR_TYPE,
+        )
+
+    # ============================================================
+    #                        EPILOGUE
+    # ============================================================
+    # Handle invalid rows: rows with no valid keys to attend to.
+    # This occurs with sliding window or causal attention (when seqlen_q > seqlen_k).
+    # For invalid rows: m_i = -inf, l_i = 0, acc = 0.
+    # We set l_i = 1.0 to avoid division by zero and ensure LSE = -inf.
+    invalid_mask = m_i == float("-inf")
+    l_i_safe = tl.where(invalid_mask, 1.0, l_i)
+    l_recip = 1 / l_i_safe[:, None]
+    acc = acc * l_recip
+
+    if ENABLE_DROPOUT:
+        dropout_scale = 1 / (1 - dropout_p)
+        acc = acc * dropout_scale
+
+    # compute log-sum-exp
+    if USE_EXP2:
+        RCP_LN2: tl.constexpr = 1.4426950408889634
+        LN2: tl.constexpr = 0.6931471824645996
+        softmax_lse = (m_i * RCP_LN2 + tl.math.log2(l_i)) * LN2
+    else:
+        softmax_lse = m_i + tl.math.log(l_i)
+
+    # Ensure invalid rows have LSE = -inf
+    softmax_lse = tl.where(invalid_mask, float("-inf"), softmax_lse)
+
+    # write back LSE(Log Sum Exponents), the log of the normalization constant
+    l_offset = (
+        LSE
+        + off_z * stride_lse_z
+        + off_h_q * stride_lse_h
+        + cu_seqlens_q_start * stride_lse_m
+    )
+    l_ptrs = l_offset + offs_m * stride_lse_m
+
+    # If seqlen_q not multiple of BLOCK_M, we need to mask out the last few rows.
+    # This is only true for the last Q block. For others, overflow_size will be -ve
+    end_m_idx = (start_m + 1) * BLOCK_M
+    overflow_size = end_m_idx - seqlen_q
+    if overflow_size > 0:
+        boundary = tl.full((BLOCK_M,), BLOCK_M - overflow_size, dtype=tl.int32)
+        l_ptrs_mask = tl.arange(0, BLOCK_M) < boundary
+        tl.store(l_ptrs, softmax_lse, mask=l_ptrs_mask)
+    else:
+        tl.store(l_ptrs, softmax_lse)
+
+    # write back O
+    o_offset = (
+        Out + off_z * stride_oz + off_h_q * stride_oh + cu_seqlens_q_start * stride_om
+    )
+    o_ptrs = o_offset + offs_m[:, None] * stride_om + offs_d_v[None, :] * stride_on
+    o_ptrs_mask = tl.full([BLOCK_M, BLOCK_DMODEL_V], 1, dtype=tl.int1)
+    if overflow_size > 0:
+        o_ptrs_mask = o_ptrs_mask & (offs_m[:, None] < seqlen_q)
+    if PADDED_HEAD_V:
+        o_ptrs_mask = o_ptrs_mask & (offs_d_v[None, :] < ACTUAL_BLOCK_DMODEL_V)
+
+    tl.store(o_ptrs, acc.to(Out.dtype.element_ty), mask=o_ptrs_mask)
+
+
+def attention_forward_prefill_triton_impl(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    o: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    sd_mask: Optional[torch.Tensor],
+    sm_scale: float,
+    alibi_slopes: Optional[torch.Tensor],
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    bias: Optional[torch.Tensor],
+    layout: Literal["bshd", "bhsd", "thd"],
+    # varlen
+    cu_seqlens_q: Optional[torch.Tensor],
+    cu_seqlens_k: Optional[torch.Tensor],
+    max_seqlens_q: int,
+    max_seqlens_k: int,
+    # dropout
+    dropout_p: float,
+    philox_seed: Optional[int],
+    philox_offset: Optional[int],
+    # misc
+    return_scores: bool,
+    use_exp2: bool,
+    # fp8
+    q_descale: Optional[torch.Tensor],
+    k_descale: Optional[torch.Tensor],
+    v_descale: Optional[torch.Tensor],
+    # seqused for FA v3
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    # rotary (optional)
+    rotary_cos: Optional[torch.Tensor] = None,
+    rotary_sin: Optional[torch.Tensor] = None,
+    rotary_interleaved: bool = False,
+    seqlens_rotary: Optional[torch.Tensor] = None,
+):
+    # get params, strides and shape
+    IS_VARLEN = layout == "thd"
+
+    # common assertions
+    assert (
+        0.0 <= dropout_p <= 1.0
+    ), f"dropout_p must be between 0 and 1, got {dropout_p}"
+    assert (
+        q.device == k.device == v.device == o.device
+    ), f"All tensors must be on the same device. Got: q={q.device}, k={k.device}, v={v.device}, o={o.device}"
+    assert q.dtype == k.dtype == v.dtype, "q, k, v must have the same dtype"
+    current_device = torch.cuda.current_device()
+    assert (
+        q.is_cuda and q.device.index == current_device
+    ), f"Device mismatch: Kernel will launch on cuda:{current_device}, but tensors are on {q.device}"
+
+    # get shapes and strides
+    if IS_VARLEN:
+        # shape
+        total_seqlen_q, nheads_q, head_size_q = q.shape
+        total_seqlen_k, nheads_k, head_size_k = k.shape
+        total_seqlen_v, nheads_v, head_size_v = v.shape
+
+        # assert shapes
+        assert (
+            cu_seqlens_q is not None
+        ), "cu_seqlens_q must be provided for varlen layout"
+        assert (
+            cu_seqlens_k is not None
+        ), "cu_seqlens_k must be provided for varlen layout"
+        assert (
+            max_seqlens_q is not None and max_seqlens_q > 0
+        ), "max_seqlens_q must be provided and positive for varlen layout"
+        assert (
+            max_seqlens_k is not None and max_seqlens_k > 0
+        ), "max_seqlens_k must be provided and positive for varlen layout"
+
+        # assert head dimensions
+        assert (
+            head_size_q == head_size_k
+        ), f"head sizes must match: q={head_size_q}, k={head_size_k}"
+        assert (
+            nheads_k == nheads_v
+        ), f"k and v must have same number of heads: k={nheads_k}, v={nheads_v}"
+        assert (
+            nheads_q % nheads_k == 0
+        ), f"nheads_q {nheads_q} must be divisible by nheads_k {nheads_k} for GQA/MQA"
+
+        # assert output shapes
+        assert o.shape == (
+            total_seqlen_q,
+            nheads_q,
+            head_size_v,
+        ), f"o shape {o.shape} != expected {(total_seqlen_q, nheads_q, head_size_v)}"
+
+        # assert cu_seqlens
+        assert (
+            cu_seqlens_q.dtype == torch.int32
+        ), f"cu_seqlens_q must be int32, got {cu_seqlens_q.dtype}"
+        assert (
+            cu_seqlens_k.dtype == torch.int32
+        ), f"cu_seqlens_k must be int32, got {cu_seqlens_k.dtype}"
+        assert cu_seqlens_q[0] == 0, "cu_seqlens_q must start with 0"
+        assert cu_seqlens_k[0] == 0, "cu_seqlens_k must start with 0"
+        assert (
+            cu_seqlens_q[-1] == total_seqlen_q
+        ), f"cu_seqlens_q[-1] {cu_seqlens_q[-1]} != total_seqlen_q {total_seqlen_q}"
+        assert (
+            cu_seqlens_k[-1] == total_seqlen_k
+        ), f"cu_seqlens_k[-1] {cu_seqlens_k[-1]} != total_seqlen_k {total_seqlen_k}"
+
+        # set vars
+        batch = len(cu_seqlens_q) - 1
+        head_size_qk = head_size_q
+
+        # Assert softmax_lse tensor is large enough
+        assert (
+            softmax_lse.shape[0] >= nheads_q
+        ), f"softmax_lse.shape[0]={softmax_lse.shape[0]} must be >= nheads_q={nheads_q}"
+        assert (
+            softmax_lse.shape[1] >= total_seqlen_q
+        ), f"softmax_lse.shape[1]={softmax_lse.shape[1]} must be >= total_seqlen_q={total_seqlen_q}"
+        assert (
+            softmax_lse.dtype == torch.float32
+        ), f"softmax_lse must be float32, got {softmax_lse.dtype}"
+        assert softmax_lse.device == q.device, "softmax_lse must be on same device as q"
+
+        # strides
+        stride_qb, stride_qh, stride_qm, stride_qd = (
+            0,
+            q.stride(1),
+            q.stride(0),
+            q.stride(2),
+        )
+        stride_kb, stride_kh, stride_kn, stride_kd = (
+            0,
+            k.stride(1),
+            k.stride(0),
+            k.stride(2),
+        )
+        stride_vb, stride_vh, stride_vn, stride_vd = (
+            0,
+            v.stride(1),
+            v.stride(0),
+            v.stride(2),
+        )
+        stride_ob, stride_oh, stride_om, stride_od = (
+            0,
+            o.stride(1),
+            o.stride(0),
+            o.stride(2),
+        )
+        stride_lse_z, stride_lse_h, stride_lse_m = (
+            0,
+            softmax_lse.stride(0),
+            softmax_lse.stride(1),
+        )
+    else:
+        # shapes
+        batch_q, seqlen_q, nheads_q, head_size_q = q.shape
+        batch_k, seqlen_k, nheads_k, head_size_k = k.shape
+        batch_v, seqlen_v, nheads_v, head_size_v = v.shape
+
+        # assert batch dimensions
+        assert (
+            batch_q == batch_k == batch_v
+        ), f"batch sizes must match: q={batch_q}, k={batch_k}, v={batch_v}"
+
+        # assert head dimensions
+        assert (
+            head_size_q == head_size_k
+        ), f"head sizes must match: q={head_size_q}, k={head_size_k}"
+        assert (
+            nheads_k == nheads_v
+        ), f"k and v must have same number of heads: k={nheads_k}, v={nheads_v}"
+        assert (
+            nheads_q % nheads_k == 0
+        ), f"nheads_q {nheads_q} must be divisible by nheads_k {nheads_k} for GQA/MQA"
+
+        # assert sequence lengths
+        assert (
+            seqlen_k == seqlen_v
+        ), f"k and v sequence lengths must match: k={seqlen_k}, v={seqlen_v}"
+
+        # assert output shapes
+        assert o.shape == (
+            batch_q,
+            seqlen_q,
+            nheads_q,
+            head_size_v,
+        ), f"o shape {o.shape} != expected {(batch_q, seqlen_q, nheads_q, head_size_v)}"
+
+        # set vars
+        batch = batch_q
+        head_size_qk = head_size_q
+        max_seqlens_q = seqlen_q
+        max_seqlens_k = seqlen_k
+
+        # Assert softmax_lse tensor is large enough
+        assert (
+            softmax_lse.shape[0] >= batch
+        ), f"softmax_lse.shape[0]={softmax_lse.shape[0]} must be >= batch={batch}"
+        assert (
+            softmax_lse.shape[1] >= nheads_q
+        ), f"softmax_lse.shape[1]={softmax_lse.shape[1]} must be >= nheads_q={nheads_q}"
+        assert (
+            softmax_lse.shape[2] >= seqlen_q
+        ), f"softmax_lse.shape[2]={softmax_lse.shape[2]} must be >= seqlen_q={seqlen_q}"
+        assert (
+            softmax_lse.dtype == torch.float32
+        ), f"softmax_lse must be float32, got {softmax_lse.dtype}"
+        assert softmax_lse.device == q.device, "softmax_lse must be on same device as q"
+
+        # strides
+        stride_qb, stride_qh, stride_qm, stride_qd = (
+            q.stride(0),
+            q.stride(2),
+            q.stride(1),
+            q.stride(3),
+        )
+        stride_kb, stride_kh, stride_kn, stride_kd = (
+            k.stride(0),
+            k.stride(2),
+            k.stride(1),
+            k.stride(3),
+        )
+        stride_vb, stride_vh, stride_vn, stride_vd = (
+            v.stride(0),
+            v.stride(2),
+            v.stride(1),
+            v.stride(3),
+        )
+        stride_ob, stride_oh, stride_om, stride_od = (
+            o.stride(0),
+            o.stride(2),
+            o.stride(1),
+            o.stride(3),
+        )
+        stride_lse_z, stride_lse_h, stride_lse_m = softmax_lse.stride()
+
+    # apply rotary embeddings
+    if rotary_cos is not None and rotary_sin is not None:
+        if IS_VARLEN:
+            raise NotImplementedError(
+                "Rotary embeddings with varlen (thd layout) prefill are not implemented yet."
+            )
+        seqlen_offsets = seqlens_rotary if seqlens_rotary is not None else 0
+        local = (window_size_left != -1) or (window_size_right != -1)
+        q, _ = apply_rotary(
+            q,
+            None,
+            rotary_cos,
+            rotary_sin,
+            causal=causal,
+            local=local,
+            interleaved=rotary_interleaved,
+            seqlen_offsets=seqlen_offsets,
+        )
+
+    # fp8 setup and assertions
+    IS_FP8 = is_fp8([q, k, v])
+    if IS_FP8:
+        arch = get_arch()
+        if not arch.supports_fp8:
+            raise RuntimeError(f"{arch.name} does not support FP8")
+        FP8_MAX = torch.finfo(q.dtype).max
+        rec_dtype = arch.recommended_fp8_dtype(q.dtype)
+        if q.dtype != rec_dtype or k.dtype != rec_dtype or v.dtype != rec_dtype:
+            warnings.warn(
+                f"Use {rec_dtype} data type on {arch}. Got q: {q.dtype}, k: {k.dtype}, v: {v.dtype}",
+                UserWarning,
+            )
+
+        if (q_descale is None) or (k_descale is None) or (v_descale is None):
+            warnings.warn(
+                "FP8 tensors detected but descale factors not provided. Using default scale of 1.0",
+                UserWarning,
+            )
+            # Create default descale tensors if not provided
+            if q_descale is None:
+                q_descale = torch.ones(
+                    batch, nheads_q, dtype=torch.float32, device=q.device
+                )
+            if k_descale is None:
+                k_descale = torch.ones(
+                    batch, nheads_k, dtype=torch.float32, device=q.device
+                )
+            if v_descale is None:
+                v_descale = torch.ones(
+                    batch, nheads_k, dtype=torch.float32, device=q.device
+                )
+        else:
+            # Enforce exact expected shapes; no reshaping or normalization.
+            assert (
+                q_descale.dim() == 2
+                and q_descale.shape[0] == batch
+                and q_descale.shape[1] == nheads_k
+            ), f"q_descale expected shape ({batch}, {nheads_k}) got {tuple(q_descale.shape)}"
+            assert (
+                k_descale.dim() == 2
+                and k_descale.shape[0] == batch
+                and k_descale.shape[1] == nheads_k
+            ), f"k_descale expected shape ({batch}, {nheads_k}) got {tuple(k_descale.shape)}"
+            assert (
+                v_descale.dim() == 2
+                and v_descale.shape[0] == batch
+                and v_descale.shape[1] == nheads_k
+            ), f"v_descale expected shape ({batch}, {nheads_k}) got {tuple(v_descale.shape)}"
+
+        # o should be fp32 or fp16/bf16
+        assert o.dtype in [
+            torch.float16,
+            torch.bfloat16,
+            torch.float32,
+        ], f"Output tensor o must be fp16, bf16, or fp32 when using fp8, got {o.dtype}"
+
+        stride_q_descale_z = q_descale.stride(0) if q_descale is not None else 0
+        stride_k_descale_z = k_descale.stride(0) if k_descale is not None else 0
+        stride_v_descale_z = v_descale.stride(0) if v_descale is not None else 0
+
+        if DEBUG:
+            print("FP8 path triggered in fwd_prefill.py")
+    else:
+        FP8_MAX = None
+        q_descale = k_descale = v_descale = None
+        stride_q_descale_z = stride_k_descale_z = stride_v_descale_z = None
+
+        # check output dtype matches input dtype when not using fp8
+        assert (
+            o.dtype == q.dtype
+        ), f"Output dtype {o.dtype} must match input dtype {q.dtype} when not using fp8"
+
+    # check features
+    use_sliding_window = window_size_left != -1 or window_size_right != -1
+    use_alibi, (stride_az, stride_ah) = (
+        (True, alibi_slopes.stride()) if alibi_slopes is not None else (False, (0, 0))
+    )
+    # NOTE: a large bias tensor leads to overflow during pointer arithmetic
+    if bias is not None:
+        assert bias.numel() < 2**31
+
+    # Get closest power of 2 over or equal to 32 for both QK and V dimensions
+    padded_d_model_qk = 1 << (head_size_qk - 1).bit_length()
+    padded_d_model_v = 1 << (head_size_v - 1).bit_length()
+    # Smallest head_dim supported is 16. If smaller, the tile in the
+    # kernel is padded - there is no padding in memory for any dims.
+    padded_d_model_qk = max(padded_d_model_qk, 16)
+    padded_d_model_v = max(padded_d_model_v, 16)
+
+    # sd_mask assertions and strides
+    if sd_mask is not None:
+        assert dropout_p > 0.0 or return_scores, "sd_mask provided but not used"
+        assert (
+            sd_mask is not None
+        ), "sd_mask must be provided when return_scores=True or dropout_p > 0"
+        # Assert sd_mask tensor is large enough
+        assert (
+            sd_mask.shape[0] >= batch
+        ), f"sd_mask.shape[0]={sd_mask.shape[0]} must be >= batch={batch}"
+        assert (
+            sd_mask.shape[1] >= nheads_q
+        ), f"sd_mask.shape[1]={sd_mask.shape[1]} must be >= nheads_q={nheads_q}"
+        assert (
+            sd_mask.shape[2] >= max_seqlens_q
+        ), f"sd_mask.shape[2]={sd_mask.shape[2]} must be >= max_seqlens_q={max_seqlens_q}"
+        assert (
+            sd_mask.shape[3] >= max_seqlens_k
+        ), f"sd_mask.shape[3]={sd_mask.shape[3]} must be >= max_seqlens_k={max_seqlens_k}"
+        assert sd_mask.device == q.device, "sd_mask must be on same device as q"
+
+        stride_sz, stride_sh, stride_sm, stride_sn = (
+            sd_mask.stride(0),
+            sd_mask.stride(1),
+            sd_mask.stride(2),
+            sd_mask.stride(3),
+        )
+    else:
+        stride_sz, stride_sh, stride_sm, stride_sn = (0, 0, 0, 0)
+
+    if bias is not None:
+        stride_bz, stride_bh, stride_bm, stride_bn = (
+            bias.stride(0),
+            bias.stride(1),
+            bias.stride(2),
+            bias.stride(3),
+        )
+    else:
+        stride_bz, stride_bh, stride_bm, stride_bn = (0, 0, 0, 0)
+
+    # Detect if we need to force masking for all blocks (required on some architectures)
+    arch = get_arch()
+    force_masking = arch.is_rdna
+
+    num_xcd = 1 if arch.is_rdna else 8
+
+    # launch kernel
+    def grid(META):
+        return (nheads_q, triton.cdiv(max_seqlens_q, META["BLOCK_M"]), batch)
+
+    attn_fwd[grid](
+        q,
+        k,
+        v,
+        bias,
+        q_descale,
+        k_descale,
+        v_descale,
+        stride_q_descale_z,
+        stride_k_descale_z,
+        stride_v_descale_z,
+        softmax_lse,
+        o,
+        sd_mask,
+        alibi_slopes,
+        stride_qb,
+        stride_qh,
+        stride_qm,
+        stride_qd,
+        stride_kb,
+        stride_kh,
+        stride_kn,
+        stride_kd,
+        stride_vb,
+        stride_vh,
+        stride_vn,
+        stride_vd,
+        stride_ob,
+        stride_oh,
+        stride_om,
+        stride_od,
+        stride_bz,
+        stride_bh,
+        stride_bm,
+        stride_bn,
+        stride_az,
+        stride_ah,
+        stride_sz,
+        stride_sh,
+        stride_sm,
+        stride_sn,
+        stride_lse_z,
+        stride_lse_h,
+        stride_lse_m,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        seqused_q,
+        seqused_k,  # Pass seqused tensors
+        dropout_p=dropout_p,
+        philox_seed=philox_seed,
+        philox_offset_base=philox_offset,
+        HQ=nheads_q,
+        HK=nheads_k,
+        ACTUAL_BLOCK_DMODEL_QK=head_size_qk,
+        ACTUAL_BLOCK_DMODEL_V=head_size_v,
+        MAX_SEQLENS_Q=max_seqlens_q,
+        MAX_SEQLENS_K=max_seqlens_k,
+        SM_SCALE=sm_scale,
+        IS_CAUSAL=causal,
+        USE_SLIDING_WINDOW=use_sliding_window,
+        WINDOW_SIZE_LEFT=window_size_left,
+        WINDOW_SIZE_RIGHT=window_size_right,
+        IS_VARLEN=IS_VARLEN,
+        BLOCK_DMODEL_QK=padded_d_model_qk,
+        BLOCK_DMODEL_V=padded_d_model_v,
+        USE_BIAS=False if bias is None else True,
+        USE_ALIBI=use_alibi,
+        ENABLE_DROPOUT=dropout_p > 0.0,
+        USE_EXP2=use_exp2,
+        RETURN_SCORES=return_scores,
+        IS_FP8=IS_FP8,
+        FP8_MAX=FP8_MAX,
+        FP8_P_DESCALE=False,
+        USE_SEQUSED=(seqused_q is not None or seqused_k is not None),
+        FORCE_MASKING=force_masking,
+        NUM_XCD=num_xcd,
+    )

build/torch-rocm/_kernels/flash_attn_triton_amd/interface_v2.py

@@ -0,0 +1,828 @@
+import torch
+from typing import Literal, Optional, Union
+from .fwd_prefill import attention_forward_prefill_triton_impl
+from .fwd_decode import attention_forward_decode_triton_impl
+from .bwd import attention_backward_triton_impl
+from .utils import (
+    DEBUG,
+    USE_EXP2,
+    BWD_MODE,
+    PHILOX_SEED,
+    PHILOX_OFFSET,
+    SHAPE_EXPECTATIONS,
+    round_multiple,
+)
+
+
+def fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: Optional[torch.Tensor],
+    alibi_slopes: Optional[torch.Tensor],
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    return_softmax: bool,
+    gen_: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], torch.Tensor]:
+
+    # Reject FP8 tensors (FA2 AMD path does not support FP8)
+    if str(q.dtype).startswith("torch.float8"):
+        raise NotImplementedError(
+            "FP8 tensors are not supported in the AMD Triton FA2 interface. Use the FA3 path instead."
+        )
+
+    # Unsupported features assertions (keep behavior explicit like v3 shim)
+    if softcap != 0.0:
+        raise NotImplementedError(
+            "softcap is not supported in the AMD Triton FA2 interface (expected 0.0)."
+        )
+
+    if DEBUG:
+        print()
+        print("flash_attn_triton_amd.py::fwd inputs")
+        print("q:", q.shape)
+        print("k:", k.shape)
+        print("v:", v.shape)
+        print("out:", out.shape if out is not None else None)
+        print("alibi_slopes:", alibi_slopes.shape if alibi_slopes is not None else None)
+        print("dropout_p:", dropout_p)
+        print("softmax_scale:", softmax_scale)
+        print("causal:", causal)
+        print("window_size_left:", window_size_left)
+        print("window_size_right:", window_size_right)
+        print("softcap:", softcap)
+        print("return_softmax:", return_softmax)
+
+    if out is None:
+        out = torch.zeros_like(q)
+    else:
+        out.zero_()
+
+    # Layout / shapes
+    layout: Literal["bshd", "bhsd", "thd"] = "bshd"
+    max_seqlen_q = q.shape[1]
+    max_seqlen_k = k.shape[1]
+    batch, _, nheads_q, _ = q.shape
+
+    # Normalize / validate alibi
+    if alibi_slopes is not None:
+        if alibi_slopes.dim() == 1:
+            alibi_slopes = alibi_slopes.unsqueeze(0).expand(batch, -1)
+        assert alibi_slopes.is_cuda and alibi_slopes.dim() == 2
+        assert alibi_slopes.shape == (batch, nheads_q)
+
+    # Dropout + RNG seed
+    philox_seed, philox_offset = PHILOX_SEED, PHILOX_OFFSET
+    rng_state = torch.as_tensor([philox_seed, philox_offset])
+
+    # argument checks
+    assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4
+    assert q.shape[-1] == k.shape[-1] == v.shape[-1]
+    assert q.dtype == k.dtype == v.dtype
+    assert out.shape[:-1] == q.shape[:-1] and out.shape[-1] == v.shape[-1]
+    nheads_k = k.shape[2]
+    assert (nheads_q % nheads_k) == 0
+
+    # Create output tensors based on shape expectations
+    if SHAPE_EXPECTATIONS == "rounded":
+        softmax_lse = torch.zeros(
+            (batch, nheads_q, round_multiple(max_seqlen_q, 128)),
+            device=q.device,
+            dtype=torch.float32,
+        )
+        if dropout_p > 0.0 or return_softmax:
+            sd_mask = torch.zeros(
+                (
+                    batch,
+                    nheads_q,
+                    round_multiple(max_seqlen_q, 128),
+                    round_multiple(max_seqlen_k, 128),
+                ),
+                device=q.device,
+                dtype=torch.float32,
+            )
+        else:
+            sd_mask = None
+    else:
+        softmax_lse = torch.zeros(
+            (batch, nheads_q, max_seqlen_q),
+            device=q.device,
+            dtype=torch.float32,
+        )
+        if dropout_p > 0.0 or return_softmax:
+            sd_mask = torch.zeros(
+                (batch, nheads_q, max_seqlen_q, max_seqlen_k),
+                device=q.device,
+                dtype=torch.float32,
+            )
+        else:
+            sd_mask = None
+
+    # call implementation
+    if DEBUG:
+        print("Using Triton implementation")
+    attention_forward_prefill_triton_impl(
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        sd_mask,
+        softmax_scale,
+        alibi_slopes,
+        causal,
+        window_size_left,
+        window_size_right,
+        None,
+        layout,
+        None,
+        None,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        philox_seed,
+        philox_offset,
+        return_softmax,
+        USE_EXP2,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+    )
+
+    if DEBUG:
+        print("flash_attn_triton_amd.py::fwd outputs")
+        print("out:", out.shape)
+        print("softmax_lse:", softmax_lse.shape)
+        print("sd_mask:", sd_mask.shape if sd_mask is not None else None)
+        print("rng_state:", rng_state)
+
+    # --- Assertions (shape + dtype contracts) ---
+    # out: (B, Sq, Hq, D)
+    assert out.shape == q.shape, f"[fwd] out shape {out.shape} != q shape {q.shape}"
+    # softmax_lse dtype
+    assert (
+        softmax_lse.dtype == torch.float32
+    ), f"[fwd] softmax_lse dtype {softmax_lse.dtype} != torch.float32"
+    # softmax_lse shape depends on SHAPE_EXPECTATIONS
+    if SHAPE_EXPECTATIONS == "rounded":
+        expected_lse_shape = (q.shape[0], q.shape[2], round_multiple(q.shape[1], 128))
+    else:
+        expected_lse_shape = (q.shape[0], q.shape[2], q.shape[1])
+    assert (
+        softmax_lse.shape == expected_lse_shape
+    ), f"[fwd] softmax_lse shape {softmax_lse.shape} != {expected_lse_shape}"
+    if return_softmax:
+        # sd_mask: (B, Hq, Sq, Sk)
+        assert sd_mask is not None, "[fwd] return_softmax=True but sd_mask is None"
+        assert sd_mask.dim() == 4, f"[fwd] sd_mask dim {sd_mask.dim()} != 4"
+        if SHAPE_EXPECTATIONS == "rounded":
+            expected_sq = round_multiple(q.shape[1], 128)
+            expected_sk = round_multiple(k.shape[1], 128)
+            assert (
+                sd_mask.shape[0] == q.shape[0]
+                and sd_mask.shape[1] == q.shape[2]
+                and sd_mask.shape[2] == expected_sq
+                and sd_mask.shape[3] == expected_sk
+            ), f"[fwd] sd_mask shape {sd_mask.shape} != (B={q.shape[0]}, Hq={q.shape[2]}, Sq={expected_sq}, Sk={expected_sk})"
+        else:
+            assert (
+                sd_mask.shape[0] == q.shape[0]
+                and sd_mask.shape[1] == q.shape[2]
+                and sd_mask.shape[2] == q.shape[1]
+            ), f"[fwd] sd_mask leading dims {sd_mask.shape[:3]} mismatch (B,Hq,Sq) {(q.shape[0], q.shape[2], q.shape[1])}"
+    else:
+        assert sd_mask is None, "[fwd] return_softmax=False but sd_mask is not None"
+
+    return out, softmax_lse, sd_mask, rng_state
+
+
+def bwd(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor],
+    dk: Optional[torch.Tensor],
+    dv: Optional[torch.Tensor],
+    alibi_slopes: Optional[torch.Tensor],
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    deterministic: bool,
+    gen_: Optional[torch.Tensor] = None,
+    rng_state: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    if softcap != 0.0:
+        raise NotImplementedError(
+            "softcap is not supported in the AMD Triton FA2 interface (expected 0.0)."
+        )
+
+    # Check for sliding window - backward doesn't support it yet
+    is_sliding_window = (window_size_left >= 0) or (window_size_right >= 0)
+    if is_sliding_window:
+        raise NotImplementedError(
+            f"Sliding window attention is not yet supported in the AMD Triton backward pass "
+            f"(window_size_left={window_size_left}, window_size_right={window_size_right}). "
+            f"Use window_size=(-1, -1) for full attention."
+        )
+
+    if DEBUG:
+        print()
+        print("flash_attn_triton_amd.py::bwd inputs")
+        print("dout:", dout.shape)
+        print("q:", q.shape)
+        print("k:", k.shape)
+        print("v:", v.shape)
+        print("out:", out.shape)
+        print("softmax_lse:", softmax_lse.shape)
+        print("dq:", dq.shape if dq is not None else None)
+        print("dk:", dk.shape if dk is not None else None)
+        print("dv:", dv.shape if dv is not None else None)
+        print("alibi_slopes:", alibi_slopes.shape if alibi_slopes is not None else None)
+        print("dropout_p:", dropout_p)
+        print("softmax_scale:", softmax_scale)
+        print("causal:", causal)
+        print("window_size_left:", window_size_left)
+        print("window_size_right:", window_size_right)
+        print("deterministic:", deterministic)
+        print("rng_state:", rng_state)
+
+    dq = torch.zeros_like(q) if dq is None else dq.zero_()
+    dk = torch.zeros_like(k) if dk is None else dk.zero_()
+    dv = torch.zeros_like(v) if dv is None else dv.zero_()
+
+    # get shape
+    batch, seqlen_q, nheads_q, _ = q.shape
+
+    # Create delta tensor with shape based on expectations
+    # delta (softmax_d) : (B, Hq, Sq) or (B, Hq, round_multiple(Sq, 128))
+    if SHAPE_EXPECTATIONS == "rounded":
+        delta = torch.zeros(
+            (batch, nheads_q, round_multiple(seqlen_q, 128)),
+            device=q.device,
+            dtype=torch.float32,
+        )
+    else:
+        delta = torch.zeros(
+            (batch, nheads_q, seqlen_q), device=q.device, dtype=torch.float32
+        )
+
+    # Upstream change: base seeding logic on provided rng_state instead of dropout probability.
+    if rng_state is not None:
+        philox_seed, philox_offset = rng_state[0].item(), rng_state[1].item()
+    else:
+        philox_seed, philox_offset = None, None
+
+    if alibi_slopes is not None:
+        if alibi_slopes.dim() == 2:
+            pass
+        elif alibi_slopes.dim() == 1:
+            alibi_slopes = alibi_slopes.unsqueeze(0).expand(batch, -1)
+        else:
+            raise ValueError("Alibi can be (nheads,) or (batch_size, nheads).")
+
+    # call implementation
+    if DEBUG:
+        print(f"Using Triton implementation in {BWD_MODE} mode")
+    attention_backward_triton_impl(
+        do=dout,
+        q=q,
+        k=k,
+        v=v,
+        o=out,
+        softmax_lse=softmax_lse,
+        dq=dq,
+        dk=dk,
+        dv=dv,
+        delta=delta,
+        sm_scale=softmax_scale,
+        alibi_slopes=alibi_slopes,
+        causal=causal,
+        layout="bshd",
+        cu_seqlens_q=None,
+        cu_seqlens_k=None,
+        max_seqlen_q=seqlen_q,
+        max_seqlen_k=k.shape[1],
+        seqused_q=None,
+        seqused_k=None,
+        dropout_p=dropout_p,
+        philox_seed=philox_seed,
+        philox_offset=philox_offset,
+        use_exp2=USE_EXP2,
+        mode=BWD_MODE,
+    )
+
+    if DEBUG:
+        print("flash_attn_triton_amd.py::bwd outputs")
+        print("dq:", dq.shape)
+        print("dk:", dk.shape)
+        print("dv:", dv.shape)
+    # --- Assertions ---
+    assert dq.shape == q.shape, f"[bwd] dq shape {dq.shape} != q shape {q.shape}"
+    assert dk.shape == k.shape, f"[bwd] dk shape {dk.shape} != k shape {k.shape}"
+    assert dv.shape == v.shape, f"[bwd] dv shape {dv.shape} != v shape {v.shape}"
+    # delta (softmax_d) : (B, Hq, Sq)
+    if SHAPE_EXPECTATIONS == "rounded":
+        expected_delta_shape = (q.shape[0], q.shape[2], round_multiple(q.shape[1], 128))
+    else:
+        expected_delta_shape = (q.shape[0], q.shape[2], q.shape[1])
+    assert (
+        delta.shape == expected_delta_shape
+    ), f"[bwd] delta shape {delta.shape} != {expected_delta_shape}"
+    return dq, dk, dv, delta
+
+
+def varlen_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: Optional[torch.Tensor],
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    seqused_k: Optional[torch.Tensor],
+    leftpad_k: Optional[torch.Tensor],
+    block_table_: Optional[torch.Tensor],
+    alibi_slopes: Optional[torch.Tensor],
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float,
+    softmax_scale: float,
+    zero_tensors: bool,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    return_softmax: bool,
+    gen_: Optional[torch.Tensor] = None,
+    num_splits: int = 0,
+) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], torch.Tensor]:
+
+    if str(q.dtype).startswith("torch.float8"):
+        raise NotImplementedError(
+            "FP8 tensors are not supported in the AMD Triton FA2 interface (varlen_fwd). Use the FA3 path instead."
+        )
+
+    if softcap != 0.0:
+        raise NotImplementedError(
+            "softcap is not supported in varlen_fwd (expected 0.0)."
+        )
+    if leftpad_k is not None:
+        raise NotImplementedError(
+            "leftpad_k is not supported in AMD Triton FA2 varlen_fwd."
+        )
+    if block_table_ is not None:
+        raise NotImplementedError(
+            "block_table / paged attention is not supported in AMD Triton FA2 varlen_fwd."
+        )
+    if seqused_k is not None:
+        raise NotImplementedError(
+            "seqused_k is not supported in AMD Triton FA2 varlen_fwd."
+        )
+    if num_splits not in (0, 1):
+        raise NotImplementedError(
+            "num_splits > 1 not supported in AMD Triton FA2 varlen_fwd."
+        )
+
+    if DEBUG:
+        print()
+        print("flash_attn_triton_amd.py::varlen_fwd")
+        print("q:", q.shape)
+        print("k:", k.shape)
+        print("v:", v.shape)
+        print("cu_seqlens_q:", cu_seqlens_q, cu_seqlens_q.shape)
+        print("cu_seqlens_k:", cu_seqlens_k, cu_seqlens_k.shape)
+        print("alibi_slopes:", alibi_slopes)
+        print("max_seqlen_q:", max_seqlen_q)
+        print("max_seqlen_k:", max_seqlen_k)
+        print("dropout_p:", dropout_p)
+        print("softmax_scale:", softmax_scale)
+        print("causal:", causal)
+        print("window_size_left:", window_size_left)
+        print("window_size_right:", window_size_right)
+        print("gen_:", gen_)
+    out = torch.zeros_like(q) if out is None else out.zero_()
+
+    # Layout and basic info for varlen
+    layout: Literal["bshd", "bhsd", "thd"] = "thd"
+    batch = len(cu_seqlens_q) - 1
+    total_q, nheads_q, _ = q.shape
+
+    # Create softmax_lse tensor - varlen always uses exact shape (Hq, Total_Q)
+    softmax_lse = torch.zeros((nheads_q, total_q), device=q.device, dtype=torch.float32)
+
+    # Create sd_mask tensor if needed
+    if return_softmax:
+        # sd_mask: (B, Hq, Sq, Sk) - shape based on expectations
+        if SHAPE_EXPECTATIONS == "rounded":
+            sd_mask = torch.zeros(
+                (
+                    batch,
+                    nheads_q,
+                    round_multiple(max_seqlen_q, 128),
+                    round_multiple(max_seqlen_k, 128),
+                ),
+                device=q.device,
+                dtype=q.dtype,
+            )
+        else:
+            sd_mask = torch.zeros(
+                (batch, nheads_q, max_seqlen_q, max_seqlen_k),
+                device=q.device,
+                dtype=q.dtype,
+            )
+    else:
+        sd_mask = None
+
+    if alibi_slopes is not None:
+        if alibi_slopes.dim() == 1:
+            alibi_slopes = alibi_slopes.unsqueeze(0).expand(batch, -1)
+        assert alibi_slopes.is_cuda and alibi_slopes.dim() == 2
+        assert alibi_slopes.shape == (batch, nheads_q)
+
+    philox_seed, philox_offset = PHILOX_SEED, PHILOX_OFFSET
+    rng_state = torch.as_tensor([philox_seed, philox_offset])
+
+    # Inline checks (subset appropriate for varlen)
+    assert q.dim() == 3 and k.dim() == 3 and v.dim() == 3
+    assert q.shape[-1] == k.shape[-1] == v.shape[-1]
+    assert q.dtype == k.dtype == v.dtype
+    assert out.shape == q.shape
+    nheads_k = k.shape[1]
+    assert (nheads_q % nheads_k) == 0
+
+    # call implementation
+    if DEBUG:
+        print("Using Triton implementation")
+    attention_forward_prefill_triton_impl(
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        sd_mask,
+        softmax_scale,
+        alibi_slopes,
+        causal,
+        window_size_left,
+        window_size_right,
+        None,
+        layout,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        philox_seed,
+        philox_offset,
+        return_softmax,
+        USE_EXP2,
+        None,
+        None,
+        None,
+    )
+
+    if DEBUG:
+        print("varlen_fwd outputs")
+        print("out:", out, out.shape)
+        print("softmax_lse:", softmax_lse, softmax_lse.shape)
+        print("sd_mask:", sd_mask, sd_mask.shape if sd_mask is not None else None)
+    # --- Assertions ---
+    # out: (Total_Q, Hq, D)
+    assert (
+        out.shape == q.shape
+    ), f"[varlen_fwd] out shape {out.shape} != q shape {q.shape}"
+    # softmax_lse: (Hq, Total_Q)
+    expected_lse_shape = (q.shape[1], q.shape[0])
+    assert (
+        softmax_lse.shape == expected_lse_shape
+    ), f"[varlen_fwd] softmax_lse shape {softmax_lse.shape} != {expected_lse_shape}"
+    assert (
+        softmax_lse.dtype == torch.float32
+    ), f"[varlen_fwd] softmax_lse dtype {softmax_lse.dtype} != torch.float32"
+    if return_softmax:
+        # sd_mask expected: (B, Hq, max_seqlen_q, max_seqlen_k)
+        assert (
+            sd_mask is not None
+        ), "[varlen_fwd] return_softmax=True but sd_mask is None"
+        assert sd_mask.dim() == 4, f"[varlen_fwd] sd_mask dim {sd_mask.dim()} != 4"
+        batch = len(cu_seqlens_q) - 1
+        assert (
+            sd_mask.shape[0] == batch
+        ), f"[varlen_fwd] sd_mask batch {sd_mask.shape[0]} != {batch}"
+        assert (
+            sd_mask.shape[1] == q.shape[1]
+        ), f"[varlen_fwd] sd_mask nheads {sd_mask.shape[1]} != {q.shape[1]}"
+        if SHAPE_EXPECTATIONS == "rounded":
+            expected_sq = round_multiple(max_seqlen_q, 128)
+            expected_sk = round_multiple(max_seqlen_k, 128)
+            assert (
+                sd_mask.shape[2] == expected_sq and sd_mask.shape[3] == expected_sk
+            ), f"[varlen_fwd] sd_mask shape {sd_mask.shape} != (B={batch}, Hq={q.shape[1]}, Sq={expected_sq}, Sk={expected_sk})"
+        else:
+            assert (
+                sd_mask.shape[2] == max_seqlen_q and sd_mask.shape[3] == max_seqlen_k
+            ), f"[varlen_fwd] sd_mask shape {sd_mask.shape} != (B={batch}, Hq={q.shape[1]}, Sq={max_seqlen_q}, Sk={max_seqlen_k})"
+    else:
+        assert (
+            sd_mask is None
+        ), "[varlen_fwd] return_softmax=False but sd_mask is not None"
+    return out, softmax_lse, sd_mask, rng_state
+
+
+def varlen_bwd(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor],
+    dk: Optional[torch.Tensor],
+    dv: Optional[torch.Tensor],
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    alibi_slopes: Optional[torch.Tensor],
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float,
+    softmax_scale: float,
+    zero_tensors: bool,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    deterministic: bool,
+    gen_: Optional[torch.Tensor] = None,
+    rng_state: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    if str(q.dtype).startswith("torch.float8"):
+        raise NotImplementedError(
+            "FP8 tensors are not supported in the AMD Triton FA2 interface (varlen_bwd). Use the FA3 path instead."
+        )
+    if softcap != 0.0:
+        raise NotImplementedError(
+            "softcap is not supported in varlen_bwd (expected 0.0)."
+        )
+
+    if DEBUG:
+        print()
+        print("varlen_bwd")
+        print("dout:", dout.shape)
+        print("q:", q.shape)
+        print("k:", k.shape)
+        print("v:", v.shape)
+        print("out:", out)
+        print("softmax_lse:", softmax_lse.shape)
+        print("dq:", dq.shape if dq is not None else None)
+        print("dk:", dk.shape if dk is not None else None)
+        print("dv:", dv.shape if dv is not None else None)
+        print("cu_seqlens_q:", cu_seqlens_q, cu_seqlens_q.shape)
+        print("cu_seqlens_k:", cu_seqlens_k, cu_seqlens_k.shape)
+        print("alibi_slopes:", alibi_slopes)
+        print("max_seqlen_q:", max_seqlen_q)
+        print("max_seqlen_k:", max_seqlen_k)
+        print("dropout_p:", dropout_p)
+        print("softmax_scale:", softmax_scale)
+        print("causal:", causal)
+        print("window_size_left:", window_size_left)
+        print("window_size_right:", window_size_right)
+        print("deterministic:", deterministic)
+        print("gen_:", gen_)
+        print("rng_state:", rng_state)
+
+    dq = torch.zeros_like(q) if dq is None else dq.zero_()
+    dk = torch.zeros_like(k) if dk is None else dk.zero_()
+    dv = torch.zeros_like(v) if dv is None else dv.zero_()
+
+    # get shape
+    batch = len(cu_seqlens_q) - 1
+    total_q, nheads_q, _ = q.shape
+
+    # Create delta tensor with shape based on expectations
+    # delta (softmax_d) : (Hq, Total_Q) or (Hq, Total_Q + 128*batch)
+    if SHAPE_EXPECTATIONS == "rounded":
+        delta = torch.zeros(
+            (nheads_q, total_q + 128 * batch), device=q.device, dtype=torch.float32
+        )
+    else:
+        delta = torch.zeros((nheads_q, total_q), device=q.device, dtype=torch.float32)
+
+    # Upstream change: base seeding logic on provided rng_state instead of dropout probability.
+    if rng_state is not None:
+        philox_seed, philox_offset = rng_state[0].item(), rng_state[1].item()
+    else:
+        philox_seed, philox_offset = None, None
+
+    if alibi_slopes is not None:
+        if alibi_slopes.dim() == 2:
+            pass
+        elif alibi_slopes.dim() == 1:
+            alibi_slopes = alibi_slopes.unsqueeze(0).expand(batch, -1)
+        else:
+            raise ValueError("Alibi can be (nheads,) or (batch_size, nheads).")
+
+    # call implementation
+    if DEBUG:
+        print(f"Using Triton implementation in {BWD_MODE} mode")
+    attention_backward_triton_impl(
+        do=dout,
+        q=q,
+        k=k,
+        v=v,
+        o=out,
+        softmax_lse=softmax_lse,
+        dq=dq,
+        dk=dk,
+        dv=dv,
+        delta=delta,
+        sm_scale=softmax_scale,
+        alibi_slopes=alibi_slopes,
+        causal=causal,
+        layout="thd",
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+        seqused_q=None,
+        seqused_k=None,
+        dropout_p=dropout_p,
+        philox_seed=philox_seed,
+        philox_offset=philox_offset,
+        use_exp2=USE_EXP2,
+        mode=BWD_MODE,
+    )
+
+    if DEBUG:
+        print("varlen_bwd outputs")
+        print("delta:", delta, delta.shape)
+        print("dv:", dv, dv.shape)
+        print("dk:", dk, dk.shape)
+        print("dq:", dq, dq.shape)
+    # --- Assertions ---
+    assert dq.shape == q.shape, f"[varlen_bwd] dq shape {dq.shape} != q shape {q.shape}"
+    assert dk.shape == k.shape, f"[varlen_bwd] dk shape {dk.shape} != k shape {k.shape}"
+    assert dv.shape == v.shape, f"[varlen_bwd] dv shape {dv.shape} != v shape {v.shape}"
+    if SHAPE_EXPECTATIONS == "rounded":
+        batch = len(cu_seqlens_q) - 1
+        expected_delta_shape = (q.shape[1], q.shape[0] + 128 * batch)
+    else:
+        expected_delta_shape = (q.shape[1], q.shape[0])  # (Hq, Total_Q)
+    assert (
+        delta.shape == expected_delta_shape
+    ), f"[varlen_bwd] delta shape {delta.shape} != {expected_delta_shape}"
+    return dq, dk, dv, delta
+
+
+def fwd_kvcache(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    k: Optional[torch.Tensor],
+    v: Optional[torch.Tensor],
+    cache_seqlens: Optional[Union[int, torch.Tensor]],
+    rotary_cos: Optional[torch.Tensor],
+    rotary_sin: Optional[torch.Tensor],
+    cache_batch_idx: Optional[torch.Tensor],
+    cache_leftpad: Optional[torch.Tensor],
+    block_table: Optional[torch.Tensor],
+    alibi_slopes: Optional[torch.Tensor],
+    out: Optional[torch.Tensor],
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    rotary_interleaved: bool,
+    num_splits: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+
+    if softcap != 0.0:
+        raise NotImplementedError(
+            "softcap is not supported in fwd_kvcache (expected 0.0)."
+        )
+    if num_splits not in (0, 1):
+        raise NotImplementedError(
+            "num_splits > 1 not supported in AMD Triton FA2 fwd_kvcache."
+        )
+
+    if DEBUG:
+        print()
+        print("flash_attn_triton_amd.py::fwd_kvcache inputs")
+        print("q:", q, q.shape)
+        print("k_cache:", k_cache, k_cache.shape)
+        print("v_cache:", v_cache, v_cache.shape)
+        print("k:", k, k.shape if k is not None else None)
+        print("v:", v, v.shape if v is not None else None)
+        print("cache_seqlens:", cache_seqlens)
+        print("rotary_cos:", rotary_cos)
+        print("rotary_sin:", rotary_sin)
+        print("cache_batch_idx:", cache_batch_idx)
+        print("cache_leftpad:", cache_leftpad)
+        print("block_table:", block_table)
+        print("alibi_slopes:", alibi_slopes)
+        print("out:", out)
+        print("softmax_scale:", softmax_scale)
+        print("causal:", causal)
+        print("window_size_left:", window_size_left)
+        print("window_size_right:", window_size_right)
+        print("softcap:", softcap)
+        print("rotary_interleaved:", rotary_interleaved)
+        print("num_splits:", num_splits)
+
+    # output
+    out = torch.zeros_like(q) if out is None else out.zero_()
+
+    # Basic layout info for decode path
+    layout: Literal["bshd"] = "bshd"
+    q.shape[1]
+    k_cache.shape[1]
+    cache_seqlens_tensor = (
+        torch.tensor(cache_seqlens, device=q.device)
+        if isinstance(cache_seqlens, int)
+        else cache_seqlens
+    )
+    window_left = (
+        int(window_size_left.item())
+        if isinstance(window_size_left, torch.Tensor)
+        else window_size_left
+    )
+    window_right = (
+        int(window_size_right.item())
+        if isinstance(window_size_right, torch.Tensor)
+        else window_size_right
+    )
+
+    k_new = k
+    v_new = v
+
+    # get shape
+    batch, seqlen_q, nheads_q, _ = q.shape
+
+    # Create softmax_lse tensor - decode always uses exact shape (B, Hq, Sq)
+    softmax_lse = torch.zeros(
+        (batch, nheads_q, seqlen_q), device=q.device, dtype=torch.float32
+    )
+
+    if alibi_slopes is not None:
+        if alibi_slopes.dim() == 1:
+            alibi_slopes = alibi_slopes.unsqueeze(0).expand(batch, -1)
+        assert alibi_slopes.is_cuda and alibi_slopes.dim() == 2
+        assert alibi_slopes.shape == (batch, nheads_q)
+
+    # launch kernel
+    if DEBUG:
+        print("Using Triton implementation")
+    attention_forward_decode_triton_impl(
+        q,
+        k_cache,
+        v_cache,
+        k_new,
+        v_new,
+        out,
+        softmax_lse,
+        softmax_scale,
+        causal,
+        window_left,
+        window_right,
+        alibi_slopes,
+        layout,
+        cache_seqlens_tensor,
+        cache_batch_idx,
+        block_table,
+        None,
+        None,
+        None,
+        rotary_cos=rotary_cos,
+        rotary_sin=rotary_sin,
+        rotary_interleaved=rotary_interleaved,
+    )
+
+    if DEBUG:
+        print("out:", out, out.shape)
+        print("softmax_lse:", softmax_lse, softmax_lse.shape)
+    # --- Assertions ---
+    assert (
+        out.shape == q.shape
+    ), f"[fwd_kvcache] out shape {out.shape} != q shape {q.shape}"
+    expected_lse_shape = (q.shape[0], q.shape[2], q.shape[1])
+    assert (
+        softmax_lse.shape == expected_lse_shape
+    ), f"[fwd_kvcache] softmax_lse shape {softmax_lse.shape} != {expected_lse_shape}"
+    assert (
+        softmax_lse.dtype == torch.float32
+    ), f"[fwd_kvcache] softmax_lse dtype {softmax_lse.dtype} != torch.float32"
+    return out, softmax_lse

build/torch-rocm/_kernels/flash_attn_triton_amd/interface_v3.py

@@ -0,0 +1,660 @@
+import torch
+from typing import Literal, Optional, Tuple
+from .fwd_prefill import attention_forward_prefill_triton_impl
+from .fwd_decode import attention_forward_decode_triton_impl
+from .bwd import attention_backward_triton_impl
+from .utils import (
+    DEBUG,
+    USE_EXP2,
+    BWD_MODE,
+    PHILOX_SEED,
+    PHILOX_OFFSET,
+    is_fp8,
+)
+
+
+def fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    k_new: Optional[torch.Tensor],
+    v_new: Optional[torch.Tensor],
+    qv: Optional[torch.Tensor],
+    out: Optional[torch.Tensor],
+    cu_seqlens_q: Optional[torch.Tensor],
+    cu_seqlens_k: Optional[torch.Tensor],
+    cu_seqlens_k_new: Optional[torch.Tensor],
+    seqused_q: Optional[torch.Tensor],
+    seqused_k: Optional[torch.Tensor],
+    max_seqlen_q: Optional[int],
+    max_seqlen_k: Optional[int],
+    page_table: Optional[torch.Tensor],
+    kv_batch_idx: Optional[torch.Tensor],
+    leftpad_k: Optional[torch.Tensor],
+    rotary_cos: Optional[torch.Tensor],
+    rotary_sin: Optional[torch.Tensor],
+    seqlens_rotary: Optional[torch.Tensor],
+    q_descale: Optional[torch.Tensor],
+    k_descale: Optional[torch.Tensor],
+    v_descale: Optional[torch.Tensor],
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    attention_chunk: int,
+    softcap: float,
+    rotary_interleaved: bool,
+    scheduler_metadata: Optional[torch.Tensor] = None,
+    num_splits: int = 1,
+    pack_gqa: Optional[bool] = None,
+    sm_margin: int = 0,
+) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
+    """
+    Flash Attention v3 forward pass compatible interface for AMD Triton implementation.
+
+    This function maps v3 parameters to the existing AMD Triton implementation.
+    """
+
+    if DEBUG:
+        print()
+        print("interface_fa_v3.py::fwd inputs")
+        print("q:", q.shape)
+        print("k:", k.shape)
+        print("v:", v.shape)
+        print("k_new:", k_new.shape if k_new is not None else None)
+        print("v_new:", v_new.shape if v_new is not None else None)
+        print("qv:", qv.shape if qv is not None else None)
+        print("out:", out.shape if out is not None else None)
+        print("cu_seqlens_q:", cu_seqlens_q.shape if cu_seqlens_q is not None else None)
+        print("cu_seqlens_k:", cu_seqlens_k.shape if cu_seqlens_k is not None else None)
+        print(
+            "cu_seqlens_k_new:",
+            cu_seqlens_k_new.shape if cu_seqlens_k_new is not None else None,
+        )
+        print("seqused_q:", seqused_q.shape if seqused_q is not None else None)
+        print("seqused_k:", seqused_k.shape if seqused_k is not None else None)
+        print("max_seqlen_q:", max_seqlen_q)
+        print("max_seqlen_k:", max_seqlen_k)
+        print("page_table:", page_table.shape if page_table is not None else None)
+        print("kv_batch_idx:", kv_batch_idx.shape if kv_batch_idx is not None else None)
+        print("leftpad_k:", leftpad_k.shape if leftpad_k is not None else None)
+        print("rotary_cos:", rotary_cos.shape if rotary_cos is not None else None)
+        print("rotary_sin:", rotary_sin.shape if rotary_sin is not None else None)
+        print(
+            "seqlens_rotary:",
+            seqlens_rotary.shape if seqlens_rotary is not None else None,
+        )
+        print("q_descale:", q_descale.shape if q_descale is not None else None)
+        print("k_descale:", k_descale.shape if k_descale is not None else None)
+        print("v_descale:", v_descale.shape if v_descale is not None else None)
+        print("softmax_scale:", softmax_scale)
+        print("causal:", causal)
+        print("window_size_left:", window_size_left)
+        print("window_size_right:", window_size_right)
+        print("attention_chunk:", attention_chunk)
+        print("softcap:", softcap)
+        print("rotary_interleaved:", rotary_interleaved)
+        print("scheduler_metadata:", scheduler_metadata)
+        print("num_splits:", num_splits)
+        print("pack_gqa:", pack_gqa)
+        print("sm_margin:", sm_margin)
+
+    # Handle qv packed input
+    if qv is not None:
+        raise NotImplementedError(
+            "QV packed input is not yet supported in the AMD Triton backend"
+        )
+
+    # Handle softcap
+    if softcap != 0.0:
+        raise NotImplementedError(
+            f"Softcap is not yet supported in the AMD Triton backend (got softcap={softcap}, expected 0.0)"
+        )
+
+    # Handle attention_chunk
+    if attention_chunk != 0 and attention_chunk != 1:
+        raise NotImplementedError(
+            f"attention_chunk is not yet supported in the AMD Triton backend (got attention_chunk={attention_chunk})"
+        )
+
+    # Handle scheduler metadata
+    if scheduler_metadata is not None:
+        raise NotImplementedError(
+            "Scheduler metadata is not yet supported in the AMD Triton backend"
+        )
+
+    # Handle pack_gqa
+    if pack_gqa is not None and pack_gqa is not False:
+        raise NotImplementedError(
+            f"pack_gqa is not yet supported in the AMD Triton backend (got pack_gqa={pack_gqa})"
+        )
+
+    # Handle num_splits
+    if num_splits != 1:
+        raise NotImplementedError(
+            f"Split attention (num_splits > 1) is not yet supported in the AMD Triton backend (got num_splits={num_splits})"
+        )
+
+    # Handle sm_margin
+    if sm_margin != 0:
+        raise NotImplementedError(
+            f"sm_margin is not yet supported in the AMD Triton backend (got sm_margin={sm_margin}, expected 0)"
+        )
+
+    # Handle leftpad_k
+    if leftpad_k is not None:
+        raise NotImplementedError(
+            "Left padding (leftpad_k) is not yet supported in the AMD Triton backend"
+        )
+
+    # Handle cu_seqlens_k_new
+    if cu_seqlens_k_new is not None:
+        raise NotImplementedError(
+            "cu_seqlens_k_new is not yet supported in the AMD Triton backend"
+        )
+
+    # establish layout / varlen & max seq lens
+    if cu_seqlens_q is not None:
+        if len(q.shape) != 3:
+            raise ValueError(
+                f"cu_seqlens_q provided but q has shape {q.shape}, expected 3D tensor for varlen"
+            )
+        layout: Literal["bshd", "thd"] = "thd"
+        cu_seqlens_q_local = cu_seqlens_q
+        assert max_seqlen_q is not None, "max_seqlen_q required for varlen mode"
+        max_seqlens_q_local = max_seqlen_q
+        if cu_seqlens_k is not None:
+            cu_seqlens_k_local = cu_seqlens_k
+            assert (
+                max_seqlen_k is not None
+            ), "max_seqlen_k required when cu_seqlens_k provided"
+            max_seqlens_k_local = max_seqlen_k
+        else:
+            cu_seqlens_k_local = None
+            if len(k.shape) == 4:
+                max_seqlens_k_local = k.shape[1]
+            else:
+                assert max_seqlen_k is not None, "max_seqlen_k required for varlen mode"
+                max_seqlens_k_local = max_seqlen_k
+    else:
+        layout = "bshd"
+        cu_seqlens_q_local = None
+        cu_seqlens_k_local = None
+        max_seqlens_q_local = q.shape[1] if max_seqlen_q is None else max_seqlen_q
+        max_seqlens_k_local = k.shape[1] if max_seqlen_k is None else max_seqlen_k
+
+    # Now determine if we should use decode or prefill kernel
+    # Decode kernel should be used for KV cache scenarios where:
+    # 1. k_new/v_new are provided - incremental KV cache update (primary KV cache indicator)
+    # 2. kv_batch_idx is provided - KV cache batch indexing (primary KV cache indicator)
+    # 3. seqused_k without seqused_q - indicates KV cache fill levels (not varlen masking)
+    # Note: In varlen, both seqused_q and seqused_k are used for sequence masking
+    #       In KV cache, only seqused_k is used to track cache fill levels
+    # Detect KV cache scenarios:
+    # - Clear KV cache indicators (k_new, v_new, kv_batch_idx)
+    # - OR seqused_k without seqused_q (KV cache fill tracking, not varlen masking)
+    use_decode = (
+        k_new is not None  # Have new KV to append (KV cache indicator)
+        or v_new is not None  # Have new KV to append (KV cache indicator)
+        or kv_batch_idx is not None  # Have KV cache batch indexing (KV cache indicator)
+        or (
+            seqused_k is not None and seqused_q is None
+        )  # KV cache fill levels (not varlen)
+    )
+
+    # Check for unsupported features with decode kernel
+    if use_decode:
+        if layout == "thd":
+            raise NotImplementedError(
+                "Varlen is not yet supported with the decode kernel in the AMD Triton backend"
+            )
+        if kv_batch_idx is not None:
+            raise NotImplementedError(
+                "kv_batch_idx is not yet supported with the decode kernel in the AMD Triton backend"
+            )
+
+    if out is None:
+        # NOTE: Using types that are lower precision than float32 such as bfloat16 for fp8 causes mismatches on a small set of tests.
+        out_dtype = torch.float32 if is_fp8([q, k, v]) else q.dtype
+        if layout == "bshd":
+            out = torch.zeros(
+                q.shape[0],
+                q.shape[1],
+                q.shape[2],
+                v.shape[-1],
+                dtype=out_dtype,
+                device=q.device,
+            )
+        elif layout == "thd":
+            out = torch.zeros(
+                q.shape[0], q.shape[1], v.shape[-1], dtype=out_dtype, device=q.device
+            )
+        else:
+            raise ValueError(
+                f"Unsupported layout: {layout}. Only 'bshd' and 'thd' layouts are supported."
+            )
+    else:
+        out = out.zero_()
+
+    # Handle causal mask
+    causal_flag = bool(causal)
+
+    # Handle alibi slopes
+    alibi_slopes = None
+
+    # Handle dropout
+    dropout_p = 0.0
+    return_softmax = False
+    philox_seed = PHILOX_SEED
+    philox_offset = PHILOX_OFFSET
+
+    # Call implementation
+    if DEBUG:
+        print("Using Triton implementation")
+
+    if use_decode:
+        if DEBUG:
+            print(
+                f"Using Decode Triton implementation (cache_seqlens={seqused_k is not None}, k_new={k_new is not None}, v_new={v_new is not None}, kv_batch_idx={kv_batch_idx is not None})"
+            )
+
+        # Create softmax_lse tensor for decode - always exact shape (B, Hq, Sq)
+        batch, seqlen_q, nheads_q, _ = q.shape
+        softmax_lse = torch.zeros(
+            (batch, nheads_q, seqlen_q), device=q.device, dtype=torch.float32
+        )
+
+        # Decode only supports bshd layout
+        assert layout == "bshd", f"decode requires bshd layout, got {layout}"
+        attention_forward_decode_triton_impl(
+            q,
+            k,
+            v,
+            k_new,
+            v_new,
+            out,
+            softmax_lse,
+            softmax_scale,
+            causal_flag,
+            window_size_left,
+            window_size_right,
+            alibi_slopes,
+            layout,
+            seqused_k,
+            kv_batch_idx,
+            page_table,
+            q_descale,
+            k_descale,
+            v_descale,
+            rotary_cos=rotary_cos,
+            rotary_sin=rotary_sin,
+            rotary_interleaved=rotary_interleaved,
+            seqlens_rotary=seqlens_rotary,
+        )
+    else:
+        if DEBUG:
+            print("Using Prefill Triton implementation")
+
+        # Create softmax_lse tensor - FA3 always uses exact shapes
+        if layout == "thd":
+            # varlen: (Hq, Total_Q)
+            total_q, nheads_q, _ = q.shape
+            softmax_lse = torch.zeros(
+                (nheads_q, total_q), device=q.device, dtype=torch.float32
+            )
+        else:
+            # bshd: (B, Hq, Sq)
+            batch, seqlen_q, nheads_q, _ = q.shape
+            softmax_lse = torch.zeros(
+                (batch, nheads_q, seqlen_q), device=q.device, dtype=torch.float32
+            )
+
+        # sd_mask is not returned in v3 interface
+        sd_mask = None
+
+        attention_forward_prefill_triton_impl(
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            sd_mask,
+            softmax_scale,
+            alibi_slopes,
+            causal_flag,
+            window_size_left,
+            window_size_right,
+            None,
+            layout,
+            cu_seqlens_q_local,
+            cu_seqlens_k_local,
+            max_seqlens_q_local,
+            max_seqlens_k_local,
+            dropout_p,
+            philox_seed,
+            philox_offset,
+            return_softmax,
+            USE_EXP2,
+            q_descale,
+            k_descale,
+            v_descale,
+            seqused_q,
+            seqused_k,
+            rotary_cos=rotary_cos,
+            rotary_sin=rotary_sin,
+            rotary_interleaved=rotary_interleaved,
+            seqlens_rotary=seqlens_rotary,
+        )
+
+    if DEBUG:
+        print("interface_fa_v3.py::fwd outputs")
+        print("out:", out.shape)
+        print("softmax_lse:", softmax_lse.shape)
+
+    # --- Assertions (FA3 always expects exact shapes) ---
+    # out: same shape as q except last dim is v's head_dim
+    if layout == "thd":
+        # varlen: (Total_Q, Hq, Dv)
+        assert (
+            out.shape[0] == q.shape[0]
+        ), f"[fwd_v3] out.shape[0] {out.shape[0]} != q.shape[0] {q.shape[0]}"
+        assert (
+            out.shape[1] == q.shape[1]
+        ), f"[fwd_v3] out.shape[1] {out.shape[1]} != q.shape[1] {q.shape[1]}"
+        assert (
+            out.shape[2] == v.shape[-1]
+        ), f"[fwd_v3] out.shape[2] {out.shape[2]} != v.shape[-1] {v.shape[-1]}"
+    else:
+        # bshd: (B, Sq, Hq, Dv)
+        assert (
+            out.shape[0] == q.shape[0]
+        ), f"[fwd_v3] out.shape[0] {out.shape[0]} != q.shape[0] {q.shape[0]}"
+        assert (
+            out.shape[1] == q.shape[1]
+        ), f"[fwd_v3] out.shape[1] {out.shape[1]} != q.shape[1] {q.shape[1]}"
+        assert (
+            out.shape[2] == q.shape[2]
+        ), f"[fwd_v3] out.shape[2] {out.shape[2]} != q.shape[2] {q.shape[2]}"
+        assert (
+            out.shape[3] == v.shape[-1]
+        ), f"[fwd_v3] out.shape[3] {out.shape[3]} != v.shape[-1] {v.shape[-1]}"
+
+    # softmax_lse dtype
+    assert (
+        softmax_lse.dtype == torch.float32
+    ), f"[fwd_v3] softmax_lse dtype {softmax_lse.dtype} != torch.float32"
+    # softmax_lse shape depends on layout
+    expected_lse_shape: tuple[int, ...]
+    if layout == "thd":
+        # varlen: (Hq, Total_Q)
+        expected_lse_shape = (q.shape[1], q.shape[0])
+    else:
+        # bshd: (B, Hq, Sq)
+        expected_lse_shape = (q.shape[0], q.shape[2], q.shape[1])
+    assert (
+        softmax_lse.shape == expected_lse_shape
+    ), f"[fwd_v3] softmax_lse shape {softmax_lse.shape} != {expected_lse_shape}"
+
+    # Return format compatible with v3
+    # V3 returns (out, softmax_lse, out_accum, softmax_lse_accum)
+    # out_accum and softmax_lse_accum are None for Triton AMD (no split-k accumulation)
+    return out, softmax_lse, None, None
+
+
+def bwd(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor],
+    dk: Optional[torch.Tensor],
+    dv: Optional[torch.Tensor],
+    cu_seqlens_q: Optional[torch.Tensor],
+    cu_seqlens_k: Optional[torch.Tensor],
+    seqused_q: Optional[torch.Tensor],
+    seqused_k: Optional[torch.Tensor],
+    max_seqlen_q: Optional[int],
+    max_seqlen_k: Optional[int],
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    deterministic: bool,
+    sm_margin: int = 0,
+    *,
+    q_descale: Optional[torch.Tensor] = None,
+    k_descale: Optional[torch.Tensor] = None,
+    v_descale: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor]:
+    """
+    Flash Attention v3 backward pass compatible interface for AMD Triton implementation.
+
+    This function maps v3 parameters to the existing AMD Triton implementation.
+    Gradients are written in-place to dq, dk, dv if provided (or allocated internally).
+    Returns (delta,) matching the upstream FA3 API.
+
+    Args:
+        q_descale: Optional descale tensor for FP8 query (batch, nheads_k)
+        k_descale: Optional descale tensor for FP8 key (batch, nheads_k)
+        v_descale: Optional descale tensor for FP8 value (batch, nheads_k)
+    """
+
+    if DEBUG:
+        print()
+        print("interface_fa_v3.py::bwd inputs")
+        print("dout:", dout.shape)
+        print("q:", q.shape)
+        print("k:", k.shape)
+        print("v:", v.shape)
+        print("out:", out.shape)
+        print("softmax_lse:", softmax_lse.shape)
+        print("dq:", dq.shape if dq is not None else None)
+        print("dk:", dk.shape if dk is not None else None)
+        print("dv:", dv.shape if dv is not None else None)
+        print("cu_seqlens_q:", cu_seqlens_q.shape if cu_seqlens_q is not None else None)
+        print("cu_seqlens_k:", cu_seqlens_k.shape if cu_seqlens_k is not None else None)
+        print("seqused_q:", seqused_q.shape if seqused_q is not None else None)
+        print("seqused_k:", seqused_k.shape if seqused_k is not None else None)
+        print("max_seqlen_q:", max_seqlen_q)
+        print("max_seqlen_k:", max_seqlen_k)
+        print("softmax_scale:", softmax_scale)
+        print("causal:", causal)
+        print("window_size_left:", window_size_left)
+        print("window_size_right:", window_size_right)
+        print("softcap:", softcap)
+        print("deterministic:", deterministic)
+        print("sm_margin:", sm_margin)
+
+    # Check for unsupported features in backward pass
+
+    # Handle sliding window - backward doesn't support it yet
+    is_sliding_window = (window_size_left >= 0) or (window_size_right >= 0)
+    if is_sliding_window:
+        raise NotImplementedError(
+            f"Sliding window attention is not yet supported in the AMD Triton backward pass "
+            f"(window_size_left={window_size_left}, window_size_right={window_size_right}). "
+            f"Use window_size=(-1, -1) for full attention."
+        )
+
+    # Handle softcap
+    if softcap != 0.0:
+        raise NotImplementedError(
+            f"Softcap is not yet supported in the AMD Triton backend backward pass (got softcap={softcap}, expected 0.0)"
+        )
+
+    # Handle sm_margin
+    if sm_margin != 0:
+        raise NotImplementedError(
+            f"sm_margin is not yet supported in the AMD Triton backend backward pass (got sm_margin={sm_margin}, expected 0)"
+        )
+
+    # Initialize gradient tensors if not provided
+    # NOTE: Using types that are lower precision than float32 such as bfloat16 for fp8 causes mismatches on a small set of tests.
+    grad_dtype = torch.float32 if is_fp8([q, k, v]) else q.dtype
+    dq = torch.zeros_like(q, dtype=grad_dtype) if dq is None else dq.zero_()
+    dk = torch.zeros_like(k, dtype=grad_dtype) if dk is None else dk.zero_()
+    dv = torch.zeros_like(v, dtype=grad_dtype) if dv is None else dv.zero_()
+
+    # Determine layout and set max sequence lengths
+    layout: Literal["bshd", "bhsd", "thd"]
+    if cu_seqlens_q is not None and cu_seqlens_k is not None:
+        # Variable length sequence mode
+        layout = "thd"
+        batch = len(cu_seqlens_q) - 1
+        total_q, nheads_q, _ = q.shape
+        # Create delta tensor - varlen: (Hq, Total_Q)
+        delta = torch.zeros((nheads_q, total_q), device=q.device, dtype=torch.float32)
+    else:
+        # Regular batch mode
+        layout = "bshd"
+        batch, seqlen_q, nheads_q, _ = q.shape
+        # Create delta tensor - bshd: (B, Hq, Sq)
+        delta = torch.zeros(
+            (batch, nheads_q, seqlen_q), device=q.device, dtype=torch.float32
+        )
+
+    # Set max sequence lengths (infer from tensor shapes if not provided)
+    max_seqlen_q = q.shape[1] if max_seqlen_q is None else max_seqlen_q
+    max_seqlen_k = k.shape[1] if max_seqlen_k is None else max_seqlen_k
+
+    # V3 backward doesn't have dropout or alibi slopes
+    dropout_p = 0.0
+    philox_seed, philox_offset = None, None
+    alibi_slopes = None
+
+    # Call implementation
+    if DEBUG:
+        print(f"Using Triton implementation in {BWD_MODE} mode")
+    attention_backward_triton_impl(
+        do=dout,
+        q=q,
+        k=k,
+        v=v,
+        o=out,
+        softmax_lse=softmax_lse,
+        dq=dq,
+        dk=dk,
+        dv=dv,
+        delta=delta,
+        sm_scale=softmax_scale,
+        alibi_slopes=alibi_slopes,
+        causal=causal,
+        layout=layout,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+        seqused_q=seqused_q,
+        seqused_k=seqused_k,
+        dropout_p=dropout_p,
+        philox_seed=philox_seed,
+        philox_offset=philox_offset,
+        use_exp2=USE_EXP2,
+        mode=BWD_MODE,
+        q_descale=q_descale,
+        k_descale=k_descale,
+        v_descale=v_descale,
+    )
+
+    if DEBUG:
+        print("interface_fa_v3.py::bwd outputs")
+        print("dq:", dq.shape)
+        print("dk:", dk.shape)
+        print("dv:", dv.shape)
+        print("delta:", delta.shape)
+
+    # --- Assertions (FA3 always expects exact shapes) ---
+    # Gradients should match input shapes
+    assert dq.shape == q.shape, f"[bwd_v3] dq shape {dq.shape} != q shape {q.shape}"
+    assert dk.shape == k.shape, f"[bwd_v3] dk shape {dk.shape} != k shape {k.shape}"
+    assert dv.shape == v.shape, f"[bwd_v3] dv shape {dv.shape} != v shape {v.shape}"
+    # delta (softmax_d) should match softmax_lse shape
+    assert (
+        delta.dtype == torch.float32
+    ), f"[bwd_v3] delta dtype {delta.dtype} != torch.float32"
+    expected_delta_shape: tuple[int, ...]
+    if layout == "thd":
+        # varlen: (Hq, Total_Q)
+        expected_delta_shape = (q.shape[1], q.shape[0])
+    else:
+        # bshd: (B, Hq, Sq)
+        expected_delta_shape = (q.shape[0], q.shape[2], q.shape[1])
+    assert (
+        delta.shape == expected_delta_shape
+    ), f"[bwd_v3] delta shape {delta.shape} != {expected_delta_shape}"
+
+    # Return only delta to match upstream FA3 API
+    # dq, dk, dv are mutated in-place by attention_backward_triton_impl
+    return (delta,)
+
+
+def fwd_combine(
+    out_partial: torch.Tensor,
+    lse_partial: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+    out_dtype: Optional[torch.dtype] = None,
+) -> torch.Tensor:
+    """
+    Combine partial outputs from split attention computation.
+
+    This is used when num_splits > 1 to combine the partial results.
+
+    Args:
+        out_partial: Partial output tensor from split computation
+        lse_partial: Partial log-sum-exp tensor
+        out: Optional output tensor to write to
+        out_dtype: Optional dtype for output
+
+    Returns:
+        Combined output tensor
+    """
+    raise NotImplementedError(
+        "fwd_combine is not yet implemented in the AMD Triton backend"
+    )
+
+
+def get_scheduler_metadata(
+    batch_size: int,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    num_heads_q: int,
+    num_heads_kv: int,
+    headdim: int,
+    headdim_v: int,
+    qkv_dtype: torch.dtype,
+    cache_seqlens: torch.Tensor,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    cu_seqlens_k_new: Optional[torch.Tensor] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    cache_leftpad: Optional[torch.Tensor] = None,
+    page_size: Optional[int] = None,
+    max_seqlen_k_new: int = 0,
+    causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    attention_chunk: int = 0,
+    has_softcap: bool = False,
+    num_splits: int = 0,
+    pack_gqa: Optional[bool] = None,
+    sm_margin: int = 0,
+) -> None:
+    """
+    Get scheduler metadata for optimized kernel selection.
+
+    This function is used to precompute metadata for kernel scheduling in FA3.
+    The AMD Triton backend currently doesn't use scheduler metadata, so this
+    raises an error.
+
+    Args:
+        Various attention parameters used for scheduling decisions
+
+    Returns:
+        None - scheduler metadata is not used in AMD Triton backend
+    """
+    raise NotImplementedError(
+        "get_scheduler_metadata is not supported in the AMD Triton backend yet."
+    )

build/torch-rocm/_kernels/flash_attn_triton_amd/utils.py

@@ -0,0 +1,338 @@
+"""
+Utilities for Flash Attention Triton AMD backend.
+
+This module contains essential runtime utilities:
+- GPU architecture detection
+- Global configuration flags
+- Tensor shape/stride helpers
+- FP8 type detection
+"""
+
+import functools
+import os
+import json
+import logging
+from dataclasses import dataclass
+from typing import Literal, Optional, Union
+import triton.language as tl
+
+import torch
+import triton
+
+logger = logging.getLogger(__name__)
+
+AutotuneMode = Literal["off", "on", "sweep"]
+
+__all__ = [
+    # Runtime info
+    "get_arch",
+    "is_hip",
+    # Global config
+    "AutotuneMode",
+    "AUTOTUNE",
+    "DEBUG",
+    "USE_TRITON_ROCM",
+    "BWD_MODE",
+    "USE_EXP2",
+    "PHILOX_SEED",
+    "PHILOX_OFFSET",
+    "SHAPE_EXPECTATIONS",
+    # FP8
+    "is_fp8",
+    # Shape/stride helpers
+    "get_shape_from_layout",
+    "get_stride_from_layout",
+    "get_padded_headsize",
+    # Misc helpers
+    "round_multiple",
+]
+
+
+# -------------------------------
+# GPU Architecture
+# -------------------------------
+ArchFamily = Literal["cdna", "rdna"]
+
+CDNA_ARCHS = frozenset({"gfx908", "gfx90a", "gfx940", "gfx941", "gfx942", "gfx950"})
+RDNA_ARCHS = frozenset(
+    {
+        "gfx1030",
+        "gfx1100",
+        "gfx1101",
+        "gfx1102",
+        "gfx1150",
+        "gfx1151",
+        "gfx1200",
+        "gfx1201",
+    }
+)
+FP8_ARCHS = frozenset({"gfx942", "gfx950", "gfx1200", "gfx1201"})
+
+_RECOMMENDED_FP8_REPLACEMENTS: dict[str, dict[torch.dtype, torch.dtype]] = {
+    "gfx942": {
+        torch.float8_e4m3fn: torch.float8_e4m3fnuz,
+        torch.float8_e5m2: torch.float8_e5m2fnuz,
+    },
+}
+
+
+@dataclass(frozen=True)
+class GpuArch:
+    """GPU architecture information."""
+
+    name: str  # e.g., "gfx942", "gfx1100"
+    family: Optional[ArchFamily] = None
+
+    @property
+    def is_cdna(self) -> bool:
+        return self.family == "cdna"
+
+    @property
+    def is_rdna(self) -> bool:
+        return self.family == "rdna"
+
+    @property
+    def supports_fp8(self) -> bool:
+        """Check if this architecture supports FP8."""
+        return self.name in FP8_ARCHS
+
+    def recommended_fp8_dtype(self, dtype: torch.dtype) -> torch.dtype:
+        """Get the recommended FP8 dtype for this architecture.
+
+        Some architectures prefer different FP8 variants (e.g., fnuz vs fn).
+        Returns the input dtype unchanged if no replacement is recommended.
+        """
+        return _RECOMMENDED_FP8_REPLACEMENTS.get(self.name, {}).get(dtype, dtype)
+
+    @property
+    def cu_count(self) -> int:
+        """Get the number of compute units on the current GPU."""
+        return int(
+            torch.cuda.get_device_properties(
+                torch.cuda.current_device()
+            ).multi_processor_count
+        )
+
+
+# -------------------------------
+# Global Variables
+# -------------------------------
+USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE"
+AUTOTUNE: AutotuneMode = (
+    "on"
+    if os.environ.get("FLASH_ATTENTION_TRITON_AMD_AUTOTUNE", "1").lower()
+    in ("1", "true", "yes", "on")
+    else "off"
+)
+
+# User override config json for attn_fwd.
+# Note: Ignored if FLASH_ATTENTION_TRITON_AMD_AUTOTUNE is enabled.
+#
+# e.g. FLASH_ATTENTION_FWD_TRITON_AMD_CONFIG_JSON='{"BLOCK_M":32,"BLOCK_N":32,"waves_per_eu":1,"PRE_LOAD_V":false,"num_stages":1,"num_warps":4}'
+FWD_CONF_OVERRIDE = None
+try:
+    conf_json = os.getenv("FLASH_ATTENTION_FWD_TRITON_AMD_CONFIG_JSON")
+    if conf_json:
+        conf = json.loads(conf_json)
+        FWD_CONF_OVERRIDE = triton.Config(
+            conf,
+            num_stages=conf.pop("num_stages", 1),
+            num_warps=conf.pop("num_warps", 4),
+        )
+except Exception as e:
+    logger.warning(f"FLASH_ATTENTION_FWD_TRITON_AMD_CONFIG_JSON parse error: {e}")
+
+# Unified debug level:
+#   0 = off (default)
+#   1 = basic debug info (shapes, tensor stats, kernel params)
+#   2 = detailed debug (includes Triton interpreter prints in kernels)
+#
+# Set via: FLASH_ATTENTION_TRITON_AMD_DEBUG=0|1|2
+DEBUG: int = int(os.environ.get("FLASH_ATTENTION_TRITON_AMD_DEBUG", "0"))
+if AUTOTUNE != "off" or DEBUG > 0:
+    os.environ["TRITON_PRINT_AUTOTUNING"] = "1"
+if DEBUG >= 2:
+    os.environ["TRITON_INTERPRET"] = "1"
+BWD_MODE: Literal["fused", "fused_atomic", "split"] = "fused"
+USE_EXP2 = True
+PHILOX_SEED = 0x1BF58
+PHILOX_OFFSET = 0x1D4B49
+SHAPE_EXPECTATIONS: Literal["exact", "rounded"] = "exact"
+
+
+# -------------------------------
+# FP8
+# -------------------------------
+_FP8_DTYPES = frozenset(
+    {
+        torch.float8_e4m3fnuz,
+        torch.float8_e4m3fn,
+        torch.float8_e5m2,
+        torch.float8_e5m2fnuz,
+    }
+)
+
+
+def is_fp8(
+    x: Union[torch.dtype, torch.Tensor, list[torch.Tensor], tuple[torch.Tensor, ...]],
+) -> bool:
+    """Check if dtype/tensor(s) are FP8.
+
+    This is a pure function - it only checks dtypes, not architecture support.
+    Use `get_arch().supports_fp8` to check if the current GPU supports FP8.
+
+    Args:
+        x: A dtype, tensor, or list/tuple of tensors to check.
+
+    Returns:
+        True if FP8, False otherwise.
+
+    Rules for multiple tensors:
+        - If all tensors are FP8 -> return True.
+        - If none are FP8 -> return False.
+        - If a mix of FP8 and non-FP8 -> raise ValueError.
+
+    Empty list/tuple returns False.
+    """
+    # Handle dtype directly
+    if isinstance(x, torch.dtype):
+        return x in _FP8_DTYPES
+
+    # Handle single tensor
+    if isinstance(x, torch.Tensor):
+        return x.dtype in _FP8_DTYPES
+
+    # Handle list/tuple of tensors
+    if isinstance(x, (list, tuple)):
+        if len(x) == 0:
+            return False
+        flags = [t.dtype in _FP8_DTYPES for t in x]
+        if all(flags):
+            return True
+        if not any(flags):
+            return False
+        raise ValueError(
+            "Mixed FP8 and non-FP8 tensors provided; either all or none must be FP8."
+        )
+
+    raise TypeError(f"Expected dtype, Tensor, or sequence of Tensors, got {type(x)}")
+
+
+# -------------------------------
+# Shape/Stride Helpers
+# -------------------------------
+def get_shape_from_layout(
+    x: torch.Tensor,
+    layout: Literal["bshd", "bhsd", "thd"],
+    cu_seqlens: Optional[torch.Tensor] = None,
+    max_seqlen: Optional[int] = None,
+) -> tuple[int, int, int, int]:
+    """Extract (batch, max_seqlen, num_heads, head_dim) from tensor based on layout."""
+    if layout == "bhsd":
+        batch, num_heads, max_seqlen_final, head_dim = x.shape
+    elif layout == "bshd":
+        batch, max_seqlen_final, num_heads, head_dim = x.shape
+    elif layout == "thd":
+        total_seqlen, num_heads, head_dim = x.shape
+        if cu_seqlens is None:
+            raise ValueError("cu_seqlens must be provided for varlen (thd) layout")
+        if max_seqlen is None:
+            raise ValueError("max_seqlen must be provided for varlen (thd) layout")
+
+        batch, max_seqlen_final, num_heads, head_dim = (
+            len(cu_seqlens) - 1,
+            max_seqlen,
+            num_heads,
+            head_dim,
+        )
+    else:
+        raise ValueError(f"Got unsupported layout: {layout}")
+
+    return batch, max_seqlen_final, num_heads, head_dim
+
+
+def get_stride_from_layout(
+    x: torch.Tensor, layout: Literal["bshd", "bhsd", "thd"]
+) -> tuple[int, int, int, int]:
+    """Get strides in (batch, head, seq, dim) order for the given layout."""
+    if layout == "thd":
+        strides = (0, x.stride(1), x.stride(0), x.stride(2))
+    elif layout == "bhsd":
+        strides = (x.stride(0), x.stride(1), x.stride(2), x.stride(3))
+    elif layout == "bshd":
+        strides = (x.stride(0), x.stride(2), x.stride(1), x.stride(3))
+    else:
+        raise ValueError(f"Got unsupported layout: {layout}")
+    return strides
+
+
+def get_padded_headsize(size: int) -> int:
+    """Get closest power of 2 over or equal to 32."""
+    # Smallest head_dim supported is 16. If smaller, the tile in the
+    # kernel is padded - there is no padding in memory for any dims.
+    padded_d_model = 1 << (size - 1).bit_length()
+    padded_d_model = max(padded_d_model, 16)
+    return padded_d_model
+
+
+# -------------------------------
+# Misc helpers
+# -------------------------------
+def round_multiple(x: int, m: int) -> int:
+    """Round x up to the nearest multiple of m."""
+    return (x + m - 1) // m * m
+
+
+# -------------------------------
+# Runtime info
+# -------------------------------
+@functools.cache
+def is_hip() -> bool:
+    """Check if running on HIP (AMD) backend."""
+    return bool(triton.runtime.driver.active.get_current_target().backend == "hip")
+
+
+@functools.cache
+def get_arch() -> GpuArch:
+    """Get the current GPU architecture."""
+    try:
+        name: str = triton.runtime.driver.active.get_current_target().arch
+    except RuntimeError:
+        # No GPU available (e.g. import-only on Windows/CPU)
+        return GpuArch(name="unknown")
+    if name in CDNA_ARCHS:
+        return GpuArch(name=name, family="cdna")
+    elif name in RDNA_ARCHS:
+        return GpuArch(name=name, family="rdna")
+    else:
+        return GpuArch(name=name)
+
+
+@triton.jit
+def remap_xcd(pid, GRID_MN, NUM_XCDS: tl.constexpr = 8):
+    ## pid remapping on xcds
+    # Number of pids per XCD in the new arrangement
+    pids_per_xcd = (GRID_MN + NUM_XCDS - 1) // NUM_XCDS
+    # When GRID_MN cannot divide NUM_XCDS, some xcds will have
+    # pids_per_xcd pids, the other will have pids_per_xcd - 1 pids.
+    # We calculate the number of xcds that have pids_per_xcd pids as
+    # tall_xcds
+    tall_xcds = GRID_MN % NUM_XCDS
+    tall_xcds = NUM_XCDS if tall_xcds == 0 else tall_xcds
+    # Compute current XCD and local pid within the XCD
+    xcd = pid % NUM_XCDS
+    local_pid = pid // NUM_XCDS
+    # Calculate new pid based on the new grouping
+    # Note that we need to consider the following two cases:
+    # 1. the current pid is on a tall xcd
+    # 2. the current pid is on a short xcd
+    if xcd < tall_xcds:
+        pid = xcd * pids_per_xcd + local_pid
+    else:
+        pid = (
+            tall_xcds * pids_per_xcd
+            + (xcd - tall_xcds) * (pids_per_xcd - 1)
+            + local_pid
+        )
+
+    return pid

build/torch-rocm/_kernels/mha.py

@@ -0,0 +1,942 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+import functools
+import json
+import torch
+import triton
+import triton.language as tl
+
+from ..utils._triton import arch_info
+from ..utils.core import AITER_TRITON_CONFIGS_PATH
+from ..utils._triton.pid_preprocessing import remap_xcd
+from ..utils._triton.mha_kernel_utils import _compute_fp8_scaling_factors
+from ..utils._triton.kernel_repr import make_kernel_repr
+
+
+@triton.jit
+def _cdiv_fn(x, y):
+    return (x + y - 1) // y
+
+
+@triton.jit
+def _load_fn(ptrs, offset_first, offset_second, boundary_first, boundary_second):
+    if offset_first is not None and offset_second is not None:
+        mask = (offset_first[:, None] < boundary_first) & (
+            offset_second[None, :] < boundary_second
+        )
+        tensor = tl.load(ptrs, mask=mask, other=0.0)
+    elif offset_first is not None:
+        mask = offset_first[:, None] < boundary_first
+        tensor = tl.load(ptrs, mask=mask, other=0.0)
+    elif offset_second is not None:
+        mask = offset_second[None, :] < boundary_second
+        tensor = tl.load(ptrs, mask=mask, other=0.0)
+    else:
+        tensor = tl.load(ptrs)
+    return tensor
+
+
+@triton.jit
+def _compute_alibi_block(
+    alibi_slope, seqlen_q, seqlen_k, offs_m, offs_n, transpose=False
+):
+    # when seqlen_k and seqlen_q are different we want the diagonal to stick to the bottom right of the attention matrix
+    # for casual mask we want something like this where (1 is kept and 0 is masked)
+    # seqlen_q = 2 and seqlen_k = 5
+    #   1 1 1 1 0
+    #   1 1 1 1 1
+    # seqlen_q = 5 and seqlen_k = 2
+    #        0 0
+    #        0 0
+    #        0 0
+    #        1 0
+    #        1 1
+    # for alibi the diagonal is 0 indicating no penalty for attending to that spot and increasing penalty for attending further from the diagonal
+    # e.g. alibi_slope = 1, seqlen_q = 2, seqlen_k = 5, offs_m = [0, 1, 2, 3], offs_n = [0, 1, 2, 3, 4], transpose = False
+    # 1. offs_m[:,None] = [[0],
+    #                       [1],
+    # 2. offs_m[:,None] + seqlen_k = [[5],
+    #                                  [6],
+    # 3. offs_m[:,None] + seqlen_k - seqlen_q = [[3],
+    #                                             [4],
+    # 4. offs_m[:,None] + seqlen_k - seqlen_q - offs_n[None,:] = [[3], - [[0, 1, 2, 3, 4]] =  [[ 3, 2, 1, 0,-1],
+    #                                                            [4],                           [ 4, 3, 2, 1, 0]]
+    # 5. -1 * alibi_slope * tl.abs(relative_pos_block) = [[ -3, -2, -1, 0,-1],
+    #                                                     [ -4, -3, -2, -1, 0]],
+    relative_pos_block = offs_m[:, None] + seqlen_k - seqlen_q - offs_n[None, :]
+    alibi_block = -1 * alibi_slope * tl.abs(relative_pos_block)
+    if transpose:
+        return alibi_block.T
+    else:
+        return alibi_block
+
+
+@triton.jit
+def _attn_fwd_inner(
+    acc,
+    l_i,
+    m_i,
+    q,
+    q_pe,
+    k_ptrs,
+    k_pe_ptrs,
+    v_ptrs,
+    stride_kn,
+    stride_vk,
+    stride_sn,
+    start_m,
+    seqlen_k,
+    seqlen_q,
+    dropout_p,
+    sd_mask_ptrs,
+    dropout_mask_ptrs,
+    philox_seed,
+    philox_ptrs,
+    block_min,
+    block_max,
+    offs_n_causal,
+    masked_blocks,
+    n_extra_tokens,
+    alibi_slope,
+    descale_q,
+    descale_k,
+    descale_v,
+    OFFS_M: tl.constexpr,
+    OFFS_N: tl.constexpr,
+    PRELOAD_V: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_DMODEL_POW2: tl.constexpr,
+    BLOCK_DMODEL_PE: tl.constexpr,  # it's zero or a power of 2
+    SM_SCALE: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+    MASK_STEPS: tl.constexpr,
+    ENABLE_DROPOUT: tl.constexpr,
+    RETURN_SCORES: tl.constexpr,
+    PADDED_HEAD: tl.constexpr,
+    IS_FP8: tl.constexpr,
+    FP8_MAX: tl.constexpr,
+    ENABLE_PIPELINING: tl.constexpr,
+    SLIDING_WINDOW: tl.constexpr,
+):
+    RCP_LN2: tl.constexpr = 1.4426950408889634
+    HAS_PE: tl.constexpr = BLOCK_DMODEL_PE > 0
+
+    # loop over k, v, and update accumulator
+
+    num_stages: tl.constexpr = (
+        None if ENABLE_PIPELINING else 1
+    )  # Set num_stages==1 if we want to disable pipelining
+    for start_n in tl.range(block_min, block_max, BLOCK_N, num_stages=num_stages):
+        # For padded blocks, we will overrun the tensor size if
+        # we load all BLOCK_N. For others, the blocks are all within range.
+        if MASK_STEPS:
+            k_offs_n = start_n + tl.arange(0, BLOCK_N)
+        else:
+            k_offs_n = None
+        k_offs_k = None if not PADDED_HEAD else tl.arange(0, BLOCK_DMODEL_POW2)
+        k = _load_fn(k_ptrs, k_offs_k, k_offs_n, BLOCK_DMODEL, seqlen_k)
+        if HAS_PE:
+            k_pe = _load_fn(
+                k_pe_ptrs,
+                None,
+                k_offs_n,
+                (BLOCK_DMODEL + BLOCK_DMODEL_PE),
+                seqlen_k,
+            )
+        if PRELOAD_V:
+            v = _load_fn(v_ptrs, k_offs_n, k_offs_k, seqlen_k, BLOCK_DMODEL)
+
+        # We start from end of seqlen_k so only the first iteration would need
+        # to be checked for padding if it is not a multiple of block_n
+        # TODO: This can be optimized to only be true for the padded block.
+        mask = tl.full([BLOCK_M, BLOCK_N], True, dtype=tl.int1)
+        if MASK_STEPS:
+            # If this is the last block / iteration, we want to
+            # mask if the sequence length is not a multiple of block size
+            # a solution is to always do BLOCK_M // BLOCK_N + 1 steps if not is_modulo_mn.
+            # last step might get wasted but that is okay. check if this masking works For
+            # that case.
+
+            # remove the old if condition
+            # if (start_n + BLOCK_N == block_max) and (n_extra_tokens != 0):
+            # Though this will unconditionally compute mask_partial at runtime,
+            # the causal for loop does not have the if-else block any more, which
+            # helps instruction scheduling and register pressure.
+            bound_cond = (start_n + BLOCK_N == block_max) and (n_extra_tokens != 0)
+            size_n = start_n + OFFS_N[None, :]
+            mask_partial = size_n < seqlen_k
+            mask = tl.where(bound_cond, mask_partial, mask)
+
+        # compute masks
+        q_mask = OFFS_M[:, None] < seqlen_q
+        k_mask = (start_n + tl.arange(0, BLOCK_N))[None, :] < seqlen_k
+        p_mask = q_mask & k_mask
+        qk_scale = SM_SCALE * RCP_LN2
+        # -- compute qk ----
+        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+        if HAS_PE:
+            qk = tl.dot(q_pe, k_pe, acc=qk)
+        qk = tl.dot(q, k, acc=qk)
+        if IS_FP8:
+            qk = qk * (qk_scale * descale_q * descale_k)
+        else:
+            qk = qk * qk_scale
+        if IS_CAUSAL:
+            causal_boundary = start_n + offs_n_causal
+            causal_mask = OFFS_M[:, None] >= causal_boundary[None, :]
+            mask = mask & causal_mask
+
+        if SLIDING_WINDOW > 0:
+            k_pos = start_n + tl.arange(0, BLOCK_N)
+            q_adj = OFFS_M + seqlen_k - seqlen_q
+            window_mask = k_pos[None, :] >= (q_adj[:, None] - SLIDING_WINDOW)
+            mask = mask & window_mask
+
+        qk = tl.where(mask, qk, float("-inf"))
+
+        if alibi_slope is not None:
+            # Compute the global position of each token within the sequence
+            global_m_positions = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+            global_n_positions = start_n + tl.arange(0, BLOCK_N)
+            alibi_block = _compute_alibi_block(
+                alibi_slope, seqlen_q, seqlen_k, global_m_positions, global_n_positions
+            )
+            qk += alibi_block * RCP_LN2
+        # get max scores so far
+        m_ij = tl.maximum(m_i, tl.max(qk, 1))
+
+        # Compute scaled QK and softmax probabilities
+        p = tl.math.exp2(qk - m_ij[:, None])
+
+        if SLIDING_WINDOW > 0:
+            # When all qk in a row are -inf (fully out-of-window block) and m_i was -inf,
+            # exp2(-inf - (-inf)) = NaN. Sanitize by zeroing masked elements.
+            p = tl.where(mask, p, 0.0)
+
+        # CAVEAT: Must update l_ij before applying dropout
+        l_ij = tl.sum(p, 1)
+        if ENABLE_DROPOUT:
+            rng_output = tl.rand(
+                philox_seed, philox_ptrs
+            )  # TODO: use tl.randint for better performance
+            dropout_mask = rng_output > dropout_p
+            tl.store(dropout_mask_ptrs, dropout_mask, mask=p_mask)
+
+            # return scores with negative values for dropped vals
+            sd_mask = tl.where(dropout_mask, p, -p)
+            tl.store(sd_mask_ptrs, sd_mask, mask=p_mask)
+
+            # apply dropout mask in place
+            p = tl.where(dropout_mask, p, 0.0)
+        elif RETURN_SCORES:
+            # NOTE: the returned score is not the same as the reference because we need to adjust as we find new maxes per block. We are not doing that
+            tl.store(sd_mask_ptrs, p, mask=p_mask)
+
+        # -- update output accumulator --
+        # alpha is an adjustment factor for acc and li as we loop and find new maxes
+        # store the diff in maxes to adjust acc and li as we discover new maxes
+        alpha = tl.math.exp2(m_i - m_ij)
+        if SLIDING_WINDOW > 0:
+            # When m_i == m_ij == -inf, exp2(-inf - (-inf)) = NaN. alpha should be 1.0
+            # (no rescaling needed since max didn't change).
+            alpha = tl.where(m_i == m_ij, 1.0, alpha)
+        acc = acc * alpha[:, None]
+        # -- update m_i and l_i
+        l_i = l_i * alpha + l_ij
+        # update m_i and l_i
+        m_i = m_ij
+        if not PRELOAD_V:
+            v = _load_fn(v_ptrs, k_offs_n, k_offs_k, seqlen_k, BLOCK_DMODEL)
+        if IS_FP8:
+            scale_p, descale_p = _compute_fp8_scaling_factors(p, FP8_MAX)
+            acc += (
+                tl.dot((p * scale_p).to(v.type.element_ty), v) * descale_p * descale_v
+            )
+        else:
+            acc = tl.dot(p.to(v.type.element_ty), v, acc=acc)
+
+        k_ptrs += BLOCK_N * stride_kn
+        if HAS_PE:
+            k_pe_ptrs += BLOCK_N * stride_kn
+        v_ptrs += BLOCK_N * stride_vk
+        if RETURN_SCORES:
+            sd_mask_ptrs += BLOCK_N * stride_sn
+
+        if ENABLE_DROPOUT:
+            dropout_mask_ptrs += BLOCK_N * stride_sn
+            philox_ptrs += BLOCK_N * stride_sn
+
+    return acc, l_i, m_i
+
+
+_attn_fwd_repr = make_kernel_repr(
+    "_attn_fwd",
+    [
+        "IS_CAUSAL",
+        "NUM_Q_HEADS",
+        "NUM_K_HEADS",
+        "BLOCK_M",
+        "BLOCK_N",
+        "BLOCK_DMODEL",
+        "RETURN_SCORES",
+        "ENABLE_DROPOUT",
+        "IS_FP8",
+        "VARLEN",
+        "NUM_XCD",
+        "USE_INT64_STRIDES",
+        "ENABLE_SINK",
+        "SLIDING_WINDOW",
+    ],
+)
+
+
+@triton.jit(repr=_attn_fwd_repr)
+def _attn_fwd(
+    q_ptr: torch.Tensor,
+    k_ptr: torch.Tensor,
+    v_ptr: torch.Tensor,
+    descale_q_ptr: torch.Tensor,
+    descale_k_ptr: torch.Tensor,
+    descale_v_ptr: torch.Tensor,
+    out_ptr: torch.Tensor,
+    alibi_slopes_ptr: torch.Tensor,
+    s_dmask_ptr: torch.Tensor,
+    dropout_mask_ptr: torch.Tensor,
+    softmax_lse_ptr: torch.Tensor,
+    sink_ptr: torch.Tensor,
+    stride_qz_in,
+    stride_qh_in,
+    stride_qm_in,
+    stride_qk_in,
+    stride_kz_in,
+    stride_kh_in,
+    stride_kn_in,
+    stride_kk_in,
+    stride_vz_in,
+    stride_vh_in,
+    stride_vn_in,
+    stride_vk_in,
+    stride_descale_q_z_in,
+    stride_descale_k_z_in,
+    stride_descale_v_z_in,
+    stride_oz_in,
+    stride_oh_in,
+    stride_om_in,
+    stride_on_in,
+    stride_alibi_z_in,
+    stride_alibi_h_in,
+    stride_sd_z_in,
+    stride_sd_h_in,
+    stride_sd_m_in,
+    stride_sd_n_in,
+    stride_lse_z_in,
+    stride_lse_h_in,
+    stride_lse_m_in,
+    sm_scale,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    dropout_p,
+    philox_seed,
+    philox_offset_base_in,
+    SEQLEN_Q,
+    SEQLEN_K,
+    IS_CAUSAL: tl.constexpr,
+    NUM_Q_HEADS: tl.constexpr,
+    NUM_K_HEADS: tl.constexpr,
+    PRELOAD_V: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_DMODEL_POW2: tl.constexpr,
+    BLOCK_DMODEL_PE: tl.constexpr,  # it's zero or a power of 2
+    RETURN_SCORES: tl.constexpr,
+    ENABLE_DROPOUT: tl.constexpr,
+    IS_FP8: tl.constexpr,
+    FP8_MAX: tl.constexpr,
+    VARLEN: tl.constexpr,
+    BATCH,
+    NUM_XCD: tl.constexpr,
+    USE_INT64_STRIDES: tl.constexpr,
+    ENABLE_SINK: tl.constexpr,
+    SLIDING_WINDOW: tl.constexpr,
+):
+    NUM_BLOCKS = (SEQLEN_Q + BLOCK_M - 1) // BLOCK_M
+    # calculate offsets
+    wid = tl.program_id(
+        0
+    )  # workgroup id ranging: 0,1,2,...., (BATCH * NUM_Q_HEADS * NUM_BLOCKS - 1)
+    # num blocks along seqlen
+
+    off_q_head = wid % NUM_Q_HEADS
+    off_q_head = remap_xcd(off_q_head, NUM_Q_HEADS, NUM_XCD)
+    start_m = (wid // NUM_Q_HEADS) % NUM_BLOCKS
+    off_z = (wid // (NUM_BLOCKS * NUM_Q_HEADS)) % BATCH
+
+    # offsets
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = tl.arange(0, BLOCK_N)
+    offs_d = tl.arange(0, BLOCK_DMODEL_POW2)
+    HAS_PE: tl.constexpr = BLOCK_DMODEL_PE > 0
+    if HAS_PE:
+        offs_pe = BLOCK_DMODEL + tl.arange(0, BLOCK_DMODEL_PE)
+
+    # NOTE:
+    # Workaround for int64 strides, In the absence of strides being int64, parts of the offset
+    # computation is done in 32 bit and overflows resulting in segfaults
+    # If input strides are defined as int64, it disables vectorized loads which drops perf
+    # If we define new strides as stride_x = stride_x_in.to(tl.int64), that does not work
+    # because strides are tl.constexpr and cannot be upcasted
+    # If we define new strides as stride_x: tl.int64 = stride_x_in, segfault remains
+    # The permanent solution is to enable upcasting of tl.constexpr
+    # In the meantime, the following workaround provides correctness and does not drop perf
+    if USE_INT64_STRIDES:
+        stride_qz = tl.cast(stride_qz_in, tl.int64)
+        stride_qh = tl.cast(stride_qh_in, tl.int64)
+        stride_qm = tl.cast(stride_qm_in, tl.int64)
+        stride_qk = tl.cast(stride_qk_in, tl.int64)
+        stride_kz = tl.cast(stride_kz_in, tl.int64)
+        stride_kh = tl.cast(stride_kh_in, tl.int64)
+        stride_kn = tl.cast(stride_kn_in, tl.int64)
+        stride_kk = tl.cast(stride_kk_in, tl.int64)
+        stride_vz = tl.cast(stride_vz_in, tl.int64)
+        stride_vh = tl.cast(stride_vh_in, tl.int64)
+        stride_vn = tl.cast(stride_vn_in, tl.int64)
+        stride_vk = tl.cast(stride_vk_in, tl.int64)
+        if IS_FP8:
+            stride_descale_q_z = tl.cast(stride_descale_q_z_in, tl.int64)
+            stride_descale_k_z = tl.cast(stride_descale_k_z_in, tl.int64)
+            stride_descale_v_z = tl.cast(stride_descale_v_z_in, tl.int64)
+        stride_oz = tl.cast(stride_oz_in, tl.int64)
+        stride_oh = tl.cast(stride_oh_in, tl.int64)
+        stride_om = tl.cast(stride_om_in, tl.int64)
+        stride_on = tl.cast(stride_on_in, tl.int64)
+        stride_alibi_z = tl.cast(stride_alibi_z_in, tl.int64)
+        stride_alibi_h = tl.cast(stride_alibi_h_in, tl.int64)
+
+        # NOTE: philox offset is need in dropout pointer calculations
+        philox_offset_base = tl.cast(philox_offset_base_in, tl.int64)
+        stride_sd_z = tl.cast(stride_sd_z_in, tl.int64)
+        stride_sd_h = tl.cast(stride_sd_h_in, tl.int64)
+        stride_sd_m = tl.cast(stride_sd_m_in, tl.int64)
+        stride_sd_n = tl.cast(stride_sd_n_in, tl.int64)
+        stride_lse_z = tl.cast(stride_lse_z_in, tl.int64)
+        stride_lse_h = tl.cast(stride_lse_h_in, tl.int64)
+        stride_lse_m = tl.cast(stride_lse_m_in, tl.int64)
+    else:
+        stride_qz = stride_qz_in
+        stride_qm = stride_qm_in
+        stride_qk = stride_qk_in
+        stride_qh = stride_qh_in
+        stride_kz = stride_kz_in
+        stride_kh = stride_kh_in
+        stride_kn = stride_kn_in
+        stride_kk = stride_kk_in
+        stride_vz = stride_vz_in
+        stride_vh = stride_vh_in
+        stride_vn = stride_vn_in
+        stride_vk = stride_vk_in
+        stride_descale_q_z = stride_descale_q_z_in
+        stride_descale_k_z = stride_descale_k_z_in
+        stride_descale_v_z = stride_descale_v_z_in
+        stride_oz = stride_oz_in
+        stride_oh = stride_oh_in
+        stride_om = stride_om_in
+        stride_on = stride_on_in
+        stride_alibi_z = stride_alibi_z_in
+        stride_alibi_h = stride_alibi_h_in
+        philox_offset_base = philox_offset_base_in
+        stride_sd_z = stride_sd_z_in
+        stride_sd_h = stride_sd_h_in
+        stride_sd_m = stride_sd_m_in
+        stride_sd_n = stride_sd_n_in
+        stride_lse_z = stride_lse_z_in
+        stride_lse_h = stride_lse_h_in
+        stride_lse_m = stride_lse_m_in
+
+    tl.assume(stride_qz_in >= 0)
+    tl.assume(stride_qh_in >= 0)
+    tl.assume(stride_qm_in >= 0)
+    tl.assume(stride_qk_in >= 0)
+    tl.assume(stride_kz_in >= 0)
+    tl.assume(stride_kh_in >= 0)
+    tl.assume(stride_kn_in >= 0)
+    tl.assume(stride_kk_in >= 0)
+    tl.assume(stride_vz_in >= 0)
+    tl.assume(stride_vh_in >= 0)
+    tl.assume(stride_vn_in >= 0)
+    tl.assume(stride_vk_in >= 0)
+    if IS_FP8:
+        tl.assume(stride_descale_q_z_in >= 0)
+        tl.assume(stride_descale_k_z_in >= 0)
+        tl.assume(stride_descale_v_z_in >= 0)
+        tl.assume(stride_oz_in >= 0)
+        tl.assume(stride_oh_in >= 0)
+        tl.assume(stride_om_in >= 0)
+        tl.assume(stride_on_in >= 0)
+        tl.assume(stride_alibi_z_in >= 0)
+        tl.assume(stride_alibi_h_in >= 0)
+    # NOTE: philox offset is need in dropout pointer calculations
+    tl.assume(philox_offset_base_in >= 0)
+    tl.assume(stride_sd_z_in >= 0)
+    tl.assume(stride_sd_h_in >= 0)
+    tl.assume(stride_sd_m_in >= 0)
+    tl.assume(stride_sd_n_in >= 0)
+    tl.assume(stride_lse_z_in >= 0)
+    tl.assume(stride_lse_h_in >= 0)
+    tl.assume(stride_lse_m_in >= 0)
+
+    if VARLEN:
+        cu_seqlens_q_start = tl.load(cu_seqlens_q + off_z)
+        cu_seqlens_q_end = tl.load(cu_seqlens_q + off_z + 1)
+
+        seqlen_q = cu_seqlens_q_end - cu_seqlens_q_start
+        # We have a one-size-fits-all grid in id(0). Some seqlens might be too
+        # small for all start_m so for those we return early.
+        if start_m * BLOCK_M > seqlen_q:
+            return
+        cu_seqlens_k_start = tl.load(cu_seqlens_k + off_z)
+        cu_seqlens_k_end = tl.load(cu_seqlens_k + off_z + 1)
+        seqlen_k = cu_seqlens_k_end - cu_seqlens_k_start
+    else:
+        cu_seqlens_q_start = 0
+        cu_seqlens_k_start = 0
+        seqlen_q = SEQLEN_Q
+        seqlen_k = SEQLEN_K
+
+    n_blocks = _cdiv_fn(seqlen_k, BLOCK_N)
+
+    # Now we compute whether we need to exit early due to causal masking.
+    # This is because for seqlen_q > seqlen_k, M rows of the attn scores
+    # are completely masked, resulting in 0s written to the output, and
+    # inf written to LSE. We don't need to do any GEMMs in this case.
+    # This block of code determines what N is, and if this WG is operating
+    # on those M rows.
+    if IS_CAUSAL:
+        # If seqlen_q == seqlen_k, the attn scores are a square matrix.
+        # If seqlen_q != seqlen_k, attn scores are rectangular which means
+        # the causal mask boundary is bottom right aligned, and ends at either
+        # the top edge (seqlen_q < seqlen_k) or left edge.
+
+        # This captures the decrease in n_blocks if we have a rectangular attn matrix
+        n_blocks_seqlen = _cdiv_fn(
+            (start_m + 1) * BLOCK_M + seqlen_k - seqlen_q, BLOCK_N
+        )
+
+        # This is what adjusts the block_max for the current WG, only
+        # if IS_CAUSAL. Otherwise we want to always iterate through all n_blocks
+        n_blocks = min(n_blocks, n_blocks_seqlen)
+
+        # If we have no blocks after adjusting for seqlen deltas, this WG is part of
+        # the blocks that are all 0. We exit early.
+        if n_blocks <= 0:
+            offs_out = (
+                off_z * stride_oz
+                + off_q_head * stride_oh
+                + cu_seqlens_q_start * stride_om
+                + offs_m[:, None] * stride_om
+                + offs_d[None, :] * stride_on
+            )
+            acc = tl.zeros([BLOCK_M, BLOCK_DMODEL_POW2], dtype=out_ptr.type.element_ty)
+            out_mask = (offs_m[:, None] < seqlen_q) & (offs_d[None, :] < BLOCK_DMODEL)
+            tl.store(out_ptr + offs_out, acc, mask=out_mask)
+
+            if softmax_lse_ptr is not None:
+                offs_lse = (
+                    off_z * stride_lse_z
+                    + off_q_head * stride_lse_h
+                    + cu_seqlens_q_start * stride_lse_m
+                    + offs_m * stride_lse_m
+                )
+                lse_mask = offs_m < SEQLEN_Q
+                lse = tl.full([BLOCK_M], value=0.0, dtype=tl.float32)
+                tl.store(softmax_lse_ptr + offs_lse, lse, mask=lse_mask)
+                # TODO: Should dropout and return encoded softmax be handled here too?
+
+            return
+
+    grp_sz: tl.constexpr = NUM_Q_HEADS // NUM_K_HEADS
+    if grp_sz != 1:  # Grouped Query Attention
+        off_k_head = off_q_head // grp_sz
+    else:
+        off_k_head = off_q_head
+
+    # q,k,v offsets
+    q_offs = (
+        off_z * stride_qz
+        + off_q_head * stride_qh
+        + cu_seqlens_q_start * stride_qm
+        + offs_m[:, None] * stride_qm
+        + offs_d[None, :] * stride_qk
+    )
+    q_ptrs = q_ptr + q_offs
+    if HAS_PE:
+        q_pe_offs = (
+            off_z * stride_qz
+            + off_q_head * stride_qh
+            + cu_seqlens_q_start * stride_qm
+            + offs_m[:, None] * stride_qm
+            + offs_pe[None, :] * stride_qk
+        )
+        q_pe_ptrs = q_ptr + q_pe_offs
+    else:
+        q_pe_ptrs = None
+
+    k_offs = (
+        off_z * stride_kz
+        + off_k_head * stride_kh
+        + cu_seqlens_k_start * stride_kn
+        + offs_d[:, None] * stride_kk
+        + offs_n[None, :] * stride_kn
+    )
+    k_ptrs = k_ptr + k_offs
+    if HAS_PE:
+        k_pe_offs = (
+            off_z * stride_kz
+            + off_k_head * stride_kh
+            + cu_seqlens_k_start * stride_kn
+            + offs_pe[:, None] * stride_kk
+            + offs_n[None, :] * stride_kn
+        )
+        k_pe_ptrs = k_ptr + k_pe_offs
+    else:
+        k_pe_ptrs = None
+
+    v_offs = (
+        off_z * stride_vz
+        + off_k_head * stride_vh
+        + cu_seqlens_k_start * stride_vn
+        + offs_n[:, None] * stride_vn
+        + offs_d[None, :] * stride_vk
+    )
+    v_ptrs = v_ptr + v_offs
+
+    # alibi slopes
+    if alibi_slopes_ptr is not None:
+        alibi_offs = off_z * stride_alibi_z + off_q_head * stride_alibi_h
+        alibi_slope = tl.load(alibi_slopes_ptr + alibi_offs)
+    else:
+        alibi_slope = None
+
+    # s_dmask (return_scores)
+    if s_dmask_ptr is not None:
+        s_dmask_offs = (
+            off_z * stride_sd_z
+            + off_q_head * stride_sd_h
+            + offs_m[:, None] * stride_sd_m
+            + offs_n[None, :] * stride_sd_n
+        )
+        s_dmask_ptrs = s_dmask_ptr + s_dmask_offs
+    else:
+        s_dmask_ptrs = None
+
+    # dropout
+    if dropout_mask_ptr is not None:
+        dropout_mask_offs = (
+            off_z * stride_sd_z
+            + off_q_head * stride_sd_h
+            + offs_m[:, None] * stride_sd_m
+            + offs_n[None, :] * stride_sd_n
+        )
+        dropout_mask_ptrs = dropout_mask_ptr + dropout_mask_offs
+        philox_ptrs = (
+            philox_offset_base
+            + off_z * stride_sd_z
+            + off_q_head * stride_sd_h
+            + offs_m[:, None] * stride_sd_m
+            + offs_n[None, :] * stride_sd_n
+        )
+    else:
+        dropout_mask_ptrs = None
+        philox_ptrs = None
+
+    if ENABLE_SINK:
+        RCP_LN2: tl.constexpr = 1.4426950408889634
+        m_i_value = tl.load(sink_ptr + off_q_head).to(tl.float32) * RCP_LN2
+    else:
+        m_i_value = float("-inf")
+
+    m_i = tl.full([BLOCK_M], m_i_value, dtype=tl.float32)
+    l_i = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
+    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL_POW2], dtype=tl.float32)
+    if BLOCK_DMODEL == BLOCK_DMODEL_POW2:
+        q_mask = offs_m[:, None] < seqlen_q
+    else:
+        q_mask = (offs_m[:, None] < seqlen_q) & (offs_d[None, :] < BLOCK_DMODEL)
+
+    if BLOCK_M >= NUM_Q_HEADS:
+        q_cache_mod: tl.constexpr = ".cg"
+    else:
+        q_cache_mod: tl.constexpr = ""
+
+    if HAS_PE:
+        q_pe = tl.load(q_pe_ptrs, mask=q_mask, other=0.0, cache_modifier=q_cache_mod)
+    else:
+        q_pe = None
+    q = tl.load(q_ptrs, mask=q_mask, other=0.0, cache_modifier=q_cache_mod)
+    if IS_FP8:
+        descale_q = tl.load(descale_q_ptr + off_z * stride_descale_q_z + off_q_head)
+        descale_k = tl.load(descale_k_ptr + off_z * stride_descale_k_z + off_k_head)
+        descale_v = tl.load(descale_v_ptr + off_z * stride_descale_v_z + off_k_head)
+    else:
+        descale_q, descale_k, descale_v = 1.0, 1.0, 1.0
+
+    n_extra_tokens = 0
+    if seqlen_k < BLOCK_N:
+        n_extra_tokens = BLOCK_N - seqlen_k
+    elif seqlen_k % BLOCK_N:
+        n_extra_tokens = seqlen_k % BLOCK_N
+
+    # if CAUSAL, then determine masked_blocks and full blocks
+    # Here we compute how many full and masked blocks we have.
+    padded_block_k = n_extra_tokens != 0
+    is_modulo_mn = not padded_block_k and (seqlen_q % BLOCK_M == 0)
+    skipped_blocks = 0
+    if SLIDING_WINDOW > 0:
+        # Skip K blocks that are fully left of the earliest key position
+        # reachable by this Q block. The first retained block can still be
+        # partially outside the window, so we keep the per-element mask below.
+        window_start_n = start_m * BLOCK_M + seqlen_k - seqlen_q - SLIDING_WINDOW
+        skipped_blocks = tl.maximum(window_start_n, 0) // BLOCK_N
+        skipped_blocks = tl.minimum(skipped_blocks, n_blocks)
+    if IS_CAUSAL:
+        # There are always at least BLOCK_M // BLOCK_N masked blocks.
+        # Additionally there might be one more due to dissimilar seqlens.
+        masked_blocks = BLOCK_M // BLOCK_N + (not is_modulo_mn)
+    else:
+        # Padding on Q does not need to be masked in the FA loop.
+        masked_blocks = padded_block_k
+    # if IS_CAUSAL, not is_modulo_mn does not always result in an additional block.
+    # In this case we might exceed n_blocks so pick the min.
+    visible_blocks = n_blocks - skipped_blocks
+    masked_blocks = min(masked_blocks, visible_blocks)
+    n_full_blocks = visible_blocks - masked_blocks
+    block_min = skipped_blocks * BLOCK_N
+    block_max = n_blocks * BLOCK_N
+    if skipped_blocks > 0:
+        k_ptrs += skipped_blocks * BLOCK_N * stride_kn
+        if HAS_PE:
+            k_pe_ptrs += skipped_blocks * BLOCK_N * stride_kn
+        v_ptrs += skipped_blocks * BLOCK_N * stride_vn
+        if RETURN_SCORES:
+            s_dmask_ptrs += skipped_blocks * BLOCK_N * stride_sd_n
+        if ENABLE_DROPOUT:
+            dropout_mask_ptrs += skipped_blocks * BLOCK_N * stride_sd_n
+            philox_ptrs += skipped_blocks * BLOCK_N * stride_sd_n
+    # Compute for full blocks. Here we set causal to false regardless of its actual
+    # value because there is no masking. Similarly we do not need padding.
+    if n_full_blocks > 0:
+        block_max = block_min + n_full_blocks * BLOCK_N
+        acc, l_i, m_i = _attn_fwd_inner(
+            acc,
+            l_i,
+            m_i,
+            q,
+            q_pe,
+            k_ptrs,
+            k_pe_ptrs,
+            v_ptrs,
+            stride_kn,
+            stride_vn,
+            stride_sd_n,
+            start_m,
+            seqlen_k,
+            seqlen_q,
+            dropout_p,
+            s_dmask_ptrs,
+            dropout_mask_ptrs,
+            philox_seed,
+            philox_ptrs,
+            block_min,
+            block_max,
+            0,
+            0,
+            0,
+            alibi_slope,
+            descale_q,
+            descale_k,
+            descale_v,
+            offs_m,
+            offs_n,
+            PRELOAD_V,
+            BLOCK_M,
+            BLOCK_N,
+            BLOCK_DMODEL,
+            BLOCK_DMODEL_POW2,
+            BLOCK_DMODEL_PE,
+            sm_scale,
+            False,
+            MASK_STEPS=False,
+            ENABLE_DROPOUT=ENABLE_DROPOUT,
+            RETURN_SCORES=RETURN_SCORES,
+            PADDED_HEAD=BLOCK_DMODEL != BLOCK_DMODEL_POW2,
+            IS_FP8=IS_FP8,
+            FP8_MAX=FP8_MAX,
+            ENABLE_PIPELINING=True,
+            SLIDING_WINDOW=SLIDING_WINDOW,
+        )
+        block_min = block_max
+        block_max = n_blocks * BLOCK_N
+
+    # Remaining blocks, if any, are full / not masked.
+    if masked_blocks > 0:
+        if IS_CAUSAL:
+            offs_n_causal = offs_n + (seqlen_q - seqlen_k)
+        else:
+            offs_n_causal = 0
+        k_ptrs += n_full_blocks * BLOCK_N * stride_kn
+        if HAS_PE:
+            k_pe_ptrs += n_full_blocks * BLOCK_N * stride_kn
+        v_ptrs += n_full_blocks * BLOCK_N * stride_vn
+        if RETURN_SCORES:
+            s_dmask_ptrs += n_full_blocks * BLOCK_N * stride_sd_n
+        if ENABLE_DROPOUT:
+            dropout_mask_ptrs += n_full_blocks * BLOCK_N * stride_sd_n
+        acc, l_i, m_i = _attn_fwd_inner(
+            acc,
+            l_i,
+            m_i,
+            q,
+            q_pe,
+            k_ptrs,
+            k_pe_ptrs,
+            v_ptrs,
+            stride_kn,
+            stride_vn,
+            stride_sd_n,
+            start_m,
+            seqlen_k,
+            seqlen_q,
+            dropout_p,
+            s_dmask_ptrs,
+            dropout_mask_ptrs,
+            philox_seed,
+            philox_ptrs,
+            block_min,
+            block_max,
+            offs_n_causal,
+            masked_blocks,
+            n_extra_tokens,
+            alibi_slope,
+            descale_q,
+            descale_k,
+            descale_v,
+            offs_m,
+            offs_n,
+            PRELOAD_V,
+            BLOCK_M,
+            BLOCK_N,
+            BLOCK_DMODEL,
+            BLOCK_DMODEL_POW2,
+            BLOCK_DMODEL_PE,
+            sm_scale,
+            IS_CAUSAL,
+            MASK_STEPS=True,
+            ENABLE_DROPOUT=ENABLE_DROPOUT,
+            RETURN_SCORES=RETURN_SCORES,
+            PADDED_HEAD=BLOCK_DMODEL != BLOCK_DMODEL_POW2,
+            IS_FP8=IS_FP8,
+            FP8_MAX=FP8_MAX,
+            ENABLE_PIPELINING=False,
+            SLIDING_WINDOW=SLIDING_WINDOW,
+        )
+    # epilogue
+    # This helps the compiler do Newton Raphson on l_i vs on acc which is much larger.
+    l_recip = 1 / l_i[:, None]
+    acc = acc * l_recip
+    if ENABLE_DROPOUT:
+        dropout_scale = 1 / (1 - dropout_p)
+        acc = acc * dropout_scale
+    # If seqlen_q > seqlen_k but the delta is not a multiple of BLOCK_M,
+    # then we have one block with a row of all NaNs which come from computing
+    # softmax over a row of all -infs (-inf - inf = NaN). We check for that here
+    # and store 0s where there are NaNs as these rows should've been zeroed out.
+    end_m_idx = (start_m + 1) * BLOCK_M
+    start_m_idx = start_m * BLOCK_M
+    causal_start_idx = seqlen_q - seqlen_k
+    if IS_CAUSAL:
+        if causal_start_idx > start_m_idx and causal_start_idx < end_m_idx:
+            out_mask_boundary = tl.full(
+                (BLOCK_DMODEL_POW2,), causal_start_idx, dtype=tl.int32
+            )
+            mask_m_offsets = start_m_idx + tl.arange(0, BLOCK_M)
+            out_ptrs_mask = mask_m_offsets[:, None] >= out_mask_boundary[None, :]
+            z = 0.0
+            acc = tl.where(out_ptrs_mask, acc, z.to(acc.type.element_ty))
+
+    # write back LSE(Log Sum Exponents), the log of the normalization constant
+    overflow_size = end_m_idx - seqlen_q
+    if softmax_lse_ptr is not None:
+        LN2: tl.constexpr = 0.6931471824645996
+        # compute log-sum-exp in base 2 units
+        softmax_lse = m_i + tl.math.log2(l_i)
+        # convert back to natural units
+        softmax_lse *= LN2
+
+        if IS_CAUSAL:
+            # zero out nans caused by -infs when doing causal
+            lse_causal_mask = (start_m_idx + tl.arange(0, BLOCK_M)) < causal_start_idx
+            softmax_lse = tl.where(lse_causal_mask, 0.0, softmax_lse)
+
+        # If seqlen_q not multiple of BLOCK_M, we need to mask out the last few rows.
+        # This is only true for the last M block. For others, overflow_size will be -ve
+        offs_lse = (
+            off_z * stride_lse_z
+            + off_q_head * stride_lse_h
+            + cu_seqlens_q_start * stride_lse_m
+            + offs_m * stride_lse_m
+        )
+        if overflow_size > 0:
+            boundary = tl.full((BLOCK_M,), BLOCK_M - overflow_size, dtype=tl.int32)
+            lse_mask = tl.arange(0, BLOCK_M) < boundary
+            tl.store(
+                softmax_lse_ptr + offs_lse, softmax_lse, mask=lse_mask
+            )  # the log of the normalization constant
+        else:
+            tl.store(
+                softmax_lse_ptr + offs_lse, softmax_lse
+            )  # the log of the normalization constant
+
+    # write back O
+    offs_out = (
+        off_z * stride_oz
+        + off_q_head * stride_oh
+        + cu_seqlens_q_start * stride_om
+        + offs_m[:, None] * stride_om
+        + offs_d[None, :] * stride_on
+    )
+    out_mask = tl.full([BLOCK_M, 1], 1, dtype=tl.int1)
+    if overflow_size > 0:
+        out_mask = out_mask & (offs_m[:, None] < seqlen_q)
+    if BLOCK_DMODEL != BLOCK_DMODEL_POW2:
+        out_mask = out_mask & (offs_d[None, :] < BLOCK_DMODEL)
+    op = acc.to(out_ptr.dtype.element_ty)
+    tl.store(out_ptr + offs_out, op, mask=out_mask)
+
+
+@functools.lru_cache(maxsize=1024)
+def _get_config(
+    enable_dropout: bool,
+    dtype: torch.dtype,
+    has_pe: bool = False,
+):
+    if not hasattr(_get_config, "_config_dict"):
+        dev = arch_info.get_arch()
+        _get_config._config_dict = {}
+        fpath = f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHA-DEFAULT.json"
+        with open(fpath, "r") as file:
+            config = json.load(file)
+        _get_config._config_dict["default"] = config
+    fwd_cfg = _get_config._config_dict["default"]["fwd"]
+    has_dropout_or_fp32 = enable_dropout or dtype == torch.float32
+    # TODO: pe + dropout is not tuned
+    if has_pe and has_dropout_or_fp32 and "pe_dropout_or_fp32" in fwd_cfg:
+        return fwd_cfg["pe_dropout_or_fp32"]
+    elif has_pe and "pe" in fwd_cfg:
+        return fwd_cfg["pe"]
+    elif enable_dropout or dtype == torch.float32:
+        return fwd_cfg["dropout_or_fp32"]
+    else:
+        return fwd_cfg["default"]

build/torch-rocm/_kernels/mha_fused_bwd.py

@@ -0,0 +1,1073 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+import functools
+import json
+import triton
+import triton.language as tl
+
+
+from ..utils._triton import arch_info
+from ..utils.core import AITER_TRITON_CONFIGS_PATH
+from ..utils._triton.pid_preprocessing import remap_xcd
+from ..utils._triton.mha_kernel_utils import _compute_fp8_scaling_factors
+from ..utils._triton.kernel_repr import make_kernel_repr
+
+# This function computes delta given output Out and gradient DO
+# Here is the I/O shape:
+# Out: (batch, nhead_q, max_seqlens_q, headDim)
+# DO: (batch, nhead_q, max_seqlens_q, headDim)
+# Delta: (batch, nheads_q, max_seqlens_q), same as softmax_lse defined at
+_bwd_preprocess_repr = make_kernel_repr(
+    "_bwd_preprocess",
+    [
+        "BLOCK_M",
+        "BLOCK_D_MODEL",
+        "IS_VARLEN",
+        "IS_FP8",
+    ],
+)
+
+
+@triton.jit(repr=_bwd_preprocess_repr)
+def _bwd_preprocess(
+    o_ptr,
+    do_ptr,  # noqa: E741
+    delta_ptr,
+    stride_o_b,
+    stride_o_h,
+    stride_o_m,
+    stride_o_k,
+    stride_do_b,
+    stride_do_h,
+    stride_do_m,
+    stride_do_k,
+    stride_delta_b,
+    stride_delta_h,
+    stride_delta_m,
+    stride_descale_do_z,
+    cu_seqlens_q,
+    max_seqlen_q,
+    descale_do_ptr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_D_MODEL: tl.constexpr,
+    BLOCK_D_MODEL_POW2: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    IS_FP8: tl.constexpr,
+):
+    pid_m = tl.program_id(0)  # seqlen
+    bid = tl.program_id(1)  # batch
+    hid = tl.program_id(2)  # head
+
+    # Handle varlen
+    q_start = 0
+    seqlen_q = max_seqlen_q
+    if IS_VARLEN:
+        q_start = tl.load(cu_seqlens_q + bid)
+        q_end = tl.load(cu_seqlens_q + bid + 1)
+        seqlen_q = q_end - q_start
+    else:
+        q_start = 0
+        seqlen_q = max_seqlen_q
+
+    # Compute offsets
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_k = tl.arange(0, BLOCK_D_MODEL_POW2)
+
+    # O and DO may have different strides (e.g. BSHD vs SBHD memory layout),
+    # so address each with its own strides.
+    offs_o = (
+        bid * stride_o_b
+        + hid * stride_o_h
+        + q_start * stride_o_m
+        + offs_m[:, None] * stride_o_m
+        + offs_k[None, :] * stride_o_k
+    )
+    offs_do = (
+        bid * stride_do_b
+        + hid * stride_do_h
+        + q_start * stride_do_m
+        + offs_m[:, None] * stride_do_m
+        + offs_k[None, :] * stride_do_k
+    )
+
+    # create masks
+    mask_m = offs_m < seqlen_q
+    mask = mask_m[:, None]
+    PADDED_HEAD: tl.constexpr = BLOCK_D_MODEL != BLOCK_D_MODEL_POW2
+    if PADDED_HEAD:
+        mask &= offs_k[None, :] < BLOCK_D_MODEL
+
+    # load [BLOCK_M, BLOCK_D_MODEL_POW2]
+    o = tl.load(o_ptr + offs_o, mask=mask, other=0.0)
+    do = tl.load(do_ptr + offs_do, mask=mask, other=0.0)
+
+    # compute and write-back to delta
+    if IS_FP8:
+        descale_do = tl.load(descale_do_ptr + bid * stride_descale_do_z + hid)
+
+        # NOTE: do is in the fp8 range and o is not in fp8
+        delta = tl.sum(o.to(tl.float32) * (do.to(tl.float32) * descale_do), axis=1)
+    else:
+        delta = tl.sum(o.to(tl.float32) * do.to(tl.float32), axis=1)
+
+    offs_delta = (
+        bid * stride_delta_b
+        + hid * stride_delta_h
+        + q_start * stride_delta_m
+        + offs_m * stride_delta_m
+    )
+    tl.store(delta_ptr + offs_delta, delta, mask=mask_m)
+
+
+@triton.jit
+def _bwd_dkdvdq_inner(
+    dk,
+    dv,
+    Q,
+    k,
+    v,
+    DO,
+    DQ,
+    M,
+    D,
+    sm_scale,
+    stride_q_m,
+    stride_q_k,
+    stride_dq_m,
+    stride_dq_k,
+    stride_do_m,
+    stride_do_k,
+    stride_dropout_m,
+    stride_dropout_n,
+    stride_deltam,
+    dropout_p,
+    philox_seed,
+    batch_philox_offset,
+    dropout_offset,
+    seqlen_q,
+    seqlen_k,
+    start_n,
+    start_m,
+    num_steps,
+    descale_q,
+    descale_k,
+    descale_v,
+    descale_do,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_D_MODEL: tl.constexpr,
+    BLOCK_D_MODEL_POW2: tl.constexpr,
+    MASK: tl.constexpr,
+    ENABLE_DROPOUT: tl.constexpr,
+    IS_FP8: tl.constexpr,
+    FP8_MAX: tl.constexpr,
+    workgroup_id,
+):
+    tl.assume(stride_q_m >= 0)
+    tl.assume(stride_q_k >= 0)
+    tl.assume(stride_dq_m >= 0)
+    tl.assume(stride_dq_k >= 0)
+    tl.assume(stride_do_m >= 0)
+    tl.assume(stride_do_k >= 0)
+    tl.assume(stride_deltam >= 0)
+
+    PADDED_HEAD: tl.constexpr = BLOCK_D_MODEL != BLOCK_D_MODEL_POW2
+    delta_qk = seqlen_q - seqlen_k
+    offs_m = start_m + tl.arange(0, BLOCK_M)
+    offs_n = start_n + tl.arange(0, BLOCK_N)
+    offs_k = tl.arange(0, BLOCK_D_MODEL_POW2)
+
+    # mask to make sure not OOB of seqlen_q
+    mask_n = offs_n < seqlen_k
+
+    qT_ptrs_start = (
+        Q + offs_m[None, :] * stride_q_m + offs_k[:, None] * stride_q_k
+    )  # [BLOCK_D_MODEL_POW2, BLOCK_M]
+    dq_ptrs_start = (
+        DQ + offs_m[:, None] * stride_dq_m + offs_k[None, :] * stride_dq_k
+    )  # [BLOCK_M, BLOCK_D_MODEL_POW2]
+
+    do_ptrs_start = DO + offs_m[:, None] * stride_do_m + offs_k[None, :] * stride_do_k
+    curr_m = start_m
+    step_m = BLOCK_M
+    curr_philox_offset = batch_philox_offset
+
+    # Iterate over blocks(BLOCK_M size) of Q while calculating
+    # a fixed block(BLOCK_N) of dk and dv. Note, during backward
+    # pass P has to be recomputed. However, this kernel computes
+    # dV and dK, so we compute we need P^T and S^T. See backward pass
+    # equations
+    #
+    # From Flash Attention Paper:
+    # ForwardPass: S = QkT, P=softmax(S), O=PV
+    #
+    # BackwardPass equations
+    # dV = P^TdO
+    # dP = dOV^T
+    # dS = dsoftmax(dP)
+    # dQ = dSK
+    # dK = QdS^T
+
+    for iter in range(num_steps):
+        # Permute the iteration order to reduce the probability that concurrent workgroups (that share the same q head idx and batch idx) are at the same iteration
+        blk_idx = (iter + workgroup_id) % num_steps
+
+        curr_m = start_m + blk_idx * step_m
+        qT_ptrs = qT_ptrs_start + blk_idx * step_m * stride_q_m
+        dq_ptrs = dq_ptrs_start + blk_idx * step_m * stride_dq_m
+        do_ptrs = do_ptrs_start + blk_idx * step_m * stride_do_m
+
+        offs_m = curr_m + tl.arange(0, BLOCK_M)
+        mask_m = offs_m < seqlen_q
+        mask_qT = mask_m[None, :]
+        mask_do = mask_m[:, None]
+        mask_nm = mask_n[:, None] & (offs_m[None, :] < seqlen_q)
+
+        if PADDED_HEAD:
+            mask_qT &= offs_k[:, None] < BLOCK_D_MODEL
+            mask_do &= offs_k[None, :] < BLOCK_D_MODEL
+
+        # load qT
+        qT = tl.load(qT_ptrs, mask=mask_qT, other=0.0)
+
+        # dropout
+        if ENABLE_DROPOUT:
+            # NOTE: dropout is transposed because it is used to mask pT
+            philox_offs = (
+                curr_philox_offset
+                + offs_m[None, :] * stride_dropout_m
+                + offs_n[:, None] * stride_dropout_n
+            )
+            rand_vals = tl.rand(philox_seed, philox_offs)
+            dropout_mask = rand_vals > dropout_p
+            dropout_scale = 1.0 / (1 - dropout_p)
+
+        # Load M
+        m = tl.load(M + offs_m * stride_deltam, mask=mask_m, other=0.0)
+
+        # Compute qkT
+        if IS_FP8:
+            qkT = tl.dot(k, qT) * descale_q * descale_k
+        else:
+            qkT = tl.dot(k, qT)
+
+        # Compute pT(use m and also apply sm_scale)
+        pT = tl.math.exp(qkT * sm_scale - m[None, :])
+
+        if MASK:
+            causal_mask = (offs_m[None, :] - delta_qk) >= (offs_n[:, None])
+            mask = causal_mask & mask_nm
+            pT = tl.where(mask, pT, 0.0)
+
+        # load DO
+        do = tl.load(do_ptrs, mask=mask_do, other=0.0)
+
+        # dV
+        if ENABLE_DROPOUT:
+            pT_dropout = tl.where(dropout_mask, pT, 0.0) * dropout_scale
+            if IS_FP8:
+                scale_p_dropout, descale_p_dropout = _compute_fp8_scaling_factors(
+                    pT_dropout, FP8_MAX
+                )
+                dv += (
+                    tl.dot((pT_dropout * scale_p_dropout).to(do.type.element_ty), do)
+                    * descale_p_dropout
+                    * descale_do
+                )
+            else:
+                dv = tl.dot(pT_dropout.to(do.type.element_ty), do, acc=dv)
+        else:
+            if IS_FP8:
+                scale_pT, descale_pT = _compute_fp8_scaling_factors(pT, FP8_MAX)
+                dv += (
+                    tl.dot((pT * scale_pT).to(do.type.element_ty), do)
+                    * descale_pT
+                    * descale_do
+                )
+            else:
+                dv = tl.dot(pT.to(do.type.element_ty), do, acc=dv)
+
+        # Load delta
+        Di = tl.load(D + offs_m * stride_deltam, mask=mask_m)
+
+        # Compute dP and dS
+        if IS_FP8:
+            dpT = tl.dot(v, tl.trans(do)) * descale_v * descale_do
+        else:
+            dpT = tl.dot(v, tl.trans(do))
+
+        if ENABLE_DROPOUT:
+            dpT = tl.where(dropout_mask, dpT, 0.0) * dropout_scale
+
+        delta_i = Di[None, :]
+        dsT = pT * (dpT - delta_i)
+
+        # compute dk
+        if IS_FP8:
+            scale_dsT, descale_dsT = _compute_fp8_scaling_factors(dsT, FP8_MAX)
+            dk += (
+                tl.dot((dsT * scale_dsT).to(qT.type.element_ty), tl.trans(qT))
+                * descale_dsT
+                * descale_q
+            )
+        else:
+            dk = tl.dot(dsT.to(qT.type.element_ty), tl.trans(qT), acc=dk)
+
+        # We can compute the dq_partial here and do a atomic add to the correct memory location
+        # NOTE: Possible problems with the atomic add: contention, is inside a loop which has achieved bad perf before
+        # (BLOCK_M, BLOCK_N) x (BLOCK_N, D)
+        if IS_FP8:
+            dq_partial = (
+                tl.dot((dsT * scale_dsT).to(k.dtype).T, k) * descale_dsT * descale_k
+            )
+        else:
+            dq_partial = tl.dot(dsT.to(k.dtype).T, k)
+        tl.atomic_add(
+            dq_ptrs,
+            dq_partial * sm_scale,
+            mask=mask_m[:, None] & (offs_k[None, :] < BLOCK_D_MODEL),
+            sem="relaxed",
+        )
+
+    return dk, dv
+
+
+_bwd_kernel_dkdvdq_causal_repr = make_kernel_repr(
+    "_bwd_kernel_dkdvdq_causal",
+    [
+        "NUM_Q_HEADS",
+        "NUM_K_HEADS",
+        "BLOCK_M",
+        "BLOCK_N",
+        "BLK_SLICE_FACTOR",
+        "BLOCK_D_MODEL",
+        "ENABLE_DROPOUT",
+        "IS_VARLEN",
+        "IS_FP8",
+        "USE_INT64_STRIDES",
+        "NUM_XCD",
+    ],
+)
+
+
+@triton.jit(repr=_bwd_kernel_dkdvdq_causal_repr)
+def _bwd_kernel_dkdvdq_causal(
+    q_ptr,
+    k_ptr,
+    v_ptr,
+    sm_scale,
+    do_ptr,
+    dk_ptr,
+    dv_ptr,
+    dq_ptr,
+    m_ptr,
+    delta_ptr,
+    stride_q_b_in,
+    stride_q_h_in,
+    stride_q_m_in,
+    stride_q_k_in,
+    stride_k_b_in,
+    stride_k_h_in,
+    stride_k_n_in,
+    stride_k_k_in,
+    stride_v_b_in,
+    stride_v_h_in,
+    stride_v_n_in,
+    stride_v_k_in,
+    stride_dk_b_in,
+    stride_dk_h_in,
+    stride_dk_n_in,
+    stride_dk_k_in,
+    stride_dq_b_in,
+    stride_dq_h_in,
+    stride_dq_m_in,
+    stride_dq_k_in,
+    stride_delta_b_in,
+    stride_delta_h_in,
+    stride_delta_m_in,
+    stride_do_b_in,
+    stride_do_h_in,
+    stride_do_m_in,
+    stride_do_k_in,
+    stride_dropout_b_in,
+    stride_dropout_h_in,
+    stride_dropout_m_in,
+    stride_dropout_n_in,
+    stride_descale_q_z_in,
+    stride_descale_k_z_in,
+    stride_descale_v_z_in,
+    stride_descale_do_z_in,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q,
+    max_seqlen_k,
+    dropout_mask,
+    dropout_p,
+    philox_seed,
+    philox_offset_base_in,
+    descale_q_ptr,
+    descale_k_ptr,
+    descale_v_ptr,
+    descale_do_ptr,
+    NUM_Q_HEADS: tl.constexpr,
+    NUM_K_HEADS: tl.constexpr,
+    BATCH,
+    NUM_K_PIDS,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLK_SLICE_FACTOR: tl.constexpr,
+    BLOCK_D_MODEL: tl.constexpr,
+    BLOCK_D_MODEL_POW2: tl.constexpr,
+    ENABLE_DROPOUT: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    IS_FP8: tl.constexpr,
+    FP8_MAX: tl.constexpr,
+    USE_INT64_STRIDES: tl.constexpr,
+    NUM_XCD: tl.constexpr,
+):
+    if USE_INT64_STRIDES:
+        stride_q_b = tl.cast(stride_q_b_in, tl.int64)
+        stride_q_h = tl.cast(stride_q_h_in, tl.int64)
+        stride_q_m = tl.cast(stride_q_m_in, tl.int64)
+        stride_q_k = tl.cast(stride_q_k_in, tl.int64)
+        stride_k_b = tl.cast(stride_k_b_in, tl.int64)
+        stride_k_h = tl.cast(stride_k_h_in, tl.int64)
+        stride_k_n = tl.cast(stride_k_n_in, tl.int64)
+        stride_k_k = tl.cast(stride_k_k_in, tl.int64)
+        stride_v_b = tl.cast(stride_v_b_in, tl.int64)
+        stride_v_h = tl.cast(stride_v_h_in, tl.int64)
+        stride_v_n = tl.cast(stride_v_n_in, tl.int64)
+        stride_v_k = tl.cast(stride_v_k_in, tl.int64)
+        stride_dk_b = tl.cast(stride_dk_b_in, tl.int64)
+        stride_dk_h = tl.cast(stride_dk_h_in, tl.int64)
+        stride_dk_n = tl.cast(stride_dk_n_in, tl.int64)
+        stride_dk_k = tl.cast(stride_dk_k_in, tl.int64)
+        stride_dq_b = tl.cast(stride_dq_b_in, tl.int64)
+        stride_dq_h = tl.cast(stride_dq_h_in, tl.int64)
+        stride_dq_m = tl.cast(stride_dq_m_in, tl.int64)
+        stride_dq_k = tl.cast(stride_dq_k_in, tl.int64)
+        stride_delta_b = tl.cast(stride_delta_b_in, tl.int64)
+        stride_delta_h = tl.cast(stride_delta_h_in, tl.int64)
+        stride_delta_m = tl.cast(stride_delta_m_in, tl.int64)
+        stride_do_b = tl.cast(stride_do_b_in, tl.int64)
+        stride_do_h = tl.cast(stride_do_h_in, tl.int64)
+        stride_do_m = tl.cast(stride_do_m_in, tl.int64)
+        stride_do_k = tl.cast(stride_do_k_in, tl.int64)
+        stride_dropout_b = tl.cast(stride_dropout_b_in, tl.int64)
+        stride_dropout_h = tl.cast(stride_dropout_h_in, tl.int64)
+        stride_dropout_m = tl.cast(stride_dropout_m_in, tl.int64)
+        stride_dropout_n = tl.cast(stride_dropout_n_in, tl.int64)
+        philox_offset_base = tl.cast(philox_offset_base_in, tl.int64)
+        if IS_FP8:
+            stride_descale_q_z = tl.cast(stride_descale_q_z_in, tl.int64)
+            stride_descale_k_z = tl.cast(stride_descale_k_z_in, tl.int64)
+            stride_descale_v_z = tl.cast(stride_descale_v_z_in, tl.int64)
+            stride_descale_do_z = tl.cast(stride_descale_do_z_in, tl.int64)
+    else:
+        stride_q_b = stride_q_b_in
+        stride_q_h = stride_q_h_in
+        stride_q_m = stride_q_m_in
+        stride_q_k = stride_q_k_in
+        stride_k_b = stride_k_b_in
+        stride_k_h = stride_k_h_in
+        stride_k_n = stride_k_n_in
+        stride_k_k = stride_k_k_in
+        stride_v_b = stride_v_b_in
+        stride_v_h = stride_v_h_in
+        stride_v_n = stride_v_n_in
+        stride_v_k = stride_v_k_in
+        stride_dk_b = stride_dk_b_in
+        stride_dk_h = stride_dk_h_in
+        stride_dk_n = stride_dk_n_in
+        stride_dk_k = stride_dk_k_in
+        stride_dq_b = stride_dq_b_in
+        stride_dq_h = stride_dq_h_in
+        stride_dq_m = stride_dq_m_in
+        stride_dq_k = stride_dq_k_in
+        stride_delta_b = stride_delta_b_in
+        stride_delta_h = stride_delta_h_in
+        stride_delta_m = stride_delta_m_in
+        stride_do_b = stride_do_b_in
+        stride_do_h = stride_do_h_in
+        stride_do_m = stride_do_m_in
+        stride_do_k = stride_do_k_in
+        stride_dropout_b = stride_dropout_b_in
+        stride_dropout_h = stride_dropout_h_in
+        stride_dropout_m = stride_dropout_m_in
+        stride_dropout_n = stride_dropout_n_in
+        philox_offset_base = philox_offset_base_in
+        stride_descale_q_z = stride_descale_q_z_in
+        stride_descale_k_z = stride_descale_k_z_in
+        stride_descale_v_z = stride_descale_v_z_in
+        stride_descale_do_z = stride_descale_do_z_in
+
+    GROUP_SIZE = NUM_Q_HEADS // NUM_K_HEADS
+    wid = tl.program_id(0)  # workgoup id: 0, ..., NUM_Q_PIDS * BATCH * NUM_K_HEADS - 1
+
+    head_q_idx = wid % NUM_Q_HEADS
+    head_q_idx = remap_xcd(head_q_idx, NUM_Q_HEADS, NUM_XCD)
+    seq_k_blk_idx = (wid // NUM_Q_HEADS) % NUM_K_PIDS
+    batch_idx = (wid // (NUM_K_PIDS * NUM_Q_HEADS)) % BATCH
+
+    # In the backward we dont want concurrent workgroups to handle consecutive heads or blocks, so remap them to be far apart.
+    head_q_idx = (head_q_idx * 29) % NUM_Q_HEADS
+    # seq_k_blk_idx = (seq_k_blk_idx * 29) % NUM_K_PIDS
+
+    head_k_idx = head_q_idx // GROUP_SIZE
+
+    # Determine q and k start along with seqlen_q and seqlen_k
+    q_start = 0
+    k_start = 0
+    seqlen_q = max_seqlen_q
+    seqlen_k = max_seqlen_k
+    if IS_VARLEN:
+        q_start = tl.load(cu_seqlens_q + batch_idx)
+        q_end = tl.load(cu_seqlens_q + batch_idx + 1)
+        k_start = tl.load(cu_seqlens_k + batch_idx)
+        k_end = tl.load(cu_seqlens_k + batch_idx + 1)
+        seqlen_q = q_end - q_start
+        seqlen_k = k_end - k_start
+
+    dk = tl.zeros([BLOCK_N, BLOCK_D_MODEL_POW2], dtype=tl.float32)
+    dv = tl.zeros([BLOCK_N, BLOCK_D_MODEL_POW2], dtype=tl.float32)
+
+    # Figure out causal starting block since we have seqlen_q >=< seqlen_k.
+    # Unlike forward pass where we tile on M dim and iterate on N dim, so that
+    # we can skip some M blocks, in backward pass, we tile on the N dim for kv
+    # and iterate over the M. In this way, we cannot skip N blocks, but only to
+    # determine the starting M blocks to skip some initial blocks masked by
+    # causal.
+    delta_qk = seqlen_q - seqlen_k
+
+    # q < k: some blocks will have no Masked block, other needs to re-calc
+    # starting position
+    # delta_qk is negative so flip it, only multiple of BLOCK_N can skip the
+    # masked op
+    num_blocks_skip = -delta_qk // BLOCK_N
+    delta_aligned = (num_blocks_skip + 1) * BLOCK_N + delta_qk
+    start_delta_q_lt_k = delta_aligned // BLOCK_M * BLOCK_M
+    if delta_qk >= 0:
+        start_delta = delta_qk
+    else:
+        start_delta = start_delta_q_lt_k
+
+    start_n = seq_k_blk_idx * BLOCK_N
+
+    offs_k = tl.arange(0, BLOCK_D_MODEL_POW2)
+    offs_n = start_n + tl.arange(0, BLOCK_N)
+    # Mask for loading K and V
+    mask_kv = offs_n[:, None] < seqlen_k
+    PADDED_HEAD: tl.constexpr = BLOCK_D_MODEL != BLOCK_D_MODEL_POW2
+    if PADDED_HEAD:
+        mask_k = offs_k < BLOCK_D_MODEL
+        mask_kv &= mask_k[None, :]
+
+    GROUP_SIZE = NUM_Q_HEADS // NUM_K_HEADS
+    adj_k = (
+        batch_idx * stride_k_b
+        + head_k_idx * stride_k_h
+        + k_start * stride_k_n
+        + offs_n[:, None] * stride_k_n
+        + offs_k[None, :] * stride_k_k
+    )
+    adj_v = (
+        batch_idx * stride_v_b
+        + head_k_idx * stride_v_h
+        + k_start * stride_v_n
+        + offs_n[:, None] * stride_v_n
+        + offs_k[None, :] * stride_v_k
+    )
+    # load K and V: they stay in SRAM throughout the inner loop.
+    k = tl.load(k_ptr + adj_k, mask=mask_kv, other=0.0)
+    v = tl.load(v_ptr + adj_v, mask=mask_kv, other=0.0)
+
+    # If MQA / GQA, set the K and V head offsets appropriately.
+    # for head_q_idx in range(head_k_idx * GROUP_SIZE, head_k_idx * GROUP_SIZE + GROUP_SIZE):
+    if delta_qk >= 0:
+        start_m = start_n + start_delta
+        len_m = BLOCK_N
+    else:
+        start_m = max(start_n + delta_qk, 0)
+        start_m = (start_m // BLOCK_M) * BLOCK_M
+        # because we might shift the masked blocks up, we are deeper into
+        # the masked out region, so we would potentially increase the total
+        # steps with masked operation to get out of it
+        residue_m = max(start_n + delta_qk - start_m, 0)
+        len_m = BLOCK_N + residue_m
+
+    # offset input and output tensor by batch and Q/K heads
+    adj_q = batch_idx * stride_q_b + head_q_idx * stride_q_h + q_start * stride_q_m
+    adj_dq = batch_idx * stride_dq_b + head_q_idx * stride_dq_h + q_start * stride_dq_m
+
+    q_ptr_adj = q_ptr + adj_q
+    dq_ptr_adj = dq_ptr + adj_dq
+
+    adj_do = batch_idx * stride_do_b + head_q_idx * stride_do_h + q_start * stride_do_m
+    do_ptr_adj = do_ptr + adj_do
+    adj_delta = (
+        batch_idx * stride_delta_b
+        + head_q_idx * stride_delta_h
+        + q_start * stride_delta_m
+    )
+    m_ptr_adj = m_ptr + adj_delta
+    delta_ptr_adj = delta_ptr + adj_delta
+
+    # batch_philox_offset is the ACTUALLY dropout offset
+    # dropout_offset is for debug purpose and will be removed later
+    batch_philox_offset = 0
+    dropout_offset = 0
+    if ENABLE_DROPOUT:
+        batch_philox_offset = (
+            philox_offset_base
+            + batch_idx * stride_dropout_b
+            + head_q_idx * stride_dropout_h
+        )
+        dropout_offset = (
+            dropout_mask + batch_idx * stride_dropout_b + head_q_idx * stride_dropout_h
+        )
+
+    MASK_BLOCK_M: tl.constexpr = BLOCK_M // BLK_SLICE_FACTOR
+    # bound the masked operation to q len so it does not have to wast cycles
+    len_m = min(len_m, seqlen_q)
+    num_steps = tl.cdiv(len_m, MASK_BLOCK_M)
+
+    # when q < k, we may skip the initial masked op
+    if seq_k_blk_idx < num_blocks_skip:
+        num_steps = 0
+
+    if IS_FP8:
+        descale_q = tl.load(descale_q_ptr + batch_idx * stride_descale_q_z + head_q_idx)
+        descale_k = tl.load(descale_k_ptr + batch_idx * stride_descale_k_z + head_k_idx)
+        descale_v = tl.load(descale_v_ptr + batch_idx * stride_descale_v_z + head_k_idx)
+        descale_do = tl.load(
+            descale_do_ptr + batch_idx * stride_descale_do_z + head_q_idx
+        )
+    else:
+        descale_q, descale_k, descale_v, descale_do = 1.0, 1.0, 1.0, 1.0
+
+    # if unaligned start_m is negative, the current N-tile has no block on the
+    #   diagonal of causal mask, so everything have no causal mask
+    dk, dv = _bwd_dkdvdq_inner(
+        dk,
+        dv,  # output tensors
+        q_ptr_adj,
+        k,
+        v,
+        do_ptr_adj,
+        dq_ptr_adj,
+        m_ptr_adj,
+        delta_ptr_adj,
+        sm_scale,  # input tensors
+        stride_q_m,
+        stride_q_k,  # strides for q
+        stride_dq_m,
+        stride_dq_k,  # strides for q
+        stride_do_m,
+        stride_do_k,  # strides for o
+        stride_dropout_m,
+        stride_dropout_n,  # strides for dropout
+        stride_delta_m,
+        dropout_p,
+        philox_seed,
+        batch_philox_offset,
+        dropout_offset,  #
+        seqlen_q,
+        seqlen_k,  # max sequence length for q and k
+        start_n,
+        start_m,
+        num_steps,  # iteration numbers
+        descale_q,
+        descale_k,
+        descale_v,
+        descale_do,  # fp8 descale factors from user
+        MASK_BLOCK_M,
+        BLOCK_N,  # block dim
+        BLOCK_D_MODEL,
+        BLOCK_D_MODEL_POW2,  # head dim
+        MASK=True,  # causal masking
+        ENABLE_DROPOUT=ENABLE_DROPOUT,  # activate dropout
+        IS_FP8=IS_FP8,
+        FP8_MAX=FP8_MAX,
+        workgroup_id=seq_k_blk_idx,
+    )
+
+    start_m += num_steps * MASK_BLOCK_M
+    num_steps = tl.cdiv(seqlen_q - start_m, BLOCK_M)
+
+    dk, dv = _bwd_dkdvdq_inner(
+        dk,
+        dv,  # output tensors
+        q_ptr_adj,
+        k,
+        v,
+        do_ptr_adj,
+        dq_ptr_adj,
+        m_ptr_adj,
+        delta_ptr_adj,
+        sm_scale,  # input tensors
+        stride_q_m,
+        stride_q_k,  # strides for q
+        stride_dq_m,
+        stride_dq_k,  # strides for dq
+        stride_do_m,
+        stride_do_k,  # strides for o
+        stride_dropout_m,
+        stride_dropout_n,  # strides for dropout
+        stride_delta_m,
+        dropout_p,
+        philox_seed,
+        batch_philox_offset,
+        dropout_offset,  #
+        seqlen_q,
+        seqlen_k,  # max sequence length for q and k
+        start_n,
+        start_m,
+        num_steps,  # iteration numbers
+        descale_q,
+        descale_k,
+        descale_v,
+        descale_do,  # fp8 descale factors from user
+        BLOCK_M,
+        BLOCK_N,  # block dim
+        BLOCK_D_MODEL,
+        BLOCK_D_MODEL_POW2,  # head dim
+        MASK=False,  # causal masking
+        ENABLE_DROPOUT=ENABLE_DROPOUT,  # activate dropout
+        IS_FP8=IS_FP8,
+        FP8_MAX=FP8_MAX,
+        workgroup_id=seq_k_blk_idx,
+    )
+
+    # Write back dV and dK.
+    offs_dkdv = (
+        batch_idx * stride_dk_b
+        + head_k_idx * stride_dk_h
+        + k_start * stride_dk_n
+        + offs_n[:, None] * stride_dk_n
+        + offs_k[None, :] * stride_dk_k
+    )
+    tl.atomic_add(dv_ptr + offs_dkdv, dv, mask=mask_kv, sem="relaxed")
+    dk *= sm_scale
+    tl.atomic_add(dk_ptr + offs_dkdv, dk, mask=mask_kv, sem="relaxed")
+
+
+_bwd_kernel_dkdvdq_noncausal_repr = make_kernel_repr(
+    "_bwd_kernel_dkdvdq_noncausal",
+    [
+        "NUM_Q_HEADS",
+        "NUM_K_HEADS",
+        "BLOCK_M",
+        "BLOCK_N",
+        "BLOCK_D_MODEL",
+        "ENABLE_DROPOUT",
+        "IS_VARLEN",
+        "IS_FP8",
+        "USE_INT64_STRIDES",
+    ],
+)
+
+
+@triton.jit(repr=_bwd_kernel_dkdvdq_noncausal_repr)
+def _bwd_kernel_dkdvdq_noncausal(
+    Q,
+    K,
+    V,
+    sm_scale,
+    DO,
+    DK,
+    DV,
+    DQ,
+    M,
+    Delta,
+    stride_qb_in,
+    stride_qh_in,
+    stride_qm_in,
+    stride_qk_in,
+    stride_kb_in,
+    stride_kh_in,
+    stride_kn_in,
+    stride_kk_in,
+    stride_vb_in,
+    stride_vh_in,
+    stride_vn_in,
+    stride_vk_in,
+    stride_dkb_in,
+    stride_dkh_in,
+    stride_dkn_in,
+    stride_dkk_in,
+    stride_dqb_in,
+    stride_dqh_in,
+    stride_dqm_in,
+    stride_dqk_in,
+    stride_deltab_in,
+    stride_deltah_in,
+    stride_deltam_in,
+    stride_dob_in,
+    stride_doh_in,
+    stride_dom_in,
+    stride_dok_in,
+    stride_dropoutb_in,
+    stride_dropouth_in,
+    stride_dropoutm_in,
+    stride_dropoutn_in,
+    stride_descale_q_z_in,
+    stride_descale_k_z_in,
+    stride_descale_v_z_in,
+    stride_descale_do_z_in,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q,
+    max_seqlen_k,
+    dropout_mask,
+    dropout_p,
+    philox_seed,
+    philox_offset,
+    descale_q_ptr,
+    descale_k_ptr,
+    descale_v_ptr,
+    descale_do_ptr,
+    NUM_Q_HEADS: tl.constexpr,
+    NUM_K_HEADS: tl.constexpr,
+    BATCH,
+    NUM_K_PIDS,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLK_SLICE_FACTOR: tl.constexpr,
+    BLOCK_D_MODEL: tl.constexpr,
+    BLOCK_D_MODEL_POW2: tl.constexpr,
+    ENABLE_DROPOUT: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    IS_FP8: tl.constexpr,
+    FP8_MAX: tl.constexpr,
+    USE_INT64_STRIDES: tl.constexpr,
+):
+    if USE_INT64_STRIDES:
+        stride_qb = tl.cast(stride_qb_in, tl.int64)
+        stride_qh = tl.cast(stride_qh_in, tl.int64)
+        stride_qm = tl.cast(stride_qm_in, tl.int64)
+        stride_qk = tl.cast(stride_qk_in, tl.int64)
+        stride_kb = tl.cast(stride_kb_in, tl.int64)
+        stride_kh = tl.cast(stride_kh_in, tl.int64)
+        stride_kn = tl.cast(stride_kn_in, tl.int64)
+        stride_kk = tl.cast(stride_kk_in, tl.int64)
+        stride_vb = tl.cast(stride_vb_in, tl.int64)
+        stride_vh = tl.cast(stride_vh_in, tl.int64)
+        stride_vn = tl.cast(stride_vn_in, tl.int64)
+        stride_vk = tl.cast(stride_vk_in, tl.int64)
+        stride_dkb = tl.cast(stride_dkb_in, tl.int64)
+        stride_dkh = tl.cast(stride_dkh_in, tl.int64)
+        stride_dkn = tl.cast(stride_dkn_in, tl.int64)
+        stride_dkk = tl.cast(stride_dkk_in, tl.int64)
+        stride_dqb = tl.cast(stride_dqb_in, tl.int64)
+        stride_dqh = tl.cast(stride_dqh_in, tl.int64)
+        stride_dqm = tl.cast(stride_dqm_in, tl.int64)
+        stride_dqk = tl.cast(stride_dqk_in, tl.int64)
+        stride_deltab = tl.cast(stride_deltab_in, tl.int64)
+        stride_deltah = tl.cast(stride_deltah_in, tl.int64)
+        stride_deltam = tl.cast(stride_deltam_in, tl.int64)
+        stride_dob = tl.cast(stride_dob_in, tl.int64)
+        stride_doh = tl.cast(stride_doh_in, tl.int64)
+        stride_dom = tl.cast(stride_dom_in, tl.int64)
+        stride_dok = tl.cast(stride_dok_in, tl.int64)
+        stride_dropoutb = tl.cast(stride_dropoutb_in, tl.int64)
+        stride_dropouth = tl.cast(stride_dropouth_in, tl.int64)
+        stride_dropoutm = tl.cast(stride_dropoutm_in, tl.int64)
+        stride_dropoutn = tl.cast(stride_dropoutn_in, tl.int64)
+        if IS_FP8:
+            stride_descale_q_z = tl.cast(stride_descale_q_z_in, tl.int64)
+            stride_descale_k_z = tl.cast(stride_descale_k_z_in, tl.int64)
+            stride_descale_v_z = tl.cast(stride_descale_v_z_in, tl.int64)
+            stride_descale_do_z = tl.cast(stride_descale_do_z_in, tl.int64)
+    else:
+        stride_qb = stride_qb_in
+        stride_qh = stride_qh_in
+        stride_qm = stride_qm_in
+        stride_qk = stride_qk_in
+        stride_kb = stride_kb_in
+        stride_kh = stride_kh_in
+        stride_kn = stride_kn_in
+        stride_kk = stride_kk_in
+        stride_vb = stride_vb_in
+        stride_vh = stride_vh_in
+        stride_vn = stride_vn_in
+        stride_vk = stride_vk_in
+        stride_dkb = stride_dkb_in
+        stride_dkh = stride_dkh_in
+        stride_dkn = stride_dkn_in
+        stride_dkk = stride_dkk_in
+        stride_dqb = stride_dqb_in
+        stride_dqh = stride_dqh_in
+        stride_dqm = stride_dqm_in
+        stride_dqk = stride_dqk_in
+        stride_deltab = stride_deltab_in
+        stride_deltah = stride_deltah_in
+        stride_deltam = stride_deltam_in
+        stride_dob = stride_dob_in
+        stride_doh = stride_doh_in
+        stride_dom = stride_dom_in
+        stride_dok = stride_dok_in
+        stride_dropoutb = stride_dropoutb_in
+        stride_dropouth = stride_dropouth_in
+        stride_dropoutm = stride_dropoutm_in
+        stride_dropoutn = stride_dropoutn_in
+        stride_descale_q_z = stride_descale_q_z_in
+        stride_descale_k_z = stride_descale_k_z_in
+        stride_descale_v_z = stride_descale_v_z_in
+        stride_descale_do_z = stride_descale_do_z_in
+
+    # workgroup id
+    wid = tl.program_id(0)  # 0, ..., NUM_K_PIDS * BATCH * NUM_K_HEADS - 1
+
+    # Workgroups get launched first along batch dim, then in head_k dim, and then in seq k block dim
+    # This is in order to avoid contention for the tl.atomic_add (inside _bwd_dkdvdq_inner) that happens between workgroups that share the same batch and head_k.
+    bid = wid % BATCH
+    hkid = wid // BATCH % NUM_K_HEADS
+    pid = wid // (BATCH * NUM_K_HEADS) % NUM_K_PIDS
+
+    q_start = 0
+    k_start = 0
+    seqlen_q = max_seqlen_q
+    seqlen_k = max_seqlen_k
+
+    if IS_VARLEN:
+        q_start = tl.load(cu_seqlens_q + bid)
+        q_end = tl.load(cu_seqlens_q + bid + 1)
+        k_start = tl.load(cu_seqlens_k + bid)
+        k_end = tl.load(cu_seqlens_k + bid + 1)
+        seqlen_q = q_end - q_start
+        seqlen_k = k_end - k_start
+
+    dk = tl.zeros([BLOCK_N, BLOCK_D_MODEL_POW2], dtype=tl.float32)
+    dv = tl.zeros([BLOCK_N, BLOCK_D_MODEL_POW2], dtype=tl.float32)
+
+    start_n = pid * BLOCK_N
+
+    offs_k = tl.arange(0, BLOCK_D_MODEL_POW2)
+    offs_n = start_n + tl.arange(0, BLOCK_N)
+    mask_kv = offs_n[:, None] < seqlen_k
+    PADDED_HEAD: tl.constexpr = BLOCK_D_MODEL != BLOCK_D_MODEL_POW2
+    if PADDED_HEAD:
+        mask_kv &= offs_k < BLOCK_D_MODEL
+
+    GROUP_SIZE = NUM_Q_HEADS // NUM_K_HEADS
+    adj_k = (
+        bid * stride_kb
+        + hkid * stride_kh
+        + k_start * stride_kn
+        + offs_n[:, None] * stride_kn
+        + offs_k[None, :] * stride_kk
+    )
+    adj_v = (
+        bid * stride_vb
+        + hkid * stride_vh
+        + k_start * stride_vn
+        + offs_n[:, None] * stride_vn
+        + offs_k[None, :] * stride_vk
+    )
+
+    k = tl.load(K + adj_k, mask=mask_kv, other=0.0)
+    v = tl.load(V + adj_v, mask=mask_kv, other=0.0)
+
+    for hqid in range(hkid * GROUP_SIZE, hkid * GROUP_SIZE + GROUP_SIZE):
+        adj_q = bid * stride_qb + hqid * stride_qh + q_start * stride_qm
+        adj_dq = bid * stride_dqb + hqid * stride_dqh + q_start * stride_dqm
+
+        Q_ptr = Q + adj_q
+        DQ_ptr = DQ + adj_dq
+
+        adj_do = bid * stride_dob + hqid * stride_doh + q_start * stride_dom
+        DO_ptr = DO + adj_do
+        adj_delta = bid * stride_deltab + hqid * stride_deltah + q_start * stride_deltam
+        M_ptr = M + adj_delta
+        Delta_ptr = Delta + adj_delta
+
+        # dropout
+        batch_philox_offset = 0
+        dropout_offset = 0
+        if ENABLE_DROPOUT:
+            batch_philox_offset = (
+                philox_offset + bid * stride_dropoutb + hqid * stride_dropouth
+            )
+            dropout_offset = (
+                dropout_mask + bid * stride_dropoutb + hqid * stride_dropouth
+            )
+
+        if IS_FP8:
+            descale_q = tl.load(descale_q_ptr + bid * stride_descale_q_z + hqid)
+            descale_k = tl.load(descale_k_ptr + bid * stride_descale_k_z + hkid)
+            descale_v = tl.load(descale_v_ptr + bid * stride_descale_v_z + hkid)
+            descale_do = tl.load(descale_do_ptr + bid * stride_descale_do_z + hqid)
+        else:
+            descale_q, descale_k, descale_v, descale_do = 1.0, 1.0, 1.0, 1.0
+
+        start_m = 0
+        num_steps = tl.cdiv(seqlen_q, BLOCK_M)
+
+        dk, dv = _bwd_dkdvdq_inner(
+            dk,
+            dv,
+            Q_ptr,
+            k,
+            v,
+            DO_ptr,
+            DQ_ptr,
+            M_ptr,
+            Delta_ptr,
+            sm_scale,
+            stride_qm,
+            stride_qk,
+            stride_dqm,
+            stride_dqk,
+            stride_dom,
+            stride_dok,
+            stride_dropoutm,
+            stride_dropoutn,
+            stride_deltam,
+            dropout_p,
+            philox_seed,
+            batch_philox_offset,
+            dropout_offset,
+            seqlen_q,
+            seqlen_k,
+            start_n,
+            start_m,
+            num_steps,
+            descale_q,
+            descale_k,
+            descale_v,
+            descale_do,
+            BLOCK_M,
+            BLOCK_N,
+            BLOCK_D_MODEL,
+            BLOCK_D_MODEL_POW2,
+            MASK=False,
+            ENABLE_DROPOUT=ENABLE_DROPOUT,
+            IS_FP8=IS_FP8,
+            FP8_MAX=FP8_MAX,
+            workgroup_id=wid,
+        )
+
+    adj_dkdv = (
+        bid * stride_dkb
+        + hkid * stride_dkh
+        + k_start * stride_dkn
+        + offs_n[:, None] * stride_dkn
+        + offs_k[None, :] * stride_dkk
+    )
+    tl.store(DV + adj_dkdv, dv, mask=mask_kv)
+    dk *= sm_scale
+    tl.store(DK + adj_dkdv, dk, mask=mask_kv)
+
+
+@functools.lru_cache(maxsize=1024)
+def _get_config():
+    if not hasattr(_get_config, "_config_dict"):
+        dev = arch_info.get_arch()
+        _get_config._config_dict = {}
+        fpath = f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHA-DEFAULT.json"
+        with open(fpath, "r") as file:
+            config = json.load(file)
+        _get_config._config_dict = config
+
+    return _get_config._config_dict["bkwd_fused"]

build/torch-rocm/_kernels/mha_onekernel_bwd.py

@@ -0,0 +1,1783 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+import functools
+import json
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+from ..utils._triton import arch_info
+from ..utils.core import AITER_TRITON_CONFIGS_PATH
+from ..utils._triton.mha_kernel_utils import _compute_fp8_scaling_factors
+from ..utils._triton.kernel_repr import make_kernel_repr
+
+# NOTE: triton fails to import tl.constexprs so create them here for the file
+DROPOUT_USE_PYTORCH = False
+DROPOUT_DUMP = False
+
+tl_DROPOUT_USE_PYTORCH: tl.constexpr = triton.language.constexpr(DROPOUT_USE_PYTORCH)
+tl_DROPOUT_DUMP: tl.constexpr = triton.language.constexpr(DROPOUT_DUMP)
+
+
+# This function computes delta given output Out and gradient DO
+# Here is the I/O shape:
+# Out: (batch, nhead_q, max_seqlens_q, headDim)
+# DO: (batch, nhead_q, max_seqlens_q, headDim)
+# Delta: (batch, nheads_q, max_seqlens_q), same as softmax_lse defined at
+_bwd_preprocess_repr = make_kernel_repr(
+    "_bwd_preprocess",
+    [
+        "BLOCK_M",
+        "BLOCK_D_MODEL",
+        "IS_VARLEN",
+        "IS_FP8",
+    ],
+)
+
+
+@triton.jit(repr=_bwd_preprocess_repr)
+def _bwd_preprocess(
+    o_ptr,
+    do_ptr,  # noqa: E741
+    delta_ptr,
+    stride_o_b,
+    stride_o_h,
+    stride_o_m,
+    stride_o_k,
+    stride_do_b,
+    stride_do_h,
+    stride_do_m,
+    stride_do_k,
+    stride_delta_b,
+    stride_delta_h,
+    stride_delta_m,
+    stride_descale_do_z,
+    cu_seqlens_q,
+    max_seqlen_q,
+    descale_do_ptr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_D_MODEL: tl.constexpr,
+    BLOCK_D_MODEL_POW2: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    IS_FP8: tl.constexpr,
+):
+    pid_m = tl.program_id(0)  # seqlen
+    bid = tl.program_id(1)  # batch
+    hid = tl.program_id(2)  # head
+
+    # Handle varlen
+    q_start = 0
+    seqlen_q = max_seqlen_q
+    if IS_VARLEN:
+        q_start = tl.load(cu_seqlens_q + bid)
+        q_end = tl.load(cu_seqlens_q + bid + 1)
+        seqlen_q = q_end - q_start
+    else:
+        q_start = 0
+        seqlen_q = max_seqlen_q
+
+    # Compute offsets
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_k = tl.arange(0, BLOCK_D_MODEL_POW2)
+
+    # O and DO may have different strides (e.g. BSHD vs SBHD memory layout),
+    # so address each with its own strides.
+    offs_o = (
+        bid * stride_o_b
+        + hid * stride_o_h
+        + q_start * stride_o_m
+        + offs_m[:, None] * stride_o_m
+        + offs_k[None, :] * stride_o_k
+    )
+    offs_do = (
+        bid * stride_do_b
+        + hid * stride_do_h
+        + q_start * stride_do_m
+        + offs_m[:, None] * stride_do_m
+        + offs_k[None, :] * stride_do_k
+    )
+
+    # create masks
+    mask_m = offs_m < seqlen_q
+    mask = mask_m[:, None]
+    PADDED_HEAD: tl.constexpr = BLOCK_D_MODEL != BLOCK_D_MODEL_POW2
+    if PADDED_HEAD:
+        mask &= offs_k[None, :] < BLOCK_D_MODEL
+
+    # load [BLOCK_M, BLOCK_D_MODEL_POW2]
+    o = tl.load(o_ptr + offs_o, mask=mask, other=0.0)
+    do = tl.load(do_ptr + offs_do, mask=mask, other=0.0)
+
+    # compute and write-back to delta
+    if IS_FP8:
+        descale_do = tl.load(descale_do_ptr + bid * stride_descale_do_z + hid)
+
+        # NOTE: do is in the fp8 range and o is not in fp8
+        delta = tl.sum(o.to(tl.float32) * (do.to(tl.float32) * descale_do), axis=1)
+    else:
+        delta = tl.sum(o.to(tl.float32) * do.to(tl.float32), axis=1)
+
+    offs_delta = (
+        bid * stride_delta_b
+        + hid * stride_delta_h
+        + q_start * stride_delta_m
+        + offs_m * stride_delta_m
+    )
+    tl.store(delta_ptr + offs_delta, delta, mask=mask_m)
+
+
+# The main inner-loop logic for computing dK and dV.
+@triton.jit
+def _bwd_dkdv_inner(
+    dk,  # output
+    dk_pe,  # optional output, pass None for non-PE case
+    dv,  # output
+    Q,
+    k,
+    k_pe,
+    v,
+    DO,
+    M,
+    D,
+    sm_scale,  # input tensor
+    stride_qm,
+    stride_qk,
+    stride_dom,
+    stride_dok,
+    stride_dropoutm,
+    stride_dropoutn,
+    stride_deltam,
+    BLOCK_M: tl.constexpr,  # 16
+    BLOCK_N: tl.constexpr,  # 128
+    HEAD_DIM: tl.constexpr,  #
+    ACTUAL_HEAD_DIM: tl.constexpr,  #
+    PE_HEAD_DIM: tl.constexpr,
+    dropout_p,
+    philox_seed,
+    batch_philox_offset,
+    dropout_offset,
+    alibi_slope,
+    seqlen_q,
+    seqlen_k,  # max sequence length for q and k
+    # Filled in by the wrapper.
+    start_n,
+    start_m,
+    num_steps,  # iteration numbers
+    descale_q,
+    descale_k,
+    descale_v,
+    descale_do,  # fp8 descale factors from user
+    MASK: tl.constexpr,  # causal masking, only apply to tiles on mask diagonal
+    ENABLE_DROPOUT: tl.constexpr,  # activate dropout
+    USE_ALIBI: tl.constexpr,
+    USE_EXP2: tl.constexpr,  # activate exp2
+    IS_FP8: tl.constexpr,
+    FP8_MAX: tl.constexpr,
+    DEBUG_TRITON: tl.constexpr,
+    DEBUG_TRITON_DETAIL: tl.constexpr,
+    SLIDING_WINDOW: tl.constexpr,
+):
+    # if HEAD_DIM is padded
+    PADDED_HEAD: tl.constexpr = ACTUAL_HEAD_DIM != HEAD_DIM
+    HAS_PE: tl.constexpr = PE_HEAD_DIM > 0
+    delta_qk = seqlen_q - seqlen_k
+    offs_m = start_m + tl.arange(0, BLOCK_M)  # start_m + (0, 15)
+    offs_n = start_n + tl.arange(0, BLOCK_N)  # start_m + (0, 127)
+    offs_k = tl.arange(0, HEAD_DIM)
+    if HAS_PE:
+        offs_k_pe = HEAD_DIM + tl.arange(0, PE_HEAD_DIM)
+    # mask to make sure not OOB of seqlen_q
+    mask_n = offs_n < seqlen_k
+    # Q and DO are (seqlen_q, head_dim)
+    # qT_ptrs = (1, BLOCK_M) + (HEAD_DIM, 1), transpose of q
+    qT_ptrs = Q + offs_m[None, :] * stride_qm + offs_k[:, None] * stride_qk
+    if HAS_PE:
+        qT_pe_ptrs = Q + offs_m[None, :] * stride_qm + offs_k_pe[:, None] * stride_qk
+    # do_ptrs = (BLOCK_M, 1) + (1, HEAD_DIM), NOT transposed
+    do_ptrs = DO + offs_m[:, None] * stride_dom + offs_k[None, :] * stride_dok
+    # BLOCK_N must be a multiple of BLOCK_M, otherwise the code wouldn't work.
+    tl.static_assert(BLOCK_N % BLOCK_M == 0)
+    curr_m = start_m
+    step_m = BLOCK_M
+    curr_philox_offset = batch_philox_offset
+    curr_dropout_offset = dropout_offset
+    RCP_LN2: tl.constexpr = 1.4426950408889634  # = 1.0 / ln(2)
+
+    for blk_idx in range(num_steps):
+        if DEBUG_TRITON:
+            print(f"iter {blk_idx}: curr_m = {curr_m}")  # noqa: E701
+        offs_m = curr_m + tl.arange(0, BLOCK_M)
+        # update the mask because offs_m advanced
+        mask_m = offs_m < seqlen_q
+        mask_qT = mask_m[None, :]
+        mask_do = mask_m[:, None]
+        mask_nm = mask_n[:, None] & (offs_m[None, :] < seqlen_q)
+        if PADDED_HEAD:
+            mask_qT &= offs_k[:, None] < ACTUAL_HEAD_DIM
+            mask_do &= offs_k[None, :] < ACTUAL_HEAD_DIM
+        qT = tl.load(qT_ptrs, mask=mask_qT, other=0.0)
+        if HAS_PE:
+            qT_pe = tl.load(qT_pe_ptrs, mask=mask_qT, other=0.0)
+        # generate dropout mask
+        if ENABLE_DROPOUT:
+            # NOTE: dropout is transposed because it is used to mask pT
+            philox_offs = (
+                curr_philox_offset
+                + offs_m[None, :] * stride_dropoutm
+                + offs_n[:, None] * stride_dropoutn
+            )
+            if tl_DROPOUT_USE_PYTORCH:
+                dropout_offs = (
+                    offs_m[None, :] * stride_dropoutm
+                    + offs_n[:, None] * stride_dropoutn
+                )
+                dropout_mask = tl.load(curr_dropout_offset + dropout_offs, mask=mask_nm)
+            else:
+                rand_vals = tl.rand(philox_seed, philox_offs)
+                dropout_mask = rand_vals > dropout_p
+            dropout_scale = 1.0 / (1 - dropout_p)
+        # Load m before computing qk to reduce pipeline stall.
+        m = tl.load(M + offs_m * stride_deltam, mask=mask_m, other=0.0)
+        if IS_FP8:
+            qkT = tl.dot(k, qT) * descale_q * descale_k
+        else:
+            qkT = tl.dot(k, qT)
+            if HAS_PE:
+                qkT = tl.dot(k_pe, qT_pe, acc=qkT)
+        qkT_scaled = qkT * sm_scale
+
+        if USE_ALIBI:
+            relative_pos_block = offs_n[:, None] + seqlen_q - seqlen_k - offs_m[None, :]
+            alibi_block = -1 * alibi_slope * tl.abs(relative_pos_block)
+            qkT_scaled += alibi_block
+
+        if DEBUG_TRITON_DETAIL:
+            if start_n == 256:
+                print(f"qT: {qT.shape}\n", qT)
+                print(f"k: {k.shape}\n", k)
+                print(f"qkT scaled: {qkT.shape}\n", qkT_scaled)
+        # TODO: remove the scaling of m later when we removed re-scaling in fwd
+        if USE_EXP2:
+            pT = tl.math.exp2(qkT_scaled * RCP_LN2 - m[None, :] * RCP_LN2)
+        else:
+            pT = tl.math.exp(qkT_scaled - m[None, :])
+
+        # Autoregressive masking.
+        if MASK:
+            # offset offs_m with delta_qk since the causal mask starts at
+            # bottom right of the (seqlen_q, seqlen_k) matrix
+            causal_mask = (offs_m[None, :] - delta_qk) >= offs_n[:, None]
+            mask = causal_mask & mask_nm
+            if DEBUG_TRITON_DETAIL:
+                if start_n == 256:
+                    print(f"causal_mask: {causal_mask.shape}\n", causal_mask)
+                    print(
+                        f"qkT after causal: {qkT.shape}\n",
+                        tl.where(causal_mask, qkT * sm_scale, 0.0),
+                    )
+            pT = tl.where(mask, pT, 0.0)
+        if SLIDING_WINDOW > 0:
+            window_mask = offs_n[:, None] >= (
+                offs_m[None, :] - delta_qk - SLIDING_WINDOW
+            )
+            pT = tl.where(window_mask, pT, 0.0)
+        do = tl.load(do_ptrs, mask=mask_do, other=0.0)
+        # Compute dV.
+        if ENABLE_DROPOUT:
+            pT_dropout = tl.where(dropout_mask, pT, 0.0) * dropout_scale
+            if IS_FP8:
+                scale_p_dropout, descale_p_dropout = _compute_fp8_scaling_factors(
+                    pT_dropout, FP8_MAX
+                )
+                dv += (
+                    tl.dot((pT_dropout * scale_p_dropout).to(do.type.element_ty), do)
+                    * descale_p_dropout
+                    * descale_do
+                )
+            else:
+                dv = tl.dot(pT_dropout.to(do.type.element_ty), do, acc=dv)
+        else:
+            if IS_FP8:
+                scale_pT, descale_pT = _compute_fp8_scaling_factors(pT, FP8_MAX)
+                dv += (
+                    tl.dot((pT * scale_pT).to(do.type.element_ty), do)
+                    * descale_pT
+                    * descale_do
+                )
+            else:
+                dv = tl.dot(pT.to(do.type.element_ty), do, acc=dv)
+
+        if DEBUG_TRITON_DETAIL:
+            if start_n == 256:
+                print(f"pT: {pT.shape}\n", pT)
+        # D (= delta) is pre-divided by ds_scale.
+        Di = tl.load(D + offs_m * stride_deltam, mask=mask_m)
+        # Compute dP and dS.
+        if IS_FP8:
+            dpT = tl.dot(v, tl.trans(do)) * descale_v * descale_do
+        else:
+            dpT = tl.dot(v, tl.trans(do))
+        if ENABLE_DROPOUT:
+            dpT = tl.where(dropout_mask, dpT, 0.0) * dropout_scale
+        delta_i = Di[None, :]
+        dsT = pT * (dpT - delta_i)
+        if IS_FP8:
+            scale_dsT, descale_dsT = _compute_fp8_scaling_factors(dsT, FP8_MAX)
+            dk += (
+                tl.dot((dsT * scale_dsT).to(qT.type.element_ty), tl.trans(qT))
+                * descale_dsT
+                * descale_q
+            )
+        else:
+            dk = tl.dot(dsT.to(qT.type.element_ty), tl.trans(qT), acc=dk)
+            if HAS_PE:
+                dk_pe = tl.dot(
+                    dsT.to(qT_pe.type.element_ty), tl.trans(qT_pe), acc=dk_pe
+                )
+        # Increment pointers.
+        curr_m += step_m
+        qT_ptrs += step_m * stride_qm
+        if HAS_PE:
+            qT_pe_ptrs += step_m * stride_qm
+        do_ptrs += step_m * stride_dom
+    return dk, dk_pe, dv
+
+
+# the main inner-loop logic for computing dQ
+@triton.jit
+def _bwd_dq_inner(
+    dq,  # output
+    dq_pe,  # optional output, pass None for non-PE case
+    q,
+    q_pe,
+    K,
+    V,
+    do,
+    m,
+    Di,  # D (= delta) is pre-divided by ds_scale.
+    sm_scale,  # input
+    # shared by Q/K/V.
+    stride_qm,
+    stride_qk,
+    stride_kn,
+    stride_kk,
+    stride_vn,
+    stride_vk,
+    stride_dropoutm,
+    stride_dropoutn,  # stride for dropout
+    seqlen_q,
+    seqlen_k,  #
+    BLOCK_M2: tl.constexpr,  #
+    BLOCK_N2: tl.constexpr,  #
+    HEAD_DIM: tl.constexpr,
+    ACTUAL_HEAD_DIM: tl.constexpr,  #
+    PE_HEAD_DIM: tl.constexpr,
+    dropout_p,
+    philox_seed,
+    batch_philox_offset,
+    dropout_offset,
+    alibi_slope,
+    # Filled in by the wrapper.
+    start_m,
+    start_n,
+    end_n,
+    num_steps,  #
+    descale_q,
+    descale_k,
+    descale_v,
+    descale_do,  # fp8 descale factors from user
+    MASK: tl.constexpr,
+    ENABLE_DROPOUT: tl.constexpr,
+    USE_ALIBI: tl.constexpr,
+    USE_EXP2: tl.constexpr,
+    IS_FP8: tl.constexpr,
+    FP8_MAX: tl.constexpr,
+    DEBUG_TRITON: tl.constexpr,
+    DEBUG_TRITON_DETAIL: tl.constexpr,
+    SLIDING_WINDOW: tl.constexpr,
+    ENABLE_SINK: tl.constexpr,
+):
+    # if HEAD_DIM is padded
+    PADDED_HEAD: tl.constexpr = ACTUAL_HEAD_DIM != HEAD_DIM
+    HAS_PE: tl.constexpr = PE_HEAD_DIM > 0
+    delta_qk = seqlen_q - seqlen_k
+    offs_m = start_m + tl.arange(0, BLOCK_M2)
+    offs_n = start_n + tl.arange(0, BLOCK_N2)
+    offs_k = tl.arange(0, HEAD_DIM)
+    if HAS_PE:
+        offs_k_pe = HEAD_DIM + tl.arange(0, PE_HEAD_DIM)
+
+    # mask to make sure not OOB of seqlen_q
+    mask_m = offs_m < seqlen_q
+
+    kT_ptrs = K + offs_n[None, :] * stride_kn + offs_k[:, None] * stride_kk
+    if HAS_PE:
+        kT_pe_ptrs = K + offs_n[None, :] * stride_kn + offs_k_pe[:, None] * stride_kk
+    vT_ptrs = V + offs_n[None, :] * stride_vn + offs_k[:, None] * stride_vk
+    # BLOCK_M2 must be a multiple of BLOCK_N2, otherwise the code wouldn't work.
+    tl.static_assert(BLOCK_M2 % BLOCK_N2 == 0)
+    curr_n = start_n
+    step_n = BLOCK_N2
+    curr_philox_offset = batch_philox_offset
+    curr_dropout_offset = dropout_offset
+    RCP_LN2: tl.constexpr = 1.4426950408889634  # = 1.0 / ln(2)
+    delta_recomp = tl.zeros([BLOCK_M2], dtype=tl.float32)
+    for blk_idx in range(num_steps):
+        if DEBUG_TRITON:
+            print(f"iter {blk_idx}: curr_n = {curr_n}")  # noqa: E701
+        offs_n = curr_n + tl.arange(0, BLOCK_N2)
+        # end_n is needed because the end of causal True might not be perfectly
+        # aligned with the end of the block
+        mask_n = offs_n < end_n
+        if DEBUG_TRITON_DETAIL:
+            print(
+                f"start_n = {start_n}, end_n = {end_n}, offs_n: {offs_n.shape}\n{offs_n}"
+            )  # noqa: E701
+        if DEBUG_TRITON_DETAIL:
+            print(f"mask_n: {mask_n.shape}\n{mask_n}")  # noqa: E701
+        mask_kT = mask_n[None, :]
+        mask_mn = mask_m[:, None] & (offs_n[None, :] < end_n)
+        if PADDED_HEAD:
+            mask_kT &= offs_k[:, None] < ACTUAL_HEAD_DIM
+
+        kT = tl.load(kT_ptrs, mask=mask_kT, other=0.0)
+        if HAS_PE:
+            kT_pe = tl.load(kT_pe_ptrs, mask=mask_kT, other=0.0)
+        vT = tl.load(vT_ptrs, mask=mask_kT, other=0.0)
+
+        if ENABLE_DROPOUT:
+            # NOTE: dropout is transposed because it is used to mask pT
+            philox_offs = (
+                curr_philox_offset
+                + offs_m[:, None] * stride_dropoutm
+                + offs_n[None, :] * stride_dropoutn
+            )
+            if tl_DROPOUT_USE_PYTORCH:
+                dropout_offs = (
+                    offs_m[:, None] * stride_dropoutm
+                    + offs_n[None, :] * stride_dropoutn
+                )
+                dropout_mask = tl.load(curr_dropout_offset + dropout_offs, mask=mask_mn)
+            else:
+                rand_vals = tl.rand(philox_seed, philox_offs)
+                dropout_mask = rand_vals > dropout_p
+            dropout_scale = 1 / (1 - dropout_p)
+
+        if IS_FP8:
+            qk = tl.dot(q, kT) * descale_q * descale_k
+        else:
+            qk = tl.dot(q, kT)
+            if HAS_PE:
+                qk = tl.dot(q_pe, kT_pe, acc=qk)
+        qk_scaled = qk * sm_scale
+
+        if USE_ALIBI:
+            relative_pos_block = offs_m[:, None] + seqlen_k - seqlen_q - offs_n[None, :]
+            alibi_block = -1 * alibi_slope * tl.abs(relative_pos_block)
+            qk_scaled += alibi_block
+
+        if DEBUG_TRITON_DETAIL:
+            print(f"qk scaled: {qk.shape}\n", qk_scaled)  # noqa: E701
+        if USE_EXP2:
+            p = tl.math.exp2(qk_scaled * RCP_LN2 - m * RCP_LN2)
+        else:
+            p = tl.math.exp(qk_scaled - m)
+
+        # Autoregressive masking.
+        if MASK:
+            causal_mask = (offs_m[:, None] - delta_qk) >= offs_n[None, :]
+            mask = causal_mask & mask_mn
+            p = tl.where(mask, p, 0.0)
+        if SLIDING_WINDOW > 0:
+            window_mask = offs_n[None, :] >= (
+                offs_m[:, None] - delta_qk - SLIDING_WINDOW
+            )
+            p = tl.where(window_mask, p, 0.0)
+        # Compute dP and dS.
+        if IS_FP8:
+            dp = tl.dot(do, vT) * descale_do * descale_v
+        else:
+            dp = tl.dot(do, vT)
+        if ENABLE_DROPOUT:
+            dp = tl.where(dropout_mask, dp, 0.0) * dropout_scale
+        if ENABLE_SINK:
+            # Equal to do_i . o_i but skips the bf16 round-trip through forward `o`.
+            delta_recomp += tl.sum(p * dp, axis=1)
+        delta_i = Di[:, None]
+        ds = p * (dp - delta_i)
+        # Compute dQ.
+        # NOTE: We need to de-scale dq in the end, because kT was pre-scaled.
+        if IS_FP8:
+            scale_ds, descale_ds = _compute_fp8_scaling_factors(ds, FP8_MAX)
+            dq += (
+                tl.dot((ds * scale_ds).to(kT.type.element_ty), tl.trans(kT))
+                * descale_ds
+                * descale_k
+            )
+        else:
+            dq = tl.dot(ds.to(kT.type.element_ty), tl.trans(kT), acc=dq)
+            if HAS_PE:
+                dq_pe = tl.dot(ds.to(kT_pe.type.element_ty), tl.trans(kT_pe), acc=dq_pe)
+        # Increment pointers.
+        curr_n += step_n
+        kT_ptrs += step_n * stride_kn
+        if HAS_PE:
+            kT_pe_ptrs += step_n * stride_kn
+        vT_ptrs += step_n * stride_vn
+    return dq, dq_pe, delta_recomp
+
+
+_bwd_kernel_causal_repr = make_kernel_repr(
+    "bwd_kernel_causal",
+    [
+        "BLOCK_M1",
+        "BLOCK_N1",
+        "BLOCK_M2",
+        "BLOCK_N2",
+        "BLK_SLICE_FACTOR",
+        "HEAD_DIM",
+        "ENABLE_DROPOUT",
+        "IS_VARLEN",
+        "USE_ALIBI",
+        "USE_EXP2",
+        "IS_FP8",
+        "USE_INT64_STRIDES",
+        "ENABLE_SINK",
+        "SLIDING_WINDOW",
+    ],
+)
+
+
+@triton.jit(repr=_bwd_kernel_causal_repr)
+def bwd_kernel_causal(  # grid = (tl.cdiv(max_seqlen_q // BLOCK_M2), batch, nheads_q)
+    Q,
+    K,
+    V,
+    Sink,
+    sm_scale,
+    DO,
+    DQ,
+    DK,
+    DV,
+    DSink,
+    M,
+    Delta,
+    stride_qb_in,
+    stride_qh_in,
+    stride_qm_in,
+    stride_qd_in,
+    stride_kb_in,
+    stride_kh_in,
+    stride_kn_in,
+    stride_kd_in,
+    stride_vb_in,
+    stride_vh_in,
+    stride_vn_in,
+    stride_vd_in,
+    stride_dqb_in,
+    stride_dqh_in,
+    stride_dqm_in,
+    stride_dqd_in,
+    stride_dkb_in,
+    stride_dkh_in,
+    stride_dkn_in,
+    stride_dkd_in,
+    stride_dvb_in,
+    stride_dvh_in,
+    stride_dvn_in,
+    stride_dvd_in,
+    stride_deltab_in,
+    stride_deltah_in,
+    stride_deltam_in,
+    stride_dob_in,
+    stride_doh_in,
+    stride_dom_in,
+    stride_dod_in,
+    stride_dropoutb_in,
+    stride_dropouth_in,
+    stride_dropoutm_in,
+    stride_dropoutn_in,
+    stride_descale_q_z_in,
+    stride_descale_k_z_in,
+    stride_descale_v_z_in,
+    stride_descale_do_z_in,
+    stride_az_in,
+    stride_ah_in,
+    HQ,
+    HK,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q,
+    max_seqlen_k,
+    Dropout_mask,
+    dropout_p,
+    philox_seed,
+    philox_offset_base_in,
+    Alibi_slopes,
+    Descale_q,
+    Descale_k,
+    Descale_v,
+    Descale_do,
+    BLOCK_M1: tl.constexpr,
+    BLOCK_N1: tl.constexpr,
+    BLOCK_M2: tl.constexpr,
+    BLOCK_N2: tl.constexpr,
+    BLK_SLICE_FACTOR: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    ACTUAL_HEAD_DIM: tl.constexpr,
+    PE_HEAD_DIM: tl.constexpr,
+    ENABLE_DROPOUT: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    USE_ALIBI: tl.constexpr,
+    USE_EXP2: tl.constexpr,
+    IS_FP8: tl.constexpr,
+    FP8_MAX: tl.constexpr,
+    DEBUG_TRITON: tl.constexpr,
+    DEBUG_TRITON_DETAIL: tl.constexpr,
+    USE_INT64_STRIDES: tl.constexpr,
+    ENABLE_SINK: tl.constexpr,
+    SLIDING_WINDOW: tl.constexpr,
+):
+    if USE_INT64_STRIDES:
+        stride_qb = tl.cast(stride_qb_in, tl.int64)
+        stride_qh = tl.cast(stride_qh_in, tl.int64)
+        stride_qm = tl.cast(stride_qm_in, tl.int64)
+        stride_qd = tl.cast(stride_qd_in, tl.int64)
+        stride_kb = tl.cast(stride_kb_in, tl.int64)
+        stride_kh = tl.cast(stride_kh_in, tl.int64)
+        stride_kn = tl.cast(stride_kn_in, tl.int64)
+        stride_kd = tl.cast(stride_kd_in, tl.int64)
+        stride_vb = tl.cast(stride_vb_in, tl.int64)
+        stride_vh = tl.cast(stride_vh_in, tl.int64)
+        stride_vn = tl.cast(stride_vn_in, tl.int64)
+        stride_vd = tl.cast(stride_vd_in, tl.int64)
+        stride_dqb = tl.cast(stride_dqb_in, tl.int64)
+        stride_dqh = tl.cast(stride_dqh_in, tl.int64)
+        stride_dqm = tl.cast(stride_dqm_in, tl.int64)
+        stride_dqd = tl.cast(stride_dqd_in, tl.int64)
+        stride_dkb = tl.cast(stride_dkb_in, tl.int64)
+        stride_dkh = tl.cast(stride_dkh_in, tl.int64)
+        stride_dkn = tl.cast(stride_dkn_in, tl.int64)
+        stride_dkd = tl.cast(stride_dkd_in, tl.int64)
+        stride_dvb = tl.cast(stride_dvb_in, tl.int64)
+        stride_dvh = tl.cast(stride_dvh_in, tl.int64)
+        stride_dvn = tl.cast(stride_dvn_in, tl.int64)
+        stride_dvd = tl.cast(stride_dvd_in, tl.int64)
+        stride_deltab = tl.cast(stride_deltab_in, tl.int64)
+        stride_deltah = tl.cast(stride_deltah_in, tl.int64)
+        stride_deltam = tl.cast(stride_deltam_in, tl.int64)
+        stride_dob = tl.cast(stride_dob_in, tl.int64)
+        stride_doh = tl.cast(stride_doh_in, tl.int64)
+        stride_dom = tl.cast(stride_dom_in, tl.int64)
+        stride_dod = tl.cast(stride_dod_in, tl.int64)
+        philox_offset_base = tl.cast(philox_offset_base_in, tl.int64)
+        stride_dropoutb = tl.cast(stride_dropoutb_in, tl.int64)
+        stride_dropouth = tl.cast(stride_dropouth_in, tl.int64)
+        stride_dropoutm = tl.cast(stride_dropoutm_in, tl.int64)
+        stride_dropoutn = tl.cast(stride_dropoutn_in, tl.int64)
+        if IS_FP8:
+            stride_descale_q_z = tl.cast(stride_descale_q_z_in, tl.int64)
+            stride_descale_k_z = tl.cast(stride_descale_k_z_in, tl.int64)
+            stride_descale_v_z = tl.cast(stride_descale_v_z_in, tl.int64)
+            stride_descale_do_z = tl.cast(stride_descale_do_z_in, tl.int64)
+        stride_az = tl.cast(stride_az_in, tl.int64)
+        stride_ah = tl.cast(stride_ah_in, tl.int64)
+    else:
+        stride_qb = stride_qb_in
+        stride_qh = stride_qh_in
+        stride_qm = stride_qm_in
+        stride_qd = stride_qd_in
+        stride_kb = stride_kb_in
+        stride_kh = stride_kh_in
+        stride_kn = stride_kn_in
+        stride_kd = stride_kd_in
+        stride_vb = stride_vb_in
+        stride_vh = stride_vh_in
+        stride_vn = stride_vn_in
+        stride_vd = stride_vd_in
+        stride_dqb = stride_dqb_in
+        stride_dqh = stride_dqh_in
+        stride_dqm = stride_dqm_in
+        stride_dqd = stride_dqd_in
+        stride_dkb = stride_dkb_in
+        stride_dkh = stride_dkh_in
+        stride_dkn = stride_dkn_in
+        stride_dkd = stride_dkd_in
+        stride_dvb = stride_dvb_in
+        stride_dvh = stride_dvh_in
+        stride_dvn = stride_dvn_in
+        stride_dvd = stride_dvd_in
+        stride_deltab = stride_deltab_in
+        stride_deltah = stride_deltah_in
+        stride_deltam = stride_deltam_in
+        stride_dob = stride_dob_in
+        stride_doh = stride_doh_in
+        stride_dom = stride_dom_in
+        stride_dod = stride_dod_in
+        philox_offset_base = philox_offset_base_in
+        stride_dropoutb = stride_dropoutb_in
+        stride_dropouth = stride_dropouth_in
+        stride_dropoutm = stride_dropoutm_in
+        stride_dropoutn = stride_dropoutn_in
+        stride_descale_q_z = stride_descale_q_z_in
+        stride_descale_k_z = stride_descale_k_z_in
+        stride_descale_v_z = stride_descale_v_z_in
+        stride_descale_do_z = stride_descale_do_z_in
+        stride_az = stride_az_in
+        stride_ah = stride_ah_in
+
+    # program ids
+    hkid = tl.program_id(0)
+    pid = tl.program_id(1)
+    bid = tl.program_id(2)
+    if DEBUG_TRITON:
+        print(f"\npid: {pid}, bid: {bid}, hkid: {hkid}")  # noqa: E701
+    # figure out varlen start and end
+    q_start = 0
+    k_start = 0
+    seqlen_q = max_seqlen_q
+    seqlen_k = max_seqlen_k
+    if IS_VARLEN:
+        # Compute actual sequence lengths
+        q_start = tl.load(cu_seqlens_q + bid)
+        q_end = tl.load(cu_seqlens_q + bid + 1)
+        k_start = tl.load(cu_seqlens_k + bid)
+        k_end = tl.load(cu_seqlens_k + bid + 1)
+        seqlen_q = q_end - q_start
+        seqlen_k = k_end - k_start
+
+    delta_qk = seqlen_q - seqlen_k
+    if DEBUG_TRITON:
+        print(f"delta_qk = {delta_qk}")  # noqa: E701
+    PADDED_HEAD: tl.constexpr = ACTUAL_HEAD_DIM != HEAD_DIM
+    HAS_PE: tl.constexpr = PE_HEAD_DIM > 0
+    offs_d = tl.arange(0, HEAD_DIM)
+    if HAS_PE:
+        offs_d_pe = HEAD_DIM + tl.arange(0, PE_HEAD_DIM)
+    GROUP_SIZE: tl.constexpr = HQ // HK
+
+    # align the delta_qk
+    start_n = pid * BLOCK_N1
+    if start_n < seqlen_k:
+        # This section does dk and dv
+        dk = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
+        if HAS_PE:
+            dk_pe = tl.zeros([BLOCK_N1, PE_HEAD_DIM], dtype=tl.float32)
+        else:
+            # Couldn't assign None to dk_pe because _bwd_dkdv_inner can't return None.
+            dk_pe = dk
+        dv = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
+
+        # q > k: diretcly skip all the way until the start of causal block
+        start_delta_q_gt_k = delta_qk
+        # q < k: some blocks will have no Masked block, other needs to re-calc
+        # starting position
+        # delta_qk is negative so flip it, only multiple of BLOCK_N can skip the
+        # masked op
+        num_blocks_skip = -delta_qk // BLOCK_N1
+        delta_aligned = (num_blocks_skip + 1) * BLOCK_N1 + delta_qk
+        start_delta_q_lt_k = delta_aligned // BLOCK_M1 * BLOCK_M1
+        if delta_qk >= 0:
+            start_delta = delta_qk
+            if DEBUG_TRITON:
+                print(
+                    f"q >= k: start_delta = delta_qk aligned to BLOCK_M = {start_delta_q_gt_k}"
+                )  # noqa: E701
+        else:
+            start_delta = start_delta_q_lt_k
+            if DEBUG_TRITON:
+                print(
+                    f"q < k: start_delta = residue btw multiple BLOCK_N and delta_qk = {delta_aligned} = aligned to BLOCK_M = {start_delta_q_lt_k}"
+                )  # noqa: E701
+
+        offs_n = start_n + tl.arange(0, BLOCK_N1)
+        # Mask for loading K and V
+        mask_kv = offs_n[:, None] < seqlen_k
+        if PADDED_HEAD:
+            mask_d = offs_d < ACTUAL_HEAD_DIM
+            mask_kv &= mask_d[None, :]
+
+        # K/V tensors not changed for the group
+        adj_k = (
+            bid * stride_kb
+            + hkid * stride_kh
+            + k_start * stride_kn
+            + offs_n[:, None] * stride_kn
+            + offs_d[None, :] * stride_kd
+        )
+        if HAS_PE:
+            adj_k_pe = (
+                bid * stride_kb
+                + hkid * stride_kh
+                + k_start * stride_kn
+                + offs_n[:, None] * stride_kn
+                + offs_d_pe[None, :] * stride_kd
+            )
+        adj_v = (
+            bid * stride_vb
+            + hkid * stride_vh
+            + k_start * stride_vn
+            + offs_n[:, None] * stride_vn
+            + offs_d[None, :] * stride_vd
+        )
+        # load K and V: they stay in SRAM throughout the inner loop.
+        k = tl.load(K + adj_k, mask=mask_kv, other=0.0)
+        if HAS_PE:
+            k_pe = tl.load(K + adj_k_pe, mask=mask_kv, other=0.0)
+        else:
+            k_pe = None
+        v = tl.load(V + adj_v, mask=mask_kv, other=0.0)
+        # If MQA / GQA, set the K and V head offsets appropriately.
+        # hqid = hkid
+        for hqid in range(hkid * GROUP_SIZE, hkid * GROUP_SIZE + GROUP_SIZE):
+            if delta_qk >= 0:
+                start_m = start_n + start_delta
+                len_m = BLOCK_N1
+            else:
+                start_m = max(start_n + delta_qk, 0)
+                start_m = start_m // BLOCK_M1 * BLOCK_M1
+                # because we might shift the masked blocks up, we are deeper into
+                # the masked out region, so we would potentially increase the total
+                # steps with masked operation to get out of it
+                residue_m = max(start_n + delta_qk - start_m, 0)
+                len_m = BLOCK_N1 + residue_m
+                if DEBUG_TRITON:
+                    print(f"residue_m = {residue_m}")  # noqa: E701
+
+            # offset input and output tensor by batch and Q/K heads
+            adj_q = bid * stride_qb + hqid * stride_qh + q_start * stride_qm
+            Q_ptr = Q + adj_q
+            adj_do = bid * stride_dob + hqid * stride_doh + q_start * stride_dom
+            DO_ptr = DO + adj_do
+            adj_delta = (
+                bid * stride_deltab + hqid * stride_deltah + q_start * stride_deltam
+            )
+            M_ptr = M + adj_delta
+            Delta_ptr = Delta + adj_delta
+
+            if USE_ALIBI:
+                alibi_offset = bid * stride_az + hqid * stride_ah
+                alibi_slope = tl.load(Alibi_slopes + alibi_offset)
+            else:
+                alibi_slope = None
+
+            # batch_philox_offset is the ACTUALLY dropout offset
+            # dropout_offset is for debug purpose and will be removed later
+            batch_philox_offset = 0
+            dropout_offset = 0
+            if ENABLE_DROPOUT:
+                batch_philox_offset = (
+                    philox_offset_base + bid * stride_dropoutb + hqid * stride_dropouth
+                )
+                dropout_offset = (
+                    Dropout_mask + bid * stride_dropoutb + hqid * stride_dropouth
+                )
+
+            if IS_FP8:
+                descale_q = tl.load(Descale_q + bid * stride_descale_q_z + hqid)
+                descale_k = tl.load(Descale_k + bid * stride_descale_k_z + hkid)
+                descale_v = tl.load(Descale_v + bid * stride_descale_v_z + hkid)
+                descale_do = tl.load(Descale_do + bid * stride_descale_do_z + hqid)
+            else:
+                descale_q, descale_k, descale_v, descale_do = 1.0, 1.0, 1.0, 1.0
+
+            MASK_BLOCK_M1: tl.constexpr = BLOCK_M1 // BLK_SLICE_FACTOR
+            # bound the masked operation to q len so it does not have to wast cycles
+            len_m = min(len_m, seqlen_q)
+            num_steps = tl.cdiv(len_m, MASK_BLOCK_M1)
+            # when q < k, we may skip the initial masked op
+            if pid < num_blocks_skip:
+                num_steps = 0
+
+            # if start_m is negative, the current N-tile has no block on the
+            #   diagonal of causal mask, so everything have no causal mask
+            if DEBUG_TRITON:
+                print(
+                    f"Masked: start_n: {start_n}; start_m: {start_m}, num_steps: {num_steps}"
+                )  # noqa: E701
+            dk, dk_pe, dv = _bwd_dkdv_inner(
+                dk,  # output tensor
+                dk_pe,  # optional output tensor
+                dv,  # output tensor
+                Q_ptr,
+                k,
+                k_pe,
+                v,
+                DO_ptr,
+                M_ptr,
+                Delta_ptr,
+                sm_scale,  # input tensors
+                stride_qm,
+                stride_qd,  # strides for q
+                stride_dom,
+                stride_dod,  # strides for o
+                stride_dropoutm,
+                stride_dropoutn,  # strides for dropout
+                stride_deltam,
+                MASK_BLOCK_M1,
+                BLOCK_N1,  # block dim
+                HEAD_DIM,
+                ACTUAL_HEAD_DIM,  # head dim
+                PE_HEAD_DIM,
+                dropout_p,
+                philox_seed,
+                batch_philox_offset,
+                dropout_offset,
+                alibi_slope,
+                seqlen_q,
+                seqlen_k,  # max sequence length for q and k
+                start_n,
+                start_m,
+                num_steps,  # iteration numbers
+                descale_q,
+                descale_k,
+                descale_v,
+                descale_do,
+                MASK=True,  # causal masking
+                ENABLE_DROPOUT=ENABLE_DROPOUT,  # activate dropout
+                USE_ALIBI=USE_ALIBI,
+                USE_EXP2=USE_EXP2,
+                IS_FP8=IS_FP8,
+                FP8_MAX=FP8_MAX,
+                DEBUG_TRITON=DEBUG_TRITON,
+                DEBUG_TRITON_DETAIL=DEBUG_TRITON_DETAIL,
+                SLIDING_WINDOW=SLIDING_WINDOW,
+            )
+            start_m += num_steps * MASK_BLOCK_M1
+            end_m = seqlen_q
+            if SLIDING_WINDOW > 0:
+                end_m = min(start_n + BLOCK_N1 + delta_qk + SLIDING_WINDOW, seqlen_q)
+            num_steps = tl.cdiv(max(end_m - start_m, 0), BLOCK_M1)
+            end_m = start_m + num_steps * BLOCK_M1
+
+            if DEBUG_TRITON:
+                print(
+                    f"start_m after Masked step: {start_m}; num_steps: {num_steps}"
+                )  # noqa: E701
+            if DEBUG_TRITON:
+                print(
+                    f"unMasked: start_n: {start_n}, start_m: {start_m}, end_m: {end_m}, num_steps: {num_steps}"
+                )  # noqa: E701
+            if DEBUG_TRITON:
+                print("unMasked")  # noqa: E701
+            dk, dk_pe, dv = _bwd_dkdv_inner(
+                dk,  # output tensor
+                dk_pe,  # optional output tensor
+                dv,  # output tensor
+                Q_ptr,
+                k,
+                k_pe,
+                v,
+                DO_ptr,
+                M_ptr,
+                Delta_ptr,
+                sm_scale,  # input tensors
+                stride_qm,
+                stride_qd,  # strides for q
+                stride_dom,
+                stride_dod,  # strides for o
+                stride_dropoutm,
+                stride_dropoutn,  # strides for dropout
+                stride_deltam,
+                BLOCK_M1,
+                BLOCK_N1,  # block dim
+                HEAD_DIM,
+                ACTUAL_HEAD_DIM,  # head dim
+                PE_HEAD_DIM,
+                dropout_p,
+                philox_seed,
+                batch_philox_offset,
+                dropout_offset,
+                alibi_slope,
+                seqlen_q,
+                seqlen_k,  # max sequence length for q and k
+                start_n,
+                start_m,
+                num_steps,  # iteration numbers
+                descale_q,
+                descale_k,
+                descale_v,
+                descale_do,
+                MASK=False,  # causal masking
+                ENABLE_DROPOUT=ENABLE_DROPOUT,  # activate dropout
+                USE_ALIBI=USE_ALIBI,
+                USE_EXP2=USE_EXP2,
+                IS_FP8=IS_FP8,
+                FP8_MAX=FP8_MAX,
+                DEBUG_TRITON=DEBUG_TRITON,
+                DEBUG_TRITON_DETAIL=DEBUG_TRITON_DETAIL,
+                SLIDING_WINDOW=SLIDING_WINDOW,
+            )
+        # end of GQA/MQA of dkdv
+        # Write back dV
+        adj_dv = bid * stride_dvb + hkid * stride_dvh + k_start * stride_dvn
+        offs_dv = offs_n[:, None] * stride_dvn + offs_d[None, :] * stride_dvd
+        tl.store(DV + adj_dv + offs_dv, dv, mask=mask_kv)
+        # write back dk
+        adj_dk = bid * stride_dkb + hkid * stride_dkh + k_start * stride_dkn
+        offs_dk = offs_n[:, None] * stride_dkn + offs_d[None, :] * stride_dkd
+        dk *= sm_scale
+        tl.store(DK + adj_dk + offs_dk, dk, mask=mask_kv)
+        if HAS_PE:
+            offs_dk_pe = offs_n[:, None] * stride_dkn + offs_d_pe[None, :] * stride_dkd
+            dk_pe *= sm_scale
+            tl.store(DK + adj_dk + offs_dk_pe, dk_pe, mask=mask_kv)
+
+    # This part does dq
+    start_m = pid * BLOCK_M2
+    if start_m < seqlen_q:
+        # seqlen_q > seqlen_k, no need to process these tile for dq
+        if DEBUG_TRITON:
+            print(
+                f"end_n = start_m + BLOCK_M = {start_m} + {BLOCK_M2} = {start_m + BLOCK_M2}"
+            )  # noqa: E701
+        if start_m + BLOCK_M2 < delta_qk:
+            if DEBUG_TRITON:
+                print(
+                    f"start_m + BLOCK_M2 = {start_m} + {BLOCK_M2} = {start_m + BLOCK_M2} < delta_qk of {delta_qk}"
+                )  # noqa: E701
+            return
+
+        offs_m = start_m + tl.arange(0, BLOCK_M2)
+        # Mask for loading K and V
+        mask_q = offs_m[:, None] < seqlen_q
+        if PADDED_HEAD:
+            mask_d = offs_d < ACTUAL_HEAD_DIM
+            mask_q &= mask_d[None, :]
+        offs_q = offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd
+        if HAS_PE:
+            offs_q_pe = offs_m[:, None] * stride_qm + offs_d_pe[None, :] * stride_qd
+        offs_do = offs_m[:, None] * stride_dom + offs_d[None, :] * stride_dod
+        # NOTE: don't assume that the strides for k and v are the same!
+        K += bid * stride_kb + hkid * stride_kh + k_start * stride_kn
+        V += bid * stride_vb + hkid * stride_vh + k_start * stride_vn
+
+        # If MQA / GQA, set the K and V head offsets appropriately.
+        for hqid in range(hkid * GROUP_SIZE, hkid * GROUP_SIZE + GROUP_SIZE):
+            # seqlen_q < seqlen_k: delta_qk more kv tokens are added at the front
+            #   for every M-tile
+            end_n = start_m + BLOCK_M2 - delta_qk
+            # clamp end_n at [0, seqlen_k]
+            end_n = max(min(end_n, seqlen_k), 0)
+            if DEBUG_TRITON:
+                print(f"delta_qk: {delta_qk}; end_n: {end_n}")  # noqa: E701
+            # offset input and output tensor by batch and Q/K heads
+            adj_q = bid * stride_qb + hqid * stride_qh + q_start * stride_qm
+            adj_do = bid * stride_dob + hqid * stride_doh + q_start * stride_dom
+            adj_delta = (
+                bid * stride_deltab + hqid * stride_deltah + q_start * stride_deltam
+            )
+            Delta_ptr = Delta + adj_delta
+
+            if USE_ALIBI:
+                alibi_offset = bid * stride_az + hqid * stride_ah
+                alibi_slope = tl.load(Alibi_slopes + alibi_offset)
+            else:
+                alibi_slope = None
+
+            # batch_philox_offset is the ACTUALLY dropout offset
+            # dropout_offset is for debug purpose and will be removed later
+            batch_philox_offset = 0
+            dropout_offset = 0
+            if ENABLE_DROPOUT:
+                batch_philox_offset = (
+                    philox_offset_base + bid * stride_dropoutb + hqid * stride_dropouth
+                )
+                dropout_offset = (
+                    Dropout_mask + bid * stride_dropoutb + hqid * stride_dropouth
+                )
+            q = tl.load(Q + adj_q + offs_q, mask=mask_q, other=0.0)
+            if HAS_PE:
+                q_pe = tl.load(Q + adj_q + offs_q_pe, mask=mask_q, other=0.0)
+            else:
+                q_pe = None
+            do = tl.load(DO + adj_do + offs_do, mask=mask_q, other=0.0)
+            mask_m = offs_m < seqlen_q
+            m = tl.load(M + adj_delta + offs_m * stride_deltam, mask=mask_m, other=0.0)
+            m = m[:, None]
+            delta = tl.load(Delta_ptr + offs_m * stride_deltam, mask=mask_m, other=0.0)
+
+            MASK_BLOCK_N2: tl.constexpr = BLOCK_N2 // BLK_SLICE_FACTOR
+            # start can only be 0 at minimum
+            start_n = max(end_n - BLOCK_M2, 0)
+            num_steps = tl.cdiv(end_n - start_n, MASK_BLOCK_N2)
+
+            if IS_FP8:
+                descale_q = tl.load(Descale_q + bid * stride_descale_q_z + hqid)
+                descale_k = tl.load(Descale_k + bid * stride_descale_k_z + hkid)
+                descale_v = tl.load(Descale_v + bid * stride_descale_v_z + hkid)
+                descale_do = tl.load(Descale_do + bid * stride_descale_do_z + hqid)
+            else:
+                descale_q, descale_k, descale_v, descale_do = 1.0, 1.0, 1.0, 1.0
+
+            dq = tl.zeros([BLOCK_M2, HEAD_DIM], dtype=tl.float32)
+            if HAS_PE:
+                dq_pe = tl.zeros([BLOCK_M2, PE_HEAD_DIM], dtype=tl.float32)
+            else:
+                dq_pe = dq  # Couldn't assign None to dq_pe because _bwd_dq_inner can't return None.
+            dq, dq_pe, delta_recomp_masked = _bwd_dq_inner(
+                dq,  # output tensor
+                dq_pe,  # optional output tensor
+                q,
+                q_pe,
+                K,
+                V,
+                do,
+                m,
+                delta,
+                sm_scale,
+                stride_qm,
+                stride_qd,
+                stride_kn,
+                stride_kd,
+                stride_vn,
+                stride_vd,
+                stride_dropoutm,
+                stride_dropoutn,
+                seqlen_q,
+                seqlen_k,
+                BLOCK_M2,
+                MASK_BLOCK_N2,
+                HEAD_DIM,
+                ACTUAL_HEAD_DIM,
+                PE_HEAD_DIM,
+                dropout_p,
+                philox_seed,
+                batch_philox_offset,
+                dropout_offset,
+                alibi_slope,
+                start_m,
+                start_n,
+                end_n,
+                num_steps,
+                descale_q,
+                descale_k,
+                descale_v,
+                descale_do,
+                MASK=True,  #
+                ENABLE_DROPOUT=ENABLE_DROPOUT,
+                USE_ALIBI=USE_ALIBI,
+                USE_EXP2=USE_EXP2,
+                IS_FP8=IS_FP8,
+                FP8_MAX=FP8_MAX,
+                DEBUG_TRITON=DEBUG_TRITON,
+                DEBUG_TRITON_DETAIL=DEBUG_TRITON_DETAIL,
+                SLIDING_WINDOW=SLIDING_WINDOW,
+                ENABLE_SINK=ENABLE_SINK,
+            )
+            end_n -= num_steps * MASK_BLOCK_N2
+            window_start_n = 0
+            if SLIDING_WINDOW > 0:
+                window_start_n = max(start_m - delta_qk - SLIDING_WINDOW, 0)
+            start_n = window_start_n // BLOCK_N2 * BLOCK_N2
+            num_steps = tl.cdiv(max(end_n - start_n, 0), BLOCK_N2)
+            if DEBUG_TRITON:
+                print(
+                    f"unMasked: start_m: {start_m}, start_n: {start_n}, end_n: {end_n}, num_steps: {num_steps}"
+                )  # noqa: E701
+            dq, dq_pe, delta_recomp_unmasked = _bwd_dq_inner(
+                dq,  # output tensor
+                dq_pe,  # optional output tensor
+                q,
+                q_pe,
+                K,
+                V,
+                do,
+                m,
+                delta,
+                sm_scale,
+                stride_qm,
+                stride_qd,
+                stride_kn,
+                stride_kd,
+                stride_vn,
+                stride_vd,
+                stride_dropoutm,
+                stride_dropoutn,
+                seqlen_q,
+                seqlen_k,
+                BLOCK_M2,
+                BLOCK_N2,
+                HEAD_DIM,
+                ACTUAL_HEAD_DIM,
+                PE_HEAD_DIM,
+                dropout_p,
+                philox_seed,
+                batch_philox_offset,
+                dropout_offset,
+                alibi_slope,
+                start_m,
+                start_n,
+                end_n,
+                num_steps,
+                descale_q,
+                descale_k,
+                descale_v,
+                descale_do,
+                MASK=False,
+                ENABLE_DROPOUT=ENABLE_DROPOUT,
+                USE_ALIBI=USE_ALIBI,
+                USE_EXP2=USE_EXP2,
+                IS_FP8=IS_FP8,
+                FP8_MAX=FP8_MAX,
+                DEBUG_TRITON=DEBUG_TRITON,
+                DEBUG_TRITON_DETAIL=DEBUG_TRITON_DETAIL,
+                SLIDING_WINDOW=SLIDING_WINDOW,
+                ENABLE_SINK=ENABLE_SINK,
+            )
+            if ENABLE_SINK:
+                sink = tl.load(Sink + hqid).to(tl.float32)
+                if USE_EXP2:
+                    RCP_LN2: tl.constexpr = 1.4426950408889634
+                    psink = tl.math.exp2(sink * RCP_LN2 - m * RCP_LN2)
+                else:
+                    psink = tl.math.exp(sink - m)
+                delta_recomp = delta_recomp_masked + delta_recomp_unmasked
+                dsink = tl.sum(-psink * delta_recomp[:, None])
+                tl.atomic_add(DSink + hqid, dsink, sem="relaxed")
+            # Write back dQ.
+            adj_dq = bid * stride_dqb + hqid * stride_dqh + q_start * stride_dqm
+            offs_dq = offs_m[:, None] * stride_dqm + offs_d[None, :] * stride_dqd
+            dq *= sm_scale
+            tl.store(DQ + adj_dq + offs_dq, dq, mask=mask_q)
+            if HAS_PE:
+                offs_dq_pe = (
+                    offs_m[:, None] * stride_dqm + offs_d_pe[None, :] * stride_dqd
+                )
+                dq_pe *= sm_scale
+                tl.store(DQ + adj_dq + offs_dq_pe, dq_pe, mask=mask_q)
+            # end of GQA/MQA of dq
+
+
+_bwd_kernel_noncausal_repr = make_kernel_repr(
+    "bwd_kernel_noncausal",
+    [
+        "BLOCK_M1",
+        "BLOCK_N1",
+        "BLOCK_M2",
+        "BLOCK_N2",
+        "BLK_SLICE_FACTOR",
+        "HEAD_DIM",
+        "ENABLE_DROPOUT",
+        "IS_VARLEN",
+        "USE_ALIBI",
+        "USE_EXP2",
+        "IS_FP8",
+        "USE_INT64_STRIDES",
+        "ENABLE_SINK",
+        "SLIDING_WINDOW",
+    ],
+)
+
+
+@triton.jit(repr=_bwd_kernel_noncausal_repr)
+def bwd_kernel_noncausal(
+    Q,
+    K,
+    V,
+    Sink,
+    sm_scale,
+    DO,
+    DQ,
+    DK,
+    DV,
+    DSink,
+    M,
+    Delta,
+    stride_qb_in,
+    stride_qh_in,
+    stride_qm_in,
+    stride_qd_in,
+    stride_kb_in,
+    stride_kh_in,
+    stride_kn_in,
+    stride_kd_in,
+    stride_vb_in,
+    stride_vh_in,
+    stride_vn_in,
+    stride_vd_in,
+    stride_dqb_in,
+    stride_dqh_in,
+    stride_dqm_in,
+    stride_dqd_in,
+    stride_dkb_in,
+    stride_dkh_in,
+    stride_dkn_in,
+    stride_dkd_in,
+    stride_dvb_in,
+    stride_dvh_in,
+    stride_dvn_in,
+    stride_dvd_in,
+    stride_deltab_in,
+    stride_deltah_in,
+    stride_deltam_in,
+    stride_dob_in,
+    stride_doh_in,
+    stride_dom_in,
+    stride_dod_in,
+    stride_dropoutb_in,
+    stride_dropouth_in,
+    stride_dropoutm_in,
+    stride_dropoutn_in,
+    stride_descale_q_z_in,
+    stride_descale_k_z_in,
+    stride_descale_v_z_in,
+    stride_descale_do_z_in,
+    stride_az_in,
+    stride_ah_in,
+    HQ,
+    HK,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q,
+    max_seqlen_k,
+    Dropout_mask,
+    dropout_p,
+    philox_seed,
+    philox_offset_base_in,
+    Alibi_slopes,
+    Descale_q,
+    Descale_k,
+    Descale_v,
+    Descale_do,
+    BLOCK_M1: tl.constexpr,  # 32
+    BLOCK_N1: tl.constexpr,  # 128
+    BLOCK_M2: tl.constexpr,  # 128
+    BLOCK_N2: tl.constexpr,  # 32
+    BLK_SLICE_FACTOR: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    ACTUAL_HEAD_DIM: tl.constexpr,
+    PE_HEAD_DIM: tl.constexpr,
+    ENABLE_DROPOUT: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    USE_ALIBI: tl.constexpr,
+    USE_EXP2: tl.constexpr,
+    IS_FP8: tl.constexpr,
+    FP8_MAX: tl.constexpr,
+    DEBUG_TRITON: tl.constexpr,
+    DEBUG_TRITON_DETAIL: tl.constexpr,
+    USE_INT64_STRIDES: tl.constexpr,
+    ENABLE_SINK: tl.constexpr,
+    SLIDING_WINDOW: tl.constexpr,
+):
+    if USE_INT64_STRIDES:
+        stride_qb = tl.cast(stride_qb_in, tl.int64)
+        stride_qh = tl.cast(stride_qh_in, tl.int64)
+        stride_qm = tl.cast(stride_qm_in, tl.int64)
+        stride_qd = tl.cast(stride_qd_in, tl.int64)
+        stride_kb = tl.cast(stride_kb_in, tl.int64)
+        stride_kh = tl.cast(stride_kh_in, tl.int64)
+        stride_kn = tl.cast(stride_kn_in, tl.int64)
+        stride_kd = tl.cast(stride_kd_in, tl.int64)
+        stride_vb = tl.cast(stride_vb_in, tl.int64)
+        stride_vh = tl.cast(stride_vh_in, tl.int64)
+        stride_vn = tl.cast(stride_vn_in, tl.int64)
+        stride_vd = tl.cast(stride_vd_in, tl.int64)
+        stride_dqb = tl.cast(stride_dqb_in, tl.int64)
+        stride_dqh = tl.cast(stride_dqh_in, tl.int64)
+        stride_dqm = tl.cast(stride_dqm_in, tl.int64)
+        stride_dqd = tl.cast(stride_dqd_in, tl.int64)
+        stride_dkb = tl.cast(stride_dkb_in, tl.int64)
+        stride_dkh = tl.cast(stride_dkh_in, tl.int64)
+        stride_dkn = tl.cast(stride_dkn_in, tl.int64)
+        stride_dkd = tl.cast(stride_dkd_in, tl.int64)
+        stride_dvb = tl.cast(stride_dvb_in, tl.int64)
+        stride_dvh = tl.cast(stride_dvh_in, tl.int64)
+        stride_dvn = tl.cast(stride_dvn_in, tl.int64)
+        stride_dvd = tl.cast(stride_dvd_in, tl.int64)
+        stride_deltab = tl.cast(stride_deltab_in, tl.int64)
+        stride_deltah = tl.cast(stride_deltah_in, tl.int64)
+        stride_deltam = tl.cast(stride_deltam_in, tl.int64)
+        stride_dob = tl.cast(stride_dob_in, tl.int64)
+        stride_doh = tl.cast(stride_doh_in, tl.int64)
+        stride_dom = tl.cast(stride_dom_in, tl.int64)
+        stride_dod = tl.cast(stride_dod_in, tl.int64)
+        philox_offset_base = tl.cast(philox_offset_base_in, tl.int64)
+        stride_dropoutb = tl.cast(stride_dropoutb_in, tl.int64)
+        stride_dropouth = tl.cast(stride_dropouth_in, tl.int64)
+        stride_dropoutm = tl.cast(stride_dropoutm_in, tl.int64)
+        stride_dropoutn = tl.cast(stride_dropoutn_in, tl.int64)
+        if IS_FP8:
+            stride_descale_q_z = tl.cast(stride_descale_q_z_in, tl.int64)
+            stride_descale_k_z = tl.cast(stride_descale_k_z_in, tl.int64)
+            stride_descale_v_z = tl.cast(stride_descale_v_z_in, tl.int64)
+            stride_descale_do_z = tl.cast(stride_descale_do_z_in, tl.int64)
+        stride_az = tl.cast(stride_az_in, tl.int64)
+        stride_ah = tl.cast(stride_ah_in, tl.int64)
+    else:
+        stride_qb = stride_qb_in
+        stride_qh = stride_qh_in
+        stride_qm = stride_qm_in
+        stride_qd = stride_qd_in
+        stride_kb = stride_kb_in
+        stride_kh = stride_kh_in
+        stride_kn = stride_kn_in
+        stride_kd = stride_kd_in
+        stride_vb = stride_vb_in
+        stride_vh = stride_vh_in
+        stride_vn = stride_vn_in
+        stride_vd = stride_vd_in
+        stride_dqb = stride_dqb_in
+        stride_dqh = stride_dqh_in
+        stride_dqm = stride_dqm_in
+        stride_dqd = stride_dqd_in
+        stride_dkb = stride_dkb_in
+        stride_dkh = stride_dkh_in
+        stride_dkn = stride_dkn_in
+        stride_dkd = stride_dkd_in
+        stride_dvb = stride_dvb_in
+        stride_dvh = stride_dvh_in
+        stride_dvn = stride_dvn_in
+        stride_dvd = stride_dvd_in
+        stride_deltab = stride_deltab_in
+        stride_deltah = stride_deltah_in
+        stride_deltam = stride_deltam_in
+        stride_dob = stride_dob_in
+        stride_doh = stride_doh_in
+        stride_dom = stride_dom_in
+        stride_dod = stride_dod_in
+        philox_offset_base = philox_offset_base_in
+        stride_dropoutb = stride_dropoutb_in
+        stride_dropouth = stride_dropouth_in
+        stride_dropoutm = stride_dropoutm_in
+        stride_dropoutn = stride_dropoutn_in
+        stride_descale_q_z = stride_descale_q_z_in
+        stride_descale_k_z = stride_descale_k_z_in
+        stride_descale_v_z = stride_descale_v_z_in
+        stride_descale_do_z = stride_descale_do_z_in
+        stride_az = stride_az_in
+        stride_ah = stride_ah_in
+
+    # program ids
+    hkid = tl.program_id(0)
+    pid = tl.program_id(1)
+    bid = tl.program_id(2)
+    if DEBUG_TRITON:
+        print(f"\npid: {pid}, bid: {bid}, hkid: {hkid}")  # noqa: E701
+    # figure out varlen start and end
+    q_start = 0
+    k_start = 0
+    seqlen_q = max_seqlen_q
+    seqlen_k = max_seqlen_k
+    if IS_VARLEN:
+        # Compute actual sequence lengths
+        q_start = tl.load(cu_seqlens_q + bid)
+        q_end = tl.load(cu_seqlens_q + bid + 1)
+        k_start = tl.load(cu_seqlens_k + bid)
+        k_end = tl.load(cu_seqlens_k + bid + 1)
+        seqlen_q = q_end - q_start
+        seqlen_k = k_end - k_start
+
+    PADDED_HEAD: tl.constexpr = ACTUAL_HEAD_DIM != HEAD_DIM
+    HAS_PE: tl.constexpr = PE_HEAD_DIM > 0
+    offs_d = tl.arange(0, HEAD_DIM)
+    if HAS_PE:
+        offs_d_pe = HEAD_DIM + tl.arange(0, PE_HEAD_DIM)
+    GROUP_SIZE: tl.constexpr = HQ // HK
+
+    start_n = pid * BLOCK_N1
+    if start_n < seqlen_k:
+        dk = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
+        if HAS_PE:
+            dk_pe = tl.zeros([BLOCK_N1, PE_HEAD_DIM], dtype=tl.float32)
+        else:
+            # Couldn't assign None to dk_pe because _bwd_dkdv_inner can't return None.
+            dk_pe = dk
+        dv = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
+
+        offs_n = start_n + tl.arange(0, BLOCK_N1)
+        # Mask for loading K and V
+        mask_kv = offs_n[:, None] < seqlen_k
+        if PADDED_HEAD:
+            mask_d = offs_d < ACTUAL_HEAD_DIM
+            mask_kv &= mask_d[None, :]
+        # NOTE: don't assume that the strides for k and v are the same!
+        # K/V tensors not changed for the group
+        adj_k = (
+            bid * stride_kb
+            + hkid * stride_kh
+            + k_start * stride_kn
+            + offs_n[:, None] * stride_kn
+            + offs_d[None, :] * stride_kd
+        )
+        if HAS_PE:
+            adj_k_pe = (
+                bid * stride_kb
+                + hkid * stride_kh
+                + k_start * stride_kn
+                + offs_n[:, None] * stride_kn
+                + offs_d_pe[None, :] * stride_kd
+            )
+        adj_v = (
+            bid * stride_vb
+            + hkid * stride_vh
+            + k_start * stride_vn
+            + offs_n[:, None] * stride_vn
+            + offs_d[None, :] * stride_vd
+        )
+        # load K and V: they stay in SRAM throughout the inner loop.
+        k = tl.load(K + adj_k, mask=mask_kv, other=0.0)
+        if HAS_PE:
+            k_pe = tl.load(K + adj_k_pe, mask=mask_kv, other=0.0)
+        else:
+            k_pe = None
+        v = tl.load(V + adj_v, mask=mask_kv, other=0.0)
+        # If MQA / GQA, set the K and V head offsets appropriately.
+        for hqid in range(hkid * GROUP_SIZE, hkid * GROUP_SIZE + GROUP_SIZE):
+            # offset input and output tensor by batch and Q/K heads
+            adj_q = bid * stride_qb + hqid * stride_qh + q_start * stride_qm
+            Q_ptr = Q + adj_q
+            adj_do = bid * stride_dob + hqid * stride_doh + q_start * stride_dom
+            DO_ptr = DO + adj_do
+            adj_delta = (
+                bid * stride_deltab + hqid * stride_deltah + q_start * stride_deltam
+            )
+            M_ptr = M + adj_delta
+            Delta_ptr = Delta + adj_delta
+
+            if USE_ALIBI:
+                alibi_offset = bid * stride_az + hqid * stride_ah
+                alibi_slope = tl.load(Alibi_slopes + alibi_offset)
+            else:
+                alibi_slope = None
+
+            # batch_philox_offset is the ACTUALLY dropout offset
+            # dropout_offset is for debug purpose and will be removed later
+            batch_philox_offset = 0
+            dropout_offset = 0
+            if ENABLE_DROPOUT:
+                batch_philox_offset = (
+                    philox_offset_base + bid * stride_dropoutb + hqid * stride_dropouth
+                )
+                dropout_offset = (
+                    Dropout_mask + bid * stride_dropoutb + hqid * stride_dropouth
+                )
+
+            if IS_FP8:
+                descale_q = tl.load(Descale_q + bid * stride_descale_q_z + hqid)
+                descale_k = tl.load(Descale_k + bid * stride_descale_k_z + hkid)
+                descale_v = tl.load(Descale_v + bid * stride_descale_v_z + hkid)
+                descale_do = tl.load(Descale_do + bid * stride_descale_do_z + hqid)
+            else:
+                descale_q, descale_k, descale_v, descale_do = 1.0, 1.0, 1.0, 1.0
+
+            # because there is no causal, we always start from the beginning
+            start_m = 0
+            num_steps = tl.cdiv(seqlen_q, BLOCK_M1)
+            dk, dk_pe, dv = _bwd_dkdv_inner(
+                dk,  # output tensor
+                dk_pe,  # optional output tensor
+                dv,  # output tensor
+                Q_ptr,
+                k,
+                k_pe,
+                v,
+                DO_ptr,
+                M_ptr,
+                Delta_ptr,
+                sm_scale,  # input tensors
+                stride_qm,
+                stride_qd,  # strides for q
+                stride_dom,
+                stride_dod,  # strides for o
+                stride_dropoutm,
+                stride_dropoutn,  # strides for dropout
+                stride_deltam,
+                BLOCK_M1,
+                BLOCK_N1,  # block dim
+                HEAD_DIM,
+                ACTUAL_HEAD_DIM,  # head dim
+                PE_HEAD_DIM,
+                dropout_p,
+                philox_seed,
+                batch_philox_offset,
+                dropout_offset,  #
+                alibi_slope,
+                seqlen_q,
+                seqlen_k,  # max sequence length for q and k
+                start_n,
+                start_m,
+                num_steps,  # iteration numbers
+                descale_q,
+                descale_k,
+                descale_v,
+                descale_do,  # fp8 descale factors from user
+                MASK=False,  # causal masking
+                ENABLE_DROPOUT=ENABLE_DROPOUT,  # activate dropout
+                USE_ALIBI=USE_ALIBI,
+                USE_EXP2=USE_EXP2,
+                IS_FP8=IS_FP8,
+                FP8_MAX=FP8_MAX,
+                DEBUG_TRITON=DEBUG_TRITON,
+                DEBUG_TRITON_DETAIL=DEBUG_TRITON_DETAIL,
+                SLIDING_WINDOW=SLIDING_WINDOW,
+            )
+
+        # Write back dV
+        adj_dv = bid * stride_dvb + hkid * stride_dvh + k_start * stride_dvn
+        offs_dv = offs_n[:, None] * stride_dvn + offs_d[None, :] * stride_dvd
+        tl.store(DV + adj_dv + offs_dv, dv, mask=mask_kv)
+        # write back dk
+        adj_dk = bid * stride_dkb + hkid * stride_dkh + k_start * stride_dkn
+        offs_dk = offs_n[:, None] * stride_dkn + offs_d[None, :] * stride_dkd
+        dk *= sm_scale
+        tl.store(DK + adj_dk + offs_dk, dk, mask=mask_kv)
+        if HAS_PE:
+            offs_dk_pe = offs_n[:, None] * stride_dkn + offs_d_pe[None, :] * stride_dkd
+            dk_pe *= sm_scale
+            tl.store(DK + adj_dk + offs_dk_pe, dk_pe, mask=mask_kv)
+
+    # THIS PART DOES DQ
+    start_m = pid * BLOCK_M2
+    if start_m < seqlen_q:
+        offs_m = start_m + tl.arange(0, BLOCK_M2)
+        # Mask for loading K and V
+        mask_q = offs_m[:, None] < seqlen_q
+        if PADDED_HEAD:
+            mask_d = offs_d < ACTUAL_HEAD_DIM
+            mask_q &= mask_d[None, :]
+        offs_q = offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd
+        if HAS_PE:
+            offs_q_pe = offs_m[:, None] * stride_qm + offs_d_pe[None, :] * stride_qd
+        offs_do = offs_m[:, None] * stride_dom + offs_d[None, :] * stride_dod
+        K += bid * stride_kb + hkid * stride_kh + k_start * stride_kn
+        V += bid * stride_vb + hkid * stride_vh + k_start * stride_vn
+        # If MQA / GQA, set the K and V head offsets appropriately.
+        for hqid in range(hkid * GROUP_SIZE, hkid * GROUP_SIZE + GROUP_SIZE):
+            # offset input and output tensor by batch and Q/K heads
+            adj_q = bid * stride_qb + hqid * stride_qh + q_start * stride_qm
+            adj_do = bid * stride_dob + hqid * stride_doh + q_start * stride_dom
+            adj_delta = (
+                bid * stride_deltab + hqid * stride_deltah + q_start * stride_deltam
+            )
+            Delta_ptr = Delta + adj_delta
+
+            if USE_ALIBI:
+                alibi_offset = bid * stride_az + hqid * stride_ah
+                alibi_slope = tl.load(Alibi_slopes + alibi_offset)
+            else:
+                alibi_slope = None
+
+            # batch_philox_offset is the ACTUALLY dropout offset
+            # dropout_offset is for debug purpose and will be removed later
+            batch_philox_offset = 0
+            dropout_offset = 0
+            if ENABLE_DROPOUT:
+                batch_philox_offset = (
+                    philox_offset_base + bid * stride_dropoutb + hqid * stride_dropouth
+                )
+                dropout_offset = (
+                    Dropout_mask + bid * stride_dropoutb + hqid * stride_dropouth
+                )
+
+            q = tl.load(Q + adj_q + offs_q, mask=mask_q, other=0.0)
+            if HAS_PE:
+                q_pe = tl.load(Q + adj_q + offs_q_pe, mask=mask_q, other=0.0)
+            else:
+                q_pe = None
+            do = tl.load(DO + adj_do + offs_do, mask=mask_q, other=0.0)
+            mask_m = offs_m < seqlen_q
+            m = tl.load(M + adj_delta + offs_m * stride_deltam, mask=mask_m, other=0.0)
+            m = m[:, None]
+            delta = tl.load(Delta_ptr + offs_m * stride_deltam, mask=mask_m, other=0.0)
+
+            if IS_FP8:
+                descale_q = tl.load(Descale_q + bid * stride_descale_q_z + hqid)
+                descale_k = tl.load(Descale_k + bid * stride_descale_k_z + hkid)
+                descale_v = tl.load(Descale_v + bid * stride_descale_v_z + hkid)
+                descale_do = tl.load(Descale_do + bid * stride_descale_do_z + hqid)
+            else:
+                descale_q, descale_k, descale_v, descale_do = 1.0, 1.0, 1.0, 1.0
+
+            # start can only be 0 at minimum
+            start_n = 0
+            end_n = seqlen_k
+            num_steps = tl.cdiv(seqlen_k, BLOCK_N2)
+
+            dq = tl.zeros([BLOCK_M2, HEAD_DIM], dtype=tl.float32)
+            if HAS_PE:
+                dq_pe = tl.zeros([BLOCK_M2, PE_HEAD_DIM], dtype=tl.float32)
+            else:
+                dq_pe = dq  # Couldn't assign None to dq_pe because _bwd_dq_inner can't return None.
+            dq, dq_pe, delta_recomp = _bwd_dq_inner(
+                dq,  # output tensor
+                dq_pe,  # optional output tensor
+                q,
+                q_pe,
+                K,
+                V,
+                do,
+                m,
+                delta,
+                sm_scale,
+                stride_qm,
+                stride_qd,
+                stride_kn,
+                stride_kd,
+                stride_vn,
+                stride_vd,
+                stride_dropoutm,
+                stride_dropoutn,
+                seqlen_q,
+                seqlen_k,
+                BLOCK_M2,
+                BLOCK_N2,
+                HEAD_DIM,
+                ACTUAL_HEAD_DIM,
+                PE_HEAD_DIM,
+                dropout_p,
+                philox_seed,
+                batch_philox_offset,
+                dropout_offset,
+                alibi_slope,
+                start_m,
+                start_n,
+                end_n,
+                num_steps,
+                descale_q,
+                descale_k,
+                descale_v,
+                descale_do,
+                MASK=False,
+                ENABLE_DROPOUT=ENABLE_DROPOUT,
+                USE_ALIBI=USE_ALIBI,
+                USE_EXP2=USE_EXP2,
+                IS_FP8=IS_FP8,
+                FP8_MAX=FP8_MAX,
+                DEBUG_TRITON=DEBUG_TRITON,
+                DEBUG_TRITON_DETAIL=DEBUG_TRITON_DETAIL,
+                SLIDING_WINDOW=SLIDING_WINDOW,
+                ENABLE_SINK=ENABLE_SINK,
+            )
+            if ENABLE_SINK:
+                sink = tl.load(Sink + hqid).to(tl.float32)
+                if USE_EXP2:
+                    RCP_LN2: tl.constexpr = 1.4426950408889634
+                    psink = tl.math.exp2(sink * RCP_LN2 - m * RCP_LN2)
+                else:
+                    psink = tl.math.exp(sink - m)
+                dsink = tl.sum(-psink * delta_recomp[:, None])
+                tl.atomic_add(DSink + hqid, dsink, sem="relaxed")
+            # Write back dQ.
+            adj_dq = bid * stride_dqb + hqid * stride_dqh + q_start * stride_dqm
+            offs_dq = offs_m[:, None] * stride_dqm + offs_d[None, :] * stride_dqd
+            dq *= sm_scale
+            tl.store(DQ + adj_dq + offs_dq, dq, mask=mask_q)
+            if HAS_PE:
+                offs_dq_pe = (
+                    offs_m[:, None] * stride_dqm + offs_d_pe[None, :] * stride_dqd
+                )
+                dq_pe *= sm_scale
+                tl.store(DQ + adj_dq + offs_dq_pe, dq_pe, mask=mask_q)
+
+
+@functools.lru_cache(maxsize=1024)
+def _get_config():
+    if not hasattr(_get_config, "_config_dict"):
+        dev = arch_info.get_arch()
+        _get_config._config_dict = {}
+        fpath = f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHA-DEFAULT.json"
+        with open(fpath, "r") as file:
+            config = json.load(file)
+        _get_config._config_dict = config
+
+    return _get_config._config_dict["bkwd_onekernel"]

build/torch-rocm/_ops.py

@@ -0,0 +1,38 @@
+import torch
+
+def get_backend() -> str:
+    """Detect the backend by inspecting torch."""
+    import torch
+
+    if hasattr(torch, "neuron"):
+        # Needs to be sorted before specific Torch builds, since Neuron
+        # extension can be loaded into e.g. CUDA Torch builds.
+        return "neuron"
+    elif torch.version.cuda is not None:
+        return "cuda"
+    elif torch.version.hip is not None:
+        return "rocm"
+    elif torch.backends.mps.is_available():
+        return "metal"
+    elif hasattr(torch.version, "xpu") and torch.version.xpu is not None:
+        return "xpu"
+    else:
+        return "cpu"
+
+
+def _find_ops_name() -> str:
+    kernel_name = "aiter_flash_attn"
+    unique_id = "27f5cac"
+    backend = get_backend()
+    return f"_{kernel_name}_{backend}_{unique_id}"
+
+
+_OPS_NAME = _find_ops_name()
+
+ops = getattr(torch.ops, _OPS_NAME)
+
+def add_op_namespace_prefix(op_name: str) -> str:
+    """
+    Prefix op by namespace.
+    """
+    return f"{_OPS_NAME}::{op_name}"

build/torch-rocm/aiter_flash_attn/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch-rocm/metadata.json

@@ -0,0 +1,10 @@
+{
+  "name": "aiter-flash-attn",
+  "id": "_aiter_flash_attn_rocm_27f5cac",
+  "version": 1,
+  "license": "MIT",
+  "python-depends": [],
+  "backend": {
+    "type": "rocm"
+  }
+}

build/torch-rocm/mha.py

@@ -0,0 +1,1000 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+from typing import Literal, Optional, Tuple, Union
+import torch
+import triton
+import triton.language as tl
+
+from .utils import types
+from .mha_onekernel_bwd import flash_attn_onekernel_backward
+from .mha_fused_bwd import flash_attn_fused_backward
+from .utils.logger import AiterTritonLogger
+from .utils.device_info import get_num_xcds
+from ._kernels.mha import _attn_fwd, _get_config
+from ._kernels.flash_attn_triton_amd import flash_attn_2
+
+_LOGGER = AiterTritonLogger()
+
+global _USE_FUSED_BWD_KERNEL
+_USE_FUSED_BWD_KERNEL = False
+
+
+def mha_set_use_fused_bwd_kernel(value: bool):
+    """
+    Set whether to use fused backward kernel (with atomics) or one-kernel backward (without atomics).
+    Fused backward is faster but doesn't support positional encoding.
+    """
+    global _USE_FUSED_BWD_KERNEL
+    _USE_FUSED_BWD_KERNEL = value
+
+
+_MHA_IMPL: Literal["default", "dao_ai"] = "default"
+
+
+def mha_set_impl(impl: Literal["default", "dao_ai"]):
+    """Set MHA forward implementation: 'default' (_attn_fwd) or 'dao_ai' (flash_attn_triton_amd)."""
+    global _MHA_IMPL
+    _MHA_IMPL = impl
+
+
+_USE_INT64_STRIDES = True
+
+
+def mha_set_use_int64_strides(value: bool):
+    """Use 64-bit integer strides to prevent integer overflows with very large tensors."""
+    global _USE_INT64_STRIDES
+    _USE_INT64_STRIDES = value
+
+
+def _get_sliding_window_size(window_size: Tuple[int, int]) -> int:
+    return int(window_size[0]) if int(window_size[0]) >= 0 else 0
+
+
+def _flash_attn_forward(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    bias: Optional[torch.Tensor],
+    alibi_slopes: Optional[torch.Tensor],
+    return_lse: bool,  # Not used
+    return_softmax: bool,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    descale_q: Optional[torch.Tensor] = None,
+    descale_k: Optional[torch.Tensor] = None,
+    descale_v: Optional[torch.Tensor] = None,
+    sink: Optional[torch.Tensor] = None,
+    config: Optional[dict[str, any]] = None,
+) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], int, int]:
+
+    if bias is not None:
+        raise ValueError("Bias is not supported yet in the Triton Backend")
+    if window_size_right != -1:
+        raise ValueError("window_size_right is not supported yet in the Triton Backend")
+    sliding_window = window_size_left if window_size_left >= 0 else 0
+
+    # Triton cannot specialize on numpy scalar types; ensure native Python int
+    max_seqlen_q = int(max_seqlen_q)
+    max_seqlen_k = int(max_seqlen_k)
+
+    # FP8
+    IS_FP8 = types._is_fp8(q)
+    FP8_MAX: tl.constexpr = torch.finfo(q.dtype).max
+    is_varlen = True if cu_seqlens_q is not None else False
+
+    if IS_FP8:
+        o = torch.zeros(
+            (q.shape[:-1] + v.shape[-1:]), dtype=torch.float32, device=q.device
+        )
+    else:
+        o = torch.zeros((q.shape[:-1] + v.shape[-1:]), dtype=q.dtype, device=q.device)
+    if is_varlen:
+        # Layout is thd.
+        # q and k are [total_tokens, num_head, head_dim_qk].
+        # v is [total_tokens, num_head, head_dim_v].
+        batch, seqlen_q, num_q_heads = (
+            len(cu_seqlens_q) - 1,
+            max_seqlen_q,
+            q.shape[1],
+        )
+        num_k_heads = k.shape[1]
+        q_strides = (0, q.stride(1), q.stride(0), q.stride(2))
+        k_strides = (0, k.stride(1), k.stride(0), k.stride(2))
+        v_strides = (0, v.stride(1), v.stride(0), v.stride(2))
+        o_strides = (0, o.stride(1), o.stride(0), o.stride(2))
+    else:
+        # Layout is bshd.
+        # q and k are [batch, seq_len, num_head, head_dim_qk].
+        # v is [batch, seq_len, num_head, head_dim_v].
+        batch, seqlen_q, num_q_heads = (int(x) for x in q.shape[:-1])
+        num_k_heads = k.shape[2]
+        q_strides = (q.stride(0), q.stride(2), q.stride(1), q.stride(3))
+        k_strides = (k.stride(0), k.stride(2), k.stride(1), k.stride(3))
+        v_strides = (v.stride(0), v.stride(2), v.stride(1), v.stride(3))
+        o_strides = (o.stride(0), o.stride(2), o.stride(1), o.stride(3))
+
+    qk_head_dim = q.shape[-1]
+    v_head_dim = v.shape[-1]
+    pe_head_dim = qk_head_dim - v_head_dim
+    # padding for head_dim. Power of 2 or 16
+    BLOCK_DMODEL_POW2 = max(triton.next_power_of_2(v_head_dim), 16)
+    BLOCK_DMODEL_PE_POW2 = (
+        0 if pe_head_dim == 0 else max(triton.next_power_of_2(pe_head_dim), 16)
+    )
+    assert (pe_head_dim == 0 and BLOCK_DMODEL_PE_POW2 == 0) or (
+        v_head_dim == BLOCK_DMODEL_POW2 and pe_head_dim == BLOCK_DMODEL_PE_POW2
+    ), "Positional encoding support requires NOPE and PE head sizes to be unpadded powers of 2."
+    assert (not IS_FP8) or (
+        IS_FP8 and pe_head_dim == 0
+    ), "Positional encoding doesn't support FP8."
+
+    assert (sink is None) or (
+        sink is not None and sink.dim() == 1 and sink.shape[0] == num_q_heads
+    ), "Sink must be 1D and have one element per query head."
+
+    # softmax_lse [batch, num_q_heads, seqlen_q]
+    if is_varlen:
+        softmax_lse = torch.zeros(
+            (q.shape[0], num_q_heads), device=q.device, dtype=torch.float32
+        )
+        stride_lse_z, stride_lse_h, stride_lse_m = (
+            0,
+            softmax_lse.stride(1),
+            softmax_lse.stride(0),
+        )
+    else:
+        softmax_lse = torch.zeros(
+            (batch, num_q_heads, max_seqlen_q), device=q.device, dtype=torch.float32
+        )
+        stride_lse_z, stride_lse_h, stride_lse_m = softmax_lse.stride()
+
+    # exp_scores [batch, num_q_heads, seqlen_q, seqlen_k]
+    enable_dropout = dropout_p > 0.0
+    if enable_dropout:
+        philox_seed = torch.randint(0, 0xFFFFFF, (1,))[
+            0
+        ].item()  # No specific reason to restrict range to 0xffffff
+        philox_offset = torch.randint(0, 0xFFFFFF, (1,))[
+            0
+        ].item()  # Pass in an int, not Tensor
+    else:
+        philox_seed = 0
+        philox_offset = 0
+    if return_softmax or enable_dropout:
+        s_dmask = torch.zeros(
+            (batch, num_q_heads, max_seqlen_q, max_seqlen_k),
+            device=q.device,
+            dtype=torch.float32,
+        )
+        dropout_mask = torch.zeros(
+            (batch, num_q_heads, max_seqlen_q, max_seqlen_k),
+            device=q.device,
+            dtype=torch.float32,
+        )
+    else:
+        s_dmask = None
+        dropout_mask = None
+
+    if _MHA_IMPL == "dao_ai":
+        assert sink is None, "dao_ai impl does not support attention sink."
+        assert (
+            pe_head_dim == 0
+        ), "dao_ai impl does not support positional encoding (pe_head_dim > 0)."
+        assert (
+            not IS_FP8
+        ), "dao_ai impl does not support FP8. Use the default impl or FA3 path."
+        assert (
+            window_size_left == -1 and window_size_right == -1
+        ), "dao_ai impl does not support sliding window attention."
+        if is_varlen:
+            o, softmax_lse, s_dmask, _ = flash_attn_2.varlen_fwd(
+                q,
+                k,
+                v,
+                o,
+                cu_seqlens_q,
+                cu_seqlens_k,
+                seqused_k=None,
+                leftpad_k=None,
+                block_table_=None,
+                alibi_slopes=alibi_slopes,
+                max_seqlen_q=max_seqlen_q,
+                max_seqlen_k=max_seqlen_k,
+                dropout_p=dropout_p,
+                softmax_scale=softmax_scale,
+                zero_tensors=False,
+                causal=causal,
+                window_size_left=-1,
+                window_size_right=-1,
+                softcap=0.0,
+                return_softmax=return_softmax,
+            )
+        else:
+            o, softmax_lse, s_dmask, _ = flash_attn_2.fwd(
+                q,
+                k,
+                v,
+                o,
+                alibi_slopes,
+                dropout_p,
+                softmax_scale,
+                causal,
+                window_size_left=-1,
+                window_size_right=-1,
+                softcap=0.0,
+                return_softmax=return_softmax,
+            )
+        # Verify softmax_lse shape contract:
+        #   non-varlen: (batch, nheads_q, seqlen_q)
+        #   varlen:     (nheads_q, total_q)  — transposed vs default impl
+        if is_varlen:
+            assert softmax_lse.shape == (
+                num_q_heads,
+                q.shape[0],
+            ), f"dao_ai varlen softmax_lse shape {softmax_lse.shape} != expected ({num_q_heads}, {q.shape[0]})"
+        else:
+            assert (
+                softmax_lse.shape[0] == batch and softmax_lse.shape[1] == num_q_heads
+            ), f"dao_ai softmax_lse shape {softmax_lse.shape} != expected (batch={batch}, nheads={num_q_heads}, ...)"
+    else:
+        if config is None:
+            config = _get_config(enable_dropout, q.dtype, has_pe=pe_head_dim > 0)
+
+        grid = lambda META: (  # noqa: E731
+            batch * num_q_heads * triton.cdiv(seqlen_q, META["BLOCK_M"]),
+        )
+
+        _attn_fwd[grid](
+            q,
+            k,
+            v,
+            descale_q,
+            descale_k,
+            descale_v,
+            o,
+            alibi_slopes,
+            s_dmask,
+            dropout_mask,
+            softmax_lse,
+            sink,
+            *q_strides,
+            *k_strides,
+            *v_strides,
+            descale_q.stride(0) if descale_q is not None else 0,
+            descale_k.stride(0) if descale_k is not None else 0,
+            descale_v.stride(0) if descale_v is not None else 0,
+            *o_strides,
+            alibi_slopes.stride(0) if alibi_slopes is not None else 0,
+            alibi_slopes.stride(1) if alibi_slopes is not None else 0,
+            s_dmask.stride(0) if s_dmask is not None else 0,
+            s_dmask.stride(1) if s_dmask is not None else 0,
+            s_dmask.stride(2) if s_dmask is not None else 0,
+            s_dmask.stride(3) if s_dmask is not None else 0,
+            stride_lse_z if softmax_lse is not None else 0,
+            stride_lse_h if softmax_lse is not None else 0,
+            stride_lse_m if softmax_lse is not None else 0,
+            softmax_scale,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            dropout_p,
+            philox_seed,
+            philox_offset,
+            SEQLEN_Q=max_seqlen_q,
+            SEQLEN_K=max_seqlen_k,
+            IS_CAUSAL=causal,
+            NUM_Q_HEADS=num_q_heads,
+            NUM_K_HEADS=num_k_heads,
+            BLOCK_DMODEL=v_head_dim,
+            BLOCK_DMODEL_POW2=BLOCK_DMODEL_POW2,
+            BLOCK_DMODEL_PE=pe_head_dim,
+            RETURN_SCORES=return_softmax,
+            ENABLE_DROPOUT=enable_dropout,
+            IS_FP8=IS_FP8,
+            FP8_MAX=FP8_MAX,
+            VARLEN=is_varlen,
+            BATCH=batch,
+            NUM_XCD=get_num_xcds(),
+            USE_INT64_STRIDES=_USE_INT64_STRIDES,
+            ENABLE_SINK=sink is not None,
+            SLIDING_WINDOW=sliding_window,
+            **config,
+        )
+
+    return o, softmax_lse, s_dmask, philox_seed, philox_offset
+
+
+class _FlashAttnFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        k,
+        v,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        bias,
+        alibi_slopes,
+        deterministic,
+        return_lse,
+        return_softmax,
+        sink,
+        is_grad_enabled,
+        config=None,
+    ):
+        is_grad = is_grad_enabled and any(
+            x is not None and x.requires_grad for x in [q, k, v, sink]
+        )
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1] ** (-0.5)
+        head_size_og = q.size(3)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, philox_seed, philox_offset = (
+            _flash_attn_forward(
+                q,
+                k,
+                v,
+                dropout_p,
+                softmax_scale,
+                causal=causal,
+                window_size_left=int(window_size[0]),
+                window_size_right=int(window_size[1]),
+                bias=bias,
+                alibi_slopes=alibi_slopes,
+                return_lse=return_lse,
+                return_softmax=return_softmax and dropout_p > 0,
+                max_seqlen_q=q.shape[1],
+                max_seqlen_k=k.shape[1],
+                sink=sink,
+                config=config,
+            )
+        )
+
+        if is_grad:
+            ctx.save_for_backward(q, k, v, out_padded, softmax_lse, sink)
+            ctx.philox_seed = philox_seed
+            ctx.philox_offset = philox_offset
+            ctx.dropout_p = dropout_p
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.bias = bias
+            ctx.window_size = window_size
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+
+        out = out_padded[..., :head_size_og]
+        result = [out]
+        if return_lse:
+            result.append(softmax_lse)
+        if return_softmax:
+            result.append(S_dmask)
+
+        return result[0] if len(result) == 1 else tuple(result)
+
+    @staticmethod
+    def backward(ctx, do, *args):
+        q, k, v, out, softmax_lse, sink = ctx.saved_tensors
+        bias = ctx.bias
+        dbias = torch.empty_like(bias) if bias is not None else None
+        dq, dk, dv = torch.zeros_like(q), torch.empty_like(k), torch.empty_like(v)
+        dsink = (
+            torch.zeros_like(sink, dtype=torch.float32) if sink is not None else None
+        )
+        head_size_v_og = do.size(3)
+        do_padded = do
+        if head_size_v_og % 8 != 0:
+            do_padded = torch.nn.functional.pad(do, [0, 8 - head_size_v_og % 8])
+        sliding_window = _get_sliding_window_size(ctx.window_size)
+
+        if _MHA_IMPL == "dao_ai":
+            assert sink is None, "dao_ai impl does not support attention sink."
+            flash_attn_2.bwd(
+                do_padded,
+                q,
+                k,
+                v,
+                out,
+                softmax_lse,
+                dq,
+                dk,
+                dv,
+                ctx.alibi_slopes,
+                ctx.dropout_p,
+                ctx.softmax_scale,
+                ctx.causal,
+                window_size_left=-1,
+                window_size_right=-1,
+                softcap=0.0,
+                deterministic=ctx.deterministic,
+            )
+        else:
+            if _USE_FUSED_BWD_KERNEL:
+                if sliding_window > 0:
+                    raise ValueError(
+                        "Fused backward doesn't support sliding window attention. "
+                        "Disable fused backward or use the one-kernel backward."
+                    )
+                assert (
+                    sink is None and dsink is None
+                ), "Fused backward doesn't support sinks."
+                flash_attn_fused_backward(
+                    do_padded,
+                    q,
+                    k,
+                    v,
+                    out,
+                    softmax_lse,
+                    dq,
+                    dk,
+                    dv,
+                    dbias,
+                    ctx.softmax_scale,
+                    ctx.alibi_slopes,
+                    ctx.causal,
+                    None,
+                    None,
+                    max_seqlen_q=q.shape[1],
+                    max_seqlen_k=k.shape[1],
+                    dropout_p=ctx.dropout_p,
+                    philox_seed=ctx.philox_seed,
+                    philox_offset=ctx.philox_offset,
+                    USE_INT64_STRIDES=_USE_INT64_STRIDES,
+                )
+            else:
+                flash_attn_onekernel_backward(
+                    do_padded,
+                    q,
+                    k,
+                    v,
+                    out,
+                    softmax_lse,
+                    dq,
+                    dk,
+                    dv,
+                    dbias,
+                    ctx.softmax_scale,
+                    ctx.alibi_slopes,
+                    ctx.causal,
+                    None,
+                    None,
+                    max_seqlen_q=q.shape[1],
+                    max_seqlen_k=k.shape[1],
+                    dropout_p=ctx.dropout_p,
+                    philox_seed=ctx.philox_seed,
+                    philox_offset=ctx.philox_offset,
+                    USE_INT64_STRIDES=_USE_INT64_STRIDES,
+                    sink=sink,
+                    dsink=dsink,
+                    sliding_window=sliding_window,
+                )
+
+        dq = dq[..., : q.shape[-1]]  # We could have padded the head dimension
+        dk = dk[..., : k.shape[-1]]
+        dv = dv[..., : v.shape[-1]]
+        return (
+            dq,
+            dk,
+            dv,
+            None,  # dropout_p
+            None,  # softmax_scale
+            None,  # causal
+            None,  # window_size
+            dbias,
+            None,  # alibi_slopes
+            None,  # deterministic
+            None,  # return_lse
+            None,  # return_softmax
+            dsink,
+            None,  # is_grad_enabled
+            None,  # config
+        )
+
+
+def flash_attn_func(
+    q,
+    k,
+    v,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    bias=None,
+    alibi_slopes=None,
+    deterministic=True,
+    return_lse=False,
+    return_attn_probs=False,
+    sink=None,
+    config: Optional[dict[str, any]] = None,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+
+    Arguments:
+        q: (batch_size, seqlen, nheads, headdim)
+        k: (batch_size, seqlen, nheads_k, headdim)
+        v: (batch_size, seqlen, nheads_k, headdim)
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim_q).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        bias: (seqlen_q, seqlen_k)
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+        sink: (nheads,), attention sink scores (one per Q head), or None
+    Return:
+        out: (batch_size, seqlen, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    _LOGGER.info(
+        f"FLASH_ATTN:  q={tuple(q.shape)}  k={tuple(k.shape)}  v={tuple(v.shape)}"
+    )
+    return _FlashAttnFunc.apply(
+        q,
+        k,
+        v,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        bias,
+        alibi_slopes,
+        deterministic,
+        return_lse,
+        return_attn_probs,
+        sink,
+        torch.is_grad_enabled(),
+        config,
+    )
+
+
+class _FlashAttnVarlenFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        k,
+        v,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        bias,
+        alibi_slopes,
+        deterministic,
+        return_lse,
+        return_softmax,
+        block_table,
+        out,
+        sink,
+        is_grad_enabled,
+        config=None,
+    ):
+        is_grad = is_grad_enabled and any(
+            x is not None and x.requires_grad for x in [q, k, v, sink]
+        )
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1] ** (-0.5)
+        head_size_og = q.size(2)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, philox_seed, philox_offset = (
+            _flash_attn_forward(
+                q,
+                k,
+                v,
+                dropout_p,
+                softmax_scale,
+                causal=causal,
+                window_size_left=int(window_size[0]),
+                window_size_right=int(window_size[1]),
+                bias=bias,
+                alibi_slopes=alibi_slopes,
+                return_lse=return_lse,
+                return_softmax=return_softmax and dropout_p > 0.0,
+                max_seqlen_q=max_seqlen_q,
+                max_seqlen_k=max_seqlen_k,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_k=cu_seqlens_k,
+                sink=sink,
+                config=config,
+            )
+        )
+        if is_grad:
+            ctx.save_for_backward(
+                q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, sink
+            )
+            ctx.max_seqlen_q = max_seqlen_q
+            ctx.max_seqlen_k = max_seqlen_k
+            ctx.philox_seed = philox_seed
+            ctx.philox_offset = philox_offset
+            ctx.dropout_p = dropout_p
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.bias = bias
+            ctx.alibi_slopes = alibi_slopes
+        out = out_padded[..., :head_size_og]
+
+        result = [out]
+        if return_lse:
+            result.append(softmax_lse)
+        if return_softmax:
+            result.append(S_dmask)
+
+        return result[0] if len(result) == 1 else tuple(result)
+
+    @staticmethod
+    def backward(ctx, do, *args):
+        q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, sink = ctx.saved_tensors
+        dq, dk, dv = torch.zeros_like(q), torch.empty_like(k), torch.empty_like(v)
+        bias = ctx.bias
+        dbias = torch.empty_like(bias) if bias is not None else None
+        dsink = (
+            torch.zeros_like(sink, dtype=torch.float32) if sink is not None else None
+        )
+        head_size_og = do.size(2)
+        do_padded = do
+        if head_size_og % 8 != 0:
+            do_padded = torch.nn.functional.pad(do, [0, 8 - head_size_og % 8])
+        sliding_window = _get_sliding_window_size(ctx.window_size)
+
+        if _MHA_IMPL == "dao_ai":
+            assert sink is None, "dao_ai impl does not support attention sink."
+            flash_attn_2.varlen_bwd(
+                do_padded,
+                q,
+                k,
+                v,
+                out,
+                softmax_lse,
+                dq,
+                dk,
+                dv,
+                cu_seqlens_q,
+                cu_seqlens_k,
+                ctx.alibi_slopes,
+                max_seqlen_q=ctx.max_seqlen_q,
+                max_seqlen_k=ctx.max_seqlen_k,
+                dropout_p=ctx.dropout_p,
+                softmax_scale=ctx.softmax_scale,
+                zero_tensors=False,
+                causal=ctx.causal,
+                window_size_left=-1,
+                window_size_right=-1,
+                softcap=0.0,
+                deterministic=False,
+            )
+        else:
+            if _USE_FUSED_BWD_KERNEL:
+                if sliding_window > 0:
+                    raise ValueError(
+                        "Fused backward doesn't support sliding window attention. "
+                        "Disable fused backward or use the one-kernel backward."
+                    )
+                assert (
+                    sink is None and dsink is None
+                ), "Fused backward doesn't support sinks."
+                flash_attn_fused_backward(
+                    do_padded,
+                    q,
+                    k,
+                    v,
+                    out,
+                    softmax_lse,
+                    dq,
+                    dk,
+                    dv,
+                    dbias,
+                    ctx.softmax_scale,
+                    ctx.alibi_slopes,
+                    ctx.causal,
+                    cu_seqlens_q,
+                    cu_seqlens_k,
+                    max_seqlen_q=ctx.max_seqlen_q,
+                    max_seqlen_k=ctx.max_seqlen_k,
+                    dropout_p=ctx.dropout_p,
+                    philox_seed=ctx.philox_seed,
+                    philox_offset=ctx.philox_offset,
+                    USE_INT64_STRIDES=_USE_INT64_STRIDES,
+                )
+            else:
+                flash_attn_onekernel_backward(
+                    do_padded,
+                    q,
+                    k,
+                    v,
+                    out,
+                    softmax_lse,
+                    dq,
+                    dk,
+                    dv,
+                    dbias,
+                    ctx.softmax_scale,
+                    ctx.alibi_slopes,
+                    ctx.causal,
+                    cu_seqlens_q,
+                    cu_seqlens_k,
+                    max_seqlen_q=ctx.max_seqlen_q,
+                    max_seqlen_k=ctx.max_seqlen_k,
+                    dropout_p=ctx.dropout_p,
+                    philox_seed=ctx.philox_seed,
+                    philox_offset=ctx.philox_offset,
+                    USE_INT64_STRIDES=_USE_INT64_STRIDES,
+                    sink=sink,
+                    dsink=dsink,
+                    sliding_window=sliding_window,
+                )
+
+        dq = dq[..., : q.shape[-1]]  # We could have padded the head dimension
+        dk = dk[..., : k.shape[-1]]
+        dv = dv[..., : v.shape[-1]]
+        return (
+            dq,
+            dk,
+            dv,
+            None,  # cu_seqlens_q,
+            None,  # cu_seqlens_k
+            None,  # max_seqlen_q
+            None,  # max_seqlen_k
+            None,  # dropout_p
+            None,  # softmax_scale
+            None,  # causal
+            None,  # window_size
+            dbias,
+            None,  # alibi_slopes
+            None,  # deterministic
+            None,  # return_lse
+            None,  # return_softmax
+            None,  # block_table
+            None,  # out
+            dsink,
+            None,  # is_grad_enabled
+            None,  # config
+        )
+
+
+def flash_attn_varlen_func(
+    q,
+    k,
+    v,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q,
+    max_seqlen_k,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    bias=None,
+    alibi_slopes=None,
+    deterministic=False,
+    return_lse=False,
+    return_attn_probs=False,
+    block_table=None,
+    out=None,
+    sink=None,
+    config: Optional[dict[str, any]] = None,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+
+    Arguments:
+        q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.
+        k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.
+        v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.
+        cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into q.
+        cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into kv.
+        max_seqlen_q: int. Maximum query sequence length in the batch.
+        max_seqlen_k: int. Maximum key sequence length in the batch.
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        bias: (seqlen_q, seqlen_k)
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+        sink: (nheads,), attention sink scores (one per Q head), or None
+    Return:
+        out: (total, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+
+    _LOGGER.info(
+        f"FLASH_ATTN_VARLEN:  q={tuple(q.shape)}  k={tuple(k.shape)}  v={tuple(v.shape)}"
+    )
+    return _FlashAttnVarlenFunc.apply(
+        q,
+        k,
+        v,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        bias,
+        alibi_slopes,
+        deterministic,
+        return_lse,
+        return_attn_probs,
+        block_table,
+        out,
+        sink,
+        torch.is_grad_enabled(),
+        config,
+    )
+
+
+def flash_attn_with_kvcache(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    k: Optional[torch.Tensor] = None,
+    v: Optional[torch.Tensor] = None,
+    cache_seqlens: Optional[Union[torch.Tensor, int]] = None,
+    softmax_scale: Optional[float] = None,
+    causal: bool = True,
+    window_size: tuple[int, int] = (-1, -1),
+    softcap: float = 0.0,
+    num_splits: int = 0,
+    rotary_cos: Optional[torch.Tensor] = None,
+    rotary_sin: Optional[torch.Tensor] = None,
+    cache_batch_idx: Optional[torch.Tensor] = None,
+    cache_leftpad: Optional[torch.Tensor] = None,
+    block_table: Optional[torch.Tensor] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    rotary_interleaved: bool = True,
+    return_softmax_lse: bool = False,
+):
+    """
+    This mirrors the public flash_attn v2 interface for KV cache using the AMD Triton backend.
+
+    Args:
+        q: (batch, seqlen_q, nheads_q, headdim)
+        k_cache / v_cache: Either contiguous (batch, seqlen_cache, nheads_k, headdim) or paged
+            (num_blocks, page_block_size, nheads_k, headdim) when block_table provided.
+        k, v: Optional incremental tokens to append in-place (appended logically after existing cache).
+        cache_seqlens: int or (batch,) current valid lengths per batch entry.
+        softmax_scale: Optional override; defaults to 1/sqrt(headdim).
+        causal: Apply causal masking.
+        window_size: (left, right) local attention window; (-1,-1) = full.
+        softcap: (float) currently must be 0.0 (backend limitation).
+        num_splits: 0 or 1 only (backend limitation >1).
+        rotary_cos/rotary_sin: Optional rotary embeddings (applied if provided) - interleaving flag unused here.
+        cache_batch_idx/cache_leftpad: Optional indexing / left padding metadata.
+            block_table: Optional paging table mapping logical blocks for paged KV cache.
+        alibi_slopes: (nheads,) or (batch,nheads) bias slopes (currently ignored if provided - placeholder).
+        rotary_interleaved: Flag kept for parity (currently forwarded as True constant to backend which ignores it).
+            return_softmax_lse: If True returns (out, lse) else out.
+
+    Returns:
+        out (and optionally softmax_lse): (batch, seqlen_q, nheads_q, headdim)
+    """
+    # Feature guards / normalization
+    if softcap != 0.0:
+        raise NotImplementedError(
+            "softcap != 0 not supported in v2 KV cache backend yet"
+        )
+    if num_splits not in (0, 1):
+        raise NotImplementedError(
+            "num_splits > 1 not supported in v2 KV cache backend yet"
+        )
+
+    if softmax_scale is None:
+        softmax_scale = q.shape[-1] ** (-0.5)
+
+    if cache_seqlens is not None and isinstance(cache_seqlens, int):
+        cache_seqlens = torch.full(
+            (k_cache.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device
+        )
+
+    # Contiguity (align last dim contiguous requirement similar to v3 path assumptions)
+    assert q.stride(-1) == 1 and k_cache.stride(-1) == 1 and v_cache.stride(-1) == 1
+
+    out, softmax_lse = flash_attn_2.fwd_kvcache(
+        q,
+        k_cache,
+        v_cache,
+        k,
+        v,
+        cache_seqlens,
+        rotary_cos,
+        rotary_sin,
+        cache_batch_idx,
+        cache_leftpad,
+        block_table,
+        alibi_slopes,
+        None,  # out tensor
+        softmax_scale,
+        causal,
+        int(window_size[0]),
+        int(window_size[1]),
+        0.0,  # softcap (guarded)
+        rotary_interleaved,
+        num_splits,
+    )
+    return (out, softmax_lse) if return_softmax_lse else out

build/torch-rocm/mha_fused_bwd.py

@@ -0,0 +1,317 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+from typing import Optional, Dict
+import torch
+import triton
+
+from .utils.types import _is_fp8
+from .utils.logger import AiterTritonLogger
+from ._kernels.mha_fused_bwd import (
+    _bwd_preprocess,
+    _bwd_kernel_dkdvdq_causal,
+    _bwd_kernel_dkdvdq_noncausal,
+    _get_config,
+)
+from .utils.device_info import get_num_xcds
+
+_LOGGER = AiterTritonLogger()
+
+
+def flash_attn_fused_backward(
+    do: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    o: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: torch.Tensor,
+    dk: torch.Tensor,
+    dv: torch.Tensor,
+    dbias: torch.Tensor,
+    sm_scale: float,
+    alibi_slopes: Optional[torch.Tensor],
+    causal: bool,
+    cu_seqlens_q: Optional[torch.Tensor],
+    cu_seqlens_k: Optional[torch.Tensor],
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float,
+    philox_seed: Optional[int] = 0,
+    philox_offset: Optional[int] = 0,
+    descale_q: Optional[torch.Tensor] = None,
+    descale_k: Optional[torch.Tensor] = None,
+    descale_v: Optional[torch.Tensor] = None,
+    descale_do: Optional[torch.Tensor] = None,
+    USE_INT64_STRIDES: Optional[bool] = False,
+    config: Optional[Dict[str, any]] = None,
+):
+    """
+    Flash Attention fused backward pass computing dQ, dK, dV in a single kernel using atomics.
+    Supports variable-length sequences, GQA, FP8 quantization, and dropout.
+
+    Args:
+        do (torch.Tensor): Output gradient. Shape (batch, seqlen_q, num_q_heads, head_dim)
+            or (total_tokens, num_q_heads, head_dim) for varlen.
+        q (torch.Tensor): Query tensor from forward pass with same shape as do.
+        k (torch.Tensor): Key tensor with shape (batch, seqlen_k, num_k_heads, head_dim)
+            or (total_tokens_k, num_k_heads, head_dim) for varlen.
+        v (torch.Tensor): Value tensor with same shape as k.
+        o (torch.Tensor): Output from forward pass with same shape as q.
+        softmax_lse (torch.Tensor): Log-sum-exp from forward pass with shape
+            (batch, num_q_heads, seqlen_q) or (total_tokens, num_q_heads) for varlen.
+        dq (torch.Tensor): Pre-allocated query gradient with same shape as q.
+        dk (torch.Tensor): Pre-allocated key gradient with same shape as k.
+        dv (torch.Tensor): Pre-allocated value gradient with same shape as v.
+        dbias (torch.Tensor): Bias gradient (not supported, must be None).
+        sm_scale (float): Softmax scale, typically 1/sqrt(head_dim).
+        alibi_slopes (Optional[torch.Tensor]): ALiBi position bias slopes.
+        causal (bool): Apply causal masking.
+        cu_seqlens_q (Optional[torch.Tensor]): Cumulative sequence lengths for query with shape
+            (batch + 1,). Enables variable-length mode.
+        cu_seqlens_k (Optional[torch.Tensor]): Cumulative sequence lengths for key with shape
+            (batch + 1,).
+        max_seqlen_q (int): Maximum query sequence length in batch.
+        max_seqlen_k (int): Maximum key sequence length in batch.
+        dropout_p (float): Dropout probability. 0.0 disables dropout.
+        philox_seed (Optional[int]): Random seed for dropout.
+        philox_offset (Optional[int]): Random offset for dropout.
+        descale_q (Optional[torch.Tensor]): FP8 descaling factor for q.
+        descale_k (Optional[torch.Tensor]): FP8 descaling factor for k.
+        descale_v (Optional[torch.Tensor]): FP8 descaling factor for v.
+        descale_do (Optional[torch.Tensor]): FP8 descaling factor for do.
+        USE_INT64_STRIDES (Optional[bool]): Use 64-bit stride indexing for large tensors.
+        config (Optional[Dict[str, any]]): Kernel tuning parameters (preprocess_kernel,
+            dkdvdq_kernel_N64, dkdvdq_kernel_N128).
+
+    Returns:
+        torch.Tensor: Delta tensor (element-wise product of do and o) with shape matching softmax_lse.
+    """
+    _LOGGER.info(
+        f"FLASH_ATTN_FUSED_BKWD: do={tuple(do.shape)} q={tuple(q.shape)}  k={tuple(k.shape)}  v={tuple(v.shape)} "
+        + f"dq={tuple(dq.shape)}  dk={tuple(dk.shape)}  dv={tuple(dv.shape)}"
+    )
+    if dbias is not None:
+        raise ValueError("Bias is not supported yet in the Triton Backend")
+    if q.shape[-1] == k.shape[-1] and k.shape[-1] > v.shape[-1]:
+        raise ValueError(
+            "'Fused' backward doesn't support Positional Encoding (PE). Please use 'one kernel' backward implementation for PE."
+        )
+
+    IS_FP8 = _is_fp8(q)
+    if IS_FP8:
+        FP8_MAX = torch.finfo(q.dtype).max
+        descale_strides = (
+            descale_q.stride(0),
+            descale_k.stride(0),
+            descale_v.stride(0),
+            descale_do.stride(0),
+        )
+    else:
+        FP8_MAX = None
+        stride_descale_q_z = stride_descale_k_z = stride_descale_v_z = (
+            stride_descale_do_z
+        ) = None
+        descale_strides = (
+            stride_descale_q_z,
+            stride_descale_k_z,
+            stride_descale_v_z,
+            stride_descale_do_z,
+        )
+
+    IS_VARLEN = True if cu_seqlens_q is not None else False
+
+    # get strides and shape
+    if IS_VARLEN:
+        # Layout for q,k,v is thd ie [total tokens, num_head, head_dim]
+        batch, seqlen_q, num_q_heads, head_sz = (
+            len(cu_seqlens_q) - 1,
+            max_seqlen_q,
+            q.shape[1],
+            q.shape[2],
+        )
+        _, num_k_heads = max_seqlen_k, k.shape[1]
+        q_strides = (0, q.stride(1), q.stride(0), q.stride(2))
+        q_strides = (0, q.stride(1), q.stride(0), q.stride(2))
+        k_strides = (0, k.stride(1), k.stride(0), k.stride(2))
+        v_strides = (0, v.stride(1), v.stride(0), v.stride(2))
+        o_strides = (0, o.stride(1), o.stride(0), o.stride(2))
+        dq_strides = (0, dq.stride(1), dq.stride(0), dq.stride(2))
+        dk_strides = (0, dk.stride(1), dk.stride(0), dk.stride(2))
+        do_strides = (0, do.stride(1), do.stride(0), do.stride(2))
+    else:
+        # Layout for q,k,v is bshd ie [batch, seq_len, num_head, head_dim]
+        batch, seqlen_q, num_q_heads, head_sz = q.shape
+        _, num_k_heads = k.shape[1], k.shape[2]
+        q_strides = (q.stride(0), q.stride(2), q.stride(1), q.stride(3))
+        k_strides = (k.stride(0), k.stride(2), k.stride(1), k.stride(3))
+        v_strides = (v.stride(0), v.stride(2), v.stride(1), v.stride(3))
+        o_strides = (o.stride(0), o.stride(2), o.stride(1), o.stride(3))
+        dq_strides = (dq.stride(0), dq.stride(2), dq.stride(1), dq.stride(3))
+        dk_strides = (dk.stride(0), dk.stride(2), dk.stride(1), dk.stride(3))
+        do_strides = (do.stride(0), do.stride(2), do.stride(1), do.stride(3))
+
+    # BLOCK_D_MODEL, BLOCK_D_MODEL_POW2
+    # padding for head_dim. Power of 2 or 16
+    BLOCK_D_MODEL_POW2 = triton.next_power_of_2(head_sz)
+    BLOCK_D_MODEL_POW2 = max(BLOCK_D_MODEL_POW2, 16)
+
+    # init delta
+    delta = torch.zeros_like(softmax_lse)
+    if IS_VARLEN:
+        # [total_tokens, num_q_heads, seqlen_q]
+        delta_strides = (0, delta.stride(1), delta.stride(0))
+    else:
+        # [batch, num_q_heads, seqlen_q]
+        delta_strides = delta.stride()
+
+    # preprocess
+    # compute D(delta) = rowsum(dO*O). Note, multiplication is element-wise.
+    if config is None:
+        config = _get_config()
+
+    pre_grid = (
+        triton.cdiv(max_seqlen_q, config["preprocess_kernel"]["PRE_BLOCK"]),
+        batch,
+        num_q_heads,
+    )
+
+    _bwd_preprocess[pre_grid](
+        o,
+        do,
+        delta,
+        *o_strides,
+        *do_strides,
+        *delta_strides,
+        descale_strides[3],
+        cu_seqlens_q,
+        max_seqlen_q,
+        descale_do,
+        BLOCK_M=config["preprocess_kernel"]["PRE_BLOCK"],
+        BLOCK_D_MODEL=head_sz,
+        BLOCK_D_MODEL_POW2=BLOCK_D_MODEL_POW2,
+        IS_VARLEN=IS_VARLEN,
+        IS_FP8=IS_FP8,
+    )
+    # dropout_mask
+    use_dropout = dropout_p > 0.0
+    if use_dropout:
+        dropout_mask = torch.zeros(
+            (batch, num_q_heads, max_seqlen_q, max_seqlen_k),
+            device=q.device,
+            dtype=torch.float32,
+        )
+        dropout_strides = dropout_mask.stride()
+    else:
+        dropout_mask = None
+        dropout_strides = (0, 0, 0, 0)
+
+    # Fuses dk,dv and dq computations into one kernel using atomics
+    if BLOCK_D_MODEL_POW2 > 160 or q.dtype == torch.float32:
+        config_dkdvdq = config["dkdvdq_kernel_N64"]
+    else:
+        config_dkdvdq = config["dkdvdq_kernel_N128"]
+
+    num_k_pids = (max_seqlen_k + config_dkdvdq["BLOCK_N"] - 1) // config_dkdvdq[
+        "BLOCK_N"
+    ]
+    # NUM_SMS = torch.cuda.get_device_properties("cuda").multi_processor_count
+    if causal:
+        grid_dkdvdq = (batch * num_q_heads * num_k_pids,)
+
+        _bwd_kernel_dkdvdq_causal[grid_dkdvdq](
+            q,
+            k,
+            v,
+            sm_scale,
+            do,
+            dk,
+            dv,
+            dq,
+            softmax_lse,
+            delta,
+            *q_strides,
+            *k_strides,
+            *v_strides,
+            *dk_strides,
+            *dq_strides,
+            *delta_strides,
+            *do_strides,
+            *dropout_strides,
+            *descale_strides,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            dropout_mask,
+            dropout_p,
+            philox_seed,
+            philox_offset,
+            descale_q,
+            descale_k,
+            descale_v,
+            descale_do,
+            NUM_Q_HEADS=num_q_heads,
+            NUM_K_HEADS=num_k_heads,
+            BATCH=batch,
+            NUM_K_PIDS=num_k_pids,
+            BLOCK_D_MODEL=head_sz,
+            BLOCK_D_MODEL_POW2=BLOCK_D_MODEL_POW2,
+            ENABLE_DROPOUT=use_dropout,
+            IS_VARLEN=IS_VARLEN,
+            IS_FP8=IS_FP8,
+            FP8_MAX=FP8_MAX,
+            USE_INT64_STRIDES=USE_INT64_STRIDES,
+            NUM_XCD=get_num_xcds(),
+            **config_dkdvdq,
+        )
+    else:
+        # in non causal inner loop over grouped q heads
+        grid_dkdvdq = (batch * num_k_heads * num_k_pids,)
+        _bwd_kernel_dkdvdq_noncausal[grid_dkdvdq](
+            q,
+            k,
+            v,
+            sm_scale,
+            do,
+            dk,
+            dv,
+            dq,
+            softmax_lse,
+            delta,
+            *q_strides,
+            *k_strides,
+            *v_strides,
+            *dk_strides,
+            *dq_strides,
+            *delta_strides,
+            *do_strides,
+            *dropout_strides,
+            *descale_strides,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            dropout_mask,
+            dropout_p,
+            philox_seed,
+            philox_offset,
+            descale_q,
+            descale_k,
+            descale_v,
+            descale_do,
+            NUM_Q_HEADS=num_q_heads,
+            NUM_K_HEADS=num_k_heads,
+            BATCH=batch,
+            NUM_K_PIDS=num_k_pids,
+            BLOCK_D_MODEL=head_sz,
+            BLOCK_D_MODEL_POW2=BLOCK_D_MODEL_POW2,
+            ENABLE_DROPOUT=use_dropout,
+            IS_VARLEN=IS_VARLEN,
+            IS_FP8=IS_FP8,
+            FP8_MAX=FP8_MAX,
+            USE_INT64_STRIDES=USE_INT64_STRIDES,
+            **config_dkdvdq,
+        )
+
+    return delta

build/torch-rocm/mha_onekernel_bwd.py

@@ -0,0 +1,380 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+from typing import Optional, Dict
+import torch
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+from .utils.types import _is_fp8
+from .utils.logger import AiterTritonLogger
+
+from ._kernels.mha_onekernel_bwd import (
+    _bwd_preprocess,
+    bwd_kernel_causal,
+    bwd_kernel_noncausal,
+    _get_config,
+)
+
+_LOGGER = AiterTritonLogger()
+
+
+# NOTE: triton fails to import tl.constexprs so create them here for the file
+DROPOUT_USE_PYTORCH = False
+DROPOUT_DUMP = False
+
+tl_DROPOUT_USE_PYTORCH: tl.constexpr = triton.language.constexpr(DROPOUT_USE_PYTORCH)
+tl_DROPOUT_DUMP: tl.constexpr = triton.language.constexpr(DROPOUT_DUMP)
+
+
+def flash_attn_onekernel_backward(
+    do: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    o: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: torch.Tensor,
+    dk: torch.Tensor,
+    dv: torch.Tensor,
+    dbias: torch.Tensor,
+    sm_scale: float,
+    alibi_slopes: Optional[torch.Tensor],
+    causal: bool,
+    cu_seqlens_q: Optional[torch.Tensor],
+    cu_seqlens_k: Optional[torch.Tensor],
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float,
+    philox_seed: Optional[int] = 0,
+    philox_offset: Optional[int] = 0,
+    descale_q: Optional[torch.Tensor] = None,
+    descale_k: Optional[torch.Tensor] = None,
+    descale_v: Optional[torch.Tensor] = None,
+    descale_do: Optional[torch.Tensor] = None,
+    USE_INT64_STRIDES: Optional[bool] = False,
+    sink: Optional[torch.Tensor] = None,
+    dsink: Optional[torch.Tensor] = None,
+    config: Optional[Dict[str, any]] = None,
+    sliding_window: int = 0,
+):
+    """
+    Flash Attention one-kernel backward pass with positional encoding support.
+    Computes dQ, dK, dV in separate passes without atomics. Supports Q/K head dimensions
+    larger than V for positional encoding.
+
+    Args:
+        do (torch.Tensor): Output gradient. Shape (batch, seqlen_q, num_q_heads, v_head_dim)
+            or (total_tokens, num_q_heads, v_head_dim) for varlen.
+        q (torch.Tensor): Query tensor with shape (batch, seqlen_q, num_q_heads, qk_head_dim).
+            qk_head_dim may be larger than v_head_dim for positional encoding.
+        k (torch.Tensor): Key tensor with shape (batch, seqlen_k, num_k_heads, qk_head_dim).
+        v (torch.Tensor): Value tensor with shape (batch, seqlen_k, num_k_heads, v_head_dim).
+        o (torch.Tensor): Output from forward pass with same shape as do.
+        softmax_lse (torch.Tensor): Log-sum-exp from forward pass with shape
+            (batch, num_q_heads, seqlen_q) or (total_tokens, num_q_heads) for varlen.
+        dq (torch.Tensor): Pre-allocated query gradient with same shape as q.
+        dk (torch.Tensor): Pre-allocated key gradient with same shape as k.
+        dv (torch.Tensor): Pre-allocated value gradient with same shape as v.
+        dbias (torch.Tensor): Bias gradient (not supported, must be None).
+        sm_scale (float): Softmax scale, typically 1/sqrt(head_dim).
+        alibi_slopes (Optional[torch.Tensor]): ALiBi position bias slopes with shape (num_q_heads,).
+        causal (bool): Apply causal masking.
+        cu_seqlens_q (Optional[torch.Tensor]): Cumulative sequence lengths for query with shape
+            (batch + 1,). Enables variable-length mode.
+        cu_seqlens_k (Optional[torch.Tensor]): Cumulative sequence lengths for key with shape
+            (batch + 1,).
+        max_seqlen_q (int): Maximum query sequence length in batch.
+        max_seqlen_k (int): Maximum key sequence length in batch.
+        dropout_p (float): Dropout probability. 0.0 disables dropout.
+        philox_seed (Optional[int]): Random seed for dropout.
+        philox_offset (Optional[int]): Random offset for dropout.
+        descale_q (Optional[torch.Tensor]): FP8 descaling factor for q.
+        descale_k (Optional[torch.Tensor]): FP8 descaling factor for k.
+        descale_v (Optional[torch.Tensor]): FP8 descaling factor for v.
+        descale_do (Optional[torch.Tensor]): FP8 descaling factor for do.
+        USE_INT64_STRIDES (Optional[bool]): Use 64-bit stride indexing for large tensors.
+        sink (Optional[torch.Tensor]): Attention sink scores (one per Q head). Shape (num_q_heads,).
+        dsink (Optional[torch.Tensor]): Pre-allocated sink gradient with same shape as sink.
+        config (Optional[Dict[str, any]]): Kernel tuning parameters (preprocess_kernel,
+            onekernel, onekernel_pe).
+
+    Returns:
+        torch.Tensor: Delta tensor (element-wise product of do and o) with shape matching softmax_lse.
+    """
+    _LOGGER.info(
+        f"FLASH_ATTN_ONEKERNEL_BKWD: do={tuple(do.shape)} q={tuple(q.shape)}  k={tuple(k.shape)}  v={tuple(v.shape)} "
+        + f"dq={tuple(dq.shape)}  dk={tuple(dk.shape)}  dv={tuple(dv.shape)}"
+    )
+    if dbias is not None:
+        raise ValueError("Bias is not supported yet in the Triton Backend")
+
+    use_alibi, (stride_az, stride_ah) = (
+        (True, alibi_slopes.stride()) if alibi_slopes is not None else (False, (0, 0))
+    )
+
+    IS_FP8 = _is_fp8(q)
+    if IS_FP8:
+        FP8_MAX = torch.finfo(q.dtype).max
+        descale_strides = (
+            descale_q.stride(0),
+            descale_k.stride(0),
+            descale_v.stride(0),
+            descale_do.stride(0),
+        )
+    else:
+        FP8_MAX = None
+        stride_descale_q_z = stride_descale_k_z = stride_descale_v_z = (
+            stride_descale_do_z
+        ) = None
+        descale_strides = (
+            stride_descale_q_z,
+            stride_descale_k_z,
+            stride_descale_v_z,
+            stride_descale_do_z,
+        )
+
+    IS_VARLEN = True if cu_seqlens_q is not None else False
+
+    # get strides and shape
+    if IS_VARLEN:
+        # Layout is thd.
+        # q and k are [total_tokens, num_head, head_dim_qk].
+        # v is [total_tokens, num_head, head_dim_v].
+        batch, seqlen_q, num_q_heads = (
+            len(cu_seqlens_q) - 1,
+            max_seqlen_q,
+            q.shape[1],
+        )
+        _, num_k_heads = max_seqlen_k, k.shape[1]
+        q_strides = (0, q.stride(1), q.stride(0), q.stride(2))
+        q_strides = (0, q.stride(1), q.stride(0), q.stride(2))
+        k_strides = (0, k.stride(1), k.stride(0), k.stride(2))
+        v_strides = (0, v.stride(1), v.stride(0), v.stride(2))
+        o_strides = (0, o.stride(1), o.stride(0), o.stride(2))
+        dq_strides = (0, dq.stride(1), dq.stride(0), dq.stride(2))
+        dk_strides = (0, dk.stride(1), dk.stride(0), dk.stride(2))
+        dv_strides = (0, dv.stride(1), dv.stride(0), dv.stride(2))
+        do_strides = (0, do.stride(1), do.stride(0), do.stride(2))
+    else:
+        # Layout is bshd.
+        # q and k are [batch, seq_len, num_head, head_dim_qk].
+        # v is [batch, seq_len, num_head, head_dim_v]
+        batch, seqlen_q, num_q_heads = q.shape[:-1]
+        _, num_k_heads = k.shape[1], k.shape[2]
+        q_strides = (q.stride(0), q.stride(2), q.stride(1), q.stride(3))
+        k_strides = (k.stride(0), k.stride(2), k.stride(1), k.stride(3))
+        v_strides = (v.stride(0), v.stride(2), v.stride(1), v.stride(3))
+        o_strides = (o.stride(0), o.stride(2), o.stride(1), o.stride(3))
+        dq_strides = (dq.stride(0), dq.stride(2), dq.stride(1), dq.stride(3))
+        dk_strides = (dk.stride(0), dk.stride(2), dk.stride(1), dk.stride(3))
+        dv_strides = (dv.stride(0), dv.stride(2), dv.stride(1), dv.stride(3))
+        do_strides = (do.stride(0), do.stride(2), do.stride(1), do.stride(3))
+
+    qk_head_dim = q.shape[-1]
+    v_head_dim = v.shape[-1]
+    pe_head_dim = qk_head_dim - v_head_dim
+    # BLOCK_D_MODEL, BLOCK_D_MODEL_POW2
+    # padding for head_dim. Power of 2 or 16
+    BLOCK_D_MODEL_POW2 = max(triton.next_power_of_2(v_head_dim), 16)
+    BLOCK_D_MODEL_PE_POW2 = (
+        0 if pe_head_dim == 0 else max(triton.next_power_of_2(pe_head_dim), 16)
+    )
+    assert (pe_head_dim == 0 and BLOCK_D_MODEL_PE_POW2 == 0) or (
+        v_head_dim == BLOCK_D_MODEL_POW2 and pe_head_dim == BLOCK_D_MODEL_PE_POW2
+    ), "Positional encoding support requires NOPE and PE head sizes to be unpadded powers of 2."
+    assert (not IS_FP8) or (
+        IS_FP8 and pe_head_dim == 0
+    ), "Positional encoding doesn't support FP8."
+
+    assert (sink is None) or (
+        sink is not None and sink.dim() == 1 and sink.shape[0] == num_q_heads
+    ), "Sink must be 1D and have one element per query head."
+    assert (dsink is None) or (
+        dsink is not None and dsink.dim() == 1 and dsink.shape[0] == num_q_heads
+    ), "Sink gradient must be 1D and have one element per query head."
+    assert (sink is None) == (
+        dsink is None
+    ), "Sink and its gradient must be both present or absent."
+
+    # Configs
+    if config is None:
+        config = _get_config()
+
+    # init delta
+    delta = torch.zeros_like(softmax_lse)
+    if IS_VARLEN:
+        # [total_tokens, num_q_heads, seqlen_q]
+        delta_strides = (0, delta.stride(1), delta.stride(0))
+    else:
+        # [batch, num_q_heads, seqlen_q]
+        delta_strides = delta.stride()
+
+    # preprocess
+    # compute D(delta) = rowsum(dO*O). Note, multiplication is element-wise.
+    pre_grid = (
+        triton.cdiv(max_seqlen_q, config["preprocess_kernel"]["PRE_BLOCK"]),
+        batch,
+        num_q_heads,
+    )
+    _bwd_preprocess[pre_grid](
+        o,
+        do,
+        delta,
+        *o_strides,
+        *do_strides,
+        *delta_strides,
+        descale_strides[3],
+        cu_seqlens_q,
+        max_seqlen_q,
+        descale_do,
+        BLOCK_M=config["preprocess_kernel"]["PRE_BLOCK"],
+        BLOCK_D_MODEL=v_head_dim,
+        BLOCK_D_MODEL_POW2=BLOCK_D_MODEL_POW2,
+        IS_VARLEN=IS_VARLEN,
+        IS_FP8=IS_FP8,
+    )
+
+    # dropout_mask
+    use_dropout = dropout_p > 0.0
+    if use_dropout:
+        dropout_mask = torch.zeros(
+            (batch, num_q_heads, max_seqlen_q, max_seqlen_k),
+            device=q.device,
+            dtype=torch.float32,
+        )
+        dropout_strides = dropout_mask.stride()
+    else:
+        dropout_mask = None
+        dropout_strides = (0, 0, 0, 0)
+
+    seqlen = max(max_seqlen_q, max_seqlen_k)
+
+    # "onekernel_pe" is for Positional Encoding (PE) causal case, it's going to be
+    # used if present. Otherwise, fallback to default "onekernel" config.
+    config_onekernel = (
+        config["onekernel_pe"]
+        if (pe_head_dim > 0 and causal and "onekernel_pe" in config)
+        else config["onekernel"]
+    )
+    grid = (
+        num_k_heads,
+        triton.cdiv(seqlen, config_onekernel["BLOCK_N1"]),
+        batch,
+    )
+
+    if causal:
+        bwd_kernel_causal[grid](
+            q,
+            k,
+            v,
+            sink,
+            sm_scale,
+            do,
+            dq,
+            dk,
+            dv,
+            dsink,
+            softmax_lse,
+            delta,
+            *q_strides,
+            *k_strides,
+            *v_strides,
+            *dq_strides,
+            *dk_strides,
+            *dv_strides,
+            *delta_strides,
+            *do_strides,
+            *dropout_strides,
+            *descale_strides,
+            stride_az,
+            stride_ah,
+            num_q_heads,
+            num_k_heads,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            dropout_mask,
+            dropout_p,
+            philox_seed,
+            philox_offset,
+            alibi_slopes,
+            descale_q,
+            descale_k,
+            descale_v,
+            descale_do,
+            HEAD_DIM=BLOCK_D_MODEL_POW2,
+            ACTUAL_HEAD_DIM=v_head_dim,
+            PE_HEAD_DIM=pe_head_dim,
+            ENABLE_DROPOUT=use_dropout,
+            IS_VARLEN=IS_VARLEN,
+            USE_ALIBI=use_alibi,
+            USE_EXP2=True,
+            IS_FP8=IS_FP8,
+            FP8_MAX=FP8_MAX,
+            DEBUG_TRITON=False,
+            DEBUG_TRITON_DETAIL=False,
+            USE_INT64_STRIDES=USE_INT64_STRIDES,
+            ENABLE_SINK=sink is not None,
+            SLIDING_WINDOW=sliding_window,
+            **config_onekernel,
+        )
+    else:
+        bwd_kernel_noncausal[grid](
+            q,
+            k,
+            v,
+            sink,
+            sm_scale,
+            do,
+            dq,
+            dk,
+            dv,
+            dsink,
+            softmax_lse,
+            delta,
+            *q_strides,
+            *k_strides,
+            *v_strides,
+            *dq_strides,
+            *dk_strides,
+            *dv_strides,
+            *delta_strides,
+            *do_strides,
+            *dropout_strides,
+            *descale_strides,
+            stride_az,
+            stride_ah,
+            num_q_heads,
+            num_k_heads,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            dropout_mask,
+            dropout_p,
+            philox_seed,
+            philox_offset,
+            alibi_slopes,
+            descale_q,
+            descale_k,
+            descale_v,
+            descale_do,
+            HEAD_DIM=BLOCK_D_MODEL_POW2,
+            ACTUAL_HEAD_DIM=v_head_dim,
+            PE_HEAD_DIM=pe_head_dim,
+            ENABLE_DROPOUT=use_dropout,
+            IS_VARLEN=IS_VARLEN,
+            USE_ALIBI=use_alibi,
+            USE_EXP2=True,
+            IS_FP8=IS_FP8,
+            FP8_MAX=FP8_MAX,
+            DEBUG_TRITON=False,
+            DEBUG_TRITON_DETAIL=False,
+            USE_INT64_STRIDES=USE_INT64_STRIDES,
+            ENABLE_SINK=sink is not None,
+            SLIDING_WINDOW=sliding_window,
+            **config_onekernel,
+        )
+
+    return delta

build/torch-rocm/utils/_triton/arch_info.py

@@ -0,0 +1,42 @@
+import triton
+
+
+_CACHED_ARCH = None
+
+
+def get_arch():
+    # Lazy-detect on first call so the module is importable in environments without an
+    # active GPU driver (e.g. the kernel-builder Nix build sandbox running
+    # `get_local_kernel` for layout checks).
+    global _CACHED_ARCH
+    if _CACHED_ARCH is None:
+        try:
+            _CACHED_ARCH = triton.runtime.driver.active.get_current_target().arch
+        except RuntimeError:
+            try:
+                from jax._src.lib import gpu_triton as triton_kernel_call_lib
+
+                _CACHED_ARCH = triton_kernel_call_lib.get_arch_details("0").split(":")[0]
+            except ImportError:
+                _CACHED_ARCH = ""
+    return _CACHED_ARCH
+
+
+def is_gluon_avail():
+    return get_arch() in ("gfx950", "gfx1250")
+
+
+def is_fp4_avail():
+    return get_arch() in ("gfx950", "gfx1250")
+
+
+def is_fp8_avail():
+    return get_arch() in ("gfx942", "gfx950", "gfx1250", "gfx1200", "gfx1201")
+
+
+def is_mx_scale_preshuffling_avail():
+    return get_arch() in ("gfx950", "gfx1250")
+
+
+def is_tdm_avail():
+    return get_arch() in ("gfx1250",)

build/torch-rocm/utils/_triton/kernel_repr.py

@@ -0,0 +1,57 @@
+def _sanitize_constexpr_value(value):
+    if value is None:
+        return "NONE"
+    if isinstance(value, bool):
+        return str(int(value))
+    if isinstance(value, int):
+        return str(value)
+    if isinstance(value, float):
+        if value.is_integer():
+            return str(int(value))
+        return str(value)
+
+    # for lists, tuples, sets - recursively join each
+    if isinstance(value, (list, tuple, set)):
+        items = sorted(value, key=str) if isinstance(value, set) else value
+        sanitized_items = [_sanitize_constexpr_value(item) for item in items]
+        joined = "_".join(sanitized_items)
+        return joined if joined else "NONE"
+
+    if isinstance(value, str):
+        cleaned_value = "".join(ch if ch.isalnum() else "_" for ch in value).strip("_")
+        return cleaned_value.upper() if cleaned_value else "NONE"
+
+    cleaned_value = "".join(ch if ch.isalnum() else "_" for ch in str(value)).strip("_")
+    return cleaned_value.upper() if cleaned_value else "NONE"
+
+
+def make_kernel_repr(base_name, config_keys, name_key=None):
+    # When name_key is set, the base name is taken from the matching constexpr
+    # kwarg at call time (falling back to base_name if missing/empty). Lets a
+    # single shared kernel produce caller-specific names in compiled artifacts.
+    def _repr(specialization):
+        constants = specialization.constants
+
+        name = base_name
+        if name_key is not None:
+            override = constants.get(name_key, None)
+            if override:
+                cleaned = "".join(
+                    ch if ch.isalnum() or ch == "_" else "_" for ch in str(override)
+                )
+                if cleaned:
+                    name = cleaned
+
+        name_parts = []
+        for key in config_keys:
+            value = constants.get(key, None)
+            symbol = _sanitize_constexpr_value(value)
+            name_parts.append(f"{key}_{symbol}")
+
+        if not name_parts:
+            return name
+
+        suffix = "_".join(name_parts)
+        return f"{name}_{suffix}"
+
+    return _repr

build/torch-rocm/utils/_triton/mha_kernel_utils.py

@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _compute_fp8_scaling_factors(x, fp8_max: tl.constexpr):
+    # compute fp8 scaling and descaling factor for a block
+    x_amax = tl.max(tl.abs(x))  # NOTE: abs deals with negative values
+    x_amax = tl.where(x_amax <= 1e-9, 1e-9, x_amax)
+    scale_x = fp8_max / x_amax
+    descale_x = x_amax / fp8_max
+    return scale_x, descale_x

build/torch-rocm/utils/_triton/pid_preprocessing.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: MIT
+
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def remap_xcd_chunked(
+    pid, GRID_MN, NUM_XCDS: tl.constexpr = 8, CHUNK_SIZE: tl.constexpr = 2
+):
+    # Compute current XCD and local PID
+    xcd = pid % NUM_XCDS
+    # distribute the modulo pids in round robin
+    if pid > (GRID_MN // (NUM_XCDS * CHUNK_SIZE)) * (NUM_XCDS * CHUNK_SIZE):
+        return pid
+    local_pid = pid // NUM_XCDS
+    # Calculate chunk index and position within chunk
+    chunk_idx = local_pid // CHUNK_SIZE
+    pos_in_chunk = local_pid % CHUNK_SIZE
+    # Calculate new PID
+    new_pid = chunk_idx * NUM_XCDS * CHUNK_SIZE + xcd * CHUNK_SIZE + pos_in_chunk
+    return new_pid
+
+
+@triton.jit
+def remap_xcd(pid, GRID_MN, NUM_XCDS: tl.constexpr = 8):
+    ## pid remapping on xcds
+    # Number of pids per XCD in the new arrangement
+    pids_per_xcd = (GRID_MN + NUM_XCDS - 1) // NUM_XCDS
+    # When GRID_MN cannot divide NUM_XCDS, some xcds will have
+    # pids_per_xcd pids, the other will have pids_per_xcd - 1 pids.
+    # We calculate the number of xcds that have pids_per_xcd pids as
+    # tall_xcds
+    tall_xcds = GRID_MN % NUM_XCDS
+    tall_xcds = NUM_XCDS if tall_xcds == 0 else tall_xcds
+    # Compute current XCD and local pid within the XCD
+    xcd = pid % NUM_XCDS
+    local_pid = pid // NUM_XCDS
+    # Calculate new pid based on the new grouping
+    # Note that we need to consider the following two cases:
+    # 1. the current pid is on a tall xcd
+    # 2. the current pid is on a short xcd
+    if xcd < tall_xcds:
+        pid = xcd * pids_per_xcd + local_pid
+    else:
+        pid = (
+            tall_xcds * pids_per_xcd
+            + (xcd - tall_xcds) * (pids_per_xcd - 1)
+            + local_pid
+        )
+
+    return pid
+
+
+@triton.jit
+def pid_grid(pid: int, num_pid_m: int, num_pid_n: int, GROUP_SIZE_M: tl.constexpr = 1):
+    """
+    Maps 1D pid to 2D grid coords (pid_m, pid_n).
+
+    Args:
+        - pid: 1D pid
+        - num_pid_m: grid m size
+        - num_pid_n: grid n size
+        - GROUP_SIZE_M: tl.constexpr: default is 1
+    """
+    if GROUP_SIZE_M == 1:
+        pid_m = pid // num_pid_n
+        pid_n = pid % num_pid_n
+    else:
+        num_pid_in_group = GROUP_SIZE_M * num_pid_n
+        group_id = pid // num_pid_in_group
+        first_pid_m = group_id * GROUP_SIZE_M
+        group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+        tl.assume(group_size_m >= 0)
+        pid_m = first_pid_m + (pid % group_size_m)
+        pid_n = (pid % num_pid_in_group) // group_size_m
+
+    return pid_m, pid_n
+
+
+@triton.jit
+def pid_grid_3d(pid: int, num_pid_m: int, num_pid_n: int, num_pid_k):
+    """
+    Maps 1D pid to 3D grid coords (pid_m, pid_n, pid_k).
+    Args:
+        - pid: 1D pid
+        - num_pid_m: grid m size
+        - num_pid_n: grid n size
+        - num_pid_k: grid k size
+
+    Returns:
+        - pid_m, pid_n, pid_k: 3D grid coordinates
+    """
+    pid_m = pid % num_pid_m
+    pid_n = (pid // num_pid_m) % num_pid_n
+    pid_k = pid // (num_pid_m * num_pid_n) % num_pid_k
+
+    return pid_m, pid_n, pid_k

build/torch-rocm/utils/core.py

@@ -0,0 +1,5 @@
+import os
+
+this_dir = os.path.dirname(os.path.abspath(__file__))
+AITER_TRITON_OPS_PATH = os.path.abspath(f"{this_dir}/../")
+AITER_TRITON_CONFIGS_PATH = os.path.abspath(f"{this_dir}/../configs")

build/torch-rocm/utils/device_info.py

@@ -0,0 +1,14 @@
+import torch
+
+
+def get_num_sms():
+    # Returns the Compute Unit count of the current device
+    current_device_index = torch.cuda.current_device()
+    current_device = torch.cuda.get_device_properties(current_device_index)
+    num_sms = current_device.multi_processor_count
+    return num_sms
+
+
+def get_num_xcds():
+    # Currently, you can't query this programmatically. For gfx942/gfx950 it's 8, so we hardcode that here.
+    return 8

build/torch-rocm/utils/logger.py

@@ -0,0 +1,47 @@
+import os
+import logging
+
+
+# AITER Triton Logger which is singleton object around python logging.
+# Note: Python logging is also a singleton object, but we want to read the
+# env var AITER_LOG_LEVEL once at the beginning. Another alternative is to do
+# this in __init__.py. In fact, that's how CK logger is setup. We can look at
+# switching to that at some point
+#
+# AITER_LOG_LEVEL follows python logging levels
+#   DEBUG
+#   INFO
+#   WARNING
+#   ERROR
+#   CRITICAL
+#
+class AiterTritonLogger(object):
+    _instance = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super(AiterTritonLogger, cls).__new__(cls)
+            log_level_str = os.getenv("AITER_TRITON_LOG_LEVEL", "WARNING").upper()
+            numeric_level = getattr(logging, log_level_str, logging.WARNING)
+            cls._instance._logger = logging.getLogger("AITER_TRITON")
+            cls._instance._logger.setLevel(numeric_level)
+
+        return cls._instance
+
+    def get_logger(self):
+        return self._logger
+
+    def debug(self, msg):
+        self._logger.debug(msg)
+
+    def info(self, msg):
+        self._logger.info(msg)
+
+    def warning(self, msg):
+        self._logger.warning(msg)
+
+    def error(self, msg):
+        self._logger.error(msg)
+
+    def critical(self, msg):
+        self._logger.critical(msg)

build/torch-rocm/utils/types.py

@@ -0,0 +1,85 @@
+import torch
+import triton.language as tl
+from ._triton import arch_info
+
+
+def get_dtype_max(dtype):
+    if torch.is_floating_point(torch.tensor([], dtype=dtype)):
+        return torch.finfo(dtype).max
+    else:
+        return torch.iinfo(dtype).max
+
+
+def get_fp8_dtypes():
+    if arch_info.get_arch() in ("gfx950", "gfx1250", "gfx1200", "gfx1201"):
+        e5m2_dtype = torch.float8_e5m2
+        e4m3_dtype = torch.float8_e4m3fn
+    else:
+        e5m2_dtype = torch.float8_e5m2fnuz
+        e4m3_dtype = torch.float8_e4m3fnuz
+
+    return e5m2_dtype, e4m3_dtype
+
+
+def get_fp8_e4m3_dtype():
+    if arch_info.get_arch() in ("gfx950", "gfx1250", "gfx1200", "gfx1201"):
+        e4m3_dtype = torch.float8_e4m3fn
+    else:
+        e4m3_dtype = torch.float8_e4m3fnuz
+
+    return e4m3_dtype
+
+
+e5m2_dtype, e4m3_dtype = get_fp8_dtypes()
+str_to_torch_dtype = {
+    "float64": torch.float64,
+    "float32": torch.float32,
+    "fp32": torch.float32,
+    "bfloat16": torch.bfloat16,
+    "bf16": torch.bfloat16,
+    "float16": torch.float16,
+    "fp16": torch.float16,
+    "float8_e5m2": e5m2_dtype,
+    "float8_e4m3fn": e4m3_dtype,
+    "e5m2fnuz": e5m2_dtype,
+    "e4m3fnuz": e4m3_dtype,
+    "fp8e4m3": e4m3_dtype,
+    "fp8e5m2": e5m2_dtype,
+    "int64": torch.int64,
+    "int32": torch.int32,
+    "int16": torch.int16,
+    "uint8": torch.uint8,
+    "int8": torch.int8,
+    "mxfp4_e2m1": torch.uint8,  # OCP MXFP4 packs two 4-bits into 8-bit
+}
+
+torch_to_triton_dtype = {
+    torch.float64: tl.float64,
+    torch.float32: tl.float32,
+    torch.float16: tl.float16,
+    torch.bfloat16: tl.bfloat16,
+    torch.float8_e4m3fn: tl.float8e4nv,
+    torch.float8_e4m3fnuz: tl.float8e4b8,
+    torch.float8_e5m2: tl.float8e5,
+    torch.float8_e5m2fnuz: tl.float8e5b16,
+    torch.int64: tl.int64,
+    torch.int32: tl.int32,
+    torch.int16: tl.int16,
+    torch.int8: tl.int8,
+    torch.uint8: tl.uint8,
+}
+
+
+def _is_fp8(x):
+    if x.dtype in {
+        torch.float8_e4m3fnuz,
+        torch.float8_e4m3fn,
+        torch.float8_e5m2,
+        torch.float8_e5m2fnuz,
+    }:
+        if arch_info.is_fp8_avail():
+            return True
+        else:
+            raise RuntimeError("This device does not support fp8")
+    else:
+        return False