[cagrikymk]

Purpose

This fix changes masked scale data from 127 (value 1.0 for e8m0) to 0 (value 2^(-127). This change is benign in terms of accuracy as the main value side already loads corresponding lanes for mfma as 0. So, scales get multiplied by 0 and has no accuracy effect.

When other=127, triton 3.8.0+gitc0a142ff leads to reading stale LDS data for the lanes corresponding to the masked scales due to a compiler bug. This bug is fixed on upstream triton but this fix resolves stale data reading issue for both compiler versions.

With other=127 and triton 3.8.0+gitc0a142ff, if LDS addresses corresponding to the masked lanes have "NaN", this leads to the result and produce inaccurate results. This is a rare case but it can still happen in practice. Reproduction relies on what LDS contains from previous runs.

This is a cleaner solution as hardware can already handle writing 0s to the masked lanes almost for free compared to writing any other value.

As a side effect, it also reduces total instruction count:

Main is upstream triton/main.
3.8.0 is triton 3.8.0+gitc0a142ff.

Instruction counts that differ between the builds (whole kernel)

instruction	3.8.0 other=127	3.8.0 other=0	main other=127	main other=0
s_mov_b32	28	10	18	10
s_xor_b64	0	0	15	0
s_and_saveexec_b64	11	0	10	0
s_or_b64	11	0	10	0
s_andn2_saveexec_b64	0	0	10	0
ds_write_b32	0	0	9	0
v_mov_b32_e32	0	0	9	0
s_cbranch_execz	8	0	12	0
v_cndmask_b32_e64	76	83	76	83
s_waitcnt	28	34	29	34
v_lshlrev_b16_e32	4	0	4	0
v_or_b32_sdwa	4	0	4	0
v_perm_b32	20	24	20	24
s_cbranch_execnz	3	0	3	0
s_movk_i32	3	1	6	1
v_bfrev_b32_e32	4	1	4	1
v_lshl_or_b32	15	18	15	18
v_or_b32_e32	9	6	9	6
s_barrier	6	6	4	4
s_load_dword	2	0	2	0
s_nop	3	4	5	4
s_or_saveexec_b64	0	0	2	0
v_add_u32_e32	100	98	100	98
v_cmp_lt_i32_e64	0	0	2	0
v_cndmask_b32_e32	1	3	1	3
s_ashr_i32	1	0	1	0
s_bfe_i32	0	1	0	1
s_branch	2	1	3	1
s_load_dwordx2	1	2	1	2
s_lshl_b32	4	3	4	3
v_add3_u32	3	4	3	4
v_and_or_b32	0	1	0	1
v_ashrrev_i32_e32	3	2	3	2
v_cmp_eq_u32_e64	3	3	4	3
v_cmp_le_i32_e64	0	0	1	0
v_lshl_add_u32	2	2	3	2
v_lshrrev_b32_e32	5	4	5	4
v_mad_u64_u32	8	7	8	7
v_mul_lo_u32	1	2	1	2
v_or3_b32	1	0	1	0
all instructions	695	645	742	643

Performance is unchanged for triton 3.8.0+gitc0a142ff.

Performance: DeepSeek-V4.1-Flash shapes

aiter gemm_afp8wfp8 (v0.1.24, shipped configs, Triton 3.8.0) on one MI355X. Time per call in us, from HIP-graph replay with weights cold in cache, median of 5-7 replays. Each cell: without fix / with other=0; * marks a point that returns NaN without the fix after the 0xFF LDS fill.

N x K	layer	M = 1	M = 16	M = 32	M = 128	M = 8192
1152x5120	shared expert gate_up (TP4)	4.78 / 4.81	5.32 / 5.20	4.97 / 4.93	7.91 / 7.92	88.49 / 87.09
1792x5120	attention wq_a|wkv	5.57* / 5.63	5.73 / 5.75	6.00 / 5.99	8.89 / 8.78	97.27 / 96.50
2304x5120	shared expert gate_up (TP2)	6.03* / 6.15	6.11 / 6.11	6.39 / 6.39	9.97 / 10.05	137.24 / 134.19
4096x1280	indexer wq_b	3.47 / 3.46	4.11 / 4.06	5.21 / 5.19	5.49 / 5.48	67.36 / 68.15
5120x576		3.10 / 3.06	3.17 / 3.09	3.58 / 3.43	4.49 / 4.51	43.79 / 42.80
5120x1152		4.19 / 3.96	4.99 / 4.78	5.03 / 4.74	6.41 / 6.39	66.22 / 65.28
5120x2048		4.79 / 4.83	4.85 / 4.80	5.50 / 5.63	7.68 / 7.66	113.94 / 113.60
5120x4096		7.60 / 7.59	8.63 / 8.08	8.54 / 8.53	12.92 / 12.88	207.88 / 205.94
8192x1280		4.48 / 4.49	5.07 / 5.06	5.17 / 5.17	7.34 / 7.39	111.54 / 110.94
16384x1280		6.47 / 6.46	7.91 / 8.03	7.44 / 7.15	10.93 / 10.68	214.24 / 213.51

Test Plan

Existing tests pass for test_rocm_mxfp8_linear.py

Test Result

Tests pass.

No similar PR exists as of writing this.

Claude is used to collect benchmarks and assembly analysis.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[github-actions]

✅ @cagrikymk, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Your branch must contain every commit currently on its upstream target branch. Merge or rebase onto the latest target branch, then rerun the command. Append --allow-stale to a run command to test an outdated branch at your own risk.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[cagrikymk]

/ci run

[github-actions]

✅ Triggered Buildkite CI #92235 for commit e525bb7405b9.

[cagrikymk]

/ci retry

[github-actions]

✅ Triggered Buildkite CI #92281 for commit 9ad97905ee41, running 4 failed step(s) from Buildkite CI #92235.