[BUG] Incorrect matmul results with ​num_warps=8​: missing cross-warp-group barrier after WGMMA wait

[lizhangyu258]

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:

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:

python python/tutorials/03-matrix-multiplication.py --only_unit_test

Observed output:

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