Skip to content

refactor(dsv4): port #734 rope_cs collapse to HCA/SWA decode sparse kernels #763

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wangqin1723-max wants to merge 1 commit into
base: main
Choose a base branch
from
Open

refactor(dsv4): port #734 rope_cs collapse to HCA/SWA decode sparse kernels #763

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 7 additions & 12 deletions models/deepseek/v4/decode_sparse_attn_hca.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
# Copyright (c) PyPTO Contributors.
# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
# CANN Open Software License Agreement Version 2.0 (the "License").
Expand Down Expand Up @@ -67,8 +67,6 @@
# (its [64,128] softmax and co-resident QK+PV L0C accumulators overflow Vec/L0C).
QK_M_TILE = 32
ATTN_K_TILE = 128
ROPE_TILE = 16
ROPE_INTERLEAVE_TILE = 2 * ROPE_TILE
# proj_a cube K-frag. 256 (not 128) keeps the B-cache-line floor: B is K-contiguous
# under b_trans, so K*2B(bf16) = 512B == the a2a3 L2 line (K=128 was 256B, half a
# line -> wasted MTE2 DMA). At 256 the cube's L0A/L0B operand staging hits 100%
Expand Down Expand Up @@ -302,21 +300,18 @@
# so it overlaps it and is off merge_norm's critical path.
rope_cos_il = pl.create_tensor([T, ROPE_DIM], dtype=pl.FP32)
rope_sin_signed = pl.create_tensor([T, ROPE_DIM], dtype=pl.FP32)
for cp in pl.spmd(HALF_ROPE // ROPE_TILE, name_hint="rope_cs"):
cp_r0 = cp * ROPE_TILE
cp_c0 = 2 * cp_r0
with pl.at(level=pl.Level.CORE_GROUP, name_hint="rope_cs"):
cs_col = pl.col_expand_mul(
pl.full([T, ROPE_INTERLEAVE_TILE], dtype=pl.FP32, value=1.0),
pl.cast(pl.arange(0, [1, ROPE_INTERLEAVE_TILE], dtype=pl.INT32), target_type=pl.FP32))
pl.full([T, ROPE_DIM], dtype=pl.FP32, value=1.0),
pl.cast(pl.arange(0, [1, ROPE_DIM], dtype=pl.INT32), target_type=pl.FP32))
cs_dup_f = pl.cast(pl.cast(pl.mul(cs_col, 0.5), target_type=pl.INT32, mode="trunc"), target_type=pl.FP32)
cs_dup_idx = pl.cast(cs_dup_f, target_type=pl.INT32) # j>>1
cs_lane = pl.sub(cs_col, pl.mul(cs_dup_f, 2.0)) # j%2
cs_sign = pl.neg(pl.sub(pl.mul(cs_lane, 2.0), 1.0)) # [+1,-1,...] (conjugate)
cs_cos = pl.cast(freqs_cos[0:T, cp_r0 : cp_r0 + ROPE_TILE], target_type=pl.FP32)
cs_sin = pl.cast(freqs_sin[0:T, cp_r0 : cp_r0 + ROPE_TILE], target_type=pl.FP32)
rope_cos_il[0:T, cp_c0 : cp_c0 + ROPE_INTERLEAVE_TILE] = pl.gather(cs_cos, dim=-1, index=cs_dup_idx)
rope_sin_signed[0:T, cp_c0 : cp_c0 + ROPE_INTERLEAVE_TILE] = pl.mul(
pl.gather(cs_sin, dim=-1, index=cs_dup_idx), cs_sign)
cs_cos = pl.cast(freqs_cos[0:T, 0:HALF_ROPE], target_type=pl.FP32)
cs_sin = pl.cast(freqs_sin[0:T, 0:HALF_ROPE], target_type=pl.FP32)
rope_cos_il[0:T, 0:ROPE_DIM] = pl.gather(cs_cos, dim=-1, index=cs_dup_idx)
rope_sin_signed[0:T, 0:ROPE_DIM] = pl.mul(pl.gather(cs_sin, dim=-1, index=cs_dup_idx), cs_sign)

# Online-softmax merge across sparse-K tiles, sink-norm, then fused inverse RoPE.
# One spmd block per (token, head-tile) -- T*(H//H_TILE) blocks -- so the merge
Expand Down
21 changes: 8 additions & 13 deletions models/deepseek/v4/decode_sparse_attn_swa.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
# Copyright (c) PyPTO Contributors.
# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
# CANN Open Software License Agreement Version 2.0 (the "License").
Expand Down Expand Up @@ -59,8 +59,6 @@
# (its [64,128] softmax and co-resident QK+PV L0C accumulators overflow Vec/L0C).
QK_M_TILE = 32
ATTN_K_TILE = 128
ROPE_TILE = 16
ROPE_INTERLEAVE_TILE = 2 * ROPE_TILE
# proj_a cube K-frag. 256 (not 128) keeps the B-cache-line floor: B is K-contiguous
# under b_trans, so K*2B(bf16) = 512B == the a2a3 L2 line (K=128 was 256B, half a
# line -> wasted MTE2 DMA). At 256 the cube's L0A/L0B operand staging hits 100%
Expand Down Expand Up @@ -177,7 +175,7 @@
# slot. Current decode tokens must be inserted into ori_kv by the caller
# before this function runs; there is no MTP overlay path here.
sparse_bias = pl.create_tensor([T, PADDED_TOPK], dtype=pl.FP32)
for vb_blk in pl.spmd(T // VALID_TOKEN_TILE, name_hint="swa_valid_bias"):
for vb_blk in pl.spmd(T // VALID_TOKEN_TILE, name_hint="swa_valid_bias", allow_early_resolve=True):
vb = vb_blk * VALID_TOKEN_TILE
v_col = pl.cast(pl.arange(0, [1, ATTN_K_TILE], dtype=pl.INT32), target_type=pl.FP32)
v_col_m = pl.col_expand(pl.full([VALID_TOKEN_TILE, ATTN_K_TILE], dtype=pl.FP32, value=0.0), v_col)
Expand Down Expand Up @@ -260,21 +258,18 @@
# so it overlaps it and is off merge_norm's critical path.
rope_cos_il = pl.create_tensor([T, ROPE_DIM], dtype=pl.FP32)
rope_sin_signed = pl.create_tensor([T, ROPE_DIM], dtype=pl.FP32)
for cp in pl.spmd(HALF_ROPE // ROPE_TILE, name_hint="rope_cs"):
cp_r0 = cp * ROPE_TILE
cp_c0 = 2 * cp_r0
with pl.at(level=pl.Level.CORE_GROUP, name_hint="rope_cs"):
cs_col = pl.col_expand_mul(
pl.full([T, ROPE_INTERLEAVE_TILE], dtype=pl.FP32, value=1.0),
pl.cast(pl.arange(0, [1, ROPE_INTERLEAVE_TILE], dtype=pl.INT32), target_type=pl.FP32))
pl.full([T, ROPE_DIM], dtype=pl.FP32, value=1.0),
pl.cast(pl.arange(0, [1, ROPE_DIM], dtype=pl.INT32), target_type=pl.FP32))
cs_dup_f = pl.cast(pl.cast(pl.mul(cs_col, 0.5), target_type=pl.INT32, mode="trunc"), target_type=pl.FP32)
cs_dup_idx = pl.cast(cs_dup_f, target_type=pl.INT32) # j>>1
cs_lane = pl.sub(cs_col, pl.mul(cs_dup_f, 2.0)) # j%2
cs_sign = pl.neg(pl.sub(pl.mul(cs_lane, 2.0), 1.0)) # [+1,-1,...] (conjugate)
cs_cos = pl.cast(freqs_cos[0:T, cp_r0 : cp_r0 + ROPE_TILE], target_type=pl.FP32)
cs_sin = pl.cast(freqs_sin[0:T, cp_r0 : cp_r0 + ROPE_TILE], target_type=pl.FP32)
rope_cos_il[0:T, cp_c0 : cp_c0 + ROPE_INTERLEAVE_TILE] = pl.gather(cs_cos, dim=-1, index=cs_dup_idx)
rope_sin_signed[0:T, cp_c0 : cp_c0 + ROPE_INTERLEAVE_TILE] = pl.mul(
pl.gather(cs_sin, dim=-1, index=cs_dup_idx), cs_sign)
cs_cos = pl.cast(freqs_cos[0:T, 0:HALF_ROPE], target_type=pl.FP32)
cs_sin = pl.cast(freqs_sin[0:T, 0:HALF_ROPE], target_type=pl.FP32)
rope_cos_il[0:T, 0:ROPE_DIM] = pl.gather(cs_cos, dim=-1, index=cs_dup_idx)
rope_sin_signed[0:T, 0:ROPE_DIM] = pl.mul(pl.gather(cs_sin, dim=-1, index=cs_dup_idx), cs_sign)

# Online-softmax merge across sparse-K tiles, sink-norm, then fused inverse RoPE.
# One spmd block per (token, head-tile) -- T*(H//H_TILE) blocks -- so the merge
Expand Down
Loading