perf(dsv4): speed up decode indexer score kernel at 8k#695
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 Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThe scoring stage in the decode indexer is refactored into a three-phase pipeline: INT8 KV quantization with dequant scales, INT8 matmul producing INT32 accumulators, and a reduce/dequant stage writing score_flat. New tiling constants are added, and TopK selection is inlined into ChangesDecode indexer scoring refactor
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
Suggested labels: 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 |
|
Warning Gemini encountered an error creating the review. You can try again by commenting |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
models/deepseek/v4/decode_indexer.py (1)
229-341: 🚀 Performance & Scalability | 🔵 TrivialOperational: recheck vector-core contention end-to-end before merge.
The
QUANT_NSPLIT/REDUCE_NSPLITfanout raises effective parallelism toT * NSPLIT = 32vector lanes contending with other decode stages. As the PR notes,decode_fwdend-to-end perf/accuracy wasn't revalidated after this change. Recommend capturing decode-path metrics (vector-core occupancy, stage latency atstart_pos=8192) under the full pipeline, not just the standalone kernel, before landing.🤖 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/decode_indexer.py` around lines 229 - 341, The fanout change in the score path increases vector-lane contention, so revalidate the full decode pipeline before merging. Re-run end-to-end `decode_fwd`/decode-path measurements with this `QUANT_NSPLIT` and `REDUCE_NSPLIT` setup, focusing on vector-core occupancy and stage latency at `start_pos=8192`. Confirm the impact under the complete pipeline, not just the standalone kernel, and only keep the change if the metrics and accuracy remain acceptable.
🤖 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.
Nitpick comments:
In `@models/deepseek/v4/decode_indexer.py`:
- Around line 229-341: The fanout change in the score path increases vector-lane
contention, so revalidate the full decode pipeline before merging. Re-run
end-to-end `decode_fwd`/decode-path measurements with this `QUANT_NSPLIT` and
`REDUCE_NSPLIT` setup, focusing on vector-core occupancy and stage latency at
`start_pos=8192`. Confirm the impact under the complete pipeline, not just the
standalone kernel, and only keep the change if the metrics and accuracy remain
acceptable.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: a9e828d7-54e5-4dc3-9ba9-3e39aa7213f9
📒 Files selected for processing (1)
models/deepseek/v4/decode_indexer.py
With decode defaulting to start_pos=8192, the indexer score SPMD scales with compressed KV length and dominated the kernel (~288us). Restructure it into three fanned stages plus intra-lane cuts, to ~42us on a2a3. - CACHE_TILE 32 -> 64; INT8 amax via a single abs tile (drop the chunked re-cast loop). Keep the abs-based max(|x|): max(row_max, -row_min) is WRONG on real signed KV (decode_attention_csa x_out fails ~27%); the standalone indexer fixture is all-positive so it never caught it. - fold the per-position dequant scale past the ReLU and head-sum, and apply the same reorder in golden_indexer so the fp32 op order matches - inline the topk helper back into indexer - split the fused mat+vec score path into three GM-handoff stages so the vector-bound work is not pinned by the cube pairing: score_kv_quant (vec) -> score_mat (cube, MAT_TILE=512) -> score_reduce (vec, pl.pipeline(stage=2), REDUCE_TILE=128) - fan score_kv_quant and score_reduce across NSPLIT=4 extra lanes (T * NSPLIT = 32), striding the cache tiles per lane Validated: standalone decode_indexer.py and decode_attention_csa.py (x_out + kv_cache) both PASS on a2a3. The extra lanes contend for vector cores, so re-check the end-to-end win under decode_fwd.
zhangqi-chen
force-pushed
the
perf-indexer
branch
from
90313bd to
870830f
Compare
1 participant
Summary
start_pos=8192the DeepSeek-V4 decode indexerscoreSPMD scales with compressed KV length and dominates the kernel (~288us). Restructure it into three fanned stages plus intra-lane cuts, down to ~42us on a2a3.score_kv_quant(vec) →score_mat(cube, MAT_TILE=512) →score_reduce(vec,pl.pipeline(stage=2), REDUCE_TILE=128).score_kv_quantandscore_reduceacross NSPLIT=4 extra lanes (T * NSPLIT = 32), striding the cache tiles per lane.CACHE_TILE32 → 64; compute the INT8 amax from a single abs tile (drop the chunked re-cast loop). Keep the abs-basedmax(|x|):max(row_max, -row_min)is wrong on real signed KV (the standalone fixture is all-positive so it never caught it).indexer.decode_indexer.pyanddecode_attention_csa.py(x_out + kv_cache) both PASS. The extra lanes contend for vector cores, so the end-to-end win underdecode_fwdshould be re-checked.