[WS2][GEMM][Forward]: implement PR4 TP+CP with FFN and collective reductions

[Flink-ddd]

resolves #239 (PR4 Forward track)

Overview

This PR builds upon the Tensor Parallel (TP) foundation established in PR3, extending the Qwen3-8B FFN orchestration to support Context Parallelism (CP). It implements the necessary 2D mesh topology and specific cross-rank communication patterns for weight gradients under CP.

Key Changes

1. FFNContext & Setup:

   • Introduced CP configuration support within the ctx.
   • Updated the test launcher to construct a 2D process mesh, establishing TP groups [0,1], [2,3] and CP groups [0,2], [1,3].

2. Communication Placement & Logic (TP + CP):

   • Forward Pass: Retained TP all_reduce semantics. CP shards token rows and does not require forward activation collectives.
   • Backward Pass (CP Weight Reductions): Because CP rank lanes maintain replicated local parameter shards across sequence chunks, the backward pass now correctly triggers a CP all_reduce(SUM) on the weight gradients (dW) for down_weight, gate_weight, and up_weight.

3. Constraints & Guardrails:

   • Maintained strict isolation from Sequence Parallelism (no SP AG/RS introduced).
   • Continued enforcement of the deterministic GEMM contract to prevent floating-point drift.

Verification & Test Results

1. Local Validation

   • tests/test_tensor_parallel_ffn.py: 8 passed (Successfully verified batch/padding invariance across the 2D TP+CP topology).
   • mypy tensor_parallel.py: Passed
   • ruff linting: Passed

2. GPU Hardware Validation (Pending)

   [!NOTE]
   Reserved section for multi-GPU hardware test metrics (Fast Path vs. Consistent Path) and execution logs:

   [To be filled after running on target GPU cluster]
   - Test Command: 
   - Output / Logs: 
   

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

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: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 8e9cb7a1-fbdf-4260-9954-87d90bce9eb6

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

---

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.}