solution.py

import torch
import time

# 尝试导入 NPU 相关依赖
try:
    import torch_npu
except ImportError:
    pass

# 尝试导入 TileLang 相关依赖
try:
    import tilelang
    from tilelang import language as T
    HAS_TILELANG = True
except ImportError:
    HAS_TILELANG = False

# =============================================================================
# CUDA 路径：TileLang 实现
# =============================================================================
if HAS_TILELANG:
    @tilelang.jit(
        out_idx=[-1],
        pass_configs={
            tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
            tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
            tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: False,
        },
    )
    def _tilelang_sparse_attention_kernel(batch, heads_q, heads_kv, seq_len, kv_len, dim, scale, block_size, top_k):
        tl_scale = scale * 1.44269504  # log2(e) for exp2
        groups = heads_q // heads_kv
        q_shape = [batch, heads_q, seq_len, dim]
        kv_shape = [batch, heads_kv, kv_len, dim]
        index_shape = [batch, heads_kv, seq_len, top_k]
        
        dtype = T.float16
        accum_dtype = T.float32
        index_dtype = T.int32
        
        BS = block_size
        BK = tilelang.math.next_power_of_2(dim) if dim <= 128 else 128
        BV = BK
        num_stages = 3 if (top_k <= 4 and block_size <= 32) else 1
        threads = 64 if (top_k <= 4 or top_k >= 32) else 32
        use_shared_acc = threads >= 64
        NV = tilelang.cdiv(dim, BV)
        
        qk_policy = T.GemmWarpPolicy.FullCol if top_k >= 32 else T.GemmWarpPolicy.FullRow
        pv_policy = T.GemmWarpPolicy.FullCol if top_k >= 32 else T.GemmWarpPolicy.FullRow

        @T.prim_func
        def sparse_attn_main(
            Q: T.Tensor(q_shape, dtype),
            K: T.Tensor(kv_shape, dtype),
            V: T.Tensor(kv_shape, dtype),
            BlockIndices: T.Tensor(index_shape, index_dtype),
            Output: T.Tensor(q_shape, dtype),
        ):
            with T.Kernel(seq_len, NV, batch * heads_kv, threads=threads) as (bx, by, bz):
                Q_shared = T.alloc_shared([groups, BK], dtype)
                K_shared = T.alloc_shared([BS, BK], dtype)
                V_shared = T.alloc_shared([BS, BV], dtype)
                if top_k >= 32:
                    O_shared = T.alloc_shared([groups, BV], dtype)
                
                acc_s = T.alloc_fragment([groups, BS], accum_dtype)
                if use_shared_acc:
                    acc_s_cast = T.alloc_shared([groups, BS], dtype)
                else:
                    acc_s_cast = T.alloc_fragment([groups, BS], dtype)
                acc_o = T.alloc_fragment([groups, BV], accum_dtype)
                scores_max = T.alloc_fragment([groups], accum_dtype)
                scores_max_prev = T.alloc_fragment([groups], accum_dtype)
                scores_scale = T.alloc_fragment([groups], accum_dtype)
                scores_sum = T.alloc_fragment([groups], accum_dtype)
                logsum = T.alloc_fragment([groups], accum_dtype)
                
                i_n = bx
                i_v = by
                i_bh = bz
                i_b = i_bh // heads_kv
                i_hk = i_bh % heads_kv
                
                T.copy(Q[i_b, i_hk * groups : (i_hk + 1) * groups, i_n, :], Q_shared)
                T.fill(acc_o, 0)
                T.fill(logsum, 0)
                T.fill(scores_max, -T.infinity(accum_dtype))
                
                for i in T.Pipelined(top_k, num_stages=num_stages):
                    block_idx = BlockIndices[i_b, i_hk, i_n, i]
                    kv_start = block_idx * BS
                    if block_idx >= 0 and kv_start <= i_n:
                        T.copy(K[i_b, i_hk, kv_start : kv_start + BS, :], K_shared)
                        if i_n >= (kv_start + BS - 1):
                            T.fill(acc_s, 0)
                        else:
                            for gi, si in T.Parallel(groups, BS):
                                acc_s[gi, si] = T.if_then_else(
                                    i_n >= (kv_start + si), 0, -T.infinity(acc_s.dtype)
                                )
                        T.gemm(Q_shared, K_shared, acc_s, transpose_B=True, policy=qk_policy)
                        if i == 0:
                            T.fill(scores_max, -T.infinity(accum_dtype))
                            T.reduce_max(acc_s, scores_max, dim=1, clear=True)
                            for gi, si in T.Parallel(groups, BS):
                                acc_s[gi, si] = T.exp2(acc_s[gi, si] * tl_scale - scores_max[gi] * tl_scale)
                            T.reduce_sum(acc_s, scores_sum, dim=1)
                            for gi in T.Parallel(groups):
                                logsum[gi] = scores_sum[gi]
                        else:
                            T.copy(scores_max, scores_max_prev)
                            T.fill(scores_max, -T.infinity(accum_dtype))
                            T.reduce_max(acc_s, scores_max, dim=1, clear=True)
                            for gi in T.Parallel(groups):
                                scores_scale[gi] = T.exp2(scores_max_prev[gi] * tl_scale - scores_max[gi] * tl_scale)
                            for gi, si in T.Parallel(groups, BS):
                                acc_s[gi, si] = T.exp2(acc_s[gi, si] * tl_scale - scores_max[gi] * tl_scale)
                            T.reduce_sum(acc_s, scores_sum, dim=1)
                            for gi in T.Parallel(groups):
                                logsum[gi] = logsum[gi] * scores_scale[gi] + scores_sum[gi]
                            for gi, vi in T.Parallel(groups, BV):
                                acc_o[gi, vi] *= scores_scale[gi]
                        if use_shared_acc and top_k <= 4 and BS <= 32:
                            for gi, si in T.Parallel(groups, BS):
                                acc_s_cast[gi, si] = acc_s[gi, si]
                        else:
                            T.copy(acc_s, acc_s_cast)
                        T.copy(V[i_b, i_hk, kv_start : kv_start + BS, i_v * BV : (i_v + 1) * BV], V_shared)
                        T.gemm(acc_s_cast, V_shared, acc_o, policy=pv_policy)
                
                for gi, vi in T.Parallel(groups, BV):
                    acc_o[gi, vi] /= logsum[gi]
                
                if top_k >= 32:
                    T.copy(acc_o, O_shared)
                    T.copy(O_shared, Output[i_b, i_hk * groups : (i_hk + 1) * groups, i_n, i_v * BV : (i_v + 1) * BV])
                else:
                    T.copy(acc_o, Output[i_b, i_hk * groups : (i_hk + 1) * groups, i_n, i_v * BV : (i_v + 1) * BV])
        return sparse_attn_main

_kernel_cache = {}

def _sparse_attention_cuda(q, k, v, index, block_size, sm_scale):
    B, H_Q, N, D_H = q.shape
    _, H_K, M, _ = k.shape
    _, _, _, TOP_K = index.shape
    cache_key = (B, H_Q, H_K, N, M, D_H, TOP_K, block_size, sm_scale)
    if cache_key not in _kernel_cache:
        kernel = _tilelang_sparse_attention_kernel(
            batch=B, heads_q=H_Q, heads_kv=H_K, seq_len=N, kv_len=M,
            dim=D_H, scale=sm_scale, block_size=block_size, top_k=TOP_K,
        )
        _kernel_cache[cache_key] = kernel
    else:
        kernel = _kernel_cache[cache_key]
    if index.dtype != torch.int32:
        index = index.to(torch.int32)
    return kernel(q, k, v, index)

# =============================================================================
# NPU 路径：PyTorch 原生实现 (优化版)
# =============================================================================
def _sparse_attention_npu(q, k, v, index, bs, sm_scale):
    B, h_q, N, d_h = q.shape
    _, h_k, M, _ = k.shape
    _, _, _, top_k = index.shape
    g = h_q // h_k 
    device = q.device
    dtype = q.dtype

    # 1. 维度重构：将 Batch 和 KV-Head 合并，提高并行度
    # [B*h_k, g, N, d_h]
    q_reshaped = q.view(B * h_k, g, N, d_h)
    index_reshaped = index.view(B * h_k, N, top_k)
    k_reshaped = k.view(B * h_k, M, d_h)
    v_reshaped = v.view(B * h_k, M, d_h)
    
    out = torch.empty_like(q_reshaped)

    # 2. 预准备常量
    offsets = torch.arange(bs, device=device).view(1, 1, bs)
    n_range = torch.arange(N, device=device).view(N, 1)

    # 3. 序列分块处理 (Large 规模下 256 是显存与性能的平衡点)
    chunk_size = 256 
    
    # 4. 遍历合并后的头维度 (B * h_k)
    # 虽然这里有循环，但内部是对 [g, chunk, top_k*bs] 的大矩阵运算，NPU 效率依然很高
    for i_bh in range(B * h_k):
        cur_q = q_reshaped[i_bh]      # [g, N, d_h]
        cur_idx = index_reshaped[i_bh] # [N, top_k]
        cur_k = k_reshaped[i_bh]      # [M, d_h]
        cur_v = v_reshaped[i_bh]      # [M, d_h]
        
        for i_start in range(0, N, chunk_size):
            i_end = min(i_start + chunk_size, N)
            cur_N = i_end - i_start
            
            # --- 索引计算 ---
            # t_idx: [cur_N, top_k * bs]
            t_idx = (cur_idx[i_start:i_end, :].unsqueeze(-1) * bs + offsets).view(cur_N, -1)
            mask_invalid = (t_idx < 0)
            t_idx_clamped = t_idx.clamp(min=0, max=M - 1).to(torch.int32)
            
            # --- Gather K, V ---
            # gather_idx: [cur_N, top_k*bs, d_h]
            # 这里的显存占用仅约 256 * 2048 * 128 * 2 = 128 MB，非常安全
            gather_idx = t_idx_clamped.unsqueeze(-1).expand(-1, -1, d_h)
            k_selected = torch.gather(cur_k.unsqueeze(0).expand(cur_N, -1, -1), 1, gather_idx)
            v_selected = torch.gather(cur_v.unsqueeze(0).expand(cur_N, -1, -1), 1, gather_idx)

            # --- Attention Score ---
            # q_chunk: [g, cur_N, d_h], k_selected: [cur_N, top_k*bs, d_h]
            # 利用广播：[g, cur_N, 1, d_h] @ [1, cur_N, d_h, top_k*bs]
            scores = torch.matmul(
                cur_q[:, i_start:i_end, :].unsqueeze(2), 
                k_selected.unsqueeze(0).transpose(-1, -2)
            ).squeeze(2) # [g, cur_N, top_k*bs]
            
            scores.mul_(sm_scale)

            # --- Mask ---
            i_global = n_range[i_start:i_end, :]
            mask_causal = (t_idx > i_global) # [cur_N, top_k*bs]
            # mask_invalid 是 [cur_N, top_k*bs], mask_causal 也是
            scores.masked_fill_(mask_invalid.unsqueeze(0) | mask_causal.unsqueeze(0), float('-inf'))

            # --- Softmax & Aggregation ---
            probs = torch.softmax(scores, dim=-1) # [g, cur_N, top_k*bs]
            
            # [g, cur_N, 1, top_k*bs] @ [1, cur_N, top_k*bs, d_h] -> [g, cur_N, d_h]
            chunk_out = torch.matmul(probs.unsqueeze(2), v_selected.unsqueeze(0)).squeeze(2)
            
            out[i_bh, :, i_start:i_end, :] = chunk_out.to(dtype)

    return out.view(B, h_q, N, d_h)

# =============================================================================
# 统一入口
# =============================================================================
def sparse_attention(q, k, v, index, bs, sm_scale):
    device_type = q.device.type   
    if device_type == "cuda":
        return _sparse_attention_cuda(q, k, v, index, bs, sm_scale)
    else:
        return _sparse_attention_npu(q, k, v, index, bs, sm_scale)