40 layers › 2 groups of 6 layers › 5 reusing layers › Compressed sparse attention › Sparse attention
Weighted row sum
×no parameterskernel.py 380
For each head, add the gathered rows by their weights.
- in
- [B batch, 64 heads, T positions, n listed rows] [B batch, T positions, n listed rows, 512 d_h]
- out
- [B batch, 64 heads, T positions, 512 d_h]
Explanation
The rows are the same for every head: the heads differ only in their queries.
Codekernel.py 310–403
@tilelang.jit(pass_configs=pass_configs)
def sparse_attn_kernel(h: int, d: int, scale=None):
"""Sparse multi-head attention via index gathering + online softmax (FlashAttention-style).
For each (batch, seq_pos), gathers top-k KV positions by index, computes attention
with numerically stable running max/sum, and includes a learnable attn_sink bias."""
b = T.symbolic("b")
m = T.symbolic("m")
n = T.symbolic("n")
topk = T.symbolic("topk")
if scale is None:
scale = (1.0 / d) ** 0.5
num_stages = 2
threads = 256
block = 64
num_blocks = tilelang.cdiv(topk, block)
@T.prim_func
def sparse_attn_kernel_(
q: T.Tensor[(b, m, h, d), BF16],
kv: T.Tensor[(b, n, d), BF16],
o: T.Tensor[(b, m, h, d), BF16],
attn_sink: T.Tensor[(h,), FP32],
topk_idxs: T.Tensor[(b, m, topk), INT32],
):
with T.Kernel(m, b, threads=threads) as (bx, by):
q_shared = T.alloc_shared((h, d), BF16)
kv_shared = T.alloc_shared((block, d), BF16)
o_shared = T.alloc_shared((h, d), BF16)
acc_s_cast = T.alloc_shared((h, block), BF16)
idxs = T.alloc_fragment(block, INT32)
acc_s = T.alloc_fragment((h, block), FP32)
acc_o = T.alloc_fragment((h, d), FP32)
scores_max = T.alloc_fragment(h, FP32)
scores_max_prev = T.alloc_fragment(h, FP32)
scores_scale = T.alloc_fragment(h, FP32)
scores_sum = T.alloc_fragment(h, FP32)
sum_exp = T.alloc_fragment(h, FP32)
T.clear(acc_o)
T.clear(sum_exp)
# Finite lower bound instead of -inf: a row with no valid index (all -1) would otherwise
# produce exp(-inf - (-inf)) = NaN. With a finite bound such rows yield an all-zero output,
# matching the training kernel's convention.
T.fill(scores_max, -1e30)
T.copy(q[by, bx, :, :], q_shared)
for t in T.Pipelined(num_blocks, num_stages=num_stages):
for i in T.Parallel(block):
idxs[i] = T.if_then_else(t * block + i < topk, topk_idxs[by, bx, t * block + i], -1)
for i, j in T.Parallel(block, d):
kv_shared[i, j] = T.if_then_else(idxs[i] != -1, kv[by, idxs[i], j], 0)
for i, j in T.Parallel(h, block):
acc_s[i, j] = T.if_then_else(idxs[j] != -1, 0, -T.infinity(FP32))
T.gemm(q_shared, kv_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow)
for i, j in T.Parallel(h, block):
acc_s[i, j] *= scale
T.copy(scores_max, scores_max_prev)
T.reduce_max(acc_s, scores_max, dim=1, clear=False)
for i in T.Parallel(h):
scores_scale[i] = T.exp(scores_max_prev[i] - scores_max[i])
for i, j in T.Parallel(h, block):
acc_s[i, j] = T.exp(acc_s[i, j] - scores_max[i])
T.reduce_sum(acc_s, scores_sum, dim=1)
for i in T.Parallel(h):
sum_exp[i] = sum_exp[i] * scores_scale[i] + scores_sum[i]
T.copy(acc_s, acc_s_cast)
for i, j in T.Parallel(h, d):
acc_o[i, j] *= scores_scale[i]
T.gemm(acc_s_cast, kv_shared, acc_o, policy=T.GemmWarpPolicy.FullRow)
for i in T.Parallel(h):
sum_exp[i] += T.exp(attn_sink[i] - scores_max[i])
for i, j in T.Parallel(h, d):
acc_o[i, j] /= sum_exp[i]
T.copy(acc_o, o_shared)
T.copy(o_shared, o[by, bx, :, :])
return sparse_attn_kernel_
def sparse_attn(q: torch.Tensor, kv: torch.Tensor, attn_sink: torch.Tensor, topk_idxs: torch.Tensor, softmax_scale: float) -> torch.Tensor:
b, s, h, d = q.size()
# Pad heads to 16 for kernel efficiency (stripped after)
if h < 16:
q = torch.cat([q, q.new_zeros(b, s, 16 - h, d)], dim=2)
attn_sink = torch.cat([attn_sink, attn_sink.new_zeros(16 - h)])
o = torch.empty_like(q)
kernel = sparse_attn_kernel(q.size(2), d, softmax_scale)
kernel(q, kv, o, attn_sink, topk_idxs)
if h < 16:
o = o.narrow(2, 0, h).contiguous()
return oSources1
Direct evidence
Sparse attention: each query head scores the gathered key–value rows, the softmax has one more term in its denominator from the head’s learned sink, and the output is the weighted sum of the same rows.
sparse-attn · deepseek-ai/DeepSeek-V4.1-Flash / inference/kernel.py
Pinned revision 2cba9e42aa02 ↗ · lines 310–403
Read code excerpt
@tilelang.jit(pass_configs=pass_configs)
def sparse_attn_kernel(h: int, d: int, scale=None):
"""Sparse multi-head attention via index gathering + online softmax (FlashAttention-style).
For each (batch, seq_pos), gathers top-k KV positions by index, computes attention
with numerically stable running max/sum, and includes a learnable attn_sink bias."""
b = T.symbolic("b")
m = T.symbolic("m")
n = T.symbolic("n")
topk = T.symbolic("topk")
if scale is None:
scale = (1.0 / d) ** 0.5
num_stages = 2
threads = 256
block = 64
num_blocks = tilelang.cdiv(topk, block)
@T.prim_func
def sparse_attn_kernel_(
q: T.Tensor[(b, m, h, d), BF16],
kv: T.Tensor[(b, n, d), BF16],
o: T.Tensor[(b, m, h, d), BF16],
attn_sink: T.Tensor[(h,), FP32],
topk_idxs: T.Tensor[(b, m, topk), INT32],
):
with T.Kernel(m, b, threads=threads) as (bx, by):
q_shared = T.alloc_shared((h, d), BF16)
kv_shared = T.alloc_shared((block, d), BF16)
o_shared = T.alloc_shared((h, d), BF16)
acc_s_cast = T.alloc_shared((h, block), BF16)
idxs = T.alloc_fragment(block, INT32)
acc_s = T.alloc_fragment((h, block), FP32)
acc_o = T.alloc_fragment((h, d), FP32)
scores_max = T.alloc_fragment(h, FP32)
scores_max_prev = T.alloc_fragment(h, FP32)
scores_scale = T.alloc_fragment(h, FP32)
scores_sum = T.alloc_fragment(h, FP32)
sum_exp = T.alloc_fragment(h, FP32)
T.clear(acc_o)
T.clear(sum_exp)
# Finite lower bound instead of -inf: a row with no valid index (all -1) would otherwise
# produce exp(-inf - (-inf)) = NaN. With a finite bound such rows yield an all-zero output,
# matching the training kernel's convention.
T.fill(scores_max, -1e30)
T.copy(q[by, bx, :, :], q_shared)
for t in T.Pipelined(num_blocks, num_stages=num_stages):
for i in T.Parallel(block):
idxs[i] = T.if_then_else(t * block + i < topk, topk_idxs[by, bx, t * block + i], -1)
for i, j in T.Parallel(block, d):
kv_shared[i, j] = T.if_then_else(idxs[i] != -1, kv[by, idxs[i], j], 0)
for i, j in T.Parallel(h, block):
acc_s[i, j] = T.if_then_else(idxs[j] != -1, 0, -T.infinity(FP32))
T.gemm(q_shared, kv_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow)
for i, j in T.Parallel(h, block):
acc_s[i, j] *= scale
T.copy(scores_max, scores_max_prev)
T.reduce_max(acc_s, scores_max, dim=1, clear=False)
for i in T.Parallel(h):
scores_scale[i] = T.exp(scores_max_prev[i] - scores_max[i])
for i, j in T.Parallel(h, block):
acc_s[i, j] = T.exp(acc_s[i, j] - scores_max[i])
T.reduce_sum(acc_s, scores_sum, dim=1)
for i in T.Parallel(h):
sum_exp[i] = sum_exp[i] * scores_scale[i] + scores_sum[i]
T.copy(acc_s, acc_s_cast)
for i, j in T.Parallel(h, d):
acc_o[i, j] *= scores_scale[i]
T.gemm(acc_s_cast, kv_shared, acc_o, policy=T.GemmWarpPolicy.FullRow)
for i in T.Parallel(h):
sum_exp[i] += T.exp(attn_sink[i] - scores_max[i])
for i, j in T.Parallel(h, d):
acc_o[i, j] /= sum_exp[i]
T.copy(acc_o, o_shared)
T.copy(o_shared, o[by, bx, :, :])
return sparse_attn_kernel_
def sparse_attn(q: torch.Tensor, kv: torch.Tensor, attn_sink: torch.Tensor, topk_idxs: torch.Tensor, softmax_scale: float) -> torch.Tensor:
b, s, h, d = q.size()
# Pad heads to 16 for kernel efficiency (stripped after)
if h < 16:
q = torch.cat([q, q.new_zeros(b, s, 16 - h, d)], dim=2)
attn_sink = torch.cat([attn_sink, attn_sink.new_zeros(16 - h)])
o = torch.empty_like(q)
kernel = sparse_attn_kernel(q.size(2), d, softmax_scale)
kernel(q, kv, o, attn_sink, topk_idxs)
if h < 16:
o = o.narrow(2, 0, h).contiguous()
return o