Perf: early-resolve the DeepSeek-V4 Flash CSA decode producers - #915
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 adjusts ChangesDeepSeek early-resolution scope tuning
Estimated code review effort: 1 (Trivial) | ~5 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e5f11c1e6e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
## Summary - Mark six CSA decode producers `allow_early_resolve=True` so the scheduler may stage them speculatively: `rmsnorm_rope_cache_write` (`decode_compressor_ratio4`), `weights_proj_reduce` (`decode_indexer`), `kv_and_cache_write` (`decode_indexer_compressor`), and `kv_touch`, `qk_pv`, and `rope_cs` (`decode_sparse_attn_csa`). - Apply the same set to both the `deepseek_v4_flash_dspark` and `deepseek_v4_flash_mtp` trees. Scheduling hints only; no kernel math, tiling, or dependency edge changes. Every touched scope is reached only through the CSA decode path. The shared `qkv_proj_rope` inline kernel is deliberately left alone: `decode_swa`, `decode_hca`, `prefill_csa`, `prefill_swa`, `prefill_hca`, and the `prefill_cp_*` flows all call it, and none of them receive these CSA hints, so changing its early-resolve state would alter unbenchmarked callers. ## Benchmark setup Measured with `PYPTO_BENCH=1` (the repository default 100 measured rounds after 5 warmup launches) on one a2a3 card, physical device 0, via `python models/deepseek_v4_flash_<tree>/decode_csa.py -p a2a3 -d 0`. The two trees run different token counts at their defaults (dspark T=32, mtp T=8), so the trees are not comparable to one another -- only each tree against itself. All twelve runs below passed golden validation for both `kv_cache` and `x_out`. Three independent before/after launch pairs were run per tree, alternating the patch in and out, because a single pair cannot separate the shift from the within-run spread (roughly 45us peak-to-peak on dspark). ## Results: median effective_us deepseek_v4_flash_dspark (T=32): | Pair | Before | After | Delta | | ---: | ---: | ---: | ---: | | 1 | 654.6 | 646.9 | -7.7 | | 2 | 652.2 | 637.1 | -15.1 | | 3 | 646.6 | 636.8 | -9.8 | | mean | 651.1 | 640.3 | -10.9 (-1.7%) | deepseek_v4_flash_mtp (T=8): | Pair | Before | After | Delta | | ---: | ---: | ---: | ---: | | 1 | 365.3 | 358.5 | -6.8 | | 2 | 369.6 | 355.6 | -14.0 | | 3 | 365.9 | 356.0 | -9.9 | | mean | 366.9 | 356.7 | -10.2 (-2.8%) | All six paired deltas are negative, and the per-run minima separate cleanly in both trees (dspark after 606.7-628.3 against before 630.3-634.9; mtp after 346.0-349.7 against before 355.0-358.3). The dspark medians do overlap in one pair (after 646.9 against before 646.6), so the dspark result rests on the paired comparison rather than on disjoint median ranges.
e5f11c1 to
3b3fe13
Compare
2 participants
Summary
allow_early_resolve=Trueso the schedulermay stage them speculatively:
rmsnorm_rope_cache_write(
decode_compressor_ratio4),weights_proj_reduce(decode_indexer),kv_and_cache_write(decode_indexer_compressor), andkv_touch,qk_pv, andrope_cs(decode_sparse_attn_csa).deepseek_v4_flash_dsparkanddeepseek_v4_flash_mtptrees. Scheduling hints only; no kernel math,tiling, or dependency edge changes.
Every touched scope is reached only through the CSA decode path. The shared
qkv_proj_ropeinline kernel is deliberately left alone:decode_swa,decode_hca,prefill_csa,prefill_swa,prefill_hca, and theprefill_cp_*flows all call it, and none of them receive these CSA hints,so changing its early-resolve state would alter unbenchmarked callers.
Benchmark setup
Measured with
PYPTO_BENCH=1(the repository default 100 measured roundsafter 5 warmup launches) on one a2a3 card, physical device 0, via
python models/deepseek_v4_flash_<tree>/decode_csa.py -p a2a3 -d 0.The two trees run different token counts at their defaults (dspark T=32,
mtp T=8), so the trees are not comparable to one another -- only each tree
against itself. All twelve runs below passed golden validation for both
kv_cacheandx_out.Three independent before/after launch pairs were run per tree, alternating
the patch in and out, because a single pair cannot separate the shift from
the within-run spread (roughly 45us peak-to-peak on dspark).
Results: median effective_us
deepseek_v4_flash_dspark (T=32):
deepseek_v4_flash_mtp (T=8):
All six paired deltas are negative, and the per-run minima separate
cleanly in both trees (dspark after 606.7-628.3 against before
630.3-634.9; mtp after 346.0-349.7 against before 355.0-358.3). The
dspark medians do overlap in one pair (after 646.9 against before 646.6),
so the dspark result rests on the paired comparison rather than on
disjoint median ranges.