Perf: fuse DSv4-flash indexer score scope, pack SWA output store#831
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThe PR reworks indexer score computation around a fused reduction loop and changes sparse attention output packing to combine no-RoPE and inverse-RoPE portions into one contiguous write. ChangesIndexer score fusion
Sparse attention output packing
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant Score as score SPMD
participant KV as kv_cache_i8_flat
participant Query as qr_full
participant Output as score_flat
Score->>KV: Read paged INT8 KV
Score->>Query: Read query tiles
Score->>Score: Compute and dequantize scores
Score->>Output: Write masked FP32 scores
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with 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.
Inline comments:
In `@models/deepseek/v4-flash/decode_indexer.py`:
- 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.
In `@models/deepseek/v4-flash/decode_sparse_attn.py`:
- Around line 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.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 0e71136f-ecb2-46bb-90e2-bd57eb20d36e
📒 Files selected for processing (2)
models/deepseek/v4-flash/decode_indexer.pymodels/deepseek/v4-flash/decode_sparse_attn.py
| 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). |
There was a problem hiding this comment.
🗄️ 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.
| # 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 |
There was a problem hiding this comment.
🎯 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()}')
PYRepository: 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:
- 1: https://pypto.gitcode.com/api/operation/pypto-concat.html
- 2: https://github.com/hw-native-sys/pypto
- 3: https://deepwiki.com/chenqi123/cann-recipes-infer/7.1-pypto-programming-model
- 4: https://docs.pola.rs/api/python/dev/reference/api/polars.concat.html
- 5: https://github.com/pola-rs/polars/blob/py-0.20.31/py-polars/polars/functions/eager.py
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.
| # 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
Two independent DeepSeek V4-flash decode optimizations. decode_indexer -- fuse the score matmul and reduce: - Merge the `score_mat` (cube) and `score_reduce` (vector) spmd scopes into one mixed `score` scope, so each page's C8 matmul feeds the relu / weight / head-sum epilogue on chip - Drop the `score_acc_gm` [T * IDX_KV_LEN, IDX_N_HEADS] INT32 scratch and its GM round-trip; cube page i+1 pipelines against vector page i - `slot_num=2` sizes the cube->vector ring at 64KB so it fits Vec alongside the FP32 epilogue; the 8-slot default would need 256KB - REDUCE_NSPLIT 4 -> 2 so T*NSPLIT = 16 mixed blocks map to 16 AIC + 32 AIV; drop the now-unused MAT_TILE and QUANT_NSPLIT - Take the relu and -inf clamp bounds as scalar `pl.maximum` operands instead of materializing two `pl.full` tiles decode_sparse_attn -- pack the merge_norm output store: - Concat the nope and inverse-RoPE halves on chip so each head row takes one HEAD_DIM-wide `o_packed` write instead of two Isolated latency (a2a3, decode B=4 S=2, mean of 100 rounds, 3 runs): indexer 93.1 -> 87.3 us, sparse_attn 189.3 -> 186.9 us.
868891b to
3ad252f
Compare
2 participants
Summary
Two independent DeepSeek V4-flash decode optimizations.
decode_indexer— fuse the score matmul and reduce:score_mat(cube) andscore_reduce(vector) spmd scopes into one mixedscorescope, so each page's C8 matmul feeds the relu / weight / head-sum epilogue on chipscore_acc_gm[T * IDX_KV_LEN, IDX_N_HEADS]INT32 scratch and its GM round-trip; cube page i+1 pipelines against vector page islot_num=2sizes the cube→vector ring at 64KB so it fits Vec alongside the FP32 epilogue; the 8-slot default would need 256KBREDUCE_NSPLIT4 → 2 soT*NSPLIT= 16 mixed blocks map to 16 AIC + 32 AIV; drop the now-unusedMAT_TILEandQUANT_NSPLITpl.maximumoperands instead of materializing twopl.fulltilesdecode_sparse_attn— pack themerge_normoutput store:HEAD_DIM-wideo_packedwrite instead of twoIsolated latency (a2a3, decode B=4 S=2, mean of 100 rounds, 3 runs): indexer 93.1 → 87.3 us, sparse_attn 189.3 → 186.9 us.