[Triton/Gluon] gemm_a16w16_atomic accumulate into an existing output buffer

[nsusanto]

Motivation

Kimi K3 fuses the outputs of shared experts and routed experts. gemm_a16w16_atomic does not have a built in accumulate for an existing accumulator buffer.

Technical Details

• accumulate=True on gemm_a16w16_atomic, which requires y and never zeroes it:
  • when no split-K, load y, add in fp32, store;
  • when split-K, atomic adds on top of y.
• Strided y (a column slice) is updated in place.
• GEMM-A16W16-ATOMIC-N=896-K=3584.json: tuned gfx950 configs for the Kimi-K3 shape.

Test Plan

Added test_gemm_a16_w16_atomic_accumulate
pytest op_tests/triton_tests/gemm/basic/test_gemm_a16w16.py -k atomic

Submission Checklist

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

[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; main branch always runs both MI35X and MI300X
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 5895 --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

Register the tuning harness and reset the accumulation buffer before each benchmark run.

Review effort: Lite
Findings: 1

Open (1)

• Register new tuning harness in HARNESS_CONFIG_NAMES · New

What changed in this PR

Adds accumulate=True support to Triton atomic A16W16 GEMM, with split-K/strided-output tests and gfx950 tuning.

Changes:

• Adds in-place accumulation API and kernel behavior.
• Adds contiguous and strided accumulation tests.
• Adds tuning harness and gfx950 configuration.

File	Description
op_tests/​triton_tests/​gemm/​basic/​test_gemm_a16w16.py	Adds accumulation tests.
aiter/​ops/​triton/​utils/​_triton/​tuning/​harness_gemm_a16w16_atomic_accumulate.py	Adds accumulation tuning harness.
aiter/​ops/​triton/​gemm/​basic/​gemm_a16w16_atomic.py	Exposes accumulation support.
aiter/​ops/​triton/​configs/​gfx950/​triton/​gemm/​gemm_a16w16_atomic/​GEMM-A16W16-ATOMIC-N=896-K=3584.json	Adds tuned gfx950 configuration.
aiter/​ops/​triton/​_triton_kernels/​gemm/​basic/​gemm_a16w16_atomic.py	Implements accumulation in the kernel.

---

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

[Boss2002n]

Another Q is - should accumulate not be in the config json? (does it have any perf improvements on other models is another way to ask this question)

what do u think is better? config json or just let the caller pass an argument directly?

[Boss2002n]

LGTM!

[azaidy]

LGTM!

[nsusanto]

MI35X Test Failures are unrelated. Safe to merge.