Add: compose DeepSeek V4 SWA decode TP output

[wangqin1723-max]

• Chain rank-local SWA heads through TP token/head exchange, sharded
  output projection, and dependency-ordered FP32 reduce-scatter.
• Select TP-derived shapes from --tp 1/2/4 at import time and default
  standalone and CI execution to TP2.
• Consolidate the TP output path and fixtures into decode_swa.py, reuse
  the shared output projection, and validate max and subcapacity layouts
  with exact inactive-tail checks.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 9da02348-be59-4fd9-901f-c3c2758f35e1

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Added a DeepSeek-V4 SWA decode output pipeline with distributed attention, sharded output projection, reduce-scatter, deterministic fixtures, PyTorch golden validation, poisoned-tail checks, and a CLI for compilation and test execution.

Changes

SWA decode output pipeline

Layer / File(s)	Summary
Distributed SWA decode kernel models/deepseek_v4_flash_dspark/decode_swa_output_cp.py	Adds decode_swa_output_cp and l3_decode_swa_output_cp. The pipeline validates dynamic shapes, performs sparse attention, all-to-all reshaping, sharded O projection, and reduce-scatter.
Fixtures and golden validation models/deepseek_v4_flash_dspark/decode_swa_output_cp.py	Adds deterministic tensor specifications, golden_decode_swa_output_cp, and build_o_local_compare. Validation covers valid rows and poisoned inactive-capacity rows.
Compilation and test harness models/deepseek_v4_flash_dspark/decode_swa_output_cp.py	Adds CLI options for device selection, maximum and subcapacity cases, compilation modes, pass dumps, distributed compilation, and golden comparison.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant CLI
  participant l3_decode_swa_output_cp
  participant decode_swa_output_cp
  participant DistributedCollectives
  CLI->>l3_decode_swa_output_cp: Launch distributed decode
  l3_decode_swa_output_cp->>decode_swa_output_cp: Invoke device kernel
  decode_swa_output_cp->>DistributedCollectives: Exchange attention outputs
  decode_swa_output_cp->>DistributedCollectives: Reduce-scatter projected output
  decode_swa_output_cp-->>CLI: Return o_local

Loading

Possibly related PRs

• hw-native-sys/pypto-lib#771: Refactors and optimizes SWA attention output projection and merge processing.
• hw-native-sys/pypto-lib#931: Adds the DSA-CP decode output communication path extended by this SWA implementation.

Poem

A rabbit hops through SWA’s flow,
With scattered sums in rows below.
Attention aligns, projections gleam,
Golden checks guard every stream.
“No poisoned tail shall pass,” I cheer—
A clean decode arrives here! 🐇

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly identifies the addition of the DeepSeek V4 SWA decode TP output composition, which matches the main change.
Description check	✅ Passed	The description accurately summarizes TP exchange, sharded projection, reduce-scatter, shape selection, and validation changes.

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[zhangqi-chen]

ci devices =2

[zhangqi-chen]

use --tp 1/2/4，like style in moe/lmhead