flagos-ai / flagos-ai/FlagTree
[BUG] Incorrect matmul results with num_warps=8: missing cross-warp-group barrier after WGMMA wait
- Dominant language
- Python
- Stars
- 350
- Forks
- 149
- Avg merge
- 2d 4h
- Merged PRs (30d)
- 81
Description
## Problem
The FP16 matrix multiplication tutorial produces substantially different results from `torch.matmul` when its unit-test uses eight warps. In one run, the maximum elementwise absolute difference was **19.605**. Some elements match exactly while others differ by several units, suggesting a synchronization or shared-memory lifetime problem rather than FP16 rounding alone.
## Reproduction
In `python/tutorials/03-matrix-multiplication.py`, set the `--only_unit_test` CUDA autotune configuration to `num_warps=8`, use FP16 inputs with shapes `a=(8192, 8200)` and `b=(8200, 8192)`, and calculate the maximum absolute difference after the Triton and PyTorch matmuls:
```python
def absolute_error(triton_output, torch_output):
return (triton_output.float() - torch_output.float()).abs()
torch.manual_seed(0)
a = torch.rand((8192, 8200), device=DEVICE, dtype=torch.float16) - 0.5
b = torch.rand((8200, 8192), device=DEVICE, dtype=torch.float16) - 0.5
triton_output = matmul(a, b)
torch_output = torch.matmul(a, b)
print(f"triton_output_with_fp16_inputs={triton_output}")
print(f"torch_output_with_fp16_inputs={torch_output}")
error = absolute_error(triton_output, torch_output)
print(f"fp16_matmul_absolute_error={error}")
print(f"fp16_matmul_max_absolute_error={error.max().item()}")
if torch.allclose(triton_output, torch_output, atol=1e-2, rtol=0):
print("✅ Triton and Torch match")
else:
print("❌ Triton and Torch differ")
```
Run:
```bash
python python/tutorials/03-matrix-multiplication.py --only_unit_test
```
Observed output:
```text
python python/tutorials/03-matrix-multiplication.py --only_unit_test
triton_output_with_fp16_inputs=tensor([[ -5.8594, 6.1523, -4.1797, ..., 8.2422, -23.2656, -3.5508],
[ 8.2891, 0.7261, -0.8872, ..., 7.4922, -0.6553, -8.9375],
[ 2.2051, -12.2188, -4.6719, ..., 10.2109, -4.2734, 7.7500],
...,
[ 2.1328, -1.5596, 8.6875, ..., -1.0107, 5.9688, -10.4531],
[ 6.3438, 7.4414, -9.6719, ..., 10.1875, -4.4336, 7.5586],
[ 10.9219, -0.9517, 4.1680, ..., -6.3086, -5.6250, -7.0820]],
device='cuda:0', dtype=torch.float16)
torch_output_with_fp16_inputs=tensor([[ -5.8594, 6.1523, -4.1797, ..., 9.2422, -23.1875, -2.8398],
[ 8.2891, 0.7261, -0.8872, ..., 6.9688, -0.2142, -9.1641],
[ 3.5098, -11.1875, -7.3672, ..., 10.2109, -4.2734, 7.7500],
...,
[ 2.1328, -1.5596, 8.6875, ..., -1.0107, 5.9688, -10.4531],
[ 8.8516, 5.4492, -10.0703, ..., 10.1875, -4.4336, 7.5586],
[ 10.4375, 0.9399, 7.0547, ..., -6.3086, -5.6250, -7.0820]],
device='cuda:0', dtype=torch.float16)
fp16_matmul_absolute_error=tensor([[0.0000, 0.0000, 0.0000, ..., 1.0000, 0.0781, 0.7109],
[0.0000, 0.0000, 0.0000, ..., 0.5234, 0.4410, 0.2266],
[1.3047, 1.0312, 2.6953, ..., 0.0000, 0.0000, 0.0000],
...,
[0.0000, 0.0000, 0.0000, ..., 0.0000, 0.0000, 0.0000],
[2.5078, 1.9922, 0.3984, ..., 0.0000, 0.0000, 0.0000],
[0.4844, 1.8916, 2.8867, ..., 0.0000, 0.0000, 0.0000]],
device='cuda:0')
fp16_matmul_max_absolute_error=19.60546875
❌ Triton and Torch differ
```
Environment: FlagTree checkout `71a8b316a`; CUDA/H20 host; FP16 inputs; unit-test autotune configuration with `num_warps=8`.
## Suspected cause
FlagTree's [`WarpGroupDotWaitOpConversion`](https://github.com/flagos-ai/FlagTree/blob/71a8b316a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM.cpp) lowers the wait to `WGMMAWaitGroupOp` and returns without a CTA barrier. A WGMMA wait synchronizes only the issuing warp group; with more than four warps, another warp group may still depend on the shared-memory operands.
Upstream Triton later fixed this in [PR #11056](https://github.com/triton-lang/triton/pull/11056): after the WGMMA wait, it inserts a local CTA barrier when `!op.getWarpGroupLocal() && lookupNumWarps(op) > 4`. The eight-warp configuration in this reproduction meets the warp-count part of that condition. The upstream PR also adds analysis to mark waits as warp-group-local when a later barrier already covers the dependency, avoiding unnecessary barriers.
Contributor guide
Research direction
Start in third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM.cpp at WarpGroupDotWaitOpConversion, then run python python/tutorials/03-matrix-multiplication.py --only_unit_test with the eight-warp FP16 configuration. Compare the Triton and PyTorch outputs and verify the maximum absolute error meets the stated allclose check without regressing other warp configurations.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 75/100