[TRITON] Fix unit tests on gfx950 - part 2

[brunomazzottiamd]

Motivation

Triton test suite isn't passing on gfx950. This PR fixes test_gemm_afp4wfp4.py and test_gemm_a8w8.py, slightly improving the situation.

Technical Details

Fix test_gemm_afp4wfp4.py

Triton commit de2ba3946b changed AMDMFMALayout.instr_shape from a 2-element [M, N] to a 3-element [M, N, K] list. Extend the previously 2-element [32, 32] to [32, 32, 64]. K=64 is the K dimension of the mfma_scale_f32_32x32x64_f8f6f4 hardware instruction used for FP4 on gfx950.

Restrict AFP4/WFP4 AOT tests to Triton 3.5. Avoid using prebuilt AOT kernels on newer Triton versions where the metadata format is incompatible.

Fix test_gemm_a8w8.py

Fix _gemm_a8w8_kernel

Same instr_shape API break (Triton de2ba3946b). The kernel uses mfma_scaled for FP8 and plain mfma for INT8, which target different hardware instructions with different K dimensions:

• FP8 mfma_scale_f32_16x16x128_f8f6f4 (K=128, K_WIDTH=32)
• INT8 mfma_i32_16x16x64_i8 (K=64, K_WIDTH=16)

SwizzledSharedLayout.vec is updated to match K_WIDTH per data type specialization.

Fix _gemm_a8w8_preshuffled_kernel

The linear_nk layout and its reshape - permute - reshape - trans unshuffle sequence were designed for K=32 / K_WIDTH=16, so applying K=128 breaks the layout conversion. Since mfma_scaled was already invoked with a_scale=None and b_scale=None (per-tensor scale applied to the accumulator separately), replace it with plain mfma, targeting the unscaled mfma_f32_16x16x32_fp8_fp8 (K=32) that the preshuffled layout was built for.

Fix test_gemm_a8w8.py

Relax absolute tolerance from 0.02 to 0.03 to accommodate the preshuffled FP8 path (unscaled dot + software accumulator scale).

Fix ff_a16w16_fused_ungated.py

The k-loop staggers each N-block's start position by k_cyclic_offset = pid_n % cdiv(K, BLOCK_SIZE_K) to reducetl.atomic_add contention on y_ptrs. The y_mask K-boundary check incorrectly used the raw loop counter k (always starting at 0) instead of k_cyclic_offset (the actual K position). When the cyclic offset is non-zero, k understates the real offset, producing a wrong mask and corrupting partial sums near the K boundary.

Compatibility fixes for older Gluon API (Triton < 3.6.0)

This PR also implements compatibility for old Gluon API, supporting Gluon of a Triton compiler older than version 3.6.

Test Plan

Run respective tests on gfx950:

for t in \
op_tests/triton_tests/gemm/basic/test_gemm_afp4wfp4.py \
op_tests/triton_tests/gemm/basic/test_gemm_a8w8.py \
op_tests/triton_tests/test_pa_decode_gluon.py \
; do echo "${t}"; pytest -q --no-header "${t}" | tail -1; done

The tests should pass on latest Triton TOT and Triton 3.5.0 (< 3.6.0).

Test Result

test_gemm_afp4wfp4.py, test_gemm_a8w8.py and test_pa_decode_gluon.py pass on gfx950.

TOT Triton - all test cases:

op_tests/triton_tests/gemm/basic/test_gemm_afp4wfp4.py
10656 passed, 3680 skipped in 570.54s (0:09:30)
op_tests/triton_tests/gemm/basic/test_gemm_a8w8.py
9382 passed, 1856 skipped in 1024.83s (0:17:04)
op_tests/triton_tests/test_pa_decode_gluon.py
4 passed, 1 warning in 618.54s (0:10:18)

Triton 3.5.0 - only test cases of Gluon kernels, to check compatibility with older API:

op_tests/triton_tests/gemm/basic/test_gemm_afp4wfp4.py
3584 passed, 3584 skipped in 223.31s (0:03:43)
op_tests/triton_tests/gemm/basic/test_gemm_a8w8.py
6414 passed, 1272 skipped in 689.21s (0::1::29)
op_tests/triton_tests/test_pa_decode_gluon.py
4 passed, 1 warning in 618.00s (0:10:18)

Execution of Gluon kernels only was achieved thought the following patch:

```diff
diff --git a/op_tests/triton_tests/gemm/basic/test_gemm_a8w8.py b/op_tests/triton_tests/gemm/basic/test_gemm_a8w8.py
index 10d8f99d8..2010a2ec0 100644
--- a/op_tests/triton_tests/gemm/basic/test_gemm_a8w8.py
+++ b/op_tests/triton_tests/gemm/basic/test_gemm_a8w8.py
@@ -175,7 +175,6 @@ def generate_gemm_a8w8_inputs(
 @pytest.mark.parametrize(
     "impl",
     [
-        "triton",
         "gluon",
         "gluon_shuffle",
     ],
diff --git a/op_tests/triton_tests/gemm/basic/test_gemm_afp4wfp4.py b/op_tests/triton_tests/gemm/basic/test_gemm_afp4wfp4.py
index 8a4953811..671f516ef 100644
--- a/op_tests/triton_tests/gemm/basic/test_gemm_afp4wfp4.py
+++ b/op_tests/triton_tests/gemm/basic/test_gemm_afp4wfp4.py
@@ -241,7 +241,7 @@ def run_triton(
     [True, False],
 )
 @pytest.mark.parametrize("skip_reduce", [True, False])
-@pytest.mark.parametrize("impl", ["triton", "gluon"])
+@pytest.mark.parametrize("impl", ["gluon"])
 def test_gemm_afp4_wfp4(
     M: int,
     N: int,

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[azaidy]

LGTM!

[brunomazzottiamd]

The only test failures are the expected ones, i.e. test_mha_backward on Triton Tests (MI35X) / Shard 0:

=========================== short test summary info ============================
FAILED op_tests/triton_tests/attention/test_mha.py::test_mha_backward[True-False-0.0-True-128-8-32-512-512-1]
FAILED op_tests/triton_tests/attention/test_mha.py::test_mha_backward[True-False-0.0-True-128-8-32-512-512-4]
FAILED op_tests/triton_tests/attention/test_mha.py::test_mha_backward[True-False-0.0-True-128-8-32-1024-1024-1]
FAILED op_tests/triton_tests/attention/test_mha.py::test_mha_backward[True-False-0.0-True-128-8-32-1024-1024-4]
FAILED op_tests/triton_tests/attention/test_mha.py::test_mha_backward[True-False-0.0-True-128-8-32-2048-2048-1]
FAILED op_tests/triton_tests/attention/test_mha.py::test_mha_backward[True-False-0.0-True-128-8-32-2048-2048-4]
FAILED op_tests/triton_tests/attention/test_mha.py::test_mha_backward[True-False-0.0-True-128-8-64-512-512-1]
FAILED op_tests/triton_tests/attention/test_mha.py::test_mha_backward[True-False-0.0-True-128-8-64-512-512-4]
FAILED op_tests/triton_tests/attention/test_mha.py::test_mha_backward[True-False-0.0-True-128-8-64-1024-1024-1]
FAILED op_tests/triton_tests/attention/test_mha.py::test_mha_backward[True-False-0.0-True-128-8-64-1024-1024-4]
FAILED op_tests/triton_tests/attention/test_mha.py::test_mha_backward[True-False-0.0-True-128-8-64-2048-2048-1]
FAILED op_tests/triton_tests/attention/test_mha.py::test_mha_backward[True-False-0.0-True-128-8-64-2048-2048-4]
==== 12 failed, 5269 passed, 2016 skipped, 6 warnings in 3920.96s (1:05:20) ====

These failures aren't in the scope of this PR. We're good to merge, everything else passed.

FYI: @gyohuangxin