Expose AQLayout as tunable parameter for CKTile blockscale 8-warp GEMM kernels

[samremes]

Summary

For CKTile blockscale FP8 GEMM 8-warp kernels on gfx950, the activation quantization scale (x_scale) is currently transposed from RowMajor to ColumnMajor on the host before every kernel launch. This PR makes the AQLayout a tunable parameter, allowing the kernel to natively read x_scale in RowMajor layout and skip the host-side transpose entirely.

Changes

• gemm_a8w8_blockscale_cktile_common.cuh: Add AQRowMajor template parameter to TileGemmConfig. When enabled, the 8-warp pipeline uses RowMajor AQLayout and the host-side x_scale transpose + allocation is skipped.
• gemm_a8w8_blockscale_cktile_instance.py: Add AQRowMajor field to TileKernelInstance with is_eight_warp property. Register a new RowMajor variant (kernel ID 12) alongside the existing ColumnMajor 8-warp kernel (ID 11), so the tuner can evaluate both.
• gen_instances_cktile.py: Propagate AQRowMajor into the generated C++ template instantiation.
• test_gemm_a8w8_blockscale_cktile_aq_rowmajor.py: New test covering instance name encoding, numerical accuracy (RowMajor vs ColumnMajor vs PyTorch reference), and padded weight stride handling (simulating vLLM's _maybe_pad_fp8_weight).

Performance

Benchmarked on gfx950 using the tuning script (gemm_a8w8_blockscale_tune.py --libtype cktile). Comparing the two 8-warp variants head-to-head:

• Shapes with large N, small K (e.g. Nx7168, K=2048): RowMajor is 3–8% faster than ColumnMajor, as the transpose overhead is significant relative to compute time. The tuner selects the RowMajor kernel as the overall best for M=1024.
• Shapes with large K (e.g. Nx2048, K=7168): ColumnMajor remains 8–21% faster, likely because the longer compute phase amortizes the transpose cost.

The tuner automatically selects the best variant per shape — no manual configuration needed.

Non-8-warp kernels

Non-8-warp kernels always use RowMajor AQLayout and are unaffected by this change. The AQRowMajor flag is only meaningful for 8-warp configurations.

Test plan

• test_instance_names: Verifies _aqrm suffix encoding and is_eight_warp property
• test_accuracy: RowMajor and ColumnMajor outputs match PyTorch FP32 reference across 4 shapes
• test_padded_weight_stride: RowMajor kernel handles non-contiguous (padded) weight tensors correctly
• The tuner picks up kernel 12 and selects it when it's the fastest

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

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

Extended tests (opt-in via labels):

Label	Tests
ci:triton-355	Run Triton tests on MI355 in addition to MI325
ci:sglang	SGLang integration tests
ci:atom	ATOM benchmark (DeepSeek-R1 + GPT-OSS)
ci:vllm	vLLM benchmark
ci:all	All of the above

Add labels via the sidebar or gh pr edit 2487 --add-label <label>

[Copilot]

Pull request overview

This PR makes activation-quantization scale layout (AQLayout) a tunable parameter for CKTile FP8 blockscale GEMM 8-warp kernels on gfx950, enabling an 8-warp RowMajor variant that can avoid the per-launch host-side x_scale transpose.

Changes:

• Add an AQRowMajor template parameter to TileGemmConfig and use it to select RowMajor vs ColumnMajor AQ layout for 8-warp kernels.
• Extend CKTile kernel instance generation to encode/propagate AQRowMajor, registering a new 8-warp RowMajor variant (kernel ID 12).
• Add a new op test validating name encoding, numerical accuracy vs reference, and padded-weight stride handling for the RowMajor AQ variant.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.

File	Description
op_tests/test_gemm_a8w8_blockscale_cktile_aq_rowmajor.py	Adds a gfx950-only test for RowMajor vs ColumnMajor AQ variants, including accuracy and padded weight stride coverage.
csrc/ck_gemm_a8w8_blockscale/include/gemm_a8w8_blockscale_cktile_common.cuh	Introduces AQRowMajor in config and adjusts host-side x_scale handling + AQ layout selection logic.
csrc/ck_gemm_a8w8_blockscale/gen_instances_cktile.py	Propagates AQRowMajor into generated C++ template instantiations.
csrc/ck_gemm_a8w8_blockscale/gemm_a8w8_blockscale_cktile_instance.py	Adds AQRowMajor to instance metadata, name suffix encoding, and registers the new RowMajor 8-warp kernel variant.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[samremes]

@valarLip Would you have time to review this PR? The CI passes with a recent merge with main.