[Triton/Gluon] Use a zero default for masked scales in gemm_afp8wfp8

[cagrikymk]

Copy PR from: vllm-project/vllm#59454

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, also the cases with LDS poisoning LDS with NaNs from another kernel.

Test Result

Tests pass.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
ci:triton-355	Run the full Triton test suite on MI35X, not only the tests the change affects
multigpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 6007 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.

[Copilot]

Copilot review overview

🟡 Changes recommended

The claimed LDS-poisoning regression is not covered by a committed test.

Review effort: Balanced
Findings: 1

Open (1)

• Add NaN-poisoning regression tests for tail paths · New

What changed in this PR

Changes masked E8M0 scale defaults to zero to prevent stale-LDS NaNs in Triton MXFP8 GEMM paths.

Changes:

• Uses zero defaults for masked activation and weight scales.
• Covers standard, packed, and preshuffled kernel variants.

File	Description
aiter/​ops/​triton/​_triton_kernels/​gemm/​basic/​gemm_afp8wfp8.py	Zero-fills masked scale lanes.

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[azaidy]

LGTM!