Add: compose DeepSeek V4 CSA decode TP output - #940
Merged
zhangqi-chen merged 1 commit intoAug 18, 2026
Merged
zhangqi-chen merged 1 commit into
zhangqi-chen merged 1 commit into
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:
📝 WalkthroughWalkthroughAdds a complete DeepSeek-V4 CSA decode output path for CP. It includes JIT execution, distributed launching, deterministic tensor fixtures, a PyTorch golden implementation, sentinel-aware comparison, and CLI cases for full and subcapacity execution. ChangesCSA CP decode output
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant CLI
participant l3_decode_csa_output_cp
participant decode_csa_output_cp
participant DistributedBuffers
participant golden_decode_csa_output_cp
CLI->>l3_decode_csa_output_cp: run selected capacity case
l3_decode_csa_output_cp->>DistributedBuffers: allocate communication buffers
l3_decode_csa_output_cp->>decode_csa_output_cp: launch rank-local decode
decode_csa_output_cp->>DistributedBuffers: exchange attention tokens and heads
decode_csa_output_cp->>DistributedBuffers: reduce-scatter output
CLI->>golden_decode_csa_output_cp: compute reference output
CLI->>golden_decode_csa_output_cp: compare valid rows and sentinel tails
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 |
wangqin1723-max
force-pushed
the
feat/compose-deepseek-v4-csa-decode-tp4-output
branch
from
August 12, 2026 02:39
46986b6 to
d1b4d63
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
models/deepseek_v4_flash_dspark/decode_csa_output_cp.py (2)
440-440: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winName the shared key dim-0 constant.
scoresdepends on the literal0.25thatinit_ori_kv(Line 249) andinit_cmp_kv(Line 276) write into key dimension 0. The value appears in three places with no link between them. If a fixture edit changes one site, the golden diverges silently and the failure looks like a kernel bug.Define one module constant, for example
FIXTURE_KEY_DIM0 = 0.25, and use it at all three sites.🤖 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_dspark/decode_csa_output_cp.py` at line 440, Define a module-level constant for the fixture key dimension-0 value and replace the duplicated 0.25 literals in init_ori_kv, init_cmp_kv, and the scores calculation with that shared constant, preserving the existing computation.
321-337: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueFactor the shared RoPE phase computation into one helper.
init_freqs_cosandinit_freqs_sincompute the samephasetensor. Only the lookup table differs. A single helper keeps the two fixtures in sync if the phase rule changes.♻️ Proposed refactor
- def init_freqs_cos(): - rank = torch.arange(SP_SIZE, dtype=torch.int32).reshape(SP_SIZE, 1, 1) - token = torch.arange(local_t, dtype=torch.int32).reshape(1, local_t, 1) - column = torch.arange(ROPE_DIM, dtype=torch.int32).reshape(1, 1, ROPE_DIM) - phase = (rank + token + column).remainder(4) - phase[:, :, HALF_ROPE:] = (phase[:, :, HALF_ROPE:] + 1).remainder(4) - values = torch.tensor((1.0, 0.0, -1.0, 0.0), dtype=torch.bfloat16) - return values[phase] - - def init_freqs_sin(): - rank = torch.arange(SP_SIZE, dtype=torch.int32).reshape(SP_SIZE, 1, 1) - token = torch.arange(local_t, dtype=torch.int32).reshape(1, local_t, 1) - column = torch.arange(ROPE_DIM, dtype=torch.int32).reshape(1, 1, ROPE_DIM) - phase = (rank + token + column).remainder(4) - phase[:, :, HALF_ROPE:] = (phase[:, :, HALF_ROPE:] + 1).remainder(4) - values = torch.tensor((0.0, 1.0, 0.0, -1.0), dtype=torch.bfloat16) - return values[phase] + def rope_phase(): + rank = torch.arange(SP_SIZE, dtype=torch.int32).reshape(SP_SIZE, 1, 1) + token = torch.arange(local_t, dtype=torch.int32).reshape(1, local_t, 1) + column = torch.arange(ROPE_DIM, dtype=torch.int32).reshape(1, 1, ROPE_DIM) + phase = (rank + token + column).remainder(4) + phase[:, :, HALF_ROPE:] = (phase[:, :, HALF_ROPE:] + 1).remainder(4) + return phase + + def init_freqs_cos(): + values = torch.tensor((1.0, 0.0, -1.0, 0.0), dtype=torch.bfloat16) + return values[rope_phase()] + + def init_freqs_sin(): + values = torch.tensor((0.0, 1.0, 0.0, -1.0), dtype=torch.bfloat16) + return values[rope_phase()]🤖 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_dspark/decode_csa_output_cp.py` around lines 321 - 337, Factor the duplicated rank/token/column RoPE phase construction from init_freqs_cos and init_freqs_sin into one shared helper, then have both functions reuse that helper while retaining their distinct lookup tables and outputs.
🤖 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_dspark/decode_csa_output_cp.py`:
- Around line 84-85: Extend the existing import-time guard beside the
divisibility check to require LOCAL_T - ROPE_CS_T_TILE to be strictly positive,
ensuring the derived subcapacity remains valid for build_tensor_specs. Preserve
the current whole-decode-request divisibility validation and raise the fixture
guard error during import rather than deferring it to CLI execution.
---
Nitpick comments:
In `@models/deepseek_v4_flash_dspark/decode_csa_output_cp.py`:
- Line 440: Define a module-level constant for the fixture key dimension-0 value
and replace the duplicated 0.25 literals in init_ori_kv, init_cmp_kv, and the
scores calculation with that shared constant, preserving the existing
computation.
- Around line 321-337: Factor the duplicated rank/token/column RoPE phase
construction from init_freqs_cos and init_freqs_sin into one shared helper, then
have both functions reuse that helper while retaining their distinct lookup
tables and outputs.
🪄 Autofix
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: c2a7a424-1c24-4a3a-9e0c-63daa94933f4
📒 Files selected for processing (1)
models/deepseek_v4_flash_dspark/decode_csa_output_cp.py
wangqin1723-max
force-pushed
the
feat/compose-deepseek-v4-csa-decode-tp4-output
branch
3 times, most recently
from
August 14, 2026 07:30
c603e96 to
ab8a116
Compare
wangqin1723-max
changed the title
- Chain CSA heads through configurable TP token/head exchange, sharded output projection, and dependency-ordered FP32 reduce-scatter. - Consolidate TP1/2/4 max and subcapacity fixtures into the CSA entrypoint. - Reuse the shared TP1 output projection and bound sparse planning by the runtime token count.
wangqin1723-max
force-pushed
the
feat/compose-deepseek-v4-csa-decode-tp4-output
branch
from
August 17, 2026 03:00
ab8a116 to
a563884
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
output projection, and dependency-ordered FP32 reduce-scatter.
entrypoint.
runtime token count.