[Triton/Gluon] [Bugfix] batched_gemm_a8w8: keep scale loads in bounds

[akii96]

Motivation

test_batched_gemm_a8w8 intermittently hits a GPU memory access fault on the MI300X Triton CI (shard 7 on main).

Technical Details

% binds tighter than +, so the scale offsets wrapped only the arange:

offs_a_scale = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) % M

When M % BLOCK_SIZE_M != 0, the last tile reads past the end of x_scale. The same happens for N and w_scale. The store is masked, so results are correct, but the load faults if the next page is unmapped. The fix wraps the full offset, as offs_am/offs_bn already do.

Test Plan

MI300X, B=16, M=320, N=1280, K=8192. The test places x_scale at the end of a hipMalloc with no mapping after it (PYTORCH_NO_HIP_MEMORY_CACHING=1):

	before	after
x_scale at end of allocation	10/10 fault	10/10 pass
x_scale mid-allocation	10/10 pass	10/10 pass

Test Result

• test_batched_gemm_a8w8.py: 196 passed
• test_fused_bmm_rope_kv_cache.py: 64 passed, 64 skipped

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 (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 6082 --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.

[akii96]

The shard 5 failure was just a runner network issue while downloading some packages, re-ran

[Copilot]

Copilot review overview

🟢 Approval recommended

The localized indexing fix bounds scale accesses without changing valid output lanes, and no unresolved issues remain.

Review effort: Balanced
Findings: None

What changed in this PR

Fixes out-of-bounds scale loads in the Triton batched A8W8 GEMM kernel for partial tiles.

Changes:

• Applies modulo to the full row and column scale offsets.
• Preserves indexing for valid output elements.

File	Description
aiter/​ops/​triton/​_triton_kernels/​gemm/​batched/​batched_gemm_a8w8.py	Corrects scale-offset wrapping for both inputs.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.