HK MLA: MI35x m16x8 retune, new m16x4 kernel, page_size + mask support

[ruanjm]

Summary

Two tracks of HK MLA work landing together:

1. MI30x kernel: feature additions

• Variable page size via template parameter.
• Causal mask support — adds nh64+mtp1, nh32+mtp3, nh16+mtp7 alongside the existing nh128+mtp0.

The existing kernel is designed for MI30x and is functional on MI35x.

2. MI35x: dedicated m16x8 + new m16x4

The new MI35x m16x8 kernel inherits from the MI30x kernel and is fine-tuned for MI35x:

• KvManager8bitsV3 to use ds_read_tr for KV loads.
• kBlockN bumped 32 → 64.
• Rescale of oaccu interleaved with the PV GEMM.
• Replace warp reduce implemented by ds_permute_b32 with v_permlane16/32_swap_b32.

The new m16x4 kernel targets num_heads * max_seqlen_qo == 64 shapes (covers (64,1), (32,2), (16,4)) on MI35x. It mirrors the m16x8 design at kBlockN=64 with Q/KV LDS overlap so it can run at occupancy=2; metadata and launch grid are plumbed end-to-end (get_mla_metadata_v1_2_device, get_mla_metadata_info_v1, and the kernel launch all agree on 2*num_cu clusters when m16x4 is dispatched). m16x4 dispatch and the cluster doubling are gated on AITER_ENABLE_EXPERIMENTAL=1.

Performance — m16x8 (-n 128,1, MI355)

Measured against the previous HK build and the ASM reference for this shape (note: ASM v3 does not exist for m16x8). Highlights:

• Mid/large ctx (≥1200) and batch ≥16 — consistent wins vs both baselines (5–17% vs ASM, 4–17% vs HK prev). This is the regime that matters for real decode serving workloads.
• Small ctx with small batch (ctx ≤ 512, batch ≤ 5) — modest regressions vs the previous HK build (typically 5–15%, occasional ~12% at ctx=21). These configurations represent very short KV caches with low concurrency and are not on the critical path for production decode (where the KV cache is what makes attention expensive in the first place); the kernel is also still within or close to ASM there. We accept this tradeoff for the substantial gains at the realistic operating points.

m16x4 vs ASM v3

The new m16x4 kernel currently runs ~20% slower than ASM v3 on the m16x4-eligible shapes. It is included here to land the kBlockN=64 + occupancy=2 plumbing and unblock further tuning; closing the gap to ASM v3 is tracked as follow-up work and m16x4 is gated behind AITER_ENABLE_EXPERIMENTAL so production traffic is unaffected.

Test Results:

The following is test results of -n 128,1 on MI355.

batch	ctx	ASM (us)	HK prev (us)	HK curr (us)	gain vs ASM (%)	gain vs prev (%)
1	23333	52.467	51.719	43.392	17.30%	16.10%
3	23333	54.562	54.211	46.625	14.55%	13.99%
5	23333	57.608	56.914	49.158	14.67%	13.63%
16	23333	109.965	117.610	102.786	6.53%	12.60%
32	23333	199.795	210.031	184.629	7.59%	12.09%
64	23333	403.895	395.417	350.695	13.17%	11.31%
128	23333	724.364	801.548	676.442	6.62%	15.61%
256	23333	1422.000	1583.000	1315.000	7.52%	16.93%
1	8192	29.806	28.896	26.926	9.66%	6.82%
3	8192	29.831	28.261	27.267	8.60%	3.52%
5	8192	31.611	29.942	28.787	8.93%	3.86%
16	8192	52.138	53.821	47.721	8.47%	11.33%
32	8192	84.845	85.508	78.311	7.70%	8.42%
64	8192	151.649	154.495	132.878	12.38%	13.99%
128	8192	267.067	280.267	246.490	7.70%	12.05%
256	8192	511.068	564.501	475.339	6.99%	15.79%
1	5200	24.606	23.708	21.910	10.96%	7.58%
3	5200	24.769	23.339	22.043	11.01%	5.55%
5	5200	26.539	24.292	22.827	13.99%	6.03%
16	5200	39.407	38.349	36.453	7.50%	4.94%
32	5200	60.707	63.077	57.899	4.63%	8.21%
64	5200	105.001	109.455	95.035	9.49%	13.17%
128	5200	182.689	191.257	171.394	6.18%	10.39%
256	5200	329.671	362.874	309.395	6.15%	14.74%
1	3200	21.044	18.331	18.795	10.69%	-2.53%
3	3200	21.343	19.749	19.303	9.56%	2.26%
5	3200	22.691	20.548	20.089	11.47%	2.23%
16	3200	31.828	29.349	29.159	8.39%	0.65%
32	3200	47.726	46.249	41.648	12.74%	9.95%
64	3200	71.196	74.426	67.318	5.45%	9.55%
128	3200	122.538	126.889	109.809	10.39%	13.46%
256	3200	212.023	222.358	192.551	9.18%	13.40%
1	1200	17.863	16.587	14.627	18.12%	11.82%
3	1200	17.625	16.581	16.488	6.45%	0.56%
5	1200	18.598	16.796	16.351	12.08%	2.65%
16	1200	23.009	20.522	21.133	8.15%	-2.98%
32	1200	28.852	25.879	24.562	14.87%	5.09%
64	1200	41.541	39.213	36.618	11.85%	6.62%
128	1200	60.574	63.209	56.015	7.53%	11.38%
256	1200	96.050	99.546	83.871	12.68%	15.75%
1	512	14.367	12.006	13.337	7.17%	-11.09%
3	512	17.986	15.767	16.958	5.72%	-7.55%
5	512	16.198	13.945	15.274	5.70%	-9.53%
16	512	20.629	17.537	16.990	17.64%	3.12%
32	512	21.617	18.975	20.076	7.13%	-5.80%
64	512	26.853	24.600	23.178	13.69%	5.78%
128	512	40.547	38.277	34.633	14.59%	9.52%
256	512	48.409	50.075	44.635	7.80%	10.86%
1	256	14.733	12.281	13.335	9.49%	-8.58%
3	256	17.175	16.085	16.893	1.64%	-5.02%
5	256	16.280	15.288	14.866	8.69%	2.76%
16	256	16.910	15.894	16.067	4.99%	-1.09%
32	256	19.423	17.752	16.931	12.83%	4.62%
64	256	22.448	21.642	19.600	12.69%	9.44%
128	256	29.061	28.159	25.505	12.24%	9.43%
256	256	33.938	33.707	30.012	11.57%	10.96%
1	64	12.518	11.960	13.218	-5.59%	-10.52%
3	64	13.032	12.640	14.278	-9.56%	-12.96%
5	64	13.718	13.205	15.193	-10.75%	-15.05%
16	64	16.577	15.058	16.565	0.07%	-10.01%
32	64	18.203	16.610	15.915	12.57%	4.18%
64	64	17.369	16.358	15.807	8.99%	3.37%
128	64	19.554	17.321	16.939	13.37%	2.21%
256	64	23.023	21.414	21.534	6.47%	-0.56%
1	21	12.549	10.615	11.832	5.71%	-11.46%
3	21	12.763	11.671	12.883	-0.94%	-10.38%
5	21	13.854	11.664	13.082	5.57%	-12.16%
16	21	14.757	13.085	14.525	1.57%	-11.00%
32	21	15.259	13.657	15.029	1.51%	-10.05%
64	21	16.051	15.173	16.011	0.25%	-5.52%
128	21	16.820	15.851	16.801	0.11%	-5.99%
256	21	21.631	19.016	20.141	6.89%	-5.92%

[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:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
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 3072 --add-label <label>

[Copilot]

Pull request overview

This PR extends and retunes the HK MLA decode path for MI3xx/MI35x by adding page-size support, adding causal-mask (“MTP”) shape coverage, introducing MI35x-specific kernels (including a new m16x4 variant), and plumbing metadata/dispatch changes end-to-end (with experimental gating via AITER_ENABLE_EXPERIMENTAL).

Changes:

• Add variable page_size support and new experimental MTP shapes (nh16/32/64 variants) in metadata + Python dispatch/tests.
• Introduce new/retuned HK decode kernels: MI3xx m16x8, MI35x m16x8 retune, and new MI35x m16x4 (occupancy=2) plus updated softmax/util helpers.
• Update hk decode entrypoint routing (hk_decode_fwd.cu) and JIT config / HipKittens pin.

Reviewed changes

Copilot reviewed 14 out of 14 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
op_tests/test_mla_persistent.py	Extends test dispatch gating and passes page_size into metadata; fixes kv_last_page_lens handling comment.
csrc/kernels/mla/metadata/v1_2_device.cuh	Adds experimental gating, cluster multiplier for m16x4 occupancy=2, and expands supported-shape logic.
csrc/kernels/mla/hk/mi3xx_v32_fwd_decode_m16x8_fp8_fp8.cuh	New MI3xx HK decode kernel (m16x8 fp8/fp8) with page-size + causal mask handling.
csrc/kernels/mla/hk/mi3xx_v32_fwd_decode_h128_fp8_fp8.cuh	Removes the prior h128-specific kernel implementation (superseded by m16x8 path).
csrc/kernels/mla/hk/mi35x_v32_fwd_decode_m16x8_fp8_fp8.cuh	New MI35x-tuned HK decode kernel (m16x8 fp8/fp8) with kBlockN=64 and PV rescale changes.
csrc/kernels/mla/hk/mi35x_v32_fwd_decode_m16x4_fp8_fp8.cuh	New MI35x HK decode kernel (m16x4 fp8/fp8), occupancy=2, product-64 shapes.
csrc/kernels/mla/hk/hk_mla_utils.cuh	Refactors shared HK traits/params/helpers; introduces HIP_CALL and expands utility functions.
csrc/kernels/mla/hk/hk_mla_softmax.cuh	Adds 16-VGPR softmax scaling helpers and switches lane/warp helpers to opus.
csrc/kernels/mla/hk_decode_fwd.cu	Updates HK decode entrypoint to route by (num_head * max_seqlen_q) and GPU arch to new kernels.
aiter/ops/attention.py	Adjusts metadata sizing for m16x4 occupancy=2 and adds experimental MTP shape recognition.
aiter/mla.py	Updates Python-side gating/dispatch to HK kernels for product-128 and product-64 shapes (experimental).
aiter/jit/optCompilerConfig.json	Updates module_hk_mla source list to the new MI3xx kernel header.
aiter/jit/core.py	Updates pinned HipKittens commit used by the JIT build system.

---

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