Perf: score DSpark CSA indexer candidates on Cube - #1379
Little-oil wants to merge 1 commit into
Conversation
|
Navigate logical layers of code changes, visualize relationships, and explore their blast radius. No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Advanced Run ID: 📒 Files selected for processing (1)
Included review availability: Your plan provides up to 1 included review per hour; 0 remain after this review. 📝 WalkthroughWalkthroughThe indexer score path now precomputes query coefficients and moves score activation, dequantization, and head reduction to Cube operations. The scoring stage uses deeper pipelining and updated tile gather and store operations. ChangesIndexer score reduction
Priority: ➖ Normal Estimated code review effort: 3 (Moderate) | ~25 minutes Change: Refactor Merge Risk: ⚪ Minimal · up to This change moves the DSpark CSA indexer's head reduction from the Vector units to the Cube, for a reported 17.9% decode speedup at long context. The scoring pipeline still waits for the quantized query to be fully written before reading it, and the earlier numeric-range and layout concerns were ruled out. No outstanding correctness or stability risk was found, so the change appears ready to merge. 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. A rabbit hops beside the score, Comment |
4526932 to
2cc3178
Compare
优化方案逐条说明（对照代码改动）打分公式是 做法参考 AscendC 几何参数： 1. 用 FIXPIPE 的 pre_relu + pre_quant 替代 Vector 的 cast + maximum改动在 score_i32 = pl.tile.matmul(
query_left,
pl.tile.move(pl.tile.transpose_view(kv_i8), target_memory=pl.Mem.Right),
)
score_relu = pl.tile.create(
[SCORE_QUERY_TILE * IDX_N_HEADS, SCORE_TILE], pl.FP16, target_memory=pl.Mem.Mat
)
score_relu = pl.tile.assemble(
score_relu, score_i32, [0, 0],
pre_quant=SCORE_FIXPIPE_SCALE,
pre_relu=True,
)
被替代的是 Vector 上的 数据位置：
关键差别是输出落在 L1 而不是 UB：中间结果不再跨 Cube/Vector 边界，也不占 UB。 2. 用 Cube 的 matmul 替代 Vector 的 row_expand_mul + col_sum紧接上一段： score_acc = pl.tile.matmul(
coefficient_left,
pl.tile.move(score_relu, target_memory=pl.Mem.Right),
)
数据位置：
第 1 点产出的 FP16 结果留在 L1，这里直接作为右操作数进 L0B，不经过 UB 也不经过 GM。 Vector 每个 tile 剩下的工作是： score_shard = pl.aiv_shard(score_acc)
for pair_token in pl.unroll(SCORE_QUERY_TILE):
score_row = pl.tile.mul(
pl.tile.slice(score_shard, [1, SCORE_LANE_ROWS], [pair_token, 0]),
kv_scale,
)取本 token 那一行，乘 key scale。原来是四条指令各处理 64×384 个元素。 3. 两个 query token 共用一块 Key tile一个请求的相邻两个 token，它们的 64 个 head 在 query_vector = pl.tile.load(
qr_hadamard_i8, [query_head_begin, 0],
[SCORE_QUERY_TILE * IDX_N_HEADS, IDX_HEAD_DIM],
target_memory=pl.Mem.Mat,
)
query_left = pl.tile.move(query_vector, target_memory=pl.Mem.Left)leaf 循环相应改成按 pair 步进： for item in pl.range(worker, query_count // SCORE_QUERY_TILE * max_leaves, TOPK_SCORE_WORKERS):
query = (item // max_leaves) * SCORE_QUERY_TILE一次 数据位置：
按「每个候选、每个 token」折算，进 L0B 的 Key 字节数从 128 降到 64。
4. head 权重排成块对角矩阵，一次 matmul 归约两个 token配对之后左操作数是 代码在 for w_pair in pl.unroll(MM_ROW_TILE // SCORE_QUERY_TILE):
w_pair_row = (w_r0 // SCORE_QUERY_TILE + w_pair) * MM_ROW_TILE
pl.tile.store(
pl.tile.full([MM_ROW_TILE, IDX_N_HEADS * SCORE_QUERY_TILE], dtype=pl.FP16, value=0.0),
[w_pair_row, 0], score_coefficient,
)
for w_token in pl.unroll(SCORE_QUERY_TILE):
...
pl.tile.store(coefficient_row, [w_pair_row + w_token, w_token * IDX_N_HEADS], score_coefficient)结果 这一点是配对能成立的前提。只做第 3 点、两个 token 各做一次 matmul 是不行的： 每个 token 的尾部单独 mask：pair 里靠后的 token 看到的候选更多，所以共用 tile 按它定大小， 5. 系数在 weights_proj_reduce 里一并算出，不新增任务改动在 score_coefficient = pl.create_tensor(
[T_PAD // SCORE_QUERY_TILE * MM_ROW_TILE, IDX_N_HEADS * SCORE_QUERY_TILE], dtype=pl.FP16
)
with pl.spmd(row_blocks, name_hint="weights_proj_reduce", deps=[qh_quant_tid], ...):
...
w_scaled = pl.tile.muls(w_sum, WEIGHTS_SCALE)
pl.tile.store(w_scaled, [w_r0, 0], weights)
...
coefficient_row = pl.tile.cast(
pl.tile.muls(
pl.tile.mul(coefficient_scale,
pl.tile.slice(w_scaled, [1, IDX_N_HEADS], [w_member, 0])),
SCORE_COEF_SCALE,
),
target_type=pl.FP16, mode="rint",
)系数是 代价是增加一条 数据位置：
系数在 leaf 里的取用路径是 GM → L1 → L0A，每个 leaf 一次。 开发过程中先写的是另起一个 数值说明
前后对照之前：一块 384 候选 tile、一个 query token。之后：一块 256 候选 tile、两个 query token。
关于 #1375 的 buffered 路径上游 #1375 新增的 它的写回是： pl.store(buf_scores, [buf_transfer_row, 0], buf_score_transfer,
pre_quant=1.0 / BUFFERED_SCORE_SCALE, pre_relu=True)
它这样做是为了把 AIC 和 AIV 完全分开：整个任务拆成一个纯 AIC 的循环和一个 把第 2 点也加进 buffered 路径是可行的后续工作：在它的 AIC 段把 但它对本 PR 的目标没有意义。 buffered 路径的启用条件是 另外该工作本身还有两个待解决项：768 候选的 Key tile 是 96 KiB，超过 L0A 的 64 KiB， |
打分路径的前后流程图范围是 之前（384 候选，1 个 token）flowchart TD
G1["idx_kv_cache（GM）"] -->|"12 次 gather_row<br/>INT8 48 KiB"| L1K["kv tile（L1）"]
G2["qr_hadamard_i8（GM）"] -->|"load<br/>INT8 8 KiB"| L1Q["query（L1）"]
L1Q -->|"move<br/>8 KiB 进 L0A"| MM["matmul<br/>q @ kv^T（Cube）"]
L1K -->|"b_trans + move<br/>48 KiB 进 L0B"| MM
MM -->|"INT32 [64,384]<br/>96 KiB 在 L0C"| CROSS["aiv_shard<br/>跨 Cube/Vector 边界"]
CROSS -->|"INT32 96 KiB<br/>L0C 到 UB"| CAST["cast 转 FP32（Vector）"]
CAST -->|"FP32 [64,384]<br/>96 KiB 在 UB"| RELU["maximum(x,0)（Vector）"]
G3["qr_hadamard_scale_dq（GM）"] -->|"load"| COEF["mul 求系数（Vector）"]
G4["weights（GM）"] -->|"load"| COEF
COEF -->|"FP32 [64,1]<br/>在 UB"| RMUL
RELU -->|"FP32 96 KiB<br/>在 UB"| RMUL["row_expand_mul（Vector）"]
RMUL -->|"FP32 96 KiB<br/>在 UB"| CSUM["col_sum（Vector）"]
CSUM -->|"FP32 [1,384]<br/>在 UB"| SMUL["mul kv_scale（Vector）"]
G5["idx_kv_scale（GM）"] -->|"6 次 gather_row<br/>FP32"| SMUL
SMUL -->|"FP32 [1,384]"| ST["store（Vector）"]
ST -->|"FP32"| G6["score_arena（GM）"]
之后（256 候选，2 个 token）flowchart TD
G1["idx_kv_cache（GM）"] -->|"8 次 gather_row<br/>INT8 32 KiB"| L1K["kv tile（L1）"]
G2["qr_hadamard_i8（GM）"] -->|"load 2 个 token<br/>INT8 16 KiB"| L1Q["query pair（L1）"]
L1Q -->|"move<br/>16 KiB 进 L0A<br/>整个 leaf 只做一次"| MM1["matmul<br/>q @ kv^T（Cube）"]
L1K -->|"transpose_view + move<br/>32 KiB 进 L0B"| MM1
MM1 -->|"INT32 [128,256]<br/>128 KiB 在 L0C"| ASM["assemble<br/>pre_relu + pre_quant（FIXPIPE）"]
ASM -->|"FP16 [128,256]<br/>64 KiB 写进 L1"| L1S["score_relu（L1）"]
G7["score_coefficient（GM）<br/>块对角"] -->|"load<br/>FP16 4 KiB"| L1C["系数（L1）"]
L1C -->|"move<br/>4 KiB 进 L0A<br/>整个 leaf 只做一次"| MM2["matmul<br/>coef @ score（Cube）"]
L1S -->|"move<br/>64 KiB 进 L0B"| MM2
MM2 -->|"FP32 [16,256]<br/>16 KiB 在 L0C"| CROSS["aiv_shard<br/>跨 Cube/Vector 边界"]
CROSS -->|"FP32 16 KiB<br/>L0C 到 UB"| SLICE["slice 取第 0、1 行（Vector）"]
SLICE -->|"每 token FP32 [1,256]<br/>在 UB"| SMUL["mul kv_scale（Vector）"]
G5["idx_kv_scale（GM）"] -->|"4 次 gather_row<br/>FP32"| SMUL
SMUL -->|"FP32 [1,256] × 2"| ST["store（Vector）"]
ST -->|"FP32"| G6["score_arena（GM）"]
两张图的差别Vector 节点从五个（ 新增的 64 head 的中间结果在改动前要出 Cube 侧进 UB，改动后停在 L1，再回到 L0B。 两个左操作数（query pair 和系数）在整个 leaf 内不变，只进一次 L0A，不随候选 tile 重复。 按「每个候选、每个 token」折算：
这个 tile 循环在整个 CSA 里的位置flowchart LR
A["qr 投影 + RoPE + Hadamard"] --> B["weights_proj"]
B --> C["weights_proj_reduce<br/>（本 PR 在这里加算块对角系数）"]
C --> D["indexer_score_topk_leaf<br/>（本 PR 改的就是它的内循环）"]
E["kv 压缩 + 写 cache"] --> D
D --> F["indexer_topk_query_merge"]
F --> G["sparse_attn_csa / qk_pv"]
G --> H["o_proj"]
128k 下 |
588775b to
3f447ad
Compare
indexer_score_topk_leaf dominates decode_csa at 128k. It kept a 64-head INT32 score on Vector and reduced it there with a cast, a ReLU, a broadcast multiply and a column sum. Move that reduction onto Cube, following the AscendC quant_lightning_indexer_v2 arch22 kernel. - FIXPIPE drains the INT32 QK accumulator from L0C straight into an FP16 Mat tile, applying ReLU on the accumulator and a fixed 1/1024 dequantization in the same writeback. This replaces Vector's cast and maximum, and the 64-head intermediate never reaches UB or GM. - A second Cube matmul contracts the heads against an FP16 coefficient, replacing Vector's row_expand_mul and col_sum. The dequantized score moves from Mat to L0B and the coefficient from Mat to L0A, so the reduction happens where the score already is. Vector keeps only the Key-scale multiply and Top-K, on a tile that is 64x narrower. - Two query tokens share one Key tile. Their coefficients sit block diagonally in one FP16 matrix, so a batch of GEMVs becomes a single GEMM and each Key tile is loaded into L0B once for two tokens instead of twice. This geometry follows hw-native-sys#1380. - The coefficient is built by a loop inside weights_proj_reduce alongside the weight that task already computes, so its preparation costs no task of its own. It carries a 1024x factor that cancels the writeback scale exactly and keeps small coefficients clear of FP16 subnormals. A2/A3, TP1, B16, S6, swimlane 0, 100 rounds after 5 warmups, A-B-B-A against 07b5d2f, one session on one device against one frozen fixture. At start_pos = 16 x 131072 the median goes 1499.6 -> 1241.3 us, -258.3 us (-17.2%). At 16 x 8192 it goes 726.5 -> 707.2 us, -19.3 us (-2.7%). On both workloads each arm's median is below both baseline arms'. The level-4 trace at 128k puts the score task at 833.6 -> 547.5 us and the whole operation at 1623.6 -> 1358.7 us over an unchanged 983 blocks. topk_idxs keeps main's comparator, topk_pair_compare, which admits an index swap only between tied scores and rejects any strictly worse pick. Contracting on Cube rounds both operands to FP16, which puts about a quarter of one percent of the diagnostic score values past the existing relative bound, so topk_scores' outlier allowance goes from 0.1% to 0.5%. That is the only gate this change touches; the decode_csa x_out comparator and every cache gate are untouched and pass as they are. a2a3sim segfaults on this kernel: the CPU model of pto.tinsert sizes the scalar pre-quant vector by rows while indexing it by columns, so a tail tile whose valid rows fall below its 64 columns reads out of bounds. Hardware never runs that code. hw-native-sys/pto-isa#334 carries a repro and a one-line fix. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This branch has not been deployed
1 participant
What this changes
indexer_score_topk_leafis half ofdecode_csaat 128k. Its head reductionmoves from Vector to Cube, and two query tokens are scored against one Key tile
instead of one, following the AscendC
quant_lightning_indexer_v2arch22kernel.
A2/A3, TP1, B16, S6, swimlane 0, 100 rounds after 5 warmups, A-B-B-A against
07b5d2f, same device and same frozen fixture.--start-pos16 x 131072--start-pos16 x 8192Per-arm medians, in run order: 128k
1501.2 / 1238.8 / 1243.8 / 1497.9, 8k720.0 / 713.2 / 701.2 / 732.9. On both workloads each PR arm's median is belowboth main arms', so the gap is larger than the drift between the repeated arms.
Main's 128k samples are bimodal and only 59 and 57 of 100 land in the fast mode;
this PR's land 100 and 99 of 100, so the median stops being a coin flip. The
means move the same way, 1769.4 / 1776.7 against 1366.5 / 1396.5.
The level-4 trace, 128k:
Block count is unchanged: the coefficient is built inside
weights_proj_reduce,so the Cube path adds no task of its own.
The buffered scoring path added by #1375 is not on this route. It needs
b_dim >= 64and a history of at least 32768, and TP1 with B16 meets neither,so 128k still runs
indexer_score_topk_leaf. Carrying the same Cube reductioninto that path is separate work.
Relation to the open PRs
This supersedes #1354 and #1300, which attack the same reduction.
#1300 kept the result bit-exact with compensated FP16 limbs and six GEMVs and
regressed. #1354 dispatches per query between the original Vector path and a
Cube path.
#1380 by @zhangqi-chen was closed in favour of this one. Its two ideas,
pairing two query tokens per Key tile and making the head weights block
diagonal, are what the third and fourth points below describe, and they are
worth 57.7 us at 128k over the version of this PR that had only the Cube
reduction. The credit for them is his.
The idea
The indexer score is
sum_h w[t,h] * relu(q[t,h] . k[c]) * kv_scale[c]. Cubealready produced the 64-head
q . k; everything after it ran on Vector as acast, a ReLU, a broadcast multiply and a column sum over a 64 x 384 FP32 tile.
pre_reluandpre_quantreplace Vector'scastandmaximum. The ReLU is applied to the INT32 accumulator and a fixed 1/1024dequantization on the way out, while L0C drains into an FP16 Mat tile. The
64-head intermediate never reaches UB or GM.
matmulreplaces Vector'srow_expand_mulandcol_sum. Thedequantized score is already on the Cube side, so the head contraction runs
there against an FP16 coefficient. Vector keeps only the Key-scale multiply.
of
qr_hadamard_i8, so one QK matmul produces both tokens' scores side byside and each Key tile is gathered and pushed through L0B once for two
tokens instead of twice.
head columns 0..63 and zeros elsewhere, row 1 carries token 1's in columns
64..127. One GEMM therefore reduces both tokens, which is what keeps the
Query and the coefficient loop-invariant across the leaf's candidate tiles.
What the arithmetic looks like before and after
Per candidate tile. Before: 384 candidates, one query token. After: 256
candidates, two query tokens.
matmul(q[64,128], kv^T[128,384])-> L0C INT32[64,384]matmul(q[128,128], kv^T[128,256])-> L0C INT32[128,256]castthenmaximum(x, 0)over 24 576 elementsrow_expand_mulover 24 576 elementscol_sumover 64 rowsmatmul(coef[16,128], R[128,256])-> L0C FP32[16,256], rows 0 and 1mulover 384 elementsmulover 2 x 256 elementsPer candidate per token, Vector goes from four passes over a 64-row column to
one multiply on one element.
Where the data sits
weights_proj_reduce, FP16 in GM, L1 -> L0A once per leaf[16,256], 16 KiB per pair across the boundaryNormalized per candidate per token, the Cube-to-Vector transfer drops from
256 bytes to 32, and the Key bytes pushed into L0B halve.
The FP16 score reaches the second matmul through
pl.tile.transpose_view, azero-copy fractal reinterpretation of the same bytes, so the reduction still
writes candidate columns and no Vector transpose is needed.
SCORE_TILEdrops from 384 to 256 because the pair doubles every Cube-sidetile: the FP16 score is
[128,256], which is all 64 KiB of L0B, and the INT32accumulator is
[128,256], which is all 128 KiB of L0C. That is also why theKey stays in L0B rather than taking L0A's faster port: with the pair there is no
L0B left to hold the Query.
The numerical contract
qk_i32 * q_scaleand only then multipliesby the head weight.
q_scaleis positive, sorelu(x*s) == s*relu(x)andapplying the ReLU to the raw accumulator is equivalent.
R = relu(qk) / 1024in FP16. The widest INT8 x 128 dot product is128 * 127^2 = 2 064 512, which scales to 2016, inside FP16's 65 504. Thesmallest nonzero product scales to 9.8e-4, far above the subnormal floor.
fp16(q_scale * weight * 1024). The 1024 cancels thewriteback scale exactly, so there is no post-scale, and it lifts small
coefficients clear of FP16 subnormals.
candidates, so the shared tile is sized from it, and every store narrows to
the storing token's own visible count.
topk_idxskeeps main's comparator on both the standalone indexer test anddecode_csa. That comparator istopk_pair_compare: it admits an index swaponly between tied scores and rejects any strictly worse pick, so a selection
that dropped a better candidate would fail it. Contracting on Cube rounds both
operands to FP16, which puts about a quarter of one percent of the diagnostic
score values past the existing relative bound, so
topk_scores' outlierallowance goes from 0.1% to 0.5%. That is the only gate this PR touches;
decode_csa'sx_outcomparator and every cache gate are untouched and passas they are.
The simulator
a2a3simsegfaults here and this PR leaves it failing rather than masking it.The CPU model of
pto.tinsertsizes the scalar pre-quant vector byGetValidRow()for a column-major source but indexes it by column, so a tailtile whose valid rows fall below its 64 columns reads past the vector and the
simulator segfaults. Hardware issues the instruction and never runs that code,
which is why every device path here passes. A one-line repro and a fix are in
hw-native-sys/pto-isa#334, which this PR waits on.
Token generation
Operator comparators only bound the score tensor; what matters is whether the
model still says the right thing. An 8192-token prompt carrying the code
BLUE-ORCHID-731asks for that code back followed by the numbers 1 to 10,greedy decode, 96 new tokens, DSpark target K0 on TP4/DP4/EP16 over 16 cards,
both arms from the same worktree, venv and devices and differing only in
decode_indexer.py.Five runs, main three times and this PR twice, across two device allocations.
Every one of them emits
BLUE-ORCHID-731then 1 through 10, consumes 8192prompt tokens and produces 96.
The free-form text that follows the instructed part differs between runs, and
that is not caused by this change: two runs of identical main code, back to
back in one allocation, diverge from each other after 66 characters. The engine
is not byte-reproducible on this prompt at 96 greedy tokens, so an arm-to-arm
text difference carries no information. What the check establishes is that the
instructed output is correct on every run of both arms.
The mechanism is visible in the comparator above.
topk_pair_compareadmits aswap between candidates whose scores are tied, FP16 rounding decides which of a
tied pair the kernel picks, and the two carry different KV. Past the instructed
part the continuation is open-ended, so a small difference in attention output
changes which token wins the argmax.
The code, hunk by hunk
All of it is in
models/deepseek_v4_flash_dspark/decode_indexer.py.SCORE_QUERY_TILEis the pair width; it falls back to 1 whenSis odd, which degenerates the block-diagonal weight to a single row andrestores per-token scoring. That fallback is exercised: the standalone
indexer test passes on device at S7,
T = 112, withtopk_idxsstill exact.SCORE_TILEgoes 384 -> 256 for the L0B and L0Cbudgets above.
SCORE_FIXPIPE_SCALEandSCORE_COEF_SCALEare the matched1/1024 and 1024 pair.
SCORE_ARENA_ROWSgrows to give each worker lane a rowper token of the pair.
weights_proj_reducealso builds the coefficient. It gainsdeps=[qh_quant_tid]for the query scale, its body moves from tensor slicingto
pl.tile.*so the coefficient can be formed from the same in-registerw_scaled, and apl.unrollwrites oneMM_ROW_TILEfractal per query pair:zeros, then token
t'sfp16(q_scale * weight * 1024)into rowt, headcolumns
t*64 .. t*64+63.indexer_score_topk_foresttakesscore_coefficientand the leaf loopstrides by
SCORE_QUERY_TILE, so each item is a (query pair, leaf) instead ofa (query, leaf).
L1 and moved to L0A once, outside the candidate-tile pipeline, because the
block-diagonal form keeps them invariant across it.
pl.create_l1becomespl.tile.create(..., target_memory=pl.Mem.Mat)so it can be moved to L0B as atransposed view.
pl.matmul:tile.matmul(query -> L0A, transpose_view(kv) -> L0B), thentile.assemble(..., pre_quant=SCORE_FIXPIPE_SCALE, pre_relu=True)into anFP16 Mat tile, then
tile.matmul(coefficient -> L0A, score -> L0B).cast,maximum,row_expand_mulandcol_sumare deleted. What remains ispl.aiv_shard, a per-token row sliceand the Key-scale multiply, wrapped in a
pl.unroll(SCORE_QUERY_TILE)thatgives each token its own tail mask and arena row. The Key-scale buffer, its
paged gather, the valid-shape narrowing and the arena store move from the
tensor-level ops to the matching
pl.tile.*ops, because the score nowarrives as a Tile and mixing the two levels is rejected.
the pair publishes its own two half-leaf rows.
Measured and rejected
This was my first attempt at pairing and it regressed: 1885.9 us at a group of
two and 1866.9 at four, against 1709.0 for one, at S8 on an earlier base.
Keeping a matmul per query
makes the Query and the coefficient vary inside the candidate-tile loop, so
they have to be re-staged per tile and that costs more than the saved Key
traffic. I concluded from this that Key traffic was not the binding
constraint. That conclusion was wrong, and Perf: accelerate DSpark TP1 long-context CSA #1380 is why: with the weights
block diagonal the group is a single GEMM, the operands stay invariant, and
the same halved Key traffic is worth 57.7 us.
at 8k, where the score task is not the critical path and the cost is therefore
pure overhead.
weights_proj_reducealready computes the weight, so thepreparation moved there and the task is gone.
16x16, so a one-row subview has no address, and
valid_shape=1on the leftoperand compiles but does not shrink the ring slot, which is sized from the
physical shape. The pair now uses two of the sixteen rows rather than one.
unpaired version of this PR, three L1 tile stages with four cross-core ring
slots measured 1542.4 us at 128k against 1706.1 for two slots, again at S8 on
an earlier base, and the Key
took L0A because it is the larger operand and L0A's transfer rate is twice
L0B's. Pairing leaves room for neither: the paired FP16 score fills all of
L0B and the paired accumulator fills all of L0C. The checked-in depths are
main's,
stage=2and one ring slot, and the Query is back in L0A.🤖 Generated with Claude Code