TLX Block Attention：为 Blackwell 架构量身定制的块稀疏注意力内核

注意力矩阵形状 (N, N)，块大小 B
只有 (i*B ~ (i+1)*B, i*B ~ (i+1)*B) 的对角块非零
其余位置 mask 为 0，不参与计算

import triton
import triton.language as tl
import torch

# ---- 编译期常量：块大小和头维度 ----
BLOCK_SIZE = 64   # B，编译期固定
HEAD_DIM = 128    # d，编译期固定

@triton.jit
def tlx_block_attn_kernel(
    Q_ptr, K_ptr, V_ptr, O_ptr,
    seq_len,       # N，运行时传入
    stride_qm, stride_qh, stride_qd,
    stride_km, stride_kh, stride_kd,
    stride_vm, stride_vh, stride_vd,
    stride_om, stride_oh, stride_od,
    HEAD_DIM: tl.constexpr,     # 编译期常量
    BLOCK_SIZE: tl.constexpr,   # 编译期常量
    NUM_BLOCKS: tl.constexpr,   # N // B，编译期常量
):
    # 当前 warp 处理的对角块序号
    block_id = tl.program_id(0)
    head_id = tl.program_id(1)

    # ---- 编译期已知：block_id 对应的行范围 ----
    # 不需要运行时 mask 判断，直接计算偏移
    q_row_start = block_id * BLOCK_SIZE
    k_row_start = block_id * BLOCK_SIZE  # 对角块：K 和 Q 同一段

    # Q 块加载：[BLOCK_SIZE, HEAD_DIM]
    q_offsets = tl.arange(0, BLOCK_SIZE) * stride_qm + q_row_start * stride_qm
    d_offsets = tl.arange(0, HEAD_DIM) * stride_qd
    q_ptrs = Q_ptr + q_offsets[:, None] * 1 + d_offsets[None, :] * 1 \
             + head_id * stride_qh
    q_block = tl.load(q_ptrs)  # [BLOCK_SIZE, HEAD_DIM]

    # K 块加载：[BLOCK_SIZE, HEAD_DIM]
    k_offsets = tl.arange(0, BLOCK_SIZE) * stride_km + k_row_start * stride_km
    k_ptrs = K_ptr + k_offsets[:, None] * 1 + d_offsets[None, :] * 1 \
             + head_id * stride_kh
    k_block = tl.load(k_ptrs)  # [BLOCK_SIZE, HEAD_DIM]

    # ---- QK^T：[BLOCK_SIZE, BLOCK_SIZE] ----
    # 因为 BLOCK_SIZE 和 HEAD_DIM 是编译期常量，
    # Triton 编译器可以精确规划 SMEM 和寄存器分配
    attn = tl.dot(q_block, tl.trans(k_block))  # [B, B]
    attn = attn * (1.0 / (HEAD_DIM ** 0.5))

    # Softmax（行内）
    attn_max = tl.max(attn, axis=1, keep_dims=True)
    attn = attn - attn_max
    attn_exp = tl.exp(attn)
    attn_sum = tl.sum(attn_exp, axis=1, keep_dims=True)
    attn_norm = attn_exp / attn_sum  # [BLOCK_SIZE, BLOCK_SIZE]

    # V 块加载：[BLOCK_SIZE, HEAD_DIM]
    v_offsets = tl.arange(0, BLOCK_SIZE) * stride_vm + k_row_start * stride_vm
    v_ptrs = V_ptr + v_offsets[:, None] * 1 + d_offsets[None, :] * 1 \
             + head_id * stride_vh
    v_block = tl.load(v_ptrs)

    # ---- Attention × V ----
    o_block = tl.dot(attn_norm, v_block)  # [BLOCK_SIZE, HEAD_DIM]

    # 写回输出
    o_offsets = tl.arange(0, BLOCK_SIZE) * stride_om + q_row_start * stride_om
    o_ptrs = O_ptr + o_offsets[:, None] * 1 + d_offsets[None, :] * 1 \
             + head_id * stride_oh
    tl.store(o_ptrs, o_block)


def tlx_block_attn(q, k, v, block_size=64):
    """块对角稀疏注意力入口函数。

    q, k, v: [batch, heads, seq_len, head_dim]
    block_size: 必须能整除 seq_len
    """
    assert q.shape[2] % block_size == 0, \
        f"seq_len {q.shape[2]} must be divisible by block_size {block_size}"

    B, H, N, D = q.shape
    num_blocks = N // block_size

    o = torch.empty_like(q)

    # 编译期常量通过 constexpr 传入 Triton
    grid = (num_blocks, H, B)

    tlx_block_attn_kernel[grid](
        q, k, v, o,
        N,
        q.stride(2), q.stride(1), q.stride(3),
        k.stride(2), k.stride(1), k.stride(3),
        v.stride(2), v.stride(1), v.stride(3),
        o.stride(2), o.stride(1), o.stride(3),
        HEAD_DIM=D,
        BLOCK_SIZE=block_size,
        NUM_BLOCKS=num_blocks,
    )
    return o


# ---- 快速验证 ----
if __name__ == "__main__":
    torch.manual_seed(42)
    B, H, N, D = 1, 2, 256, 128
    q = torch.randn(B, H, N, D, device="cuda", dtype=torch.float16)
    k = torch.randn(B, H, N, D, device="cuda", dtype=torch.float16)
    v = torch.randn(B, H, N, D, device="cuda", dtype=torch.float16)

    # TLX 块对角注意力
    o_sparse = tlx_block_attn(q, k, v, block_size=64)

    # 对照：完整注意力（只取对角块部分验证数值一致性）
    q_full = q[0, 0, 0:64]   # 第 0 个块
    k_full = k[0, 0, 0:64]
    v_full = v[0, 0, 0:64]
    attn_full = torch.matmul(q_full, k_full.T) / (D ** 0.5)
    attn_full = torch.softmax(attn_full.float(), dim=-1).half()
    o_full = torch.matmul(attn_full, v_full)

    diff = (o_sparse[0, 0, 0:64].float() - o_full.float()).abs().max().item()
    print(f"块对角注意力输出与完整注意力（同块）最大差异: {diff:.6f}")
    print(f"稀疏输出形状: {o_sparse.shape}")
    # 预期差异在 float16 精度范围内（~1e-3）

# 安装 Triton（需要 CUDA 12+ 环境）
pip install triton

# 验证 GPU 可用
python -c "import torch; print(torch.cuda.get_device_name(0))"