mslk-cuda-nightly 2026.1.19__cp310-cp310-manylinux_2_28_x86_64.whl

mslk/attention/fmha/_triton/splitk_kernels.py

@@ -0,0 +1,1534 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import functools
+import sys
+from typing import Callable, Dict, Tuple, Union
+
+import torch
+import triton
+import triton.language as tl
+
+from .vararg_kernel import unroll_varargs, VAR_ARGS_ARRAY
+
+# pyre-ignore-all-errors
+AUTOTUNER_KEY = [
+    "Z",
+    "H",
+    "G",
+    "N_CTX_Q",
+    "N_CTX_K",
+    "BLOCK_DMODEL",
+    "PACKED_PER_VAL",
+    "N_GROUPS",
+    "BLOCK_N_PER_SPLIT",
+    "PAGE_SIZE",
+]
+
+
+@triton.jit
+def _fwd_kernel_splitK(  # noqa: C901
+    Q,
+    K,
+    V,
+    sm_scale,
+    Out_splitK,  # [B, H, split_k, Mq, K]
+    LSE_splitk,  # [B, H, split_k, Mq]
+    block_tables,
+    Seq_len,
+    Seq_starts_k,
+    Seq_starts_q,
+    Seq_starts_q_multiplier,
+    additive_bias,
+    K_fp8_scale_shift,
+    V_fp8_scale_shift,
+    q_fp8_scale_shift,
+    stride_qz,
+    stride_qm,
+    stride_qg,
+    stride_qh,
+    stride_qk,
+    stride_kz,
+    stride_kn,
+    stride_kg,
+    stride_kh,
+    stride_kk,
+    stride_vz,
+    stride_vn,
+    stride_vg,
+    stride_vh,
+    stride_vk,
+    stride_osk_z,
+    stride_osk_g,
+    stride_osk_h,
+    stride_osk_s,
+    stride_osk_m,
+    stride_osk_k,
+    stride_lsek_z,
+    stride_lsek_g,
+    stride_lsek_h,
+    stride_lsek_s,
+    stride_lsek_m,
+    stride_blocktablesz,
+    stride_blocktablesl,
+    stride_bias_b,
+    stride_bias_g,
+    stride_bias_h,
+    stride_bias_qm,
+    stride_bias_km,
+    stride_k_fp8_scale_shift_z: tl.constexpr,
+    stride_k_fp8_scale_shift_n: tl.constexpr,
+    stride_k_fp8_scale_shift_g: tl.constexpr,
+    stride_k_fp8_scale_shift_h: tl.constexpr,
+    stride_v_fp8_scale_shift_z: tl.constexpr,
+    stride_v_fp8_scale_shift_n: tl.constexpr,
+    stride_v_fp8_scale_shift_g: tl.constexpr,
+    stride_v_fp8_scale_shift_h: tl.constexpr,
+    stride_q_fp8_scale_shift_z: tl.constexpr,
+    stride_q_fp8_scale_shift_m: tl.constexpr,
+    stride_q_fp8_scale_shift_g: tl.constexpr,
+    stride_q_fp8_scale_shift_h: tl.constexpr,
+    kv_cache_blocks_per_row: tl.constexpr,
+    Z: tl.constexpr,
+    N_CTX_Q: tl.constexpr,  # The number of queries
+    N_CTX_K: tl.constexpr,
+    BLOCK_N_PER_SPLIT: tl.constexpr,
+    H: tl.constexpr,
+    G: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    USE_SEQ_LEN: tl.constexpr,
+    PACKED_PER_VAL: tl.constexpr,
+    N_GROUPS: tl.constexpr,
+    # It's important that BOUNDS_CHECKS_N, BLOCK_M, BLOCK_N come at the end of
+    # the argument list, since they are provided by the heuristics/autotune decorator.
+    # Otherwise Triton throws IndexError
+    BOUNDS_CHECKS_N: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    IS_SPLITK: tl.constexpr,
+    SPLIT_K_EARLY_EXIT: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+    IS_LOCAL: tl.constexpr,
+    NUM_QUERIES_CAUSAL: tl.constexpr,  # The N_CTX_Q queries are from this many sequence positions
+    USE_PAGED_ATTENTION: tl.constexpr,
+    PAGE_SIZE: tl.constexpr,
+    WINDOW_LEFT: tl.constexpr,
+    WINDOW_RIGHT: tl.constexpr,
+    WRITE_LSE: tl.constexpr,
+    HAS_ADDITIVE_BIAS: tl.constexpr,
+    NUM_PROGRAMS_DIM2_CONST: tl.constexpr,
+    IS_HIP: tl.constexpr,
+    QUANTIZE_PV_TO_FP8: tl.constexpr,
+    QUANTIZE_QK_TO_FP8: tl.constexpr,
+    USE_FP32_SCALES: tl.constexpr,
+):
+    """This kernel can accept non-quantized or int4-quantized keys/values.
+    PACKED_PER_VAL determines the quantization type:
+        - PACKED_PER_VAL == 1 means no quantization
+        - PACKED_PER_VAL == 8 means 4-bit quantization (8 packed quantized values inside one int32)
+    For the quantized case K/V should be int32 tensors.
+    Quantization can be row-wise (when N_GROUPS = 1) or group-wise with N_GROUPS = 2, 4, or 8.
+    Quantization coefficients are stored at the beginning of the row along the last dimension of K/V
+    So K[B, H, M, :] has a form
+    [   quant_coef0, quant_coef1, ...|
+        group0_quant_value0, group0_quant_value1,... |
+        group1_quant_value0, group1_quant_value1,...]
+    where each quant_coef is an int32 which should be interpreted as 2 packed float16: scale and offset.
+
+    Note: this kernel needs to be processed by unroll_varargs
+    before compilation. That will unroll variables marked with "VAR_ARGS_ARRAY" into lists.
+    See how FwOp.apply does it below.
+
+    Set IS_SPLITK=False to indicate the MHA result should be written directly.
+    No metadata will be written.
+    """
+    internal_dtype = (
+        tl.float64 if Out_splitK.dtype.element_ty is tl.float64 else tl.float32
+    )
+    tl.static_assert(
+        (PACKED_PER_VAL == 1 and tl.constexpr(K.dtype.element_ty != tl.int32))
+        or (
+            (PACKED_PER_VAL == 4 or PACKED_PER_VAL == 8)
+            and tl.constexpr(K.dtype.element_ty == tl.int32)
+        ),
+        f"Only int4 and fp8 quantization is supported, K/V should have dtype int32 in "
+        f"the quantized case: {PACKED_PER_VAL=} {tl.constexpr(K.dtype)=} {tl.constexpr(K.dtype.element_ty)=}",
+    )
+    tl.static_assert(
+        (((N_GROUPS == 1 or N_GROUPS == 2) or N_GROUPS == 4) or N_GROUPS == 8),
+        "Number of quantization groups can be 1 (row-wise quantization), 2, 4, or 8.",
+    )
+    tl.static_assert(
+        N_GROUPS == 1 or K_fp8_scale_shift is None,
+        f"Only row-wise fp8 quantization is supported, but got {N_GROUPS=} > 1.",
+    )
+    FP8_QUANTIZED: tl.constexpr = K_fp8_scale_shift is not None
+    INT4_QUANTIZED: tl.constexpr = PACKED_PER_VAL > 1 and not FP8_QUANTIZED
+    PACKED_D_PER_GROUP: tl.constexpr = BLOCK_DMODEL // PACKED_PER_VAL // N_GROUPS
+    D_PER_GROUP: tl.constexpr = BLOCK_DMODEL // N_GROUPS
+
+    start_m = tl.program_id(0)
+    off_zhg = tl.program_id(1)
+    off_z = (off_zhg // (H * G)).to(tl.int64)
+    off_hg = off_zhg % (H * G)
+    off_h = off_hg // G
+    off_g = off_hg % G
+    splitk_idx = tl.program_id(2)
+
+    if USE_SEQ_LEN:
+        kv_len = tl.load(Seq_len + off_z)
+        if SPLIT_K_EARLY_EXIT and kv_len == 0:
+            return
+    else:
+        kv_len = N_CTX_K
+
+    if Seq_starts_k is None:
+        start_kv_idx = 0
+    else:
+        start_kv_idx = tl.load(Seq_starts_k + off_z)
+        if USE_SEQ_LEN and PAGE_SIZE > 0:
+            # gappy with paged attention stores each "end" instead of each "length"
+            # because that's what FA3 needs.
+            kv_len -= start_kv_idx
+
+    if Seq_starts_q is None:
+        q_len = N_CTX_Q
+        queries_use_batch_dim = 1
+        off_m = 0
+    else:
+        queries_use_batch_dim = 0
+        off_m = tl.load(Seq_starts_q + off_z) * Seq_starts_q_multiplier
+        q_len = tl.load(Seq_starts_q + off_z + 1) * Seq_starts_q_multiplier - off_m
+        if q_len == 0:
+            return
+
+    k_base = K + off_h * stride_kh + off_g * stride_kg
+    v_base = V + off_h * stride_vh + off_g * stride_vg
+
+    if FP8_QUANTIZED:
+        k_fp8_scale_shift_base = (
+            K_fp8_scale_shift
+            + off_h * stride_k_fp8_scale_shift_h
+            + off_g * stride_k_fp8_scale_shift_g
+        )
+        v_fp8_scale_shift_base = (
+            V_fp8_scale_shift
+            + off_h * stride_v_fp8_scale_shift_h
+            + off_g * stride_v_fp8_scale_shift_g
+        )
+    else:
+        k_fp8_scale_shift_base = None
+        v_fp8_scale_shift_base = None
+
+    # Boundaries of split-k chunk
+    chunk_hi = (splitk_idx + 1) * BLOCK_N_PER_SPLIT
+    chunk_lo = splitk_idx * BLOCK_N_PER_SPLIT
+    ignore_in_first_block = 0
+    # For paged attention case K/V_block_ptr are defined inside the loop
+    # whereas for non-paged case they are defined before the loop.
+    if PAGE_SIZE > 0:
+        # Page contains several blocks
+        BLOCKS_IN_PAGE: tl.constexpr = PAGE_SIZE // BLOCK_N
+        # Align boundaries of split-k chunk to block boundaries
+        # In the last chunk, shift hi to the right, in the other chunks, shift it to the left
+        # TODO: Replace NUM_PROGRAMS_DIM2_CONST with tl.num_programs(2) after
+        # the next Triton upgrade.
+        is_last_chunk = splitk_idx == NUM_PROGRAMS_DIM2_CONST - 1
+        shift = BLOCK_N - 1 if is_last_chunk else 0
+        lo = (tl.maximum(chunk_lo, start_kv_idx) // BLOCK_N) * BLOCK_N
+        ignore_in_first_block = tl.maximum(0, (start_kv_idx - lo))
+        hi = ((chunk_hi + shift) // BLOCK_N) * BLOCK_N
+        hi = tl.minimum(hi, kv_len + start_kv_idx)
+        block_table = block_tables + stride_blocktablesz * off_z
+        # Offset in integer blocks
+        logical_block_idx = lo // BLOCK_N
+    else:
+        lo = chunk_lo
+        hi = tl.minimum(chunk_hi, kv_len)
+        if Seq_starts_k is not None:
+            k_base += start_kv_idx * stride_kn
+            v_base += start_kv_idx * stride_vn
+        else:
+            k_base += off_z * stride_kz
+            v_base += off_z * stride_vz
+        # Additional shift by 1 along the last dimension in the quantized case, since
+        # the first element along that dim contains packed quantization coefficients.
+        K_block_ptr = tl.make_block_ptr(
+            base=k_base + stride_kk * INT4_QUANTIZED * N_GROUPS,
+            shape=(PACKED_D_PER_GROUP, hi),
+            strides=(stride_kk, stride_kn),
+            offsets=(0, lo),
+            block_shape=(PACKED_D_PER_GROUP, BLOCK_N),
+            order=(0, 1),
+        )
+        V_block_ptr = tl.make_block_ptr(
+            base=v_base + stride_vk * INT4_QUANTIZED * N_GROUPS,
+            shape=(hi, PACKED_D_PER_GROUP if not QUANTIZE_PV_TO_FP8 else D_PER_GROUP),
+            strides=(stride_vn, stride_vk),
+            offsets=(lo, 0),
+            block_shape=(
+                BLOCK_N,
+                PACKED_D_PER_GROUP if not QUANTIZE_PV_TO_FP8 else D_PER_GROUP,
+            ),
+            order=(1, 0),
+        )
+
+        if INT4_QUANTIZED:
+            # Pointers to quantization coefficients. Even those they are 1D,
+            # we use block pointers here so the pointer arithmetic is in int64,
+            # as otherwise the offsets for V_scale_shift_block_ptr may overflow.
+            K_scale_shift_block_ptr = tl.make_block_ptr(
+                base=k_base,
+                shape=(1, hi),
+                strides=(stride_kk, stride_kn),
+                offsets=(0, lo),
+                block_shape=(1, BLOCK_N),
+                order=(0, 1),
+            )
+            V_scale_shift_block_ptr = tl.make_block_ptr(
+                base=v_base,
+                shape=(hi, 1),
+                strides=(stride_vn, stride_vk),
+                offsets=(lo, 0),
+                block_shape=(BLOCK_N, 1),
+                order=(1, 0),
+            )
+        elif FP8_QUANTIZED:
+            if Seq_starts_k is not None:
+                k_fp8_scale_shift_base += start_kv_idx * stride_k_fp8_scale_shift_n
+                v_fp8_scale_shift_base += start_kv_idx * stride_v_fp8_scale_shift_n
+            else:
+                k_fp8_scale_shift_base += off_z * stride_k_fp8_scale_shift_z
+                v_fp8_scale_shift_base += off_z * stride_v_fp8_scale_shift_z
+            K_scale_shift_block_ptr = tl.make_block_ptr(
+                base=k_fp8_scale_shift_base,
+                shape=(1, hi),
+                strides=(1, stride_k_fp8_scale_shift_n),
+                offsets=(0, lo),
+                block_shape=(1, BLOCK_N),
+                order=(0, 1),
+            )
+            V_scale_shift_block_ptr = tl.make_block_ptr(
+                base=v_fp8_scale_shift_base,
+                shape=(hi, 1),
+                strides=(stride_v_fp8_scale_shift_n, 1),
+                offsets=(lo, 0),
+                block_shape=(BLOCK_N, 1),
+                order=(1, 0),
+            )
+        else:
+            K_scale_shift_block_ptr = None
+            V_scale_shift_block_ptr = None
+
+        if HAS_ADDITIVE_BIAS:
+            additive_bias_block_ptr = tl.make_block_ptr(
+                base=additive_bias
+                + off_z * stride_bias_b
+                + off_g * stride_bias_g
+                + off_h * stride_bias_h,
+                shape=(N_CTX_Q, hi),
+                strides=(stride_bias_qm, stride_bias_km),
+                offsets=(start_m * BLOCK_M, lo),
+                block_shape=(BLOCK_M, BLOCK_N),
+                order=(0, 1),
+            )
+
+    if SPLIT_K_EARLY_EXIT and lo >= hi:
+        return
+
+    Q_block_ptr = tl.make_block_ptr(
+        base=Q
+        + off_m * stride_qm
+        + off_h * stride_qh
+        + off_z * stride_qz * queries_use_batch_dim
+        + off_g * stride_qg,
+        shape=(q_len, BLOCK_DMODEL),
+        strides=(stride_qm, stride_qk),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, D_PER_GROUP),
+        order=(1, 0),
+    )
+
+    # initialize pointer to m and l
+    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
+
+    # Before compilation, this kernel will be processed by unroll_varargs.
+    # That turns tensors annotated as the one below into lists of tensors of length N_GROUPS.
+    # This is a solution for Triton native lack of support for lists of tensors.
+    acc: "VAR_ARGS_ARRAY"  # noqa: F821
+
+    for i in range(len(acc)):  # noqa: F821
+        acc[i] = tl.zeros([BLOCK_M, D_PER_GROUP], dtype=internal_dtype)  # noqa: F821
+    # 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
+    #
+    # We declare log2e as a constant with a precisely-specified type to guarantee that
+    # triton will use the exact same value in all instances below, rather than sometimes
+    # using float32 and sometimes using float64.  For more discussion see:
+    # https://github.com/triton-lang/triton/issues/5466
+    log2e = tl.full((), 1.44269504, tl.float32)
+    qk_scale = sm_scale * log2e
+    # load q: it will stay in SRAM throughout
+    q: "VAR_ARGS_ARRAY"  # noqa: F821
+
+    if QUANTIZE_QK_TO_FP8:
+        # Create a block pointer for q_scale
+        q_scale_block_ptr = tl.make_block_ptr(
+            base=q_fp8_scale_shift
+            + off_m * stride_q_fp8_scale_shift_m
+            + off_h * stride_q_fp8_scale_shift_h
+            + off_g * stride_q_fp8_scale_shift_g
+            + off_z * stride_q_fp8_scale_shift_z * queries_use_batch_dim,
+            shape=(q_len, 1),
+            strides=(stride_q_fp8_scale_shift_m, 1),
+            offsets=(start_m * BLOCK_M, 0),
+            block_shape=(BLOCK_M, 1),
+            order=(1, 0),
+        )
+
+        # For FP8 quantized query, load and dequantize
+        for i in range(len(acc)):  # noqa: F821
+            # Load quantized query
+            q_quantized = tl.load(  # noqa: F821
+                tl.advance(Q_block_ptr, (0, i * D_PER_GROUP)), boundary_check=(0,)
+            )
+
+            # Load q_scale for dequantization - q_scale is per row
+            q_scale = tl.load(
+                tl.advance(q_scale_block_ptr, (0, i * D_PER_GROUP)), boundary_check=(0,)
+            )
+            q[i] = q_quantized.to(Q.dtype.element_ty)  # noqa: F821
+    else:
+        # Regular query loading
+        for i in range(len(acc)):  # noqa: F821
+            q[i] = tl.load(  # noqa: F821
+                tl.advance(Q_block_ptr, (0, i * D_PER_GROUP)), boundary_check=(0,)
+            )
+
+    if IS_CAUSAL or IS_LOCAL:
+        # Why does the masking conditon below work as a causal mask?
+        # Assuming num_queries <= BLOCK_M:
+        # kv_pos = kv_start + range(0, BLOCK_N)
+        # q_offset = start_m * BLOCK_M + range(0, BLOCK_M)
+        # q_pos = kv_start + kv_len - num_queries + q_offset % num_queries
+        # mask = q_pos - kv_pos >= 0
+        # So the final masking condition is:
+        #   range(0, BLOCK_M) % num_queries - range(0, BLOCK_N) >= num_queries - kv_len
+
+        q_offset = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+        diag_idx = (q_offset[:, None] % NUM_QUERIES_CAUSAL) - tl.arange(0, BLOCK_N)[
+            None, :
+        ]
+        diag_idx_shifted = tl.constexpr(diag_idx - NUM_QUERIES_CAUSAL + kv_len)
+
+    # loop over k, v and update accumulator
+    for start_n in range(lo, hi, BLOCK_N):
+        if PAGE_SIZE > 0:
+            # Offset in integer blocks from the beginning of the page
+            block_offset_in_page = logical_block_idx % BLOCKS_IN_PAGE
+            # Offset in integer pages
+            logical_page_idx = logical_block_idx // BLOCKS_IN_PAGE
+            physical_page_idx = tl.load(
+                block_table + stride_blocktablesl * logical_page_idx
+            ).to(tl.int32)
+            offset = physical_page_idx * PAGE_SIZE + block_offset_in_page * BLOCK_N
+
+            current_block_size = min(hi - start_n, BLOCK_N)
+            K_block_ptr = tl.make_block_ptr(
+                base=k_base + stride_kk * INT4_QUANTIZED * N_GROUPS,
+                shape=(PACKED_D_PER_GROUP, offset + current_block_size),
+                strides=(stride_kk, stride_kn),
+                offsets=(0, offset),
+                block_shape=(PACKED_D_PER_GROUP, BLOCK_N),
+                order=(0, 1),
+            )
+            V_block_ptr = tl.make_block_ptr(
+                base=v_base + stride_vk * INT4_QUANTIZED * N_GROUPS,
+                shape=(
+                    offset + current_block_size,
+                    PACKED_D_PER_GROUP if not QUANTIZE_PV_TO_FP8 else D_PER_GROUP,
+                ),
+                strides=(stride_vn, stride_vk),
+                offsets=(offset, 0),
+                block_shape=(
+                    BLOCK_N,
+                    PACKED_D_PER_GROUP if not QUANTIZE_PV_TO_FP8 else D_PER_GROUP,
+                ),
+                order=(1, 0),
+            )
+            if INT4_QUANTIZED:
+                # Pointers to quantization coefficients. Even those they are 1D,
+                # we use block pointers here so the pointer arithmetic is in int64,
+                # as otherwise the offsets for V_scale_shift_block_ptr may overflow.
+                K_scale_shift_block_ptr = tl.make_block_ptr(
+                    base=k_base,
+                    shape=(1, offset + current_block_size),
+                    strides=(stride_kk, stride_kn),
+                    offsets=(0, offset),
+                    block_shape=(1, BLOCK_N),
+                    order=(0, 1),
+                )
+                V_scale_shift_block_ptr = tl.make_block_ptr(
+                    base=v_base,
+                    shape=(offset + current_block_size, 1),
+                    strides=(stride_vn, stride_vk),
+                    offsets=(offset, 0),
+                    block_shape=(BLOCK_N, 1),
+                    order=(1, 0),
+                )
+            elif FP8_QUANTIZED:
+                K_scale_shift_block_ptr = tl.make_block_ptr(
+                    base=k_fp8_scale_shift_base,
+                    shape=(1, offset + current_block_size),
+                    strides=(1, stride_k_fp8_scale_shift_n),
+                    offsets=(0, offset),
+                    block_shape=(1, BLOCK_N),
+                    order=(0, 1),
+                )
+                V_scale_shift_block_ptr = tl.make_block_ptr(
+                    base=v_fp8_scale_shift_base,
+                    shape=(offset + current_block_size, 1),
+                    strides=(stride_v_fp8_scale_shift_n, 1),
+                    offsets=(offset, 0),
+                    block_shape=(BLOCK_N, 1),
+                    order=(1, 0),
+                )
+            else:
+                K_scale_shift_block_ptr = None
+                V_scale_shift_block_ptr = None
+            logical_block_idx += 1
+
+        k: "VAR_ARGS_ARRAY"  # noqa: F821
+        v: "VAR_ARGS_ARRAY"  # noqa: F821
+
+        if QUANTIZE_PV_TO_FP8:
+            v_dtype = tl.float8e4nv
+        else:
+            v_dtype = Q.dtype.element_ty
+        for i in range(len(acc)):  # noqa: F821
+            # Load and dequantize K/V with appropriate return values based on quantization flags
+            result = load_dequantize_k_v_group(  # noqa: F821
+                K_block_ptr,
+                V_block_ptr,
+                K_scale_shift_block_ptr,
+                V_scale_shift_block_ptr,
+                BOUNDS_CHECKS_N,
+                PACKED_PER_VAL,
+                PACKED_D_PER_GROUP,
+                FP8_QUANTIZED,
+                Q.dtype.element_ty,
+                v_dtype,
+                i,
+                IS_HIP,
+                QUANTIZE_PV_TO_FP8,
+                QUANTIZE_QK_TO_FP8,
+                USE_FP32_SCALES,
+            )
+
+            # Unpack results based on quantization configuration
+            if QUANTIZE_PV_TO_FP8 and QUANTIZE_QK_TO_FP8:
+                k[i], v[i], v_scale, k_scale = result  # noqa: F821
+            elif QUANTIZE_PV_TO_FP8:
+                k[i], v[i], v_scale = result  # noqa: F821
+            elif QUANTIZE_QK_TO_FP8:
+                k[i], v[i], k_scale = result  # noqa: F821
+            else:
+                k[i], v[i] = result  # noqa: F821
+        # -- compute qk ---
+        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+        for i in range(len(acc)):  # noqa: F821
+            qk += tl.dot(q[i], k[i])  # noqa: F821
+
+        if QUANTIZE_QK_TO_FP8:
+            # Reshape k_scale for proper broadcasting with qk
+            # k_scale has shape (BLOCK_N,), we need to reshape it to (1, BLOCK_N)
+            # for proper broadcasting with qk of shape (BLOCK_M, BLOCK_N)
+            k_scale_reshaped = tl.reshape(k_scale, (1, BLOCK_N))
+
+            # Apply k_scale to qk
+            qk = qk * k_scale_reshaped
+            qk = qk * tl.reshape(q_scale, (BLOCK_M, 1))  # noqa: F821
+
+        # Apply qk_scale (scalar)
+        qk *= qk_scale
+
+        if start_n == lo and ignore_in_first_block > 0:
+            qk = tl.where(
+                tl.arange(0, BLOCK_N) < ignore_in_first_block, float("-inf"), qk
+            )
+
+        if HAS_ADDITIVE_BIAS:
+            loaded_bias = tl.load(
+                additive_bias_block_ptr,
+                boundary_check=(0, 1) if BOUNDS_CHECKS_N else (0,),
+            )
+            qk += loaded_bias.to(tl.float32) * log2e
+            additive_bias_block_ptr = tl.advance(additive_bias_block_ptr, (0, BLOCK_N))
+
+        # TODO: This is slow, and only needed at the last iteration.
+        # Maybe we can unroll the last iteration instead?
+        if BOUNDS_CHECKS_N:
+            qk = tl.where(tl.arange(0, BLOCK_N) < hi - start_n, qk, float("-inf"))
+        if IS_CAUSAL:
+            # -- apply the causal mask --
+            qk = tl.where(diag_idx_shifted >= start_n - start_kv_idx, qk, float("-inf"))
+        if IS_LOCAL:
+            # -- apply the local window size mask --
+            qk = tl.where(
+                diag_idx_shifted < start_n - start_kv_idx + WINDOW_LEFT + 1,
+                qk,
+                float("-inf"),
+            )
+            if not IS_CAUSAL and WINDOW_RIGHT >= 0:
+                qk = tl.where(
+                    diag_idx_shifted >= start_n - start_kv_idx - WINDOW_RIGHT,
+                    qk,
+                    float("-inf"),
+                )
+
+        # -- compute scaling constant ---
+        m_i_new = tl.maximum(m_i, tl.max(qk, 1))
+        alpha = tl.math.exp2(m_i - m_i_new)
+        p = tl.math.exp2(qk - m_i_new[:, None])
+        if HAS_ADDITIVE_BIAS or (IS_CAUSAL or IS_LOCAL):
+            # NOTE: It's possible that an entire block is masked out.
+            # if this is the case, `m_i_new=nan` and everything becomes nan
+            alpha = tl.where(m_i_new == float("-inf"), 0, alpha)
+            p = tl.where(m_i_new[:, None] == float("-inf"), 0, p)
+
+        # -- update m_i and l_i --
+        l_i = l_i * alpha + tl.sum(p, 1)
+        m_i = m_i_new
+        if not QUANTIZE_PV_TO_FP8:
+            p = p.to(v_dtype)
+        else:
+            # Apply v-scale to P
+            p = p * tl.trans(v_scale)
+
+            # Quantize P to FP8
+            MAX_FP8 = 448
+            amax = tl.max(p, axis=1)  # rowmax(P)
+            p_scale = tl.maximum(amax / MAX_FP8, 1e-9)
+            p_scaled = p / p_scale[:, None]
+            p_clamped = tl.clamp(p_scaled, 0, MAX_FP8)
+
+            # covert P to FP8
+            p = p_clamped.to(v_dtype, fp_downcast_rounding="rtne")
+
+        # -- scale and update acc --
+        for i in range(len(acc)):  # noqa: F821
+            acc[i] *= alpha[:, None]  # noqa: F821
+            if not QUANTIZE_PV_TO_FP8:
+                acc[i] += tl.dot(p, v[i])  # noqa: F821
+            else:
+                # Re-scale PV using p_scale
+                acc[i] += tl.dot(p, v[i]) * p_scale[:, None]  # noqa: F821
+
+        if not PAGE_SIZE:
+            # update pointers
+            K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
+            V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))
+            if PACKED_PER_VAL > 1:
+                K_scale_shift_block_ptr = tl.advance(
+                    K_scale_shift_block_ptr, (0, BLOCK_N)
+                )
+                V_scale_shift_block_ptr = tl.advance(
+                    V_scale_shift_block_ptr, (BLOCK_N, 0)
+                )
+
+    # write back O
+    O_block_ptr = tl.make_block_ptr(
+        base=Out_splitK
+        + off_z.to(tl.int64) * stride_osk_z * queries_use_batch_dim
+        + off_m * stride_osk_m
+        + off_g * stride_osk_g
+        + off_h * stride_osk_h
+        + splitk_idx * stride_osk_s,
+        shape=(q_len, D_PER_GROUP),
+        strides=(stride_osk_m, 1),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, D_PER_GROUP),
+        order=(1, 0),
+    )
+    for i in range(len(acc)):  # noqa: F821
+        # If for the current batch element there are no tokens in the current split-k chunk (because
+        # seqlen is too short), l_i will be 0, so we need to make sure attention is filled with zeros and not NaNs.
+        attn_out = tl.where(l_i[:, None] == 0, 0.0, acc[i] / l_i[:, None])  # noqa: F821
+        tl.store(
+            tl.advance(O_block_ptr, (0, i * D_PER_GROUP)),
+            attn_out.to(Out_splitK.dtype.element_ty),  # noqa: F821
+            boundary_check=(0,),
+        )
+    if WRITE_LSE:
+        LSE_splitk_ptr = (
+            LSE_splitk
+            + off_z * stride_lsek_z * queries_use_batch_dim
+            + off_m * stride_lsek_m
+            + off_g * stride_lsek_g
+            + off_h * stride_lsek_h
+            + splitk_idx * stride_lsek_s
+            + (start_m * BLOCK_M + tl.arange(0, BLOCK_M)) * stride_lsek_m
+        )
+        mask = start_m * BLOCK_M + tl.arange(0, BLOCK_M) < q_len
+        # Can be float64 to improve numerics
+        lse_dtype = LSE_splitk.dtype.element_ty
+        tl.store(
+            LSE_splitk_ptr,
+            (tl.math.log2(l_i.to(lse_dtype)) + m_i.to(lse_dtype)) / log2e,
+            mask=mask,
+        )
+
+
+def gen_config(
+    block_m: int,
+    block_n: int,
+    stages: int,
+    warps: int,
+) -> triton.Config:
+    """A more compact way to define a triton.Config, so it fits on one line"""
+
+    return triton.Config(
+        {
+            "BLOCK_M": block_m,
+            "BLOCK_N": block_n,
+        },
+        num_stages=stages,
+        num_warps=warps,
+    )
+
+
+def _get_splitk_kernel(num_groups):
+    """
+    Kernel _fwd_kernel_splitK needs to be post-processed by unroll_varargs
+    to specialize it for a given number of quantization groups N_GROUPS
+    before we can apply triton.heuristics and triton.autotune, so we
+    don't do them as decorators.
+    """
+
+    _fwd_kernel_splitK_unrolled = unroll_varargs(_fwd_kernel_splitK, N=num_groups)
+    kernel = triton.heuristics(
+        {
+            "BOUNDS_CHECKS_N": lambda args: bool(
+                (args["BLOCK_N_PER_SPLIT"] % args["BLOCK_N"])
+                or (
+                    args["BLOCK_N_PER_SPLIT"] > 0
+                    and args["N_CTX_K"] % args["BLOCK_N_PER_SPLIT"]
+                )
+                or args["USE_SEQ_LEN"]
+            )
+        }
+    )(_fwd_kernel_splitK_unrolled)
+    return kernel
+
+
+def early_config_prune(configs, named_args, **kwargs):
+    use_paged_attention = kwargs["USE_PAGED_ATTENTION"]
+    page_size = kwargs["PAGE_SIZE"]
+    if use_paged_attention:
+        return list(
+            filter(lambda config: page_size % config.kwargs["BLOCK_N"] == 0, configs)
+        )
+    else:
+        return configs
+
+
+@functools.lru_cache(None)
+def autotune_kernel(kernel: Callable):
+    BLOCK_M_VALUES = [16, 32, 64, 128]
+    BLOCK_N_VALUES = [16, 32, 64, 128]
+    STAGES_VALUES = [1, 2] if torch.version.hip else [1, 2, 3]
+    WARPS_VALUES = [1, 2, 4, 8]
+
+    TRITON_CONFIGS = [
+        gen_config(block_m, block_n, stages, warps)
+        for block_m in BLOCK_M_VALUES
+        for block_n in BLOCK_N_VALUES
+        for stages in STAGES_VALUES
+        for warps in WARPS_VALUES
+        if block_n >= block_m
+    ]
+
+    kernel = triton.autotune(
+        configs=TRITON_CONFIGS,
+        key=AUTOTUNER_KEY,
+        use_cuda_graph=True,
+        prune_configs_by={
+            "early_config_prune": early_config_prune,
+        },
+    )(kernel)
+    return kernel
+
+
+# This object contains forward kernels wrapped into autotuner for different number
+# of quantization groups.
+_fwd_kernel_splitK_autotune: Dict[int, triton.runtime.Autotuner] = {}
+# The loop below:
+# - transforms the jitted kernel with unroll_varargs producing a new kernel of each value of num_groups
+# - wraps the kernel into triton.heuristics
+# - wraps kernel into Triton autotuner. Autotuning itself happens the first time the kernel is called
+if sys.version_info >= (3, 9):
+    # unroll_varargs requires Python 3.9+
+    for num_groups in [1, 2, 4, 8]:
+        _fwd_kernel_splitK_autotune[num_groups] = autotune_kernel(
+            _get_splitk_kernel(num_groups)
+        )
+
+    def get_autotuner_cache(
+        num_groups: int,
+    ) -> Dict[Tuple[Union[int, str]], triton.Config]:
+        """Returns a triton.runtime.autotuner.AutoTuner.cache object, which
+        represents mappings from kernel autotune keys (tuples describing kernel inputs)
+        to triton.Config
+        """
+        return _fwd_kernel_splitK_autotune[num_groups].cache
+
+    def set_autotuner_cache(
+        cache: Dict[Tuple[Union[int, str]], triton.Config], num_groups: int
+    ) -> None:
+        _fwd_kernel_splitK_autotune[num_groups].cache = cache
+
+
+@triton.jit
+def load_dequantize_k_v_group(
+    K_block_ptr,
+    V_block_ptr,
+    K_scale_shift_block_ptr,
+    V_scale_shift_block_ptr,
+    BOUNDS_CHECKS_N: tl.constexpr,
+    PACKED_PER_VAL: tl.constexpr,
+    PACKED_D_PER_GROUP: tl.constexpr,
+    FP8_QUANTIZED: tl.constexpr,
+    q_dtype: tl.constexpr,
+    v_dtype: tl.constexpr,  # Q.dtype.element_ty
+    group_id: tl.constexpr,
+    IS_HIP: tl.constexpr,
+    QUANTIZE_PV_TO_FP8: tl.constexpr,
+    QUANTIZE_QK_TO_FP8: tl.constexpr,
+    USE_FP32_SCALES: tl.constexpr,
+):
+    """Load K/V for a given block. In case of int4/fp8-quantized K/V, dequantize them after loading.
+    If quantization is group-wise, use group_id to advance the pointers to the current group.
+
+    Returns:
+        - k, v: loaded and potentially dequantized tensors
+        - v_scale (optional): V scale factor if QUANTIZE_PV_TO_FP8 is True
+        - k_scale (optional): K scale factor if QUANTIZE_QK_TO_FP8 is True
+    """
+    # Advance to the current quantization group
+    K_block_ptr = tl.advance(K_block_ptr, (PACKED_D_PER_GROUP * group_id, 0))
+    V_block_ptr = tl.advance(
+        V_block_ptr,
+        (0, PACKED_D_PER_GROUP * group_id)
+        if not QUANTIZE_PV_TO_FP8
+        else (0, PACKED_D_PER_GROUP * PACKED_PER_VAL * group_id),
+    )
+
+    # -- load k, v --
+    k = tl.load(K_block_ptr, boundary_check=(1,) if BOUNDS_CHECKS_N else ())
+    v = tl.load(V_block_ptr, boundary_check=(0,) if BOUNDS_CHECKS_N else ())
+
+    # Initialize return values
+    # v_scale = None
+    # k_scale = None
+
+    if FP8_QUANTIZED:
+        k, v, v_scale, k_scale = _process_fp8_quantization(
+            k,
+            v,
+            K_scale_shift_block_ptr,
+            V_scale_shift_block_ptr,
+            BOUNDS_CHECKS_N,
+            PACKED_PER_VAL,
+            q_dtype,
+            v_dtype,
+            IS_HIP,
+            QUANTIZE_PV_TO_FP8,
+            QUANTIZE_QK_TO_FP8,
+            USE_FP32_SCALES,
+        )
+
+    elif PACKED_PER_VAL > 1:
+        # Int4 quantization.
+        k, v = _process_int4_quantization(
+            k,
+            v,
+            K_scale_shift_block_ptr,
+            V_scale_shift_block_ptr,
+            group_id,
+            BOUNDS_CHECKS_N,
+            PACKED_PER_VAL,
+            q_dtype,
+            IS_HIP,
+        )
+
+    # Return appropriate values based on quantization flags
+    if QUANTIZE_PV_TO_FP8 and QUANTIZE_QK_TO_FP8:
+        # Return both v_scale and k_scale for applying to P and K
+        return k, v, v_scale, k_scale
+    elif QUANTIZE_PV_TO_FP8:
+        # Return v_scale for applying v_scale to P
+        return k, v, v_scale
+    elif QUANTIZE_QK_TO_FP8:
+        # Return k_scale for applying k_scale to K
+        return k, v, k_scale
+    else:
+        return k, v
+
+
+@triton.jit
+def _process_fp8_quantization(
+    k,
+    v,
+    K_scale_shift_block_ptr,
+    V_scale_shift_block_ptr,
+    BOUNDS_CHECKS_N: tl.constexpr,
+    PACKED_PER_VAL: tl.constexpr,
+    q_dtype: tl.constexpr,
+    v_dtype: tl.constexpr,
+    IS_HIP: tl.constexpr,
+    QUANTIZE_PV_TO_FP8: tl.constexpr,
+    QUANTIZE_QK_TO_FP8: tl.constexpr,
+    USE_FP32_SCALES: tl.constexpr,
+):
+    """Process FP8 quantization for K and V tensors."""
+    # Process V tensor
+    v_scale_shift = tl.load(
+        V_scale_shift_block_ptr, boundary_check=(0,) if BOUNDS_CHECKS_N else ()
+    )
+    v_scale, v_shift = _extract_scale_shift(v_scale_shift, IS_HIP, USE_FP32_SCALES)
+    if not QUANTIZE_PV_TO_FP8:
+        v = dequantize(
+            v,
+            v_scale,
+            v_shift if not USE_FP32_SCALES else None,
+            PACKED_PER_VAL,
+            IS_HIP,
+            USE_FP32_SCALES,
+        ).to(v_dtype)
+    else:
+        # Do not dequantize V; V needs to be FP8 for PV.
+        tl.static_assert(PACKED_PER_VAL == 4, "Assert: int32 packs four FP8 values")
+
+    # Process K tensor
+    k_scale_shift = tl.load(
+        K_scale_shift_block_ptr, boundary_check=(1,) if BOUNDS_CHECKS_N else ()
+    )
+    k_scale, k_shift = _extract_scale_shift(k_scale_shift, IS_HIP, USE_FP32_SCALES)
+    if IS_HIP:
+        if not QUANTIZE_QK_TO_FP8:
+            k = dequantize_k_hip(k, k_scale, k_shift, PACKED_PER_VAL).to(q_dtype)
+        else:
+            # For QUANTIZE_QK_TO_FP8, unpack int32 to 8-bit entries and interpret as fp8
+            tl.static_assert(PACKED_PER_VAL == 4, "Assert: int32 packs four FP8 values")
+    else:
+        if not QUANTIZE_QK_TO_FP8:
+            k_t = dequantize(
+                tl.trans(k),
+                tl.trans(k_scale),
+                tl.trans(k_shift) if not USE_FP32_SCALES else None,
+                PACKED_PER_VAL,
+                IS_HIP,
+                USE_FP32_SCALES,
+            ).to(q_dtype)
+            k = tl.trans(k_t)
+        else:
+            # For QUANTIZE_QK_TO_FP8, unpack int32 to 8-bit entries and interpret as fp8
+            tl.static_assert(PACKED_PER_VAL == 4, "Assert: int32 packs four FP8 values")
+            k_t = tl.trans(k)
+            k_t = _unpack_fp8_tensor(k_t, PACKED_PER_VAL, IS_HIP)
+            k = tl.trans(k_t)
+
+    return k, v, v_scale, k_scale
+
+
+@triton.jit
+def _extract_scale_shift(
+    scale_shift, IS_HIP: tl.constexpr, USE_FP32_SCALES: tl.constexpr
+):
+    """Extract scale and shift values from packed representation."""
+    if IS_HIP:
+        return cast_uint32_to_float(scale_shift)
+    elif USE_FP32_SCALES:
+        return scale_shift.to(tl.float32, bitcast=True), 0
+    else:
+        return cast_uint32_to_half2(scale_shift)
+
+
+@triton.jit
+def _unpack_fp8_tensor(x_, PACKED_PER_VAL: tl.constexpr, IS_HIP: tl.constexpr):
+    """Unpack FP8 K/V tensor from int32 packed representation."""
+    tl.static_assert(PACKED_PER_VAL == 4, "Assert: int32 packs four FP8 values")
+
+    BLOCK_DMODEL_PACKED: tl.constexpr = x_.shape[1]
+    BLOCK_N: tl.constexpr = x_.shape[0]
+    # Create bit offsets for unpacking (0, 8, 16, 24 bits)
+    offsets = tl.arange(0, PACKED_PER_VAL) * 8
+
+    # Extract 8-bit values by right-shifting and masking
+    unpacked_values = x_[:, :, None, :] >> offsets
+    unpacked_values = tl.reshape(
+        unpacked_values, (BLOCK_N, BLOCK_DMODEL_PACKED * PACKED_PER_VAL)
+    )
+
+    # Convert to FP8 through bitcast
+    fp8_type = tl.float8e4b8 if IS_HIP else tl.float8e4nv
+    x_ = unpacked_values.to(tl.uint8).to(fp8_type, bitcast=True)
+
+    return x_
+
+
+@triton.jit
+def _process_int4_quantization(
+    k,
+    v,
+    K_scale_shift_block_ptr,
+    V_scale_shift_block_ptr,
+    group_id: tl.constexpr,
+    BOUNDS_CHECKS_N: tl.constexpr,
+    PACKED_PER_VAL: tl.constexpr,
+    dtype: tl.constexpr,
+    IS_HIP: tl.constexpr,
+):
+    """Process INT4 quantization for K and V tensors."""
+    # Advance scale/shift pointers for INT4
+    K_scale_shift_block_ptr = tl.advance(K_scale_shift_block_ptr, (group_id, 0))
+    V_scale_shift_block_ptr = tl.advance(V_scale_shift_block_ptr, (0, group_id))
+
+    k_scale_shift = tl.load(
+        K_scale_shift_block_ptr, boundary_check=(1,) if BOUNDS_CHECKS_N else ()
+    )
+    v_scale_shift = tl.load(
+        V_scale_shift_block_ptr, boundary_check=(0,) if BOUNDS_CHECKS_N else ()
+    )
+    if IS_HIP:
+        k_scale, k_shift = cast_uint32_to_float(k_scale_shift)
+        v_scale, v_shift = cast_uint32_to_float(v_scale_shift)
+        v = dequantize(v, v_scale, v_shift, PACKED_PER_VAL, IS_HIP).to(dtype)
+        k = dequantize_k_hip(k, k_scale, k_shift, PACKED_PER_VAL).to(dtype)
+    else:
+        k_scale, k_shift = cast_uint32_to_half2(k_scale_shift)
+        v_scale, v_shift = cast_uint32_to_half2(v_scale_shift)
+        v = dequantize(v, v_scale, v_shift, PACKED_PER_VAL, IS_HIP).to(dtype)
+        k_t = dequantize(
+            tl.trans(k),
+            tl.trans(k_scale),
+            tl.trans(k_shift),
+            PACKED_PER_VAL,
+            IS_HIP,
+        ).to(dtype)
+        k = tl.trans(k_t)
+
+    return k, v
+
+
+@triton.jit
+def cast_uint64_to_float2(scale_shift):
+    """Using FP32 scales, so only extract one fp32 from the packed int64"""
+    scale = scale_shift & 0xFFFFFFFF
+    scale = scale.to(tl.uint32).to(tl.float32, bitcast=True)
+    return scale, 0
+
+
+@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 cast_uint32_to_float(scale_shift):
+    """Extract two float16 packed into one int32 as float32"""
+    scale = scale_shift & 0xFFFF
+    shift = scale_shift >> 16
+    scale = scale.to(tl.uint16).to(tl.float16, bitcast=True).to(tl.float32)
+    shift = shift.to(tl.uint16).to(tl.float16, bitcast=True).to(tl.float32)
+    return scale, shift
+
+
+@triton.jit
+def dequantize_k_hip(
+    x_,
+    scale,
+    shift,
+    PACKED_PER_VAL: tl.constexpr,
+):
+    """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.
+    """
+    # x_ : (BLOCK_N, D // PACKED_PER_VAL)
+    # scale: (BLOCK_N, 1)
+    # offsets: (PACKED_PER_VAL,)
+    BLOCK_N: tl.constexpr = x_.shape[1]
+    BLOCK_DMODEL_PACKED: tl.constexpr = x_.shape[0]
+    offsets = tl.arange(0, PACKED_PER_VAL) * (32 // PACKED_PER_VAL)
+    quant_offset = (
+        x_[:, None, :, :] >> offsets[:, None]
+    )  # (BLOCK_N, D // PACKED_PER_VAL, PACKED_PER_VAL)
+
+    quant_offset = tl.reshape(
+        quant_offset, (BLOCK_DMODEL_PACKED * PACKED_PER_VAL, BLOCK_N)
+    )
+
+    if PACKED_PER_VAL == 4:
+        # FP8 quantization.
+        fp8_type = tl.float8e4b8 if torch.version.hip is not None else tl.float8e4nv
+        dequant = (
+            quant_offset.to(tl.uint8).to(fp8_type, bitcast=True).to(scale.dtype) * scale
+            + shift
+        )
+    else:
+        # Int4 quantization.
+        # 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)
+            .to(tl.float32)
+        )
+        quant_offset = quant_offset * 32768.0
+        scale_512 = scale * 512
+
+        dequant = quant_offset * scale_512 + shift
+    return dequant
+
+
+@triton.jit
+def dequantize(
+    x_,
+    scale,
+    shift,
+    PACKED_PER_VAL: tl.constexpr,
+    IS_HIP: tl.constexpr,
+    USE_FP32_SCALES: tl.constexpr = False,
+):
+    """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.
+    """
+    # x_ : (BLOCK_N, D // PACKED_PER_VAL)
+    # scale: (BLOCK_N, 1)
+    # offsets: (PACKED_PER_VAL,)
+    BLOCK_N: tl.constexpr = x_.shape[0]
+    BLOCK_DMODEL_PACKED: tl.constexpr = x_.shape[1]
+    offsets = tl.arange(0, PACKED_PER_VAL) * (32 // PACKED_PER_VAL)
+    quant_offset = (
+        x_[:, :, None, :] >> offsets
+    )  # (BLOCK_N, D // PACKED_PER_VAL, PACKED_PER_VAL)
+
+    quant_offset = tl.reshape(
+        quant_offset, (BLOCK_N, BLOCK_DMODEL_PACKED * PACKED_PER_VAL)
+    )
+    if PACKED_PER_VAL == 4:
+        # FP8 quantization.
+        fp8_type = tl.float8e4b8 if torch.version.hip is not None else tl.float8e4nv
+        dequant = (
+            quant_offset.to(tl.uint8).to(fp8_type, bitcast=True).to(scale.dtype) * scale
+        )
+        if not USE_FP32_SCALES:
+            # Use asymmetric quantization only for FP16 scales
+            dequant += shift
+    else:
+        # Int4 quantization.
+        # Trick - instead of converting int4 to float16 we view it as float16
+        # and then multiply by 32768 * 512 == 2**24
+        if IS_HIP:
+            # Do final math in float32 to avoid casting to bf16 on MI300. There
+            # no direct instructions for this so its less performant on this workload.
+            quant_offset = (
+                (quant_offset & 0xF)
+                .to(tl.uint16)
+                .to(tl.float16, bitcast=True)
+                .to(tl.float32)
+            )
+            quant_offset = quant_offset * 32768.0
+        else:
+            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
+        if not USE_FP32_SCALES:
+            # Use asymmetric quantization only for FP16 scales
+            dequant += shift
+    return dequant
+
+
+@triton.jit
+def _splitK_reduce(
+    Out_splitK,  # [B, G, H, split_k, Mq, K]
+    LSE_splitK,  # [B, G, H, split_k, Mq]
+    Out,  # [B, H, M, K]
+    LSE,  # [B, H, M]
+    split_k: tl.constexpr,
+    splitK_pow2: tl.constexpr,
+    stride_osk_z: tl.constexpr,
+    stride_osk_g: tl.constexpr,
+    stride_osk_h: tl.constexpr,
+    stride_osk_s: tl.constexpr,
+    stride_osk_m: tl.constexpr,
+    stride_osk_k: tl.constexpr,
+    stride_lsek_z: tl.constexpr,
+    stride_lsek_g: tl.constexpr,
+    stride_lsek_h: tl.constexpr,
+    stride_lsek_s: tl.constexpr,
+    stride_lsek_m: tl.constexpr,
+    stride_oz: tl.constexpr,
+    stride_og: tl.constexpr,
+    stride_oh: tl.constexpr,
+    stride_om: tl.constexpr,
+    stride_ok: tl.constexpr,
+    stride_lse_z: tl.constexpr,
+    stride_lse_g: tl.constexpr,
+    stride_lse_h: tl.constexpr,
+    stride_lse_m: tl.constexpr,
+    head_dim: tl.constexpr,
+    head_dim_pow_2: tl.constexpr,
+    H: tl.constexpr,
+    G: tl.constexpr,
+    WRITE_LSE: tl.constexpr,
+):
+    # grid = (M, B * G * H, 1)
+    off_m = tl.program_id(0).to(tl.int64)
+    off_zhg = tl.program_id(1).to(tl.int64)
+    off_z = off_zhg // (H * G)
+    off_h = (off_zhg // G) % H
+    off_g = off_zhg % G
+
+    head_dim_mask = tl.arange(0, head_dim_pow_2) < head_dim
+
+    Out_splitK_ptr = (
+        Out_splitK
+        + stride_osk_z * off_z
+        + stride_osk_g * off_g
+        + stride_osk_h * off_h
+        + stride_osk_m * off_m
+        + tl.arange(0, head_dim_pow_2)[None, :]
+        + stride_osk_s * tl.arange(0, splitK_pow2)[:, None]
+    )
+
+    LSE_splitK_ptr0 = (
+        LSE_splitK
+        + stride_lsek_z * off_z
+        + stride_lsek_g * off_g
+        + stride_lsek_h * off_h
+        + stride_lsek_m * off_m
+        + stride_lsek_s * tl.arange(0, splitK_pow2)
+    )
+
+    if splitK_pow2 > split_k:
+        mask_1d = tl.arange(0, splitK_pow2) < split_k
+        mask_2d = mask_1d[:, None] & head_dim_mask[None, :]
+        lse_splitk = tl.load(LSE_splitK_ptr0, mask=mask_1d, other=float("-inf"))
+        lse_max = tl.max(lse_splitk)
+        out_splitk = tl.load(
+            Out_splitK_ptr, mask=mask_2d, other=0
+        )  # (split_k, head_dim_pow_2)
+        lse_splitk = tl.load(
+            LSE_splitK_ptr0, mask=mask_1d, other=float("-inf")
+        )  # (split_k,)
+    else:
+        lse_splitk = tl.load(LSE_splitK_ptr0)
+        lse_max = tl.max(lse_splitk)
+        out_splitk = tl.load(Out_splitK_ptr)
+        lse_splitk = tl.load(LSE_splitK_ptr0)
+
+    sumexp_normalized_splitk = tl.math.exp2(
+        (lse_splitk - lse_max).to(tl.float32) * 1.44269504
+    )  # (split_k,)
+    sumexp_normalized = tl.sum(sumexp_normalized_splitk, axis=0)  # scalar
+    # Compute numerator
+    numerator_normalized = tl.sum(
+        out_splitk * sumexp_normalized_splitk[:, None], axis=0
+    )
+    acc = numerator_normalized / sumexp_normalized
+    acc = tl.where(lse_max == float("-inf"), 0.0, acc)
+
+    Out_ptr = (
+        Out
+        + stride_oz * off_z
+        + stride_oh * off_h
+        + stride_og * off_g
+        + stride_om * off_m
+        + tl.arange(0, head_dim_pow_2)
+    )
+    if acc.dtype is tl.float64 and Out.dtype.element_ty is not tl.float64:
+        # must avoid direct cast f64->f16
+        acc = acc.to(tl.float32)
+    tl.store(Out_ptr, acc, mask=head_dim_mask)
+
+    if WRITE_LSE:
+        l_ptrs = (
+            LSE
+            + off_z * stride_lse_z
+            + off_g * stride_lse_g
+            + off_h * stride_lse_h
+            + off_m * stride_lse_m
+        )
+        to_store = lse_max + tl.math.log2(sumexp_normalized) / 1.44269504
+        to_store = tl.where(lse_max == float("-inf"), lse_max, to_store)
+        tl.store(l_ptrs, to_store)
+
+
+@triton.jit
+def _splitK_reduce_varargs(
+    Out_splitK: "VAR_ARGS_ARRAY",  # list of [B, G, H, Mq, K];
+    LSE_splitK: "VAR_ARGS_ARRAY",  # list of [B, G, H, Mq]
+    Out,  # [B, G, H, M, K]
+    LSE,  # [B, G, H, M]
+    stride_osk_z: "VAR_ARGS_ARRAY",
+    stride_osk_g: "VAR_ARGS_ARRAY",
+    stride_osk_h: "VAR_ARGS_ARRAY",
+    stride_osk_m: "VAR_ARGS_ARRAY",
+    stride_osk_k: "VAR_ARGS_ARRAY",
+    stride_lsek_z: "VAR_ARGS_ARRAY",
+    stride_lsek_g: "VAR_ARGS_ARRAY",
+    stride_lsek_h: "VAR_ARGS_ARRAY",
+    stride_lsek_m: "VAR_ARGS_ARRAY",
+    stride_oz,
+    stride_og,
+    stride_oh,
+    stride_om,
+    stride_ok,
+    stride_lse_z,
+    stride_lse_g,
+    stride_lse_h,
+    stride_lse_m,
+    head_dim: tl.constexpr,
+    head_dim_pow_2: tl.constexpr,
+    H: tl.constexpr,
+    G: tl.constexpr,
+    WRITE_LSE: tl.constexpr,
+):
+    """
+    This version of reduce kernel takes attention and LSE of chunks as lists of tensors,
+    as opposed to _splitK_reduce, which takes each as a stacked tensor.
+    """
+    # grid = (M, B * G * H, 1)
+    off_m = tl.program_id(0).to(tl.int64)
+    off_zhg = tl.program_id(1).to(tl.int64)
+    off_z = off_zhg // (H * G)
+    off_h = (off_zhg // G) % H
+    off_g = off_zhg % G
+    head_dim_mask = tl.arange(0, head_dim_pow_2) < head_dim
+
+    out_splitk_offset: "VAR_ARGS_ARRAY"  # noqa: F821
+    for i in range(len(Out_splitK)):
+        out_splitk_offset[i] = (  # noqa: F821
+            stride_osk_z[i] * off_z  # type: ignore # noqa: F821
+            + stride_osk_g[i] * off_g
+            + stride_osk_h[i] * off_h
+            + stride_osk_m[i] * off_m
+            + tl.arange(0, head_dim_pow_2)
+        )
+    lse_splitk_offset: "VAR_ARGS_ARRAY"  # noqa: F821
+    for i in range(len(Out_splitK)):
+        lse_splitk_offset[i] = (  # noqa: F821
+            stride_lsek_z[i] * off_z  # type: ignore # noqa: F821
+            + stride_lsek_g[i] * off_g
+            + stride_lsek_h[i] * off_h
+            + stride_lsek_m[i] * off_m
+        )
+
+    lse_max = float("-inf")
+    for split_k_idx in range(len(Out_splitK)):  # type: ignore # noqa: F821
+        LSE_splitK_ptr = LSE_splitK[split_k_idx] + lse_splitk_offset[split_k_idx]  # type: ignore # noqa: F821
+        lse_splitk = tl.load(LSE_splitK_ptr)
+        lse_max = tl.maximum(lse_max, lse_splitk)
+
+    sumexp_normalized = 0.0
+    numerator_normalized = tl.zeros([head_dim_pow_2], dtype=tl.float32)
+
+    for split_k_idx in range(len(Out_splitK)):  # type: ignore # noqa: F821
+        out_splitk = tl.load(
+            Out_splitK[split_k_idx] + out_splitk_offset[split_k_idx],  # type: ignore # noqa: F821
+            mask=head_dim_mask,
+        )
+        lse_splitk = tl.load(LSE_splitK[split_k_idx] + lse_splitk_offset[split_k_idx])  # type: ignore # noqa: F821
+        # Compute denominator
+        sumexp_normalized_splitk = tl.math.exp2(
+            (lse_splitk - lse_max).to(tl.float32) * 1.44269504
+        )
+        sumexp_normalized += sumexp_normalized_splitk
+
+        # Compute numerator
+        numerator_normalized += out_splitk * sumexp_normalized_splitk
+
+    acc = numerator_normalized / sumexp_normalized
+    acc = tl.where(lse_max == float("-inf"), 0.0, acc)
+
+    Out_ptr = (
+        Out
+        + stride_oz * off_z
+        + stride_oh * off_h
+        + stride_og * off_g
+        + stride_om * off_m
+        + tl.arange(0, head_dim_pow_2)
+    )
+    if acc.dtype is tl.float64 and Out.dtype.element_ty is not tl.float64:
+        # must avoid direct cast f64->f16
+        acc = acc.to(tl.float32)
+    tl.store(Out_ptr, acc, mask=head_dim_mask)
+
+    if WRITE_LSE:
+        l_ptrs = (
+            LSE
+            + off_z * stride_lse_z
+            + off_g * stride_lse_g
+            + off_h * stride_lse_h
+            + off_m * stride_lse_m
+        )
+        to_store = lse_max + tl.math.log2(sumexp_normalized) / 1.44269504
+        to_store = tl.where(lse_max == float("-inf"), lse_max, to_store)
+        tl.store(l_ptrs, to_store)
+
+
+@triton.jit
+def _splitK_reduce_varargs_backward(
+    Out_splitK: "VAR_ARGS_ARRAY",  # list of [B, G, H, Mq, K];
+    LSE_splitK: "VAR_ARGS_ARRAY",  # list of [B, G, H, Mq]
+    Dout_splitK: "VAR_ARGS_ARRAY",  # gradients - same shape as the inputs themselves
+    DLSE_splitK: "VAR_ARGS_ARRAY",
+    Out,  # [B, G, H, M, K]
+    LSE,  # [B, G, H, M]
+    DOut,
+    DLSE,
+    # strides of chunked inputs: attention and LSE
+    stride_osk_z: "VAR_ARGS_ARRAY",
+    stride_osk_g: "VAR_ARGS_ARRAY",
+    stride_osk_h: "VAR_ARGS_ARRAY",
+    stride_osk_m: "VAR_ARGS_ARRAY",
+    stride_osk_k: "VAR_ARGS_ARRAY",
+    stride_lsek_z: "VAR_ARGS_ARRAY",
+    stride_lsek_g: "VAR_ARGS_ARRAY",
+    stride_lsek_h: "VAR_ARGS_ARRAY",
+    stride_lsek_m: "VAR_ARGS_ARRAY",
+    # strides of merged outputs: attention and LSE
+    stride_oz,
+    stride_og,
+    stride_oh,
+    stride_om,
+    stride_ok,
+    stride_lse_z,
+    stride_lse_g,
+    stride_lse_h,
+    stride_lse_m,
+    # strides of gradients
+    stride_doz,
+    stride_dog,
+    stride_doh,
+    stride_dom,
+    stride_dok,
+    stride_dlse_z,
+    stride_dlse_g,
+    stride_dlse_h,
+    stride_dlse_m,
+    BLOCK_SIZE: tl.constexpr,
+    H: tl.constexpr,
+    G: tl.constexpr,
+):
+    """
+    Backward for _splitK_reduce_varargs. Similar to forward, it takes
+    attention and LSE of chunks as lists of tensors,
+    and outputs the corresponding gradients in the same format.
+    """
+
+    # grid = (M, B * G * H, 1)
+    off_m = tl.program_id(0).to(tl.int64)
+    off_zhg = tl.program_id(1).to(tl.int64)
+    off_z = off_zhg // (H * G)
+    off_h = (off_zhg // G) % H
+    off_g = off_zhg % G
+
+    # Compute offsets inside each attention/LSE chunk.
+    # Note that each chunk can have different strides, so offsets can also be different.
+    out_splitk_offset: "VAR_ARGS_ARRAY"  # noqa: F821
+    for i in range(len(Out_splitK)):
+        out_splitk_offset[i] = (  # type: ignore # noqa: F821
+            stride_osk_z[i] * off_z
+            + stride_osk_g[i] * off_g
+            + stride_osk_h[i] * off_h
+            + stride_osk_m[i] * off_m
+            + tl.arange(0, BLOCK_SIZE)
+        )
+    lse_splitk_offset: "VAR_ARGS_ARRAY"  # noqa: F821
+    for i in range(len(Out_splitK)):
+        lse_splitk_offset[i] = (  # type: ignore # noqa: F821
+            stride_lsek_z[i] * off_z
+            + stride_lsek_g[i] * off_g
+            + stride_lsek_h[i] * off_h
+            + stride_lsek_m[i] * off_m
+        )
+
+    lse_max = float("-inf")
+    for split_k_idx in range(len(Out_splitK)):  # type: ignore # noqa: F821
+        LSE_splitK_ptr = LSE_splitK[split_k_idx] + lse_splitk_offset[split_k_idx]  # type: ignore # noqa: F821
+        lse_splitk = tl.load(LSE_splitK_ptr)
+        lse_max = tl.maximum(lse_max, lse_splitk)
+
+    # Load attention and the corresponding gradient
+    offset_out = (
+        stride_oz * off_z
+        + stride_oh * off_h
+        + stride_og * off_g
+        + stride_om * off_m
+        + tl.arange(0, BLOCK_SIZE)
+    )
+    offset_dout = (
+        stride_doz * off_z
+        + stride_doh * off_h
+        + stride_dog * off_g
+        + stride_dom * off_m
+        + tl.arange(0, BLOCK_SIZE)
+    )
+    out = tl.load(Out + offset_out)
+    dattn = tl.load(DOut + offset_dout)
+
+    # Load LSE and the corresponding gradient
+    offset_lse = (
+        stride_lse_z * off_z
+        + stride_lse_h * off_h
+        + stride_lse_g * off_g
+        + stride_lse_m * off_m
+    )
+    offset_dlse = (
+        stride_dlse_z * off_z
+        + stride_dlse_h * off_h
+        + stride_dlse_g * off_g
+        + stride_dlse_m * off_m
+    )
+    lse = tl.load(LSE + offset_lse)
+    dlse = tl.load(DLSE + offset_dlse)
+
+    for split_k_idx in range(len(Out_splitK)):  # type: ignore # noqa: F821
+        # Load attention and LSE of chunks
+        out_splitk = tl.load(Out_splitK[split_k_idx] + out_splitk_offset[split_k_idx])  # type: ignore # noqa: F821
+        lse_splitk = tl.load(LSE_splitK[split_k_idx] + lse_splitk_offset[split_k_idx])  # type: ignore # noqa: F821
+
+        # Pointers to save gradients of attention and LSE of chunks
+        dout_splitk_ptr = Dout_splitK[split_k_idx] + out_splitk_offset[split_k_idx]  # type: ignore # noqa: F821
+        dlse_splitk_ptr = DLSE_splitK[split_k_idx] + lse_splitk_offset[split_k_idx]  # type: ignore # noqa: F821
+
+        # dX/dattn_i = dX/dattn * dattn/dattn_i + dX/dlse * dlse/dattn_i, and dlse/dattn_i == 0
+        dattn_dattn_i = tl.exp(lse_splitk - lse_max) / tl.exp(lse - lse_max)
+        dX_dattn_i = dattn_dattn_i * dattn
+        tl.store(dout_splitk_ptr, dX_dattn_i)
+
+        dattn_dlse_i = (out_splitk - out) * dattn_dattn_i
+
+        # dX/dlse_i = dX/dattn * dattn/dlse_i + dX/dlse * dlse/dlse_i
+        dlse_dlse_i = dattn_dattn_i
+        dX_dlse_i = dlse_dlse_i * dlse + tl.sum(
+            dattn_dlse_i * dattn
+        )  # Sum is over the hidden dimension
+        tl.store(dlse_splitk_ptr, dX_dlse_i)