Add: support decode attention TP and SP boundaries - #925
wangqin1723-max wants to merge 1 commit into
Conversation
|
Warning Review limit reached
Next review available in: 11 minutes You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (17)
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 |
Gather SP token shards and reduce-scatter FP32 output-projection partials across TP groups. Shard Q-B heads, attention sinks, O-A groups, and O-B inputs for SWA, HCA, and CSA decode while preserving the full-head QKV entry used by prefill. Size decode metadata, compressors, and indexer paths for the native full attention batch.
374b404 to
f026dc0
Compare
|
Issue #905 now targets the vllm-ascend DSA-CP-on layout. This PR models the previous CP-off direction by sharding Q-B/attention heads and gathering tokens before attention, which conflicts with the updated contract: Q-B and attention sink must be replicated, each rank keeps its local 128 queries over all 64 heads, and attention exits through a group-major all-to-all before sharded O-A/O-B plus reduce-scatter. The replacement is split into reviewable decode-only slices. The first draft is #931, which defines the three communication seams and exposes the SWA packed-head/task dependency boundary. HCA/CSA and the sharded receive-side O projection will follow. Prefill remains unchanged. Closing this PR as superseded; its history stays available for reference. |
Part of #905. This is the first replacement slice for the old CP-off attention direction in #925. ## What - Define the TP4/SP4 decode component layout: Q-B, attention sink, and shared expert replicated; O-A, O-B, and vocabulary sharded four ways. - Add subgroup-local hidden-state all-gather, group-major attention-output all-to-all, and FP32 O-B reduce-scatter seams. - Make all three collectives capacity-static but runtime-row aware: valid rank slabs are packed at `rank * local_t`, and invalid tails remain untouched. - Add the receive-side sharded O projection: two local O-A groups, a 2048-wide O-B K shard, per-group A8 quantization, and an FP32 partial for reduce-scatter. - Split the SWA, HCA, and CSA decode kernels at the same group-major packed-head/task-ID boundary while preserving each legacy public wrapper and local output projection. - Use simulator-compatible single-buffer TPUT transfers with an 8-row communication tile. ## 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 only `group_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 compact `local_t` rows 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-files` - Four-rank A3 simulator + exact golden for AG/A2A/RS at current S=2: - full capacity `local_t=32` - poisoned subcapacity `local_t=31`, including preserved output tails - Merged #928 S=8 base: four-rank A3 simulator + exact golden at full capacity `local_t=128` and dynamic subcapacity `local_t=127` - Receive-side O projection A3 simulator + golden: - current S=2 at 128/124 group rows - merged #928 S=8 base at 512/508 group rows - strict valid-prefix check with NaN-poisoned receive padding and zero output tails - group-distinct O-A weights spanning all 4096 K columns and non-unit O-B scales - SWA seam: A3 simulator + golden at current local T=32 and temporary S=8 local T=128 - HCA seam: A3 simulator + golden at T=32, subcapacity T=8, and temporary S=8 T=128 - CSA seam: T=32, subcapacity T=8, and temporary S=8 T=128 compile and complete A3 simulator runtime. The pinned local stack has a pre-existing broad CSA golden mismatch reproduced on the untouched base; after normalizing renamed symbols, all nine generated PTO kernels are byte-identical to the base. Physical 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. Because `config.py` changes, 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 unchanged `lm_head.py` until 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
1 participant
reduce-scatter FP32 O-B partials back to their owning ranks.
SWA, HCA, and CSA decode at selectable TP1/TP2/TP4/TP8 degrees.
host wrapper for TP2/TP4/TP8. TP1 creates no DistributedConfig and
runs no SP AllGather or ReduceScatter.
64-request attention batch instead of the pre-TP batch stand-in.
while decode TP uses a separate rank-local QKV entry.
consumers so communication windows are safe to reuse across layers.