Skip to content

Perf: fuse DSv4-flash indexer score scope, pack SWA output store #831

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

Merged
zhangqi-chen merged 1 commit into from
Jul 25, 2026
+22 −40
Merged

Perf: fuse DSv4-flash indexer score scope, pack SWA output store #831

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
56 changes: 18 additions & 38 deletions models/deepseek/v4-flash/decode_indexer.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 @@ -61,12 +61,11 @@
# tiling
CACHE_TILE = 64
assert BLOCK_SIZE % CACHE_TILE == 0, "CACHE_TILE must not cross a paged idx_kv_cache block"
# matmul/reduce tile over contiguous GM scratch, not the paged KV cache
MAT_TILE = 512
# REDUCE_TILE tiles the paged C8 cache one 128-row page per fused matmul+reduce step.
REDUCE_TILE = 128
# score_kv_quant / score_reduce fan the cache-tile loop across NSPLIT extra lanes: T * NSPLIT.
QUANT_NSPLIT = 4
REDUCE_NSPLIT = 4
# the fused score scope fans the cache-page loop across REDUCE_NSPLIT extra lanes: T * NSPLIT.
# T*NSPLIT=16 mixed blocks map to 16 AIC + 32 AIV (1:2), one clean wave on the 24+48 chip.
REDUCE_NSPLIT = 2
Q_TILE = 256
# Q_OUT_TILE is the per-task N granularity (sets idx_qr_proj task count); MM_N_TILE
# is the Mat-safe cube N-tile. Q_OUT_TILE fans Q_OUT_TILE // MM_N_TILE cube ops per
Expand Down Expand Up @@ -258,31 +257,14 @@
idx_block_table_flat = pl.reshape(idx_block_table, [B * IDX_CACHE_MAX_BLOCKS])
score_flat = pl.reshape(score, [T, SCORE_LEN])

# No score_init: reduce writes the valid region; the tail is never read (topk re-masks).
# Two GM-handoff stages: matmul (cube, reads paged C8 directly) -> reduce (vec).
score_acc_gm = pl.create_tensor([T * IDX_KV_LEN, IDX_N_HEADS], dtype=pl.INT32)

# read paged C8 KV one page per tile, matmul with the per-step-quantized query
for tg in pl.spmd(T, name_hint="score_mat", allow_early_resolve=True):
b = tg // S
s = tg - b * S
clen_b = pl.read(kv_seq_lens, [b]) // COMPRESS_RATIO
cblk_b = (clen_b + BLOCK_SIZE - 1) // BLOCK_SIZE
qb = b * S * IDX_N_HEADS
qr_full = qr_hadamard_i8[qb + s * IDX_N_HEADS : qb + (s + 1) * IDX_N_HEADS, 0 : IDX_HEAD_DIM]
for cb in pl.pipeline(0, cblk_b, stage=2):
cache0 = cb * BLOCK_SIZE
idx_blk_id = pl.cast(
pl.read(idx_block_table_flat, [b * IDX_CACHE_MAX_BLOCKS + cb]),
pl.INDEX,
)
kv0 = idx_blk_id * BLOCK_SIZE
base = tg * IDX_KV_LEN + cache0
kv_i8_mat = kv_cache_i8_flat[kv0 : kv0 + BLOCK_SIZE, :]
score_acc_mat = pl.matmul(kv_i8_mat, qr_full, out_dtype=pl.INT32, b_trans=True)
score_acc_gm[base : base + BLOCK_SIZE, :] = score_acc_mat

for unit in pl.spmd(T * REDUCE_NSPLIT, name_hint="score_reduce", allow_early_resolve=True):
# No score_init: the loop writes the valid region; the tail is never read (topk re-masks).

@coderabbitai coderabbitai Bot Jul 24, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win

Preserve invalid-score initialization.

Pages after visible_len_t are never written, so score now contains uninitialized values there. golden_indexer explicitly returns FP32_NEG_INF for every invalid position (lines 513-546), and indexer_test exposes this tensor. Top-k re-masking protects the current consumer but does not preserve the score output contract. Initialize score to FP32_NEG_INF, or formally remove/relax this output contract and update all callers and reference validation.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@models/deepseek/v4-flash/decode_indexer.py` at line 260, Initialize the score
output to FP32_NEG_INF before the valid-region loop so all positions after
visible_len_t retain the required invalid sentinel, matching golden_indexer and
indexer_test expectations; keep the existing valid-score writes unchanged.

# Fused mixed cube+vector: the per-page C8 matmul feeds the relu/weight/head-sum reduce
# on chip, so cube page i+1 pipelines against vector page i (no score_acc_gm handoff).
# slot_num=2: the cube->vector pipe defaults to an 8-slot ring (8 * 32KB = 256KB > UB);
# a depth-2 ring (64KB) fits alongside the FP32 epilogue. Deeper rings (>=3) would cut the
# handoff stalls but overflow Vec.
for unit in pl.spmd(T * REDUCE_NSPLIT, name_hint="score", allow_early_resolve=True,
optimizations=[pl.split(pl.SplitMode.NONE, slot_num=2)]):
tg = unit // REDUCE_NSPLIT
split = unit - tg * REDUCE_NSPLIT
b = tg // S
Expand All @@ -293,33 +275,31 @@
cblk_t = (visible_len_t + REDUCE_TILE - 1) // REDUCE_TILE
tb = b * S
qb = b * S * IDX_N_HEADS
qr_full = qr_hadamard_i8[qb + s * IDX_N_HEADS : qb + (s + 1) * IDX_N_HEADS, 0 : IDX_HEAD_DIM]
qh_scale_s = pl.reshape(qr_hadamard_scale_dq[qb + s * IDX_N_HEADS : qb + (s + 1) * IDX_N_HEADS, :], [1, IDX_N_HEADS])
weights_row_s = pl.reshape(weights[tb + s : tb + s + 1, :], [1, IDX_N_HEADS])
lane_iters = (cblk_t - split + REDUCE_NSPLIT - 1) // REDUCE_NSPLIT
for cb_local in pl.pipeline(0, lane_iters, stage=2):
cb = split + cb_local * REDUCE_NSPLIT
cache0 = cb * REDUCE_TILE
valid_len = pl.min(REDUCE_TILE, visible_len_t - cache0)
base = tg * IDX_KV_LEN + cache0
idx_blk_id = pl.cast(
pl.read(idx_block_table_flat, [b * IDX_CACHE_MAX_BLOCKS + cb]),
pl.INDEX,
)
kv0 = idx_blk_id * BLOCK_SIZE
score_acc_red = score_acc_gm[base : base + REDUCE_TILE, :]
kv_i8_mat = kv_cache_i8_flat[kv0 : kv0 + BLOCK_SIZE, :]
score_acc_red = pl.matmul(kv_i8_mat, qr_full, out_dtype=pl.INT32, b_trans=True)
kv_dq_red = kv_scale_flat[kv0 : kv0 + REDUCE_TILE, :] # paged per-position dequant scale
score_tile_red = pl.cast(score_acc_red, target_type=pl.FP32, mode="none")
# per-position dequant kv_dq_red applied after the head-sum
score_tile_red = pl.col_expand_mul(score_tile_red, qh_scale_s)
relu_score_red = pl.maximum(score_tile_red, pl.full([REDUCE_TILE, IDX_N_HEADS], dtype=pl.FP32, value=0.0))
relu_score_red = pl.maximum(score_tile_red, 0.0)
weighted_score_red = pl.col_expand_mul(relu_score_red, weights_row_s)
# per-position dequant kv_dq_red applied after the head-sum
weighted_score_row = pl.mul(pl.row_sum(weighted_score_red), kv_dq_red)
weighted_score_s = pl.reshape(weighted_score_row, [1, REDUCE_TILE])
weighted_score_valid_s = pl.fillpad(pl.set_validshape(weighted_score_s, 1, valid_len), pad_value=pl.PadValue.min)
weighted_score_valid_s = pl.maximum(
weighted_score_valid_s,
pl.full([1, REDUCE_TILE], dtype=pl.FP32, value=FP32_NEG_INF),
)
weighted_score_valid_s = pl.maximum(weighted_score_valid_s, FP32_NEG_INF)
score_flat[tb + s : tb + s + 1, cache0 : cache0 + REDUCE_TILE] = weighted_score_valid_s

topk_idxs_flat = pl.reshape(topk_idxs, [T, SCORE_LEN])
Expand Down
6 changes: 4 additions & 2 deletions models/deepseek/v4-flash/decode_sparse_attn.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 @@ -457,8 +457,10 @@
n_hh = n_gh - n_g * HEADS_PER_GROUP
n_pack_row = n_g * T + m_t
n_col = n_hh * HEAD_DIM
o_packed[n_pack_row : n_pack_row + 1, n_col : n_col + NOPE_DIM] = n_bf16[n_hi : n_hi + 1, 0 : NOPE_DIM]
o_packed[n_pack_row : n_pack_row + 1, n_col + NOPE_DIM : n_col + HEAD_DIM] = n_rope_bf16[n_hi : n_hi + 1, 0 : ROPE_DIM]
# one HEAD_DIM-wide store per head row instead of two: concat the nope and
# inverse-RoPE halves on chip so o_packed takes a single contiguous write.
n_full_bf16 = pl.concat(n_bf16[n_hi : n_hi + 1, 0 : NOPE_DIM], n_rope_bf16[n_hi : n_hi + 1, 0 : ROPE_DIM])
o_packed[n_pack_row : n_pack_row + 1, n_col : n_col + HEAD_DIM] = n_full_bf16
Comment on lines +460 to +463

@coderabbitai coderabbitai Bot Jul 24, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/usr/bin/env bash
set -euo pipefail

rg -n -C 3 '\bpl\.concat\s*\(' .
rg -n -C 3 '(^|[[:space:]])(def|class).*concat|concat[[:space:]]*=' .

Repository: hw-native-sys/pypto-lib

Length of output: 3837

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "== files =="
git ls-files | sed -n '1,120p'

echo
echo "== imports/defs around relevant files =="
for f in models/deepseek/v4-flash/decode_sparse_attn.py models/qwen3/14b/decode_layer_a8w8.py models/qwen3/14b/prefill_fwd.py; do
  echo "-- $f --"
  wc -l "$f"
  sed -n '1,80p' "$f"
done

echo
echo "== concat occurrences with dimensions around them =="
python3 - <<'PY'
import pathlib, re
for p in pathlib.Path('.').rglob('*.py'):
    if '.git' in p.parts: continue
    text=p.read_text(errors='ignore')
    for m in re.finditer(r'pl\.concat\s*\(', text):
        line=text[:m.start()].count('\n')+1
        print(f'{p}:{line}: {text[m.start():m.start()+160].split(chr(10))[0].strip()}')
PY

Repository: hw-native-sys/pypto-lib

Length of output: 20032

🌐 Web query:

PyPTO language pl.concat documented API tensors list dim default documentation

💡 Result:

In the PyPTO programming framework, the pypto.concat API is used to concatenate multiple tensors along a specified dimension [1]. The function signature and parameter details are as follows: API Signature: pypto.concat(tensors: List[Tensor], dim: int = 0) -> Tensor [1] Parameters: 1. tensors: A list of input tensors to be concatenated. The framework requires the list to contain at least 2 tensors (and no more than 128) [1]. All input tensors must have the same data type and the same number of dimensions, and must match in all dimensions except for the dimension specified by dim [1]. 2. dim: An integer specifying the dimension along which to concatenate the tensors [1]. The default value is 0 [1]. Important Constraints: - The dim parameter must satisfy the constraint -input.dim <= dim < input.dim [1]. - The API does not automatically deduce the output tensor's shape; the user must ensure the output Tensor's validShape is correct before calling concat [1]. - Input tensors cannot be empty, and the Shape Size must not exceed 2147483647 (INT32_MAX) [1]. For context, PyPTO is a high-performance programming framework for AI accelerators that utilizes a Tile-based programming model [2][3]. It should not be confused with Polars (pl.concat), which is a library for data manipulation and concatenation of DataFrames and Series [4][5]. [1][2][4]

Citations:

Pass pl.concat a tensor list and specify dim=1.

pypto.concat requires tensors: List[Tensor] and defaults to dim=0; this call passes the slices positionally and would not concatenate them into the required [1, HEAD_DIM] row tensor.

Proposed fix 
-            n_full_bf16 = pl.concat(n_bf16[n_hi : n_hi + 1, 0 : NOPE_DIM], n_rope_bf16[n_hi : n_hi + 1, 0 : ROPE_DIM])
+            n_full_bf16 = pl.concat(
+                [
+                    n_bf16[n_hi : n_hi + 1, 0 : NOPE_DIM],
+                    n_rope_bf16[n_hi : n_hi + 1, 0 : ROPE_DIM],
+                ],
+                dim=1,
+            )
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
# one HEAD_DIM-wide store per head row instead of two: concat the nope and
# inverse-RoPE halves on chip so o_packed takes a single contiguous write.
n_full_bf16 = pl.concat(n_bf16[n_hi : n_hi + 1, 0 : NOPE_DIM], n_rope_bf16[n_hi : n_hi + 1, 0 : ROPE_DIM])
o_packed[n_pack_row : n_pack_row + 1, n_col : n_col + HEAD_DIM] = n_full_bf16
# one HEAD_DIM-wide store per head row instead of two: concat the nope and
# inverse-RoPE halves on chip so o_packed takes a single contiguous write.
n_full_bf16 = pl.concat(
[
n_bf16[n_hi : n_hi + 1, 0 : NOPE_DIM],
n_rope_bf16[n_hi : n_hi + 1, 0 : ROPE_DIM],
],
dim=1,
)
o_packed[n_pack_row : n_pack_row + 1, n_col : n_col + HEAD_DIM] = n_full_bf16
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@models/deepseek/v4-flash/decode_sparse_attn.py` around lines 460 - 463,
Update the n_full_bf16 construction in the sparse-attention packing logic to
pass the nope and inverse-RoPE slices as a tensor list to pl.concat and
explicitly concatenate along dim=1, preserving the required [1, HEAD_DIM] shape
before writing to o_packed.

Source: MCP tools


# ========================================================================
# Back-to-back grouped output projection (manual scope, PER-GROUP INT8 quant).
Expand Down
Loading