[atom-vllm] Attention CP for DSA models

[whx-sjtu]

Motivation

Advanced version of #196.

Technical Details

• For Attention part, weights are stored normally in TP. When running prefill or mixed batch, gather attention weights layer by layer to obtain full weights, and then perform full CP for attention as described in [RFC]: Parallel strategy for Deepseek v3.2 like models #196.
• For MoE part, stay full TP. When running prefill or mixed batch, replace the all-reduces before and after MoE with all gather and reduce scatter.
• As attention weights are stored in normal TP, no extra memory cost for duplicated weights. Also Decode path can run with normal TP.

Prefill performance are optimized without hurting decode.

Test Plan

GLM 5.2 600k-600

Test Result

[github-actions]

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

• ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

• ✅ Run after the PR is approved and Pre Checkin passes
• ✅ Run immediately when an approval review is submitted
• ✅ Can be requested before approval with labels

Label	Tests
ci:full	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom	Run native ATOM model accuracy tests
ci:vllm	Run ATOM vLLM OOT model accuracy tests
ci:sglang	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 1618 --add-label <label>