from typing import Optional, Tuple
import functools

import tilelang
import tilelang.language as T
import torch

tilelang.set_log_level("WARNING")
cu_count = torch.cuda.get_device_properties("cuda").multi_processor_count

pass_configs = {
    tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
    tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
    # tilelang.PassConfigKey.TL_DISABLE_FAST_MATH: True,
    tilelang.PassConfigKey.TL_ENABLE_AGGRESSIVE_SHARED_MEMORY_MERGE: True,
    tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: True,
}

BF16 = "bfloat16"
FP8 = "float8_e4m3"
FP32 = "float32"

def fast_log2_ceil(x):
    bits_x = T.reinterpret("uint32", x)
    exp_x = (bits_x >> 23) & 0xFF
    man_bits = bits_x & ((1 << 23) - 1)
    return T.Cast("int32", exp_x - 127 + T.if_then_else(man_bits != 0, 1, 0))


def fast_pow2(x):
    bits_x = (x + 127) << 23
    return T.reinterpret("float32", bits_x)


def fast_round_scale(amax, fp8_max_inv):
    return fast_pow2(fast_log2_ceil(amax * fp8_max_inv))


@tilelang.jit(pass_configs=pass_configs)
def act_quant_kernel(
    N, in_dtype=BF16, out_dtype=FP8, scale_dtype=FP32, round_scale=False
):
    M = T.symbolic("M")
    fp8_min = -448.0
    fp8_max = 448.0
    fp8_max_inv = 1 / fp8_max
    num_stages = 0 if round_scale else 2
    blk_m = 32
    group_size = 128

    @T.prim_func
    def act_quant_kernel_(
        X: T.Tensor[(M, N), in_dtype],
        Y: T.Tensor[(M, N), out_dtype],
        S: T.Tensor[(M, T.ceildiv(N, group_size)), scale_dtype],
    ):
        with T.Kernel(T.ceildiv(M, blk_m), T.ceildiv(N, group_size), threads=128) as (
            pid_m,
            pid_n,
        ):
            x_shared = T.alloc_shared((blk_m, group_size), in_dtype)
            x_local = T.alloc_fragment((blk_m, group_size), in_dtype)
            amax_local = T.alloc_fragment((blk_m,), scale_dtype)
            s_local = T.alloc_fragment((blk_m,), scale_dtype)
            y_local = T.alloc_fragment((blk_m, group_size), out_dtype)
            y_shared = T.alloc_shared((blk_m, group_size), out_dtype)

            for _ in T.Pipelined(1, num_stages=num_stages):
                T.copy(X[pid_m * blk_m, pid_n * group_size], x_shared)
                T.copy(x_shared, x_local)
                T.reduce_absmax(x_local, amax_local, dim=1)
                for i in T.Parallel(blk_m):
                    amax_local[i] = T.max(amax_local[i], 1e-4)
                    if round_scale:
                        s_local[i] = fast_round_scale(amax_local[i], fp8_max_inv)
                    else:
                        s_local[i] = amax_local[i] * fp8_max_inv
                for i, j in T.Parallel(blk_m, group_size):
                    y_local[i, j] = T.clamp(
                        x_local[i, j] / s_local[i], fp8_min, fp8_max
                    )
                for i in T.Parallel(blk_m):
                    S[pid_m * blk_m + i, pid_n] = s_local[i]
                T.copy(y_local, y_shared)
                T.copy(y_shared, Y[pid_m * blk_m, pid_n * group_size])

    return act_quant_kernel_


def act_quant(
    x: torch.Tensor, block_size: int = 128, scale_fmt: Optional[str] = None
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Quantizes the input tensor `x` using block-wise quantization.

    Args:
        x (torch.Tensor): The input tensor to be quantized. Must be contiguous and its last dimension size must be divisible by `block_size`.
        block_size (int, optional): The size of the blocks to be used for quantization. Default is 128.
        scale_fmt (Optional[str], optional): The format of the scale. Default is None.
    Returns:
        Tuple[torch.Tensor, torch.Tensor]: A tuple containing:
            - The quantized tensor with dtype `torch.float8_e4m3fn`.
            - A tensor of scaling factors with dtype `torch.float32`.
    """
    assert x.is_contiguous(), "Input tensor must be contiguous"
    assert (
        x.size(-1) % block_size == 0
    ), f"Last dimension size must be divisible by block_size (block_size={block_size})"
    N = x.size(-1)
    y = torch.empty_like(x, dtype=torch.float8_e4m3fn)
    s = x.new_empty(*x.size()[:-1], N // block_size, dtype=torch.float32)
    kernel = act_quant_kernel(N, round_scale=scale_fmt is not None)
    kernel(x.view(-1, N), y.view(-1, N), s.view(-1, N // block_size))
    return y, s

@tilelang.jit(out_idx=[4], pass_configs=pass_configs)
def fp8_index_kernel(
    h: int, d: int, m_split: int, blk_n1: int, blk_n2: int, disable_buffer_ops: bool = False, threads: int = 256, clear_accum: bool = True
):
    b, m, n = T.symbolic("b"), T.symbolic("m"), T.symbolic("n")
    # if m_split * h > 128, use Square policy to avoid register spill
    gemm_policy = T.GemmWarpPolicy.FullRow if m_split * h <= 128 else T.GemmWarpPolicy.Square

    @T.prim_func
    def fp8_index_kernel_(
        q: T.Tensor[(b, m, h, d), FP8],
        q_s: T.Tensor[(b, m, h), FP32],
        k: T.Tensor[(b, n, d), FP8],
        k_s: T.Tensor[(b, n), FP32],
        o: T.Tensor[(b, m, n), FP32],
    ) -> None:
        with T.Kernel(b, T.ceildiv(n, blk_n1), m, threads=threads) as (
            i_b, i1_n, i_m_block
        ):
            if disable_buffer_ops:
                T.disable_buffer_ops(o)
            m_start = i_m_block
            q_smem = T.alloc_shared((h, d), FP8)
            k_smem = T.alloc_shared((blk_n2, d), FP8)
            T.annotate_layout({
                q_smem: tilelang.layout.make_hcu_swizzled_layout(q_smem, major_pack=1),
                k_smem: tilelang.layout.make_hcu_swizzled_layout(k_smem, major_pack=1),
            })
            q_frag = T.alloc_fragment((h, d), FP8)
            q_s_frag = T.alloc_fragment(h, FP32)
            k_frag = T.alloc_fragment((blk_n2, d), FP8)
            k_s_frag = T.alloc_fragment(blk_n2, FP32)
            logits = T.alloc_fragment((blk_n2, h), FP32)
            logits_sum = T.alloc_fragment(blk_n2, FP32)
            T.copy(q[i_b, m_start, 0, 0], q_smem)
            T.copy(q_smem, q_frag)
            T.copy(q_s[i_b, m_start, 0], q_s_frag)
            for i2_n in T.Pipelined(blk_n1 // blk_n2, num_stages=0):
                T.copy(k[i_b, i1_n * blk_n1 + i2_n * blk_n2, 0], k_smem)
                T.copy(k_s[i_b, i1_n * blk_n1 + i2_n * blk_n2], k_s_frag)
                T.clear(logits)
                T.copy(k_smem, k_frag)
                T.gemm(k_frag, q_frag, logits, transpose_A=False, transpose_B=True, policy=gemm_policy)

                for i_h, i3_n in T.Parallel(h, blk_n2):
                    logits[i3_n, i_h] = T.max(logits[i3_n, i_h], 0) * q_s_frag[i_h]
                T.reduce_sum(logits, logits_sum, dim=1)
                for i3_n in T.Parallel(blk_n2):
                    logits_sum[i3_n] *= k_s_frag[i3_n]
                T.copy(logits_sum, o[i_b, m_start, i1_n * blk_n1 + i2_n * blk_n2])

    @T.prim_func
    def fp8_index_kernel_1(
        q: T.Tensor[(b, m, h, d), FP8],
        q_s: T.Tensor[(b, m, h), FP32],
        k: T.Tensor[(b, n, d), FP8],
        k_s: T.Tensor[(b, n), FP32],
        o: T.Tensor[(b, m, n), FP32],
    ) -> None:
        with T.Kernel(b, T.ceildiv(n, blk_n1), T.ceildiv(m, 2), threads=threads) as (
            i_b, i1_n, i_m_block
        ):
            if disable_buffer_ops:
                T.disable_buffer_ops(o)
            m_start = i_m_block * 2
            q_smem0 = T.alloc_shared((h, d), FP8)
            q_smem1 = T.alloc_shared((h, d), FP8)
            k_smem = T.alloc_shared((blk_n2, d), FP8)
            T.annotate_layout({
                q_smem0: tilelang.layout.make_hcu_swizzled_layout(q_smem0, major_pack=1),
                q_smem1: tilelang.layout.make_hcu_swizzled_layout(q_smem1, major_pack=1),
                k_smem: tilelang.layout.make_hcu_swizzled_layout(k_smem, major_pack=1),
            })
            q_frag0 = T.alloc_fragment((h, d), FP8)
            q_frag1 = T.alloc_fragment((h, d), FP8)
            q_s_frag0 = T.alloc_fragment(h, FP32)
            q_s_frag1 = T.alloc_fragment(h, FP32)
            k_frag = T.alloc_fragment((blk_n2, d), FP8)
            k_s_frag = T.alloc_fragment(blk_n2, FP32)
            logits0 = T.alloc_fragment((blk_n2, h), FP32)
            logits1 = T.alloc_fragment((blk_n2, h), FP32)
            logits_sum0 = T.alloc_fragment(blk_n2, FP32)
            logits_sum1 = T.alloc_fragment(blk_n2, FP32)
            T.copy(q[i_b, m_start, 0, 0], q_smem0)
            T.copy(q[i_b, m_start + 1, 0, 0], q_smem1)
            T.copy(q_smem0, q_frag0)
            T.copy(q_smem1, q_frag1)
            T.copy(q_s[i_b, m_start, 0], q_s_frag0)
            T.copy(q_s[i_b, m_start + 1, 0], q_s_frag1)
            for i2_n in T.Pipelined(blk_n1 // blk_n2, num_stages=0):
                T.copy(k[i_b, i1_n * blk_n1 + i2_n * blk_n2, 0], k_smem)
                T.copy(k_s[i_b, i1_n * blk_n1 + i2_n * blk_n2], k_s_frag)
                T.clear(logits0)
                T.clear(logits1)
                T.copy(k_smem, k_frag)
                T.gemm(k_frag, q_frag0, logits0, transpose_A=False, transpose_B=True, policy=gemm_policy)
                T.gemm(k_frag, q_frag1, logits1, transpose_A=False, transpose_B=True, policy=gemm_policy)

                for i_h, i3_n in T.Parallel(h, blk_n2):
                    logits0[i3_n, i_h] = T.max(logits0[i3_n, i_h], 0) * q_s_frag0[i_h]
                    logits1[i3_n, i_h] = T.max(logits1[i3_n, i_h], 0) * q_s_frag1[i_h]
                T.reduce_sum(logits0, logits_sum0, dim=1)
                T.reduce_sum(logits1, logits_sum1, dim=1)
                for i3_n in T.Parallel(blk_n2):
                    logits_sum0[i3_n] *= k_s_frag[i3_n]
                    logits_sum1[i3_n] *= k_s_frag[i3_n]
                T.copy(logits_sum0, o[i_b, m_start, i1_n * blk_n1 + i2_n * blk_n2])
                T.copy(logits_sum1, o[i_b, m_start + 1, i1_n * blk_n1 + i2_n * blk_n2])

    @T.prim_func
    def fp8_index_kernel_2(
        q: T.Tensor[(b, m, h, d), FP8],
        q_s: T.Tensor[(b, m, h), FP32],
        k: T.Tensor[(b, n, d), FP8],
        k_s: T.Tensor[(b, n), FP32],
        o: T.Tensor[(b, m, n), FP32],
    ) -> None:
        with T.Kernel(b, T.ceildiv(n, blk_n1), T.ceildiv(m, 4), threads=threads) as (
            i_b, i1_n, i_m_block
        ):
            if disable_buffer_ops:
                T.disable_buffer_ops(o)
            m_start = i_m_block * 4
            q_smem0 = T.alloc_shared((h, d), FP8)
            q_smem1 = T.alloc_shared((h, d), FP8)
            q_smem2 = T.alloc_shared((h, d), FP8)
            q_smem3 = T.alloc_shared((h, d), FP8)
            k_smem = T.alloc_shared((blk_n2, d), FP8)
            T.annotate_layout({
                q_smem0: tilelang.layout.make_hcu_swizzled_layout(q_smem0, major_pack=1),
                q_smem1: tilelang.layout.make_hcu_swizzled_layout(q_smem1, major_pack=1),
                q_smem2: tilelang.layout.make_hcu_swizzled_layout(q_smem2, major_pack=1),
                q_smem3: tilelang.layout.make_hcu_swizzled_layout(q_smem3, major_pack=1),
                k_smem: tilelang.layout.make_hcu_swizzled_layout(k_smem, major_pack=1),
            })
            q_frag0 = T.alloc_fragment((h, d), FP8)
            q_frag1 = T.alloc_fragment((h, d), FP8)
            q_frag2 = T.alloc_fragment((h, d), FP8)
            q_frag3 = T.alloc_fragment((h, d), FP8)
            q_s_frag0 = T.alloc_fragment(h, FP32)
            q_s_frag1 = T.alloc_fragment(h, FP32)
            q_s_frag2 = T.alloc_fragment(h, FP32)
            q_s_frag3 = T.alloc_fragment(h, FP32)
            T.copy(q[i_b, m_start, 0, 0], q_smem0)
            T.copy(q[i_b, m_start + 1, 0, 0], q_smem1)
            T.copy(q[i_b, m_start + 2, 0, 0], q_smem2)
            T.copy(q[i_b, m_start + 3, 0, 0], q_smem3)
            T.copy(q_smem0, q_frag0)
            T.copy(q_smem1, q_frag1)
            T.copy(q_smem2, q_frag2)
            T.copy(q_smem3, q_frag3)
            T.copy(q_s[i_b, m_start, 0], q_s_frag0)
            T.copy(q_s[i_b, m_start + 1, 0], q_s_frag1)
            T.copy(q_s[i_b, m_start + 2, 0], q_s_frag2)
            T.copy(q_s[i_b, m_start + 3, 0], q_s_frag3)
            k_frag = T.alloc_fragment((blk_n2, d), FP8)
            k_s_frag = T.alloc_fragment(blk_n2, FP32)
            logits0 = T.alloc_fragment((blk_n2, h), FP32)
            logits1 = T.alloc_fragment((blk_n2, h), FP32)
            logits2 = T.alloc_fragment((blk_n2, h), FP32)
            logits3 = T.alloc_fragment((blk_n2, h), FP32)
            logits_sum0 = T.alloc_fragment(blk_n2, FP32)
            logits_sum1 = T.alloc_fragment(blk_n2, FP32)
            logits_sum2 = T.alloc_fragment(blk_n2, FP32)
            logits_sum3 = T.alloc_fragment(blk_n2, FP32)
            for i2_n in T.Pipelined(blk_n1 // blk_n2, num_stages=0):
                T.copy(k[i_b, i1_n * blk_n1 + i2_n * blk_n2, 0], k_smem)
                T.copy(k_s[i_b, i1_n * blk_n1 + i2_n * blk_n2], k_s_frag)
                T.clear(logits0)
                T.clear(logits1)
                T.clear(logits2)
                T.clear(logits3)
                T.copy(k_smem, k_frag)
                T.gemm(k_frag, q_frag0, logits0, transpose_A=False, transpose_B=True, k_pack=1, policy=gemm_policy)
                T.gemm(k_frag, q_frag1, logits1, transpose_A=False, transpose_B=True, k_pack=1, policy=gemm_policy)
                T.gemm(k_frag, q_frag2, logits2, transpose_A=False, transpose_B=True, k_pack=1, policy=gemm_policy)
                T.gemm(k_frag, q_frag3, logits3, transpose_A=False, transpose_B=True, k_pack=1, policy=gemm_policy)

                for i_h, i3_n in T.Parallel(h, blk_n2):
                    logits0[i3_n, i_h] = T.max(logits0[i3_n, i_h], 0) * q_s_frag0[i_h]
                    logits1[i3_n, i_h] = T.max(logits1[i3_n, i_h], 0) * q_s_frag1[i_h]
                    logits2[i3_n, i_h] = T.max(logits2[i3_n, i_h], 0) * q_s_frag2[i_h]
                    logits3[i3_n, i_h] = T.max(logits3[i3_n, i_h], 0) * q_s_frag3[i_h]
                T.reduce_sum(logits0, logits_sum0, dim=1)
                T.reduce_sum(logits1, logits_sum1, dim=1)
                T.reduce_sum(logits2, logits_sum2, dim=1)
                T.reduce_sum(logits3, logits_sum3, dim=1)
                for i3_n in T.Parallel(blk_n2):
                    logits_sum0[i3_n] *= k_s_frag[i3_n]
                    logits_sum1[i3_n] *= k_s_frag[i3_n]
                    logits_sum2[i3_n] *= k_s_frag[i3_n]
                    logits_sum3[i3_n] *= k_s_frag[i3_n]
                T.copy(logits_sum0, o[i_b, m_start, i1_n * blk_n1 + i2_n * blk_n2])
                T.copy(logits_sum1, o[i_b, m_start + 1, i1_n * blk_n1 + i2_n * blk_n2])
                T.copy(logits_sum2, o[i_b, m_start + 2, i1_n * blk_n1 + i2_n * blk_n2])
                T.copy(logits_sum3, o[i_b, m_start + 3, i1_n * blk_n1 + i2_n * blk_n2])

    @T.prim_func
    def fp8_index_kernel_3(
        q: T.Tensor[(b, m, h, d), FP8],
        q_s: T.Tensor[(b, m, h), FP32],
        k: T.Tensor[(b, n, d), FP8],
        k_s: T.Tensor[(b, n), FP32],
        o: T.Tensor[(b, m, n), FP32],
    ) -> None:
        with T.Kernel(b, T.ceildiv(n, blk_n1), T.ceildiv(m, 8), threads=threads) as (
            i_b, i1_n, i_m_block
        ):
            if disable_buffer_ops:
                T.disable_buffer_ops(o)
            m_start = i_m_block * 8
            q_smem0 = T.alloc_shared((h, d), FP8)
            q_smem1 = T.alloc_shared((h, d), FP8)
            q_smem2 = T.alloc_shared((h, d), FP8)
            q_smem3 = T.alloc_shared((h, d), FP8)
            k_smem = T.alloc_shared((blk_n2, d), FP8)
            T.annotate_layout({
                q_smem0: tilelang.layout.make_hcu_swizzled_layout(q_smem0, major_pack=1),
                q_smem1: tilelang.layout.make_hcu_swizzled_layout(q_smem1, major_pack=1),
                q_smem2: tilelang.layout.make_hcu_swizzled_layout(q_smem2, major_pack=1),
                q_smem3: tilelang.layout.make_hcu_swizzled_layout(q_smem3, major_pack=1),
                k_smem: tilelang.layout.make_hcu_swizzled_layout(k_smem, major_pack=1),
            })
            q_pre_frag0 = T.alloc_fragment((h, d), FP8)
            q_pre_frag1 = T.alloc_fragment((h, d), FP8)
            q_pre_frag2 = T.alloc_fragment((h, d), FP8)
            q_pre_frag3 = T.alloc_fragment((h, d), FP8)
            q_pre_frag4 = T.alloc_fragment((h, d), FP8)
            q_pre_frag5 = T.alloc_fragment((h, d), FP8)
            q_pre_frag6 = T.alloc_fragment((h, d), FP8)
            q_pre_frag7 = T.alloc_fragment((h, d), FP8)
            q_frag0 = T.alloc_fragment((h, d), FP8)
            q_frag1 = T.alloc_fragment((h, d), FP8)
            q_frag2 = T.alloc_fragment((h, d), FP8)
            q_frag3 = T.alloc_fragment((h, d), FP8)
            q_frag4 = T.alloc_fragment((h, d), FP8)
            q_frag5 = T.alloc_fragment((h, d), FP8)
            q_frag6 = T.alloc_fragment((h, d), FP8)
            q_frag7 = T.alloc_fragment((h, d), FP8)
            q_s_frag0 = T.alloc_fragment(h, FP32)
            q_s_frag1 = T.alloc_fragment(h, FP32)
            q_s_frag2 = T.alloc_fragment(h, FP32)
            q_s_frag3 = T.alloc_fragment(h, FP32)
            q_s_frag4 = T.alloc_fragment(h, FP32)
            q_s_frag5 = T.alloc_fragment(h, FP32)
            q_s_frag6 = T.alloc_fragment(h, FP32)
            q_s_frag7 = T.alloc_fragment(h, FP32)
            k_frag = T.alloc_fragment((blk_n2, d), FP8)
            k_s_frag = T.alloc_fragment(blk_n2, FP32)
            logits0 = T.alloc_fragment((blk_n2, h), FP32)
            logits1 = T.alloc_fragment((blk_n2, h), FP32)
            logits2 = T.alloc_fragment((blk_n2, h), FP32)
            logits3 = T.alloc_fragment((blk_n2, h), FP32)
            logits4 = T.alloc_fragment((blk_n2, h), FP32)
            logits5 = T.alloc_fragment((blk_n2, h), FP32)
            logits6 = T.alloc_fragment((blk_n2, h), FP32)
            logits7 = T.alloc_fragment((blk_n2, h), FP32)
            logits_sum0 = T.alloc_fragment(blk_n2, FP32)
            logits_sum1 = T.alloc_fragment(blk_n2, FP32)
            logits_sum2 = T.alloc_fragment(blk_n2, FP32)
            logits_sum3 = T.alloc_fragment(blk_n2, FP32)
            logits_sum4 = T.alloc_fragment(blk_n2, FP32)
            logits_sum5 = T.alloc_fragment(blk_n2, FP32)
            logits_sum6 = T.alloc_fragment(blk_n2, FP32)
            logits_sum7 = T.alloc_fragment(blk_n2, FP32)

            T.copy(q[i_b, m_start, 0, 0], q_pre_frag0)
            T.copy(q[i_b, m_start + 1, 0, 0], q_pre_frag1)
            T.copy(q[i_b, m_start + 2, 0, 0], q_pre_frag2)
            T.copy(q[i_b, m_start + 3, 0, 0], q_pre_frag3)
            T.copy(q[i_b, m_start + 4, 0, 0], q_pre_frag4)
            T.copy(q[i_b, m_start + 5, 0, 0], q_pre_frag5)
            T.copy(q[i_b, m_start + 6, 0, 0], q_pre_frag6)
            T.copy(q[i_b, m_start + 7, 0, 0], q_pre_frag7)
            T.copy(q_s[i_b, m_start, 0], q_s_frag0)
            T.copy(q_s[i_b, m_start + 1, 0], q_s_frag1)
            T.copy(q_s[i_b, m_start + 2, 0], q_s_frag2)
            T.copy(q_s[i_b, m_start + 3, 0], q_s_frag3)
            T.copy(q_s[i_b, m_start + 4, 0], q_s_frag4)
            T.copy(q_s[i_b, m_start + 5, 0], q_s_frag5)
            T.copy(q_s[i_b, m_start + 6, 0], q_s_frag6)
            T.copy(q_s[i_b, m_start + 7, 0], q_s_frag7)

            T.copy(q_pre_frag0, q_smem0)
            T.copy(q_pre_frag1, q_smem1)
            T.copy(q_pre_frag2, q_smem2)
            T.copy(q_pre_frag3, q_smem3)
            T.copy(q_smem0, q_frag0)
            T.copy(q_smem1, q_frag1)
            T.copy(q_smem2, q_frag2)
            T.copy(q_smem3, q_frag3)
            T.copy(q_pre_frag4, q_smem0)
            T.copy(q_pre_frag5, q_smem1)
            T.copy(q_pre_frag6, q_smem2)
            T.copy(q_pre_frag7, q_smem3)
            T.copy(q_smem0, q_frag4)
            T.copy(q_smem1, q_frag5)
            T.copy(q_smem2, q_frag6)
            T.copy(q_smem3, q_frag7)

            for i2_n in T.Pipelined(blk_n1 // blk_n2, num_stages=0):
                T.copy(k[i_b, i1_n * blk_n1 + i2_n * blk_n2, 0], k_smem)
                T.copy(k_s[i_b, i1_n * blk_n1 + i2_n * blk_n2], k_s_frag)
                T.clear(logits0)
                T.clear(logits1)
                T.clear(logits2)
                T.clear(logits3)
                T.clear(logits4)
                T.clear(logits5)
                T.clear(logits6)
                T.clear(logits7)
                T.copy(k_smem, k_frag)
                T.gemm(k_frag, q_frag0, logits0, transpose_A=False, transpose_B=True, k_pack=1, policy=gemm_policy)
                T.gemm(k_frag, q_frag1, logits1, transpose_A=False, transpose_B=True, k_pack=1, policy=gemm_policy)
                T.gemm(k_frag, q_frag2, logits2, transpose_A=False, transpose_B=True, k_pack=1, policy=gemm_policy)
                T.gemm(k_frag, q_frag3, logits3, transpose_A=False, transpose_B=True, k_pack=1, policy=gemm_policy)
                T.gemm(k_frag, q_frag4, logits4, transpose_A=False, transpose_B=True, k_pack=1, policy=gemm_policy)
                T.gemm(k_frag, q_frag5, logits5, transpose_A=False, transpose_B=True, k_pack=1, policy=gemm_policy)
                T.gemm(k_frag, q_frag6, logits6, transpose_A=False, transpose_B=True, k_pack=1, policy=gemm_policy)
                T.gemm(k_frag, q_frag7, logits7, transpose_A=False, transpose_B=True, k_pack=1, policy=gemm_policy)
                for i_h, i3_n in T.Parallel(h, blk_n2):
                    logits0[i3_n, i_h] = T.max(logits0[i3_n, i_h], 0) * q_s_frag0[i_h]
                    logits1[i3_n, i_h] = T.max(logits1[i3_n, i_h], 0) * q_s_frag1[i_h]
                    logits2[i3_n, i_h] = T.max(logits2[i3_n, i_h], 0) * q_s_frag2[i_h]
                    logits3[i3_n, i_h] = T.max(logits3[i3_n, i_h], 0) * q_s_frag3[i_h]
                    logits4[i3_n, i_h] = T.max(logits4[i3_n, i_h], 0) * q_s_frag4[i_h]
                    logits5[i3_n, i_h] = T.max(logits5[i3_n, i_h], 0) * q_s_frag5[i_h]
                    logits6[i3_n, i_h] = T.max(logits6[i3_n, i_h], 0) * q_s_frag6[i_h]
                    logits7[i3_n, i_h] = T.max(logits7[i3_n, i_h], 0) * q_s_frag7[i_h]

                T.reduce_sum(logits0, logits_sum0, dim=1)
                T.reduce_sum(logits1, logits_sum1, dim=1)
                T.reduce_sum(logits2, logits_sum2, dim=1)
                T.reduce_sum(logits3, logits_sum3, dim=1)
                T.reduce_sum(logits4, logits_sum4, dim=1)
                T.reduce_sum(logits5, logits_sum5, dim=1)
                T.reduce_sum(logits6, logits_sum6, dim=1)
                T.reduce_sum(logits7, logits_sum7, dim=1)
                for i3_n in T.Parallel(blk_n2):
                    logits_sum0[i3_n] *= k_s_frag[i3_n]
                    logits_sum1[i3_n] *= k_s_frag[i3_n]
                    logits_sum2[i3_n] *= k_s_frag[i3_n]
                    logits_sum3[i3_n] *= k_s_frag[i3_n]
                    logits_sum4[i3_n] *= k_s_frag[i3_n]
                    logits_sum5[i3_n] *= k_s_frag[i3_n]
                    logits_sum6[i3_n] *= k_s_frag[i3_n]
                    logits_sum7[i3_n] *= k_s_frag[i3_n]

                T.copy(logits_sum0, o[i_b, m_start, i1_n * blk_n1 + i2_n * blk_n2])
                T.copy(logits_sum1, o[i_b, m_start + 1, i1_n * blk_n1 + i2_n * blk_n2])
                T.copy(logits_sum2, o[i_b, m_start + 2, i1_n * blk_n1 + i2_n * blk_n2])
                T.copy(logits_sum3, o[i_b, m_start + 3, i1_n * blk_n1 + i2_n * blk_n2])
                T.copy(logits_sum4, o[i_b, m_start + 4, i1_n * blk_n1 + i2_n * blk_n2])
                T.copy(logits_sum5, o[i_b, m_start + 5, i1_n * blk_n1 + i2_n * blk_n2])
                T.copy(logits_sum6, o[i_b, m_start + 6, i1_n * blk_n1 + i2_n * blk_n2])
                T.copy(logits_sum7, o[i_b, m_start + 7, i1_n * blk_n1 + i2_n * blk_n2])

    if m_split == 1:
        return fp8_index_kernel_
    elif m_split == 2:
        return fp8_index_kernel_1
    elif m_split == 4:
        return fp8_index_kernel_2
    else:
        return fp8_index_kernel_3


@functools.lru_cache(maxsize=64)
def _get_config_module(h: int, d: int, cu_count: int):
    """Get config module for (h, d, cu_count). Returns None if not found."""
    config_module_name = f"fp8_index_tuned_config_h{h}_d{d}_cu{cu_count}"
    try:
        return __import__(f"aiter.ops.tilelang.configs.fp8_index.{config_module_name}", fromlist=["get_tuned_config"])
    except (ImportError, AttributeError):
        return None


@functools.lru_cache(maxsize=128)
def _get_fp8_index_kernel(
    h: int,
    d: int,
    m_split: int,
    blk_n1: int,
    blk_n2: int,
    disable_buffer_ops: bool = False,
    threads: int = 256,
    clear_accum: bool = True,
):
    """Cached kernel creation. Dispatches to fp8_index_kernel_ / _1 / _2 based on m_split."""
    assert m_split in (1, 2, 4, 8), "m_split must be 1, 2, 4, or 8"
    print(
        f"[fp8_index] kernel config: h={h} d={d} m_split={m_split} blk_n1={blk_n1} blk_n2={blk_n2} "
        f"threads={threads} clear_accum={clear_accum} disable_buffer_ops={disable_buffer_ops}"
    )
    return fp8_index_kernel(
        h, d, m_split, blk_n1, blk_n2, disable_buffer_ops, threads=threads, clear_accum=clear_accum
    )


def fp8_index(
    q: torch.Tensor,
    q_s: torch.Tensor,
    k: torch.Tensor,
    k_s: torch.Tensor,
) -> torch.Tensor:
    """
    Perform index score using FP8 precision.

    Args:
        q (torch.Tensor): The Q tensor, must be contiguous.
        q_s (torch.Tensor): The scaling factor for Q (float), must be contiguous.
        k (torch.Tensor): The K tensor, must be contiguous.
        k_s (torch.Tensor): The scaling factor for K (e8m0 here), must be contiguous.

        fp8 q @ fp8 k -> fp32 logits
        relu(fp32 logits) * q_s (weights) -> fp32 logits
        fp32 logits -> fp32 logits_sum
        fp32 logits_sum * k_s (e8m0) -> fp32 index_score
    """
    b, m, h, d = q.shape
    n = k.shape[1]

    # Use tuned config; fallback to default if not found (run tune_fp8_index.py to generate)
    mod = _get_config_module(h, d, cu_count)
    if mod is not None:
        m_split, blk_n1, blk_n2 = mod.get_tuned_config(m, n)
    else:
        m_split, blk_n1, blk_n2 = 1, 512, 128

    disable_buffer_ops = False
    if b * m * n * 4 >= 4294967296:
        disable_buffer_ops = True

    kernel = _get_fp8_index_kernel(
        h, d, m_split, blk_n1, blk_n2, disable_buffer_ops, threads=256, clear_accum=False
    )
    return kernel(q, q_s, k, k_s)
