import math

import torch
import triton
import triton.language as tl

@triton.jit
def flash_attention_fwd_kernel(
    q_ptr,
    k_ptr,
    v_ptr,
    out_ptr,
    n_ctx,
    sm_scale,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_D: tl.constexpr,
):
    # program_id(0): 当前处理哪一块 Query
    # program_id(1): 当前处理哪一个 batch-head
    block_m_id = tl.program_id(0)
    batch_head_id = tl.program_id(1)

    # 把 [B, H, N, D] 的前两维合并成 B*H
    bh_offset = batch_head_id * n_ctx * BLOCK_D

    # 当前 Query 块的行下标，以及 head_dim 的列下标
    offs_m = block_m_id * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_d = tl.arange(0, BLOCK_D)

    # 加载一块 Q，形状为 [BLOCK_M, BLOCK_D]
    q_offsets = bh_offset + offs_m[:, None] * BLOCK_D + offs_d[None, :]
    q = tl.load(
        q_ptr + q_offsets,
        mask=offs_m[:, None] < n_ctx,
        other=0.0,
    )

    # 在线 Softmax 的三个状态：
    # m_i：目前见过的每一行最大值
    # l_i：目前见过的 exp(score - max) 之和
    # acc：目前累积的 Softmax(score) @ V 分子
    m_i = tl.full((BLOCK_M,), -float("inf"), tl.float32)
    l_i = tl.zeros((BLOCK_M,), tl.float32)
    acc = tl.zeros((BLOCK_M, BLOCK_D), tl.float32)

    # 不一次性生成完整的 N×N 分数矩阵，
    # 而是一块一块读取 K 和 V。
    for start_n in tl.range(0, n_ctx, BLOCK_N):
        offs_n = start_n + tl.arange(0, BLOCK_N)

        # K 按 [D, BLOCK_N] 加载，方便计算 Q @ K^T
        k_offsets = bh_offset + offs_n[None, :] * BLOCK_D + offs_d[:, None]
        k = tl.load(
            k_ptr + k_offsets,
            mask=offs_n[None, :] < n_ctx,
            other=0.0,
        )

        # V 按 [BLOCK_N, D] 加载
        v_offsets = bh_offset + offs_n[:, None] * BLOCK_D + offs_d[None, :]
        v = tl.load(
            v_ptr + v_offsets,
            mask=offs_n[:, None] < n_ctx,
            other=0.0,
        )

        # 当前小块的注意力分数，形状 [BLOCK_M, BLOCK_N]
        scores = tl.dot(q, k) * sm_scale

        # 最后一块可能越界，越界位置不能进入 Softmax
        scores = tl.where(
            offs_n[None, :] < n_ctx,
            scores,
            -float("inf"),
        )

        # -------- 在线 Softmax 更新 --------
        # 当前分数块每一行的最大值
        block_max = tl.max(scores, axis=1)

        # 合并旧最大值和当前块最大值
        m_new = tl.maximum(m_i, block_max)

        # 旧结果因为最大值变化，需要重新缩放
        alpha = tl.exp(m_i - m_new)

        # 当前块在新最大值下的指数值
        p = tl.exp(scores - m_new[:, None])

        # 更新分母
        l_new = alpha * l_i + tl.sum(p, axis=1)

        # 更新分子：
        # 旧分子先乘 alpha，再加当前块的 p @ V
        acc = acc * alpha[:, None] + tl.dot(p.to(tl.float16), v)

        m_i = m_new
        l_i = l_new

    # Softmax 分子 / Softmax 分母
    out = acc / l_i[:, None]

    out_offsets = bh_offset + offs_m[:, None] * BLOCK_D + offs_d[None, :]
    tl.store(
        out_ptr + out_offsets,
        out,
        mask=offs_m[:, None] < n_ctx,
    )

def simple_flash_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
    assert q.is_cuda and k.is_cuda and v.is_cuda
    assert q.dtype == k.dtype == v.dtype == torch.float16
    assert q.shape == k.shape == v.shape
    assert q.ndim == 4
    assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous()

    batch, heads, n_ctx, head_dim = q.shape
    assert head_dim in (16, 32, 64, 128)

    out = torch.empty_like(q)

    block_m = 32
    block_n = 32

    # 每个 program 计算一个 Query 块；
    # 第二维枚举所有 batch-head。
    grid = (triton.cdiv(n_ctx, block_m), batch * heads)

    flash_attention_fwd_kernel[grid](
        q,
        k,
        v,
        out,
        n_ctx,
        1.0 / math.sqrt(head_dim),
        BLOCK_M=block_m,
        BLOCK_N=block_n,
        BLOCK_D=head_dim,
        num_warps=4,
    )
    return out

def torch_reference(q, k, v):
    """直接写出标准 Attention，便于理解和对拍。"""
    scale = 1.0 / math.sqrt(q.shape[-1])
    scores = torch.matmul(q, k.transpose(-1, -2)) * scale
    probs = torch.softmax(scores, dim=-1)
    return torch.matmul(probs, v)

def main():
    if not torch.cuda.is_available():
        raise RuntimeError("需要 CUDA GPU 才能运行 Triton Kernel")

    torch.manual_seed(0)

    # 使用非整块序列长度 100，顺便测试边界 mask
    q = torch.randn(2, 4, 100, 64, device="cuda", dtype=torch.float16)
    k = torch.randn_like(q)
    v = torch.randn_like(q)

    actual = simple_flash_attention(q, k, v)
    expected = torch_reference(q, k, v)

    max_error = (actual - expected).abs().max().item()
    print("输出形状:", tuple(actual.shape))
    print("最大绝对误差:", max_error)

    torch.testing.assert_close(
        actual,
        expected,
        atol=2e-2,
        rtol=2e-2,
    )
    print("正确性检查通过")

if __name__ == "__main__":
    main()