NSA 算子解析（Torch/Tilelang）

Maxwell_Newton

335人浏览 · 2026-07-01 15:24:04

Maxwell_Newton · 2026-07-01 15:24:04 发布

Torch朴素实现

定义

NSA 的核心思想是将原本全局的 Attention 分解为三个优势互补的稀疏分支，以极低的 KV Cache 占用保留长文本的表征能力：

Compressed Attention（压缩注意力）
做法：将历史的 KV 序列按照固定的 Block 大小（例如每 128 个 Token 一个块）进行划分，通过一个简单的 Pooling 或小型网络进行粗粒度压缩。
作用：用极小规模的 KV 表达全局的“背景宏观语义”，大幅降低远距离历史信息的计算开销，同时会输出各个块的“重要性得分”。
Selected Attention（选择性注意力）
做法：基于 Compressed 阶段算出的块重要性得分，利用 Top-K 算法筛选出当前 Query 最相关的少数几个 KV 块。
作用：实现激活值级别的动态稀疏。只把真正重要的非局部信息加载进计算，既保证了长文本的召回准确率（Perfect Retrieval），又避免了全量计算。
Sliding Window Attention（滑动窗口注意力）
做法：对当前 Query 序列尾部最近的 KV 块采用固定大小的滑动窗口（Local Window）。
作用：捕捉最新、最细粒度的局部上下文信息，弥补前两路分支因为 Block-level 处理而丢失的局部细节（如临近词组、句法结构）。

如图以上三部分注意力合并，在经过一个gate，得到注意力输出。
在这里插入图片描述
NSA也是一种稀疏注意力，也就是一个q不会和全部kv计算注意力，而是选一些计算。但很多稀疏方法，都是计算一个权值，然后选出权值最高的topk个token，这从数学角度是可行的，但在 GPU 算子层面，普通的稀疏注意力通常伴随着大量的非连续内存访问和不规则的控制流，导致 Warp Stall 严重。

NSA 为了解决普通稀疏注意力硬件不友好的问题，用了一个简单的方法：分块，这也是为什么它叫做naive。分块是经典数据结构和高性能优化思想，当我们执行一个操作需要大量遍历时，不妨对序列分块，然后按块访问，这样访问下标是连续的，对内存更友好，且实现简单。并且，选择合适的块长，可以对齐GPU底层的Tensor Core MMA指令，最大程度释放Tensor Core的计吞吐

此外，NSA是DeepSeek提出的，而DeepSeek积极引入tilelang开发算子，因此可以说tilelang的编程范式，很大长度上也影响了NSA的设计思路，因为tilelang的基本单元就是tile（数据块），那么分块计算可以在软件上匹配tilelang编程范式，提高性能，减低编码难度。

在分块基础上，完整的NSA流程是：

每个块压缩成少数几个词向量作为代表，先计算q和每个块的代表的注意力，保存为压缩注意力。并且根据压缩注意力，找到最适合的topk个块，寻找方式有很多，这里不展开
对于最适合的topk个块，计算块内所有kv的完整注意力。作为选择注意力保存。
此外，还引入划窗，当前q token所在块附近的块，无条件计算完整注意力，离得近意味着语义更相关。保存为划窗注意力
这三个注意力再做一个加权合并，得到最终注意力输出。

今天的算子，因为要迁移到tilelang，适合高吞吐的问题，不适合分支太多的问题，所以只考虑选择注意力，选块的过程也不考虑，直接输入已经选好的块id，然后对这几个块做完整注意力，最后合并。

输入输出

先分析输入，QKV就是注意力的输入，shape里B是batch，T是seqlen，HQ是Q的头数，H是KV头数，这里是GQA所以HQ不等于H，一个KV头可能对应多个Q头。K是QK的隐藏层维度，V是V的隐藏层维度。

这里QK的T是相等的，也就是这是prefill推理或者自回归训练阶段，QKV都是同样的输入文本产生的。

Block indices是外面已经计算好的，计算完整注意力的块编号，前三个维度仍然是（B, T, H,），第四个维度S是选中的块数，一般固定16，可以不足，不足的用填充位填上，但是不能多。

Block counts是表示Block indices的实际有效块数，用于Block indices有效下标不足16个的情况

Selected block size. Default: 64.块长固定为64

o 是输出张量，对每个batch，每个token，每个q head，都输出一个维度V的注意力结果

r"""
Args:
    q (torch.Tensor):
        queries of shape `[B, T, HQ, K]` if `head_first=False` else `[B, HQ, T, K]`.
    k (torch.Tensor):
        keys of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`.
        GQA is enforced here. The ratio of query heads (HQ) to key/value heads (H) must be a power of 2 and >=16.
    v (torch.Tensor):
        values of shape `[B, T, H, V]` if `head_first=False` else `[B, H, T, V]`.
    block_indices (torch.LongTensor):
        Block indices of shape `[B, T, H, S]` if `head_first=False` else `[B, H, T, S]`.
        `S` is the maximum number of selected blocks for each query token, which is set to 16 in the paper.
    block_counts (torch.LongTensor):
        Block counts of shape `[B, T, H]` if `head_first=False` else `[B, H, T]`.
    block_size (int):
        Selected block size. Default: 64.

Returns:
    o (torch.Tensor):
        Outputs of shape `[B, T, HQ, V]` if `head_first=False` else `[B, HQ, T, V]`.
"""

初始化准备

G = HQ // H表示每个kv head，要对应多少个q head，SELECTED_BLOCKS_SIZE = S * BS块长*有效块数，等于所有要计算的KV个数

peat(x, "b t h d -> b t (h g) d", g=G) for x in (k, v, block_indices))，把k,v,bi这几个张量，把h维度复制成h,g，因为每个kv head都要对应g和q head计算，这类似于torch的广播机制

repeat(block_counts, "b t h -> b t (h g)", g=G)类似地，block_count也复制。

c = torch.arange(S).repeat_interleave(BS).unsqueeze(1).expand(-1, q.shape[2]).to(q.device)构造一个张量，大小对齐BS,S,head，每个要计算的kv head都有，构造的数据是个等差数列然后重复，表示每个kv head属于第几个块，后面和block_count对比，构造掩码，把不满16块时无效部分盖住

    scale = k.shape[-1] ** -0.5

    dtype = q.dtype
    HQ = q.shape[2]
    H = k.shape[2]
    D = k.shape[-1]
    G = HQ // H
    BS = block_size
    S = block_indices.shape[-1]
    SELECTED_BLOCKS_SIZE = S * BS
    k, v, block_indices = (repeat(x, "b t h d -> b t (h g) d", g=G) for x in (k, v, block_indices))
    block_counts = repeat(block_counts, "b t h -> b t (h g)", g=G)
    c = torch.arange(S).repeat_interleave(BS).unsqueeze(1).expand(-1, q.shape[2]).to(q.device)
    q, k, v = map(lambda x: x.float(), (q, k, v))
    o = torch.zeros_like(v)
    B, T = q.shape[:2]

batch循环

q_b, k_b, v_b, i_b, s_b = 先取出对应数据

i_b = i_b.unsqueeze(-1) * BS + i_b.new_tensor(range(BS))根据块编号，构造出token的具体编号i_b = i_b.unsqueeze(-1) * BS + i_b.new_tensor(range(BS))，比如block_id=3,block_size=4,那么就能构造出一个 $[12, 13, 14, 15]$ ，表示这个块内的具体token id

i_b = i_b.view(T, block_indices.shape[2], -1).transpose(1, 2)再reshape一下，把head维度放到最后，计算的基本单元是kv head和q head

# This variant keeps the gather and reduction flow explicit for readability.
for i in range(B):
    q_b, k_b, v_b, i_b, s_b = q[i], k[i], v[i], block_indices[i], block_counts[i]
    # [T, HQ, S, BS] -> [T, HQ, S*BS]
    i_b = i_b.unsqueeze(-1) * BS + i_b.new_tensor(range(BS))
    # [T, HQ, S*BS] -> [T, S*BS, HQ]
    i_b = i_b.view(T, block_indices.shape[2], -1).transpose(1, 2)

seqlen循环

内部再枚举每个token，得到i_i ,s_i,q_i,q_i,v_i，分别是这个q token对应的kv token下标，对应的有效kv块数，对应的qkv张量

for i_q in range(T):
    # [HQ, D]
    q_i = q_b[i_q] * scale
    # [S*BS, HQ] -> represents selected blocks for each query token
    i_i = i_b[i_q]
    # [HQ] -> represents the number of selected blocks for each query token
    s_i = s_b[i_q]

    k_i = torch.zeros((S * BS, HQ, D), device=k_b.device, dtype=k_b.dtype)
    v_i = torch.zeros((S * BS, HQ, D), device=v_b.device, dtype=v_b.dtype)

循环所有头+选中token

for h in range(HQ):枚举所有q head，前面已经把kv都复制到head数等于q head了。

for t in range(SELECTED_BLOCKS_SIZE):对于一个q head，枚举所有的选中的kv token

selected_block_index = i_i[t, h]取出对应的kv token的id，k_i[t, h] = k_b[selected_block_index, h, :]复制kv到临时数组准备计算

# Copy token-by-token so the mapping from selected indices is obvious.
for h in range(HQ):
    for t in range(SELECTED_BLOCKS_SIZE):
        selected_block_index = i_i[t, h]
        k_i[t, h] = k_b[selected_block_index, h, :]
        v_i[t, h] = v_b[selected_block_index, h, :]

计算注意力

attn = torch.einsum("h d, n h d -> n h", q_i, k_i)qk计算注意力得分

attn = attn.masked_fill((i_i > i_q) | (c >= s_i), float("-inf"))如果一个kv token下标，超过当前询问的q token id，或者所在块编号c超过当前的有效块数si,都无效，设置为负无穷，后面softmax时映射成0

attn = torch.softmax(attn, dim=0)做softmax归一化

o[i, i_q] = torch.einsum("n h, n h v -> h v", attn, v_i)注意力得分乘上V得到注意力输出

# [S*BS, HQ]
attn = torch.einsum("h d, n h d -> n h", q_i, k_i)
attn = attn.masked_fill((i_i > i_q) | (c >= s_i), float("-inf"))
attn = torch.softmax(attn, dim=0)
o[i, i_q] = torch.einsum("n h, n h v -> h v", attn, v_i)

完整代码

def naive_nsa_simple(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    block_indices: torch.LongTensor,
    block_counts: torch.LongTensor,
    block_size: int = 64,
) -> torch.Tensor:
    r"""
    Args:
        q (torch.Tensor):
            queries of shape `[B, T, HQ, K]` if `head_first=False` else `[B, HQ, T, K]`.
        k (torch.Tensor):
            keys of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`.
            GQA is enforced here. The ratio of query heads (HQ) to key/value heads (H) must be a power of 2 and >=16.
        v (torch.Tensor):
            values of shape `[B, T, H, V]` if `head_first=False` else `[B, H, T, V]`.
        block_indices (torch.LongTensor):
            Block indices of shape `[B, T, H, S]` if `head_first=False` else `[B, H, T, S]`.
            `S` is the maximum number of selected blocks for each query token, which is set to 16 in the paper.
        block_counts (torch.LongTensor):
            Block counts of shape `[B, T, H]` if `head_first=False` else `[B, H, T]`.
        block_size (int):
            Selected block size. Default: 64.

    Returns:
        o (torch.Tensor):
            Outputs of shape `[B, T, HQ, V]` if `head_first=False` else `[B, HQ, T, V]`.
    """
    scale = k.shape[-1] ** -0.5

    dtype = q.dtype
    HQ = q.shape[2]
    H = k.shape[2]
    D = k.shape[-1]
    G = HQ // H
    BS = block_size
    S = block_indices.shape[-1]
    SELECTED_BLOCKS_SIZE = S * BS
    k, v, block_indices = (repeat(x, "b t h d -> b t (h g) d", g=G) for x in (k, v, block_indices))
    block_counts = repeat(block_counts, "b t h -> b t (h g)", g=G)
    c = torch.arange(S).repeat_interleave(BS).unsqueeze(1).expand(-1, q.shape[2]).to(q.device)
    q, k, v = map(lambda x: x.float(), (q, k, v))
    o = torch.zeros_like(v)
    B, T = q.shape[:2]

    # This variant keeps the gather and reduction flow explicit for readability.
    for i in range(B):
        q_b, k_b, v_b, i_b, s_b = q[i], k[i], v[i], block_indices[i], block_counts[i]
        # [T, HQ, S, BS] -> [T, HQ, S*BS]
        i_b = i_b.unsqueeze(-1) * BS + i_b.new_tensor(range(BS))
        # [T, HQ, S*BS] -> [T, S*BS, HQ]
        i_b = i_b.view(T, block_indices.shape[2], -1).transpose(1, 2)
        for i_q in range(T):
            # [HQ, D]
            q_i = q_b[i_q] * scale
            # [S*BS, HQ] -> represents selected blocks for each query token
            i_i = i_b[i_q]
            # [HQ] -> represents the number of selected blocks for each query token
            s_i = s_b[i_q]

            k_i = torch.zeros((S * BS, HQ, D), device=k_b.device, dtype=k_b.dtype)
            v_i = torch.zeros((S * BS, HQ, D), device=v_b.device, dtype=v_b.dtype)

            # Copy token-by-token so the mapping from selected indices is obvious.
            for h in range(HQ):
                for t in range(SELECTED_BLOCKS_SIZE):
                    selected_block_index = i_i[t, h]
                    k_i[t, h] = k_b[selected_block_index, h, :]
                    v_i[t, h] = v_b[selected_block_index, h, :]

            # [S*BS, HQ]
            attn = torch.einsum("h d, n h d -> n h", q_i, k_i)
            attn = attn.masked_fill((i_i > i_q) | (c >= s_i), float("-inf"))
            attn = torch.softmax(attn, dim=0)
            o[i, i_q] = torch.einsum("n h, n h v -> h v", attn, v_i)

    return o.to(dtype)

Tilelang优化实现

分块

with T.Kernel(seq_len, NV, batch * head_kv, threads=threads) as (bx, by, bz):。tilelang最重要的就是分块模式，这里可以看到，先对seqlen分块，再对V分块，分成NV块，NV的计算是NV = tilelang.cdiv(dim, block_T)，也就是隐藏层这个维度按块长block_T分块，最后再把batcj和head两个维度合并分块。

注意这里是对输出张量分块的，也就是每个线程块负责一个batch里，一个q token的一个head的注意力输出，其中长度为block_T的一段。这一个结果，是由多个q head和多个kv head计算出来的，但这个维度是要被压缩的，所以枚举然后累加，不分块。如果分块了，类似mla的split kv实现，还需要把不同块的split结果做合并，更麻烦。

申请内存

申请QKV输入，注意力输出O，以及计算过程中矩阵乘法累加和online softmax所需的块内存。重点是他们的shape

Q是[G, BK]，BK是对KV维度的统一分块长度，BK = BV = block_T，G是一个kv head要负责多少个q head

来看KV，[BS, BK]，BS就是block_size，每次取一个kv块，BK是在隐藏层上只取一段。

来看输出O，[G, BV]，每个q最终都得到一个和V 维度相同的输出

acc_s 是注意力得分的中间变量，[G, BS]的shape就是 $QK^T$ 产生的注意力得分，每个q head对每个kv head都有一个得分。

acc_o 是临时累加O输出的寄存器变量，shape和O_shared 完全相同

剩下的scores_max 这些用来做online softmax的，每行只需要一个，所以shape都是[G]

Q_shared = T.alloc_shared([G, BK], dtype)
K_shared = T.alloc_shared([BS, BK], dtype)
V_shared = T.alloc_shared([BS, BV], dtype)
O_shared = T.alloc_shared([G, BV], dtype)

acc_s = T.alloc_fragment([G, BS], accum_dtype)
acc_s_cast = T.alloc_fragment([G, BS], dtype)
acc_o = T.alloc_fragment([G, BV], accum_dtype)
scores_max = T.alloc_fragment([G], accum_dtype)
scores_max_prev = T.alloc_fragment([G], accum_dtype)
scores_scale = T.alloc_fragment([G], accum_dtype)
scores_sum = T.alloc_fragment([G], accum_dtype)
logsum = T.alloc_fragment([G], accum_dtype)

计算下标，拷贝Q，初始化累加数组

it是q token 编号，iv是负责的输出隐藏层的区间编号，ibh是kv的batch head展平的编号。

可以用ibh计算出ib,ih，表示当前是第几个batch，batch内第几个head

ib ih都确定了，需要的q输入也确定了，这里是一个q head和多个kv head计算注意力，所以确定一个q，后面变得是kv，q不变，可以在最开始就拷贝得到。

执行累加的acc_o,logsum置零，需要维护最大值的scores_max设为负无穷

i_t, i_v, i_bh = bx, by, bz
i_b, i_h = i_bh // head_kv, i_bh % head_kv

NS = S
T.copy(Q[i_b, i_t, i_h * G : (i_h + 1) * G, :], Q_shared)

# acc_o/logsum/scores_max together implement streaming softmax so
# blocks can be accumulated incrementally.
T.fill(acc_o, 0)
T.fill(logsum, 0)
T.fill(scores_max, -T.infinity(accum_dtype))

枚举kv块

NS是需要计算的kv块数，串行枚举，每一次都拷贝，计算这样的多阶段，可以用流水线优化T.Pipelined(NS, num_stages=num_stages)，这一行就能自动开启流水线。

is是当前块的起始地址，先判断是否合法，if i_s <= i_t and i_s >= 0:小于0说明是无效位置，大于询问q token的id(it)也是无效位置。这是为了实现自回归训练中只能看到当前token前面token的要求。

T.copy(K[i_b, i_s : i_s + BS, i_h, :], K_shared)将K块复制进来

if is_causal:如果开启因果掩码，q token后面的位置都设置为-inf，前面的置零，如果不开启的全部置零。设为-inf可以在softmax时映射为0，相当于不参与注意力计算。前面的if i_s <= i_t虽然也是为了做到不看后面的token，但是is只是当前块的起点，当前块范围实际是 $[i s, i s + B S)$ ，可能起点is小于it，但块中后半部分超过it了，所以块内超过it的部分，需要因果掩码来手动置零，不让他们参与注意力计算。

读取数据，掩码都弄好了，做 $QK^T$ ，T.gemm(Q_shared, K_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow)

for i in T.Pipelined(NS, num_stages=num_stages):
    i_s = BlockIndices[i_b, i_t, i_h, i] * BS
    if i_s <= i_t and i_s >= 0:
        # [BS, BK]
        T.copy(K[i_b, i_s : i_s + BS, i_h, :], K_shared)

        if is_causal:
            # Causal masking is written into acc_s before score GEMM.
            for i, j in T.Parallel(G, BS):
                acc_s[i, j] = T.if_then_else(i_t >= (i_s + j), 0, -T.infinity(acc_s.dtype))
        else:
            T.clear(acc_s)

        T.gemm(Q_shared, K_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow)

online softmax第一次循环+乘V

对 $QK^T$ 结果做一个online softmax，统计每行的max，动态更新logsum

最后读取V块，和softmax结果相乘，累加到注意力输出

# Softmax
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 i in T.Parallel(G):
    scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale)
for i, j in T.Parallel(G, BS):
    acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale)
T.reduce_sum(acc_s, scores_sum, dim=1)
for i in T.Parallel(G):
    logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i]
T.copy(acc_s, acc_s_cast)

# Rescale
for i, j in T.Parallel(G, BV):
    acc_o[i, j] *= scores_scale[i]

# V * softmax(Q * K)
T.copy(V[i_b, i_s : i_s + BS, i_h, i_v * BV : (i_v + 1) * BV], V_shared)
T.gemm(acc_s_cast, V_shared, acc_o, policy=T.GemmWarpPolicy.FullRow)

online softmax第二次循环+复制到输出张量

acc_o[i, j] /= logsum[i]用第一次循环的sum缩放每个位置，然后拷贝到输出张量Output

for i, j in T.Parallel(G, BV):
    # Finalize the streaming softmax normalization.
    acc_o[i, j] /= logsum[i]
T.copy(acc_o, O_shared)
T.copy(O_shared, Output[i_b, i_t, i_h * G : (i_h + 1) * G, i_v * BV : (i_v + 1) * BV])

完整代码

@tilelang.jit(
    out_idx=[-1],
    pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True},
)
def native_sparse_attention(batch, heads, seq_len, dim, is_causal, scale=None, block_size=64, groups=1, selected_blocks=16):
    # This kernel implements selected sparse attention forward only.
    # block_indices already decides which blocks each query may read from.
    if scale is None:
        scale = (1.0 / dim) ** 0.5 * 1.44269504  # log2(e)
    else:
        scale = scale * 1.44269504  # log2(e)

    # groups is the number of query heads sharing one KV head in GQA/MQA mode.
    head_kv = heads // groups
    q_shape = [batch, seq_len, heads, dim]
    kv_shape = [batch, seq_len, head_kv, dim]
    block_indices_shape = [batch, seq_len, head_kv, selected_blocks]
    block_indices_dtype = T.int32
    dtype = T.float16
    accum_dtype = T.float32
    block_S = block_size
    block_T = min(128, tilelang.math.next_power_of_2(dim))

    NK = tilelang.cdiv(dim, block_T)
    NV = tilelang.cdiv(dim, block_T)
    assert NK == 1, "The key dimension can not be larger than 256"

    S = selected_blocks
    G = groups
    BS = block_S
    BK = BV = block_T
    num_stages = 2
    threads = 64

    @T.prim_func
    def native_sparse_attention(
        Q: T.Tensor(q_shape, dtype),
        K: T.Tensor(kv_shape, dtype),
        V: T.Tensor(kv_shape, dtype),
        BlockIndices: T.Tensor(block_indices_shape, block_indices_dtype),
        Output: T.Tensor(q_shape, dtype),
    ):
        with T.Kernel(seq_len, NV, batch * head_kv, threads=threads) as (bx, by, bz):
            # One program instance handles one query position, one V tile, and
            # one batch/head_kv pair.
            Q_shared = T.alloc_shared([G, BK], dtype)
            K_shared = T.alloc_shared([BS, BK], dtype)
            V_shared = T.alloc_shared([BS, BV], dtype)
            O_shared = T.alloc_shared([G, BV], dtype)

            acc_s = T.alloc_fragment([G, BS], accum_dtype)
            acc_s_cast = T.alloc_fragment([G, BS], dtype)
            acc_o = T.alloc_fragment([G, BV], accum_dtype)
            scores_max = T.alloc_fragment([G], accum_dtype)
            scores_max_prev = T.alloc_fragment([G], accum_dtype)
            scores_scale = T.alloc_fragment([G], accum_dtype)
            scores_sum = T.alloc_fragment([G], accum_dtype)
            logsum = T.alloc_fragment([G], accum_dtype)

            i_t, i_v, i_bh = bx, by, bz
            i_b, i_h = i_bh // head_kv, i_bh % head_kv

            NS = S
            T.copy(Q[i_b, i_t, i_h * G : (i_h + 1) * G, :], Q_shared)

            # acc_o/logsum/scores_max together implement streaming softmax so
            # blocks can be accumulated incrementally.
            T.fill(acc_o, 0)
            T.fill(logsum, 0)
            T.fill(scores_max, -T.infinity(accum_dtype))

            for i in T.Pipelined(NS, num_stages=num_stages):
                i_s = BlockIndices[i_b, i_t, i_h, i] * BS
                if i_s <= i_t and i_s >= 0:
                    # [BS, BK]
                    T.copy(K[i_b, i_s : i_s + BS, i_h, :], K_shared)

                    if is_causal:
                        # Causal masking is written into acc_s before score GEMM.
                        for i, j in T.Parallel(G, BS):
                            acc_s[i, j] = T.if_then_else(i_t >= (i_s + j), 0, -T.infinity(acc_s.dtype))
                    else:
                        T.clear(acc_s)

                    T.gemm(Q_shared, K_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow)

                    # Softmax
                    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 i in T.Parallel(G):
                        scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale)
                    for i, j in T.Parallel(G, BS):
                        acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale)
                    T.reduce_sum(acc_s, scores_sum, dim=1)
                    for i in T.Parallel(G):
                        logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i]
                    T.copy(acc_s, acc_s_cast)

                    # Rescale
                    for i, j in T.Parallel(G, BV):
                        acc_o[i, j] *= scores_scale[i]

                    # V * softmax(Q * K)
                    T.copy(V[i_b, i_s : i_s + BS, i_h, i_v * BV : (i_v + 1) * BV], V_shared)
                    T.gemm(acc_s_cast, V_shared, acc_o, policy=T.GemmWarpPolicy.FullRow)

            for i, j in T.Parallel(G, BV):
                # Finalize the streaming softmax normalization.
                acc_o[i, j] /= logsum[i]
            T.copy(acc_o, O_shared)
            T.copy(O_shared, Output[i_b, i_t, i_h * G : (i_h + 1) * G, i_v * BV : (i_v + 1) * BV])

    return native_sparse_attention

亚马逊云科技技术品牌专区

更多推荐

VeOps CLI：你的火山引擎可观测排障助手

火山引擎 CUA 依托豆包大模型，云端沙箱模拟键鼠操作跨系统自动办公。云原生架构安全弹性，任务成功率超 85%，免脚本，对话式操作，提供 30 分钟免费体验。

亚马逊云科技技术品牌专区

AI文章阅读思考 - 人工智能辅助如何影响编程技能的形成

AI编码工具对开发者能力的负面影响研究引发热议。行业首个RCT实验证实，长期依赖AI辅助会导致独立排障能力显著下降。专家建议核心开发环节保持自主思考：系统设计、风险预判、性能优化等关键决策需人工完成，仅将AI作为辅助工具。尤其在支付、高并发等核心链路中，应坚持"先草稿-AI辅助-边界检查"的工作流程，避免能力退化。该研究为AI时代开发者能力培养敲响警钟。