Bug: assert_allclose fails in test_ag_moe.py for large model shape (K=14336) due to numerical divergence

Hi Triton-distributed team,

I am testing the test_ag_moe.py script on a multi-GPU setup and have encountered a numerical correctness issue that seems specific to certain model configurations.

While the test passes for the smaller "Dummy-Model" and "Qwen1.5-MoE-A2.7B" configurations, it consistently fails for the third, larger configuration.

**Environment**
GPU: 4x NVIDIA RTX 6000 Ada Generation

Interconnect: PCIe 4.0 (verified with nvidia-smi topo -m)

CUDA Toolkit: 12.8

PyTorch: 2.7.0a0+79aa17489c.nv25.4 (from log)

Triton-distributed: main branch

Host Compiler: g++ (conda-forge gcc 12.4.0-1) 12.4.0

**Bug Description**
The script test_ag_moe.py fails with a RuntimeError originating from assert_allclose. This indicates that the output from the custom Triton implementation (C_triton) does not match the baseline PyTorch implementation (C_torch) within the specified tolerance (atol=1e-3, rtol=1e-3).

The failure only occurs for the third test case with the shape M=8192, N=4096, K=14336.

**Steps to Reproduce**

1. Navigate to the root of the cloned Triton-distributed repository.
2. Set up the environment and launch the test script for a 4-GPU node:

```
export NPROC_PER_NODE=4
source ./scripts/sentenv.sh
bash ./scripts/launch.sh python/triton_dist/test/nvidia/test_ag_moe.py
```
**Observed Behavior**
The first two tests pass successfully with significant speedups. The script then fails on the third test, printing detailed debug information about the numerical differences before raising a RuntimeError.

Key output from the logs:
```
# The first two tests pass successfully
✅ RANK 0 Dummy-Model pass
RANK 0 perf: ... speedup=8.44
✅ RANK 0 Qwen1.5-MoE-A2.7B pass
RANK 0 perf: ... speedup=23.03

# The third test fails
shape: M=8192, N=4096, K=14336; num experts=8, topk=2

# The numerical difference is reported before the error
diff count: 4672 (0.028%), [16384, 1024]
diff max: 0.5, atol: 0.001, rtol_abs: 0.0
diff locations:
tensor([[ 1, 148],
[ 5, 741],
...
[16383, 451]], device='cuda:2')
--------------------------------------------------------------

# The final traceback
[rank2]: Traceback (most recent call last):
[rank2]: File "/workspace/Triton-distributed/python/triton_dist/test/nvidia/test_ag_moe.py", line 261, in
[rank2]: perf_test(name, args.M, dtype, config, debug=args.debug, pg=TP_GROUP)
[rank2]: File "/workspace/Triton-distributed/python/triton_dist/test/nvidia/test_ag_moe.py", line 170, in perf_test
[rank2]: raise e
[rank2]: File "/workspace/Triton-distributed/python/triton_dist/test/nvidia/test_ag_moe.py", line 165, in perf_test
[rank2]: assert_allclose(C_torch, C_triton, atol=1e-3, rtol=1e-3, verbose=False)
[rank2]: File "/workspace/Triton-distributed/python/triton_dist/utils.py", line 896, in assert_allclose
[rank2]: raise RuntimeError
[rank2]: RuntimeError
```

**Expected Behavior**
All tests within test_ag_moe.py, including the larger configuration, should pass the assert_allclose correctness check.

**Additional Context**
The failure seems specific to this larger configuration (especially the large K dimension of 14336). This might be a floating-point precision issue, possibly related to the accumulation data type (float16 vs float32) within the custom Triton kernels, as the maximum absolute difference is quite large (0.5).

Thank you for looking into this.