[yuychang]

Motivation

SGLang's Kimi-K3 ROCm path writes its PTPC FP8 shared-expert down-projection straight into a preallocated buffer (shared_output = buf[offset:].view(M, hidden)), which is later all-reduced together with the MoE latent. gemm_a8w8_bpreshuffle always allocates its output, so today that result has to be copied. This PR adds an op that writes into a caller-owned buffer.

Technical Details

• New op gemm_a8w8_bpreshuffle_out(out, XQ, WQ, x_scale, w_scale) -> None. It writes into out and picks the same kernel as gemm_a8w8_bpreshuffle, because both share one dispatch helper (_gemm_a8w8_bpreshuffle_into).
• torch.compile:
  • The op declares only out as mutated (mutates_args=["out"]) and returns nothing. Its schema is (Tensor(a0!) out, Tensor XQ, Tensor WQ, Tensor x_scale, Tensor w_scale) -> ().
  • Compiled graphs therefore add no writebacks for the read-only inputs, and no op output aliases an input.
  • It is a separate op rather than an optional out= on gemm_a8w8_bpreshuffle: narrowing that op's mutates_args with an optional out broke inductor for its existing no-out callers (getitem is not an OpOverload).
• Inputs:
  • Rejected with ValueError before any launch, because no kernel supports them: non-2-D or non-fp8 XQ/WQ, non-fp32 scales or scales without M/N elements, operands that are not all on one CUDA/HIP device, WQ K < XQ K, and an out that is not a contiguous bf16 [M, N] tensor.
    • bf16 only: the FlyDSL 8-wave kernels, selected by many tuned rows, write bf16 only, so an fp16 out would fail at launch on those shapes.
    • Contiguous out only: 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 written past its bounds.
    • out must start 16-byte aligned. A dense view into a larger buffer is fine at an aligned offset; at an odd bf16 offset the CK split-K kernel faulted. The rule applies to every backend so that a valid buffer never depends on which tuned row is selected. torch.empty storage and SGLang's buf[M * latent:] view are always aligned.
• Arch: only gfx942 and gfx950 are accepted. Elsewhere the op raises NotImplementedError before launch, because the shared dispatch would reach backends this op was never validated on (the gfx1250 WMMA path and the CK MFMA fallback on RDNA); use gemm_a8w8_bpreshuffle there. This check runs after the argument checks, so invalid arguments always get ValueError.
  • Normalized, not rejected: a strided XQ, WQ or scale is copied to a dense tensor first, because CK and CKTile read only row strides and raw scale pointers. Dense inputs pass through without a copy.
  • Any device: when the operands are not on the current device, the op runs under torch.cuda.device(XQ.device), because the launchers use the current device's stream. On the normal path it adds no guard.
  • An empty out (M or N = 0) returns after validation. M=0 matches tuned FlyDSL rows whose launch would otherwise fail.
  • out must not overlap any input. A K-padded WQ keeps gemm_a8w8_bpreshuffle's padded-K handling unchanged.
• gemm_a8w8_bpreshuffle is unchanged: same signature, same custom-op schema, same results.

Test Plan

New file op_tests/test_gemm_a8w8_bpreshuffle_out.py:

• The schema marks only out as written and has no return value.
• out is bit-identical to gemm_a8w8_bpreshuffle on gfx950 tuned rows that select each backend: CK, CK split-K, CKTile, FlyDSL, FlyDSL split-K and FlyDSL 8-wave. Each case asserts that its row still selects that backend and fails otherwise, so coverage cannot silently drop when tables change.
• Bit-identical on an untuned shape and on the padded-K path (WQ K > XQ K). Both assert they have no tuned row.
• A dense view of a larger flat buffer gets the right result, and the rest of the buffer is untouched.
• torch.compile(fullgraph=True) with aot_eager and with inductor (caches disabled): out is filled and matches eager.
• A wrong shape, fp16 or fp32 out, a column-slice or transposed out, WQ K < XQ K, WQ, x_scale or w_scale on another device, an all-CPU call, a rank-3 XQ or WQ, an int8 XQ, and a bf16 or short scale each raise ValueError and leave the buffer untouched.
• A column-major or row-strided XQ, a row-strided WQ, strided scales and a non-contiguous [1, M] scale all give the same result as dense inputs.
• With two GPUs visible, operands on cuda:1 called from cuda:0 give the right result and leave cuda:0 current.
• An empty out (M=0 on a FlyDSL row and an untuned row, and N=0) is a no-op and leaves the device usable.
• Running the file directly (python op_tests/test_gemm_a8w8_bpreshuffle_out.py) benchmarks both ops over a spread of N/K shapes, including the Kimi-K3 TP=8 shared-expert down projection (N=7168, K=768), writing out as a dense view into a caller-owned [latent | shared] buffer as SGLang does, and prints a markdown summary.
• Existing op_tests/test_gemm_a8w8_bpreshuffle_pad_k.py still passes, and gemm_a8w8_bpreshuffle keeps main's schema and still compiles under inductor.
• out at an odd bf16 offset raises ValueError and leaves the buffer untouched.
• Every arch outside gfx942/gfx950 (gfx1250, gfx1201, unknown) raises NotImplementedError before launch, and an all-CPU call raises ValueError even when the arch probe fails.
• The per-backend tuned-row cases run only on gfx950 with 256 CUs.
• gemm_a8w8_bpreshuffle still compiles under inductor (fullgraph=True) and matches eager.

Test Result

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

• op_tests/test_gemm_a8w8_bpreshuffle_out.py: 53 passed, 1 skipped with one GPU visible (the skipped test needs two GPUs and passes with two). With no GPU visible, all skip.
• op_tests/test_gemm_a8w8_bpreshuffle_pad_k.py: 17 passed, 8 skipped (gfx942 tuned-table cases).
• GPU kernel time per call (us), from the in-file benchmark, which now also times gemm_a8w8_bpreshuffle followed by the copy into the caller's buffer (run_perftest, FlyDSL rows at every shape below, error 0 everywhere). N=7168 K=768 is the Kimi-K3 TP=8 shared-expert down projection:

M	N, K	bpreshuffle	bpreshuffle + copy	bpreshuffle_out
1	7168, 768	2.75	4.10	2.75
4	7168, 768	2.74	4.76	2.74
16	7168, 768	2.78	4.71	2.76
64	7168, 768	3.33	5.16	3.34
4	7168, 2048	3.43	4.65	3.42
64	7168, 2048	4.68	9.94	4.66

• Host time per call (M=4, best of 5 x 2000 calls, standalone script): K=768: 27.0 us with the copy vs 23.4 us for _out (22.6 us for bpreshuffle alone); K=2048: 24.9 vs 22.3 us (21.5 us alone). An earlier run of the same comparison read 33.0 -> 28.6 us host and 5.63 -> 3.92 us GPU at K=2048; the re-run above replaces it.
• The SGLang caller was proposed in [Perf] [AMD] kimi-k3: Fold decode activation FP8 quant into its producers sgl-project/sglang#43125 (ptpc_fp8_aiter_hip.py, run(..., out=)), which was closed without merging; this op has no in-tree caller until that change is resubmitted.
• The benchmark main() over M = 1 to 4096 and N/K = 4096/4096, 7168/2048, 2112/7168, 8192/1024 and 7168/768 selects CK, CKTile, FlyDSL and untuned rows. Error is 0 for both ops on every row, and kernel times match within noise.

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 6437 --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

Same-device validation is incomplete, and backend/fallback coverage can silently disappear.

2 open findings

• Validate all tensors share XQ's device before kernel dispatch · New
• Fail on backend mismatch instead of silently skipping coverage · New

What changed in this PR

Adds caller-owned BF16 output support for pre-shuffled A8W8 GEMM while sharing existing kernel dispatch.

Changes:

• Adds gemm_a8w8_bpreshuffle_out.
• Extracts shared dispatch logic.
• Adds backend, validation, buffer-view, and compilation tests.

File	Description
aiter/​ops/​gemm_op_a8w8.py	Adds output-buffer API and shared dispatch helper.
op_tests/​test_gemm_a8w8_bpreshuffle_out.py	Tests output semantics, dispatch, validation, and compilation.

🧠 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

GPU-device and input-dtype validation remain incomplete, and the test architecture and benchmark contracts need correction.

1 open finding

• Reject CPU operands before GPU dispatch · New

2 resolved since last review

• Validate all tensors share XQ's device before kernel dispatch
• Fail on backend mismatch instead of silently skipping coverage

Previously missed (4)

In code that hasn't changed since last review

Validate both quantized operand dtypes

aiter/​ops/​gemm_op_a8w8.py:893

Only WQ is type-checked. On a tuned 8-wave row, an int8 XQ passes the launcher's one-byte-storage check and is then consumed by the FP8 kernel, producing incorrect values rather than rejecting the input. Validate both quantized operands here.

Validate scale tensors are fp32

aiter/​ops/​gemm_op_a8w8.py:899

The documented fp32 scale contract is not enforced before dispatch. The standard FlyDSL path passes these tensors through unchanged but emits BufferCopy32b(..., Float32) scale loads (preshuffle_gemm.py:716), while the 8-wave path converts them, so non-fp32 scales have backend-dependent behavior and can be misread by the standard kernel. Validate both scale dtypes before calling the helper.

Add the standard runnable op-test benchmark flow

op_tests/​test_gemm_a8w8_bpreshuffle_out.py:19

This new op_tests/test_*.py suite omits the repository's standard runnable op-test entry point: there is no @benchmark/run_perftest candidate loop, checkAllclose, markdown summary, or main() gated by get_gfx(). As a result, the claimed allocation-overhead improvement has no in-repo benchmark reproducer and the suite cannot be exercised through the normal op-test workflow. Add the standard benchmark/main flow while retaining these pytest contract cases.

Use get_gfx consistently for architecture checks

op_tests/​test_gemm_a8w8_bpreshuffle_out.py:25

This test gate uses the live-device helper, but _backend() calls the tuned-table lookup whose architecture key comes from get_gfx(). When GPU_ARCHS and the live card differ, the gate and the backend assertion inspect different architecture contracts, so the intended gfx950 backend coverage can be silently skipped. Tests should use get_gfx() consistently with the lookup they verify.

🧠 Review effort: Balanced

[Copilot]

🟡 Changes recommended

Missing rank, FP8 operand, and FP32 scale validation can pass incompatible pointers or dimensions to GPU backends.

1 open finding

• Reject non-2D XQ and WQ tensors before dispatch · New

1 resolved since last review

• Reject CPU operands before GPU dispatch

🧠 Review effort: Balanced

[Copilot]

🟡 Changes recommended

Scale validation can permit invalid memory reads, and the new test suite does not follow the required op-test structure.

1 open finding

• Validate scale dtype and size before kernel dispatch · New

1 resolved since last review

• Reject non-2D XQ and WQ tensors before dispatch

Previously missed (2)

In code that hasn't changed since last review

Integrate GEMM coverage with the standard op-test contract

op_tests/​test_gemm_a8w8_bpreshuffle_out.py:8

This new op_tests/test_*.py suite does not follow the repository's op-test contract: it is pytest-only and has no @benchmark/run_perftest candidate loop, checkAllclose result reporting, markdown summary, or main() architecture gate. Consequently the caller-owned path is not exercised by the standard correctness/performance sweep. Integrate this coverage into the existing GEMM op test (or reshape this file to the standard aiter-op-test structure).

Use get_gfx for test-time architecture selection

op_tests/​test_gemm_a8w8_bpreshuffle_out.py:25

Tests and build-time architecture gates must use get_gfx(); get_gfx_runtime() is reserved for runtime dispatch. Using the runtime probe here bypasses GPU_ARCHS, so the backend rows selected by the test can disagree with the configured test target. Switch this test helper to the test/build-time architecture helper.

🧠 Review effort: Balanced

[yuychang]

Addressing the "previously missed" items from the Copilot review summaries, all in 043ebf3:

• Validate both quantized operand dtypes: XQ must now be fp8, like WQ. New test test_int8_activation_rejected_before_launch.
• Validate scale tensors are fp32: see the reply on the inline thread. Scales must be contiguous fp32 with M and N elements.
• Use get_gfx() for architecture checks: the per-backend gate now uses get_gfx(), the same function the tuned-table lookup uses.
• Standard op-test benchmark flow: the file now has a @benchmark / run_perftest / checkAllclose main() that compares gemm_a8w8_bpreshuffle and gemm_a8w8_bpreshuffle_out on the Kimi-K3 shared-expert down-projection shapes and prints a markdown summary. It is gated on get_gfx(). The pytest contract cases are unchanged. The benchmark function is named bench_out so pytest does not collect it.

[Copilot]

🟡 Changes recommended

Input-layout validation and non-current GPU dispatch must be corrected to prevent invalid kernel reads or launches.

1 open finding

• Validate XQ row and WQ preshuffle contiguity before dispatch · New

1 resolved since last review

• Validate scale dtype and size before kernel dispatch

Previously missed (1)

In code that hasn't changed since last review

Set XQ.device as current device before launching kernels

aiter/​ops/​gemm_op_a8w8.py:923

This call does not establish XQ.device as the current device. The CK launcher and the ordinary FlyDSL launcher use the ambient current HIP stream/device, so tensors that are all on (for example) cuda:1 can still launch on cuda:0 when that is current. Run the shared dispatch under an XQ.device guard so the documented same-device contract works on non-current GPUs as well.

🧠 Review effort: Balanced

[yuychang]

On the "previously missed" device-guard item, fixed in d848178: with every operand on cuda:1 and cuda:0 current, the call caused a GPU memory access fault. The existing gemm_a8w8_bpreshuffle faults the same way. The op now raises ValueError unless the operands are on the current device. I chose a check over a torch.cuda.device guard because the guard cost about 1.7 us of host time per call, while the check costs about 0.4 us. SGLang already sets the current device per rank. New test test_non_current_device_rejected_before_launch checks the rejection, that out is untouched, and that the same call succeeds under torch.cuda.device(1). Fixing gemm_a8w8_bpreshuffle itself is left out of this PR.

[Copilot]

🔵 Needs a closer look

The existing allocating operation lacks the claimed schema and inductor regression coverage.

0 open findings

1 resolved since last review

• Validate XQ row and WQ preshuffle contiguity before dispatch

Previously missed (1)

In code that hasn't changed since last review

Add regression test for original no-out op signature

op_tests/​test_gemm_a8w8_bpreshuffle_out.py:135

The compile regression that motivated keeping this as a separate op is not covered here: this test compiles only gemm_a8w8_bpreshuffle_out, and there is no repository test asserting that the existing no-out op still compiles or retains its schema. Add a fullgraph=True inductor regression test that invokes gemm_a8w8_bpreshuffle with its original signature, since the PR description notes that this exact caller path broke with the earlier design.

🧠 Review effort: Balanced

[Copilot]

🔵 Needs a closer look

The architecture guard needs correction, and multi-backend buffer writes and compiled mutation behavior warrant human confirmation of runtime validation.

1 open issue

• Reject unsupported gfx1201 and unknown targets before dispatch · New

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

GPU buffer writes and compiler mutation semantics need final human confirmation, and the hardware-specific test gates still need correction.

2 open issues

• Reject unsupported gfx1201 and unknown targets before dispatch
• Gate gfx950 backend tests on the required 256 CU count · New

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

Multi-backend GPU execution and compiler mutation behavior need final human review because the reported runtime results were not independently validated.

No open issues 🎉

2 resolved since last review

• Reject unsupported gfx1201 and unknown targets before dispatch
• Gate gfx950 backend tests on the required 256 CU count

_{Review effort: Balanced}

[yuychang]

Two follow-ups in this push. test_allocating_op_still_compiles_under_inductor now compiles gemm_a8w8_bpreshuffle with inductor (fullgraph=True) and checks it against eager, so the no-out path that motivated a separate op is covered. The arch check in gemm_a8w8_bpreshuffle_out now runs after the argument checks, so a CPU-tensor call on a host with no GPU gets the documented ValueError. The earlier gfx1250 items are covered by rejecting every arch other than gfx942/gfx950.

[Copilot]

🔵 Needs a closer look

Human assessment is needed because the standard CI runner executes only the benchmark, skipping schema, validation, aliasing, empty-output, and torch.compile contract tests.

No open issues 🎉

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

The mutable caller-owned-buffer op spans CK, CKTile, and FlyDSL, requiring GPU validation of correctness, torch.compile, and device/stream behavior on gfx942/gfx950 that this source-only review could not independently perform.

No open issues 🎉

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

GPU execution is unavailable, leaving the new mutating op’s caller-owned writes, torch.compile behavior, and non-current-device stream handling unverified on supported ROCm hardware.

No open issues 🎉

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

Caller-owned GPU writes and torch.compile mutation semantics across CK, CKTile, and FlyDSL require runtime validation on gfx942 and gfx950 that this source-only review cannot independently perform.

No open issues 🎉

_{Review effort: Balanced}

[demonsan]

LGTM