[yuychang]

Motivation

Some callers need gemm_a8w8_bpreshuffle and gemm_a4w4 to write into a buffer they already own, such as a preallocated activation or a CUDA-graph pool, rather than get back a newly allocated tensor. This PR adds out-variant ops for both.

Technical Details

• New ops, which write into out and return None:
  • gemm_a8w8_bpreshuffle_out(out, XQ, WQ, x_scale, w_scale): out is a contiguous [M, N] bf16 or fp16 tensor. Its dtype selects the output dtype.
  • gemm_a4w4_out(out, A, B, A_scale, B_scale, bias=None, alpha=1.0, beta=0.0, bpreshuffle=True): out is a contiguous bf16 [ceil(M/32)*32, N] tensor, because the kernels write whole 32-row tiles. Only bf16 is accepted: the asm kernel, which serves every shape without a CK tuned row, writes only bf16. The result is out[:M]. Supported on gfx950 only; on gfx1250 it raises NotImplementedError.
• Both ops declare only out as mutated (Tensor(a0!) out) and return nothing.
  • torch.compile therefore adds no writebacks for the read-only inputs, and no op output aliases an input.
  • The same pattern is used by mhc_post and batched_gemm_a16wfp4_.
• gemm_a8w8_bpreshuffle and gemm_a4w4 keep exactly the same signature and custom-op schema as on main.
  • Narrowing their mutates_args instead (one op with an optional out) broke inductor for the existing no-out A8W8 call (getitem is not an OpOverload). So the out path lives in separate ops.
• Each op pair shares its kernel dispatch (_gemm_a8w8_bpreshuffle_into, _gemm_a4w4_into), so both variants pick the same tuned or default kernel.
• Both out ops raise ValueError before any launch if out has the wrong shape, dtype or device, or is not contiguous.
  • CK and CKTile write a dense row-major [M, N], so a strided out (for example a column slice of a wider buffer) would be overwritten past its bounds.

Test Plan

New file op_tests/test_gemm_out.py:

• Both ops: the registered schema marks only out as mutated.
• gemm_a8w8_bpreshuffle_out:
  • At M=1, 37 and 256 (bf16), and M=37 (fp16), out is bit-identical to gemm_a8w8_bpreshuffle.
  • Under torch.compile(fullgraph=True) with aot_eager and with inductor, out is filled and matches eager.
  • A wrong shape, an fp32 dtype, a column slice and a transposed buffer each raise ValueError and leave the backing storage untouched.
• gemm_a4w4_out (gfx950):
  • At M=1, 32 and 45, out[:M] is bit-identical to gemm_a4w4.
  • Under torch.compile(backend="aot_eager", fullgraph=True), the padded buffer is filled and out[:M] matches eager.
    • Not tested under inductor: gemm_a4w4 already fails to lower there on main, with or without this PR (auto_functionalized_v2 was not removed).
  • An unpadded [M, N] buffer, an fp32 or fp16 dtype and a column slice each raise ValueError.
  • On gfx1250 (mocked), it raises NotImplementedError.
• Checked that both existing ops keep main's schema, and that gemm_a8w8_bpreshuffle still compiles under aot_eager and inductor.

Test Result

MI355X (gfx950), ROCm 7.2, torch 2.11.0+rocm7.2:

• op_tests/test_gemm_out.py: 21 passed (inductor caches disabled).
• The existing gemm_a8w8_bpreshuffle (no out) matches eager under aot_eager and inductor, as on main.

Submission Checklist

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

🤖 Generated with Claude Code

[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
ci:multi-gpu	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: DeepSeek-V4-Pro, MiniMax-M3-MXFP4
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 6412 --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]

🟡 Changes recommended

Custom-op mutation schemas need narrowing, and the A4W4 compile path lacks coverage.

3 open findings

• Restrict mutates_args to out to prevent compile-time writebacks · New
• Declare only out mutable to avoid unnecessary functionalization copies · New
• Add fullgraph a4w4 compile test for padded-buffer aliasing · New

What changed in this PR

Adds caller-owned output buffers to two quantized GEMM APIs.

Changes:

• Validates output shape, dtype, device, and contiguity.
• Supports padded output views for A4W4.
• Adds eager, rejection, and limited compile tests.

File	Description
aiter/​ops/​gemm_op_a8w8.py	Adds out= support and validation.
aiter/​ops/​gemm_op_a4w4.py	Adds padded out= support and gfx1250 rejection.
op_tests/​test_gemm_out.py	Tests output aliasing and validation.

🧠 Review effort: Balanced

---

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

[Copilot]

🟡 Changes recommended

The A4W4 API accepts fp16 output, but ASM-dispatched shapes reject all non-bf16 outputs.

1 open finding

• fp16 GEMM routes to bf16-only assembly kernel · New

3 resolved since last review

• Declare only out mutable to avoid unnecessary functionalization copies
• Restrict mutates_args to out to prevent compile-time writebacks
• Add fullgraph a4w4 compile test for padded-buffer aliasing

🧠 Review effort: Balanced

[Copilot]

🟡 Changes recommended

Some accepted dtype and optional-argument combinations fail or are silently ignored by selected dispatch paths.

1 open finding

• Honor optional parameters for non-ZN CK GEMM rows · New

1 resolved since last review

• fp16 GEMM routes to bf16-only assembly kernel

Previously missed (1)

In code that hasn't changed since last review

Avoid dispatching fp16 to bf16-only tuned rows

aiter/​ops/​gemm_op_a8w8.py:869

fp16 is accepted here for every configured shape, but a tuned FlyDSL 8-wave row is bf16-only. For example, the shipped GLM 5.2 config selects flydsl_bpreshuffle_8w_..._B16_... for M=8192, N=2048, K=2048, after which flydsl_8wave_gemm_a8 rejects the fp16 out. Either bypass bf16-only tuned rows for fp16 or stop accepting fp16 for those dispatches so the advertised out variant does not fail after validation.

🧠 Review effort: Balanced

[yuychang]

Closing: no current caller needs a caller-owned output for gemm_a8w8_bpreshuffle or gemm_a4w4, so this API would ship unused. If a consumer appears, the branch k3/gemm-out-param has the _out op design, its tests, and the review fixes (bf16-only gemm_a4w4_out, out as the only mutated argument).

[yuychang]

Superseded by #6437, which adds only gemm_a8w8_bpreshuffle_out (bf16-only, mutates only out). SGLang's Kimi-K3 shared-expert down-projection is its consumer. The gemm_a4w4 variant is dropped because it has no consumer.