Add: define decode DSA-CP communication and output seams - #931
Merged
zhangqi-chen merged 7 commits intoAug 11, 2026
Merged
zhangqi-chen merged 7 commits into
zhangqi-chen merged 7 commits into
Conversation
📝 WalkthroughWalkthroughAdded sequence-parallel decode communication, a receive-side grouped output projection, and staged CSA, HCA, and SWA attention paths. Existing top-level attention wrappers remain available. ChangesDeepSeek-V4 decode pipeline
Estimated code review effort: 5 (Critical) | ~90+ minutes Sequence Diagram(s)sequenceDiagram
participant Rank as SP rank
participant Decode as decode_attention_cp
participant Window as Distributed windows
participant Projection as decode_o_projection_cp
Rank->>Decode: Launch decode layout
Decode->>Window: Gather KV tokens
Decode->>Window: Exchange attention groups
Window-->>Projection: Provide grouped attention rows
Projection->>Rank: Write output partials
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 |
wangqin1723-max
changed the title
wangqin1723-max
changed the title
wangqin1723-max
marked this pull request as ready for review
August 10, 2026 10:32
![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 (3)
models/deepseek_v4_flash_dspark/decode_attention_cp.py (2)
276-282: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReturn the poisoned tensor explicitly.
init_attention_groupedpoisonsgrouped, then returnsvalues. The poisoning is visible only becausepl.reshape-freetorch.reshapereturns a view for this contiguous tensor. Any later change to the initialization chain that breaks contiguity silently removes the poisoned tail, and the fixture stops testing the capacity rows. Return the poisoned tensor in the required shape.♻️ Proposed change
def init_attention_grouped(): shape = (SP_SIZE, O_GROUPS * LOCAL_T_PAD, O_GROUP_IN) values = torch.arange(SP_SIZE * O_GROUPS * LOCAL_T_PAD * O_GROUP_IN, dtype=torch.int32) values = values.remainder(127).reshape(shape).to(torch.bfloat16) grouped = values.reshape(SP_SIZE, O_GROUPS, LOCAL_T_PAD, O_GROUP_IN) grouped[:, :, local_t:] = -2000.0 - return values + return grouped.reshape(shape)🤖 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_attention_cp.py` around lines 276 - 282, Update init_attention_grouped to return the poisoned grouped tensor rather than values, preserving the shape (SP_SIZE, O_GROUPS, LOCAL_T_PAD, O_GROUP_IN) so the local_t tail remains filled with -2000.0.
116-117: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueConsider tiled copies instead of single-row loops.
The three seams move the window into the output one row at a time (
group_out[group_row : group_row + 1, ...], thecopy_rowloop, and thepl.load(..., [1, D])reduction). Each iteration is a separate GM transfer. A row-tile ofCOMM_ROW_TILErows, with a clipped tail, would cut the transfer count by the tile factor. This is a fixture today, but the same seams are the template for the decode path.Also applies to: 149-154, 183-190
🤖 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_attention_cp.py` around lines 116 - 117, Replace the single-row GM copy loops in the group_out/gather_window seam, the copy_row seam, and the pl.load reduction with COMM_ROW_TILE-row transfers, clipping the tile size for the final partial block. Preserve the existing row ordering, column range, and reduction behavior while reducing the number of GM transfers.models/deepseek_v4_flash_dspark/decode_sparse_attn_csa.py (1)
124-130: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUntyped packed-head buffers weaken the new head-to-projection seam. The split introduces
o_packed_headsas a barepl.Tensorin all three head stages and in the matching local projection stages. The layout is fixed by compile-time constants at every call site, but the missing annotation removes the shape and dtype check exactly where the two new stages meet. The SWA stage uses a head-major layout while CSA and HCA use a group-major layout, so a swapped buffer would compile.
models/deepseek_v4_flash_dspark/decode_sparse_attn_csa.py#L124-L130: annotateo_packed_headsaspl.Tensor[[O_GROUPS * T_PAD, O_GROUP_IN], pl.BF16], and apply the same annotation too_packedinsparse_attn_csa_local_o_projat Line 427.models/deepseek_v4_flash_dspark/decode_sparse_attn_hca.py#L132-L138: annotateo_packed_headsaspl.Tensor[[O_GROUPS * T_PAD, O_GROUP_IN], pl.BF16]here and at Line 392.models/deepseek_v4_flash_dspark/decode_sparse_attn_swa.py#L103-L109: annotateo_packed_headsaspl.Tensor[[O_GROUPS * T_PAD * HEADS_PER_GROUP, HEAD_DIM], pl.BF16]here and at Line 294, which keeps the head-major layout explicit at the boundary and documents why Line 322 reshapes.🤖 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_sparse_attn_csa.py` around lines 124 - 130, Annotate the packed-head parameters with their fixed layout and BF16 dtype: in models/deepseek_v4_flash_dspark/decode_sparse_attn_csa.py lines 124-130 and sparse_attn_csa_local_o_proj at line 427, use the group-major shape; in models/deepseek_v4_flash_dspark/decode_sparse_attn_hca.py lines 132-138 and line 392, use the same group-major shape; and in models/deepseek_v4_flash_dspark/decode_sparse_attn_swa.py lines 103-109 and line 294, use the head-major shape. Preserve these annotations at each head-to-projection boundary.
🤖 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_o_projection_cp.py`:
- Around line 180-186: Update build_tensor_specs to validate local_t before
calculating group_t, rejecting non-positive values such as the subcapacity case
when local_t is zero. Match the range-validation behavior used by
decode_attention_cp.build_tensor_specs, while preserving valid capacity
handling.
---
Nitpick comments:
In `@models/deepseek_v4_flash_dspark/decode_attention_cp.py`:
- Around line 276-282: Update init_attention_grouped to return the poisoned
grouped tensor rather than values, preserving the shape (SP_SIZE, O_GROUPS,
LOCAL_T_PAD, O_GROUP_IN) so the local_t tail remains filled with -2000.0.
- Around line 116-117: Replace the single-row GM copy loops in the
group_out/gather_window seam, the copy_row seam, and the pl.load reduction with
COMM_ROW_TILE-row transfers, clipping the tile size for the final partial block.
Preserve the existing row ordering, column range, and reduction behavior while
reducing the number of GM transfers.
In `@models/deepseek_v4_flash_dspark/decode_sparse_attn_csa.py`:
- Around line 124-130: Annotate the packed-head parameters with their fixed
layout and BF16 dtype: in
models/deepseek_v4_flash_dspark/decode_sparse_attn_csa.py lines 124-130 and
sparse_attn_csa_local_o_proj at line 427, use the group-major shape; in
models/deepseek_v4_flash_dspark/decode_sparse_attn_hca.py lines 132-138 and line
392, use the same group-major shape; and in
models/deepseek_v4_flash_dspark/decode_sparse_attn_swa.py lines 103-109 and line
294, use the head-major shape. Preserve these annotations at each
head-to-projection boundary.
🪄 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: c55a8142-ebac-4d2e-b1c3-f32304debf5f
📒 Files selected for processing (6)
models/deepseek_v4_flash_dspark/config.pymodels/deepseek_v4_flash_dspark/decode_attention_cp.pymodels/deepseek_v4_flash_dspark/decode_o_projection_cp.pymodels/deepseek_v4_flash_dspark/decode_sparse_attn_csa.pymodels/deepseek_v4_flash_dspark/decode_sparse_attn_hca.pymodels/deepseek_v4_flash_dspark/decode_sparse_attn_swa.py
wangqin1723-max
force-pushed
the
feat/dsv4-flash-decode-dsa-cp-foundation
branch
2 times, most recently
from
August 10, 2026 12:01
db957d2 to
d5f894d
Compare
- Set the TP4/SP4 component layout for replicated Q-B, attention sinks, and shared experts plus sharded output and vocab projections - Add group-local KV gather, grouped attention all-to-all, and O-B reduce-scatter seams with a four-rank golden fixture
Split the legacy local output projection from the SWA head computation so decode DSA-CP can insert its group-major all-to-all at the exact tensor and task dependency boundary. Keep the public wrapper and numerical behavior unchanged.
- Thread the runtime local token count through AG, A2A, and RS. - Pack valid rank rows compactly inside static-capacity windows. - Cover sub-capacity rows with poisoned inputs and preserved output tails. - Use simulator-compatible one-buffer TPUT transfers with 8-row tiles.
- Project compact runtime token prefixes through sharded O-A and O-B - Keep O-B matmuls in 128-row slabs and expose FP32 completion - Cover max-capacity and poisoned-tail receive layouts in the golden fixture
Split the local output projection from HCA head computation so decode DSA-CP can insert its group-major all-to-all at the tensor and task dependency boundary. Keep the public wrapper and numerical behavior unchanged.
Collaborator
Author
|
Latest-main CI disposition for head
I am treating the simulator statuses as baseline/runtime debt, not changing DSA-CP kernels or tolerances to mask them. |
This was referenced Aug 11, 2026
|
|
||
| # Parallelism constants | ||
| TP = 4 # tensor-parallel ranks per DP group | ||
| SP = TP # sequence-parallel token owners in the TP group |
| @@ -0,0 +1,287 @@ | |||
| # Copyright (c) PyPTO Contributors. | |||
Collaborator
There was a problem hiding this comment.
in next pr, rename this to decode_o_proj.py，put all o proj related functions within it, including communication, tp=1 functions.
This was referenced Aug 18, 2026
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.
Part of #905. This is the first replacement slice for the old CP-off attention direction in #925.
What
rank * local_t, and invalid tails remain untouched.Layout contract
Each attention core publishes
[group, T_PAD, head-in-group, dim]; only the first runtime token rows in each group are valid. The all-to-all writes each source rank directly into[local_group, source_rank * local_t], producing two local O-A groups over the full TP-group token stream.The receive-side projection consumes
[LOCAL_O_GROUPS, GROUP_T_PAD, O_GROUP_IN], computes sharded O-A and O-B over onlygroup_t = SP_SIZE * local_t, and returns a rebound[GROUP_T_PAD, hidden]FP32 partial plus its completion task ID. Reduce-scatter then sends each token owner compactlocal_trows and casts only the reduced local result to BF16.Validation\n\n- Rebased without patch changes onto main at
d1cf017(#932)\n-pre-commit run --all-fileslocal_t=32local_t=31, including preserved output tailslocal_t=128and dynamic subcapacitylocal_t=127Physical A3 CI passed the four-rank communication fixture, the strict 512/508-row receive-side O projection cases, all three staged sparse-attention programs, and the unchanged LM-head, including a fresh pass after rebasing onto
d1cf017. The first A3 job had a transient unchanged LM-head S1 scheduler stall; two no-code-change A3 runs passed. Becauseconfig.pychanges, simulator CI swept all 32 runnable files in the directory: both platforms passed the new communication and O-projection fixtures, then reproduced existing directory-wide mismatches (including the CSA mismatch reproduced on the untouched base with byte-identical generated kernels) and eventually stalled in unchangedlm_head.pyuntil the 30-minute cancellation. The red simulator statuses are therefore not feature-local regressions. No physical A5 run was available. Prefill is unchanged by this PR. SWA, HCA, and CSA end-to-end output-path composition follow as decode-only slices; #913 and #923 cover the shared-expert and vocabulary boundaries.\n