[PERF] fix: Correct AR 1stage Gating Condition

[ftyghome]

Motivation

This PR addresses an issue in the AllReduce 1-stage / 2-stage dispatch threshold introduced in #3458.

In that PR, the condition was written as:

total_bytes <= self.world_size * 32 * 1024

but the intended condition should be:

total_bytes * self.world_size <= 128 * 7168 * 2

The threshold comes from the measurements in #3458. For inference workloads on models such as Kimi, the AllReduce path should use the 1-stage implementation for M <= 32 under TP4, and for M <= 16 under TP8. Larger shapes should use the 2-stage implementation.

Because the condition was computed incorrectly, the actual dispatch policy did not change as intended. In particular, TP4 was still only using the 1-stage path for M <= 8.

This PR fixes the threshold calculation and has been tested on MI350X.

Technical Details

See #3458

Test Plan

See #3458

Test Result

See #3458

[Copilot]

Copilot wasn't able to review any files in this pull request.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 3880 --add-label <label>