[rahulbatra85]

Motivation

Cache-warm inputs. Most AITER Triton benchmarks reuse a single set of inputs across benchmark iterations. Depending on their size, those inputs can stay resident in L2, so later iterations read from cache instead of HBM. This overstates performance, especially for memory-bound and small-shape kernels, and can make autotuning pick the wrong configs. In real workloads, other layers’ weights and activations pass through the cache between calls to a kernel, so its inputs usually have to be read from HBM. This PR allocates multiple input sets whose combined size is larger than the L2 cache and rotates through them, so each launch sees cold-cache behavior.
triton.testing.do_bench approximates this by flushing L2 before each run. However, the flush can’t be placed inside a captured graph without being timed along with the kernel, and it adds gaps between launches. Rotating buffers avoid both problems.

Graph capture. Most benchmarks launch kernels eagerly, which includes CPU launch overhead in the measurement. That overhead can dominate for small kernels. Production inference stacks (e.g. vLLM and SGLang decode paths) typically run under HIP graphs. The new do_bench captures the rotated launches in a graph, which makes the measured time closer to what these kernels see in deployment.

[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:windows	Windows JIT op smoke test on the gfx1201 ROCm runner
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 6838 --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

Graph capture can omit pool entries, and pool sizing counts unused buffers, so measurements do not enforce the requested active footprint.

5 open issues

• Exclude unused state and output tensors from prepare pool sizing · New
• Pool sizing overestimates each timing arm's active footprint · New
• Validate graph lengths and replay counts as positive integers · New
• Format callable values without assuming a func attribute · New
• Replay rotates only a prefix of oversized buffer pools · New

What changed in this PR

This PR moves KDA benchmarks to shared rotating-buffer graph timing. It aims to reduce cache-warm bias and measure execution without CPU launch overhead.

Changes:

• Add shared input-pool generation and graph-replay timing helpers.
• Add command-line controls for pool size, warmup, and graph replay.
• Migrate four KDA benchmarks while retaining legacy timing options as unused compatibility flags.

File	Description
op_tests/​op_benchmarks/​triton/​utils/​benchmark_utils.py	Add rotating input pools and graph timing.
op_tests/​op_benchmarks/​triton/​utils/​argparse.py	Add shared pool and graph controls.
op_tests/​op_benchmarks/​triton/​bench_kda.py	Migrate recurrent KDA timing and variability reporting.
op_tests/​op_benchmarks/​triton/​bench_fused_kda_decode.py	Migrate decode and speculative-decode comparisons.
op_tests/​op_benchmarks/​triton/​bench_flash_kda.py	Migrate FlashKDA timing to the shared helpers.
op_tests/​op_benchmarks/​triton/​bench_chunk_kda.py	Add rotating pools for full, prepare, and walk measurements.

_{Review effort: Balanced}

[aiter-bot]

NOT READY FOR HUMAN REVIEW

PR Self-Review

Revision 9d5b2676e8b0 · mode shadow

Waiting for the PR author. The readiness report will appear after the self-review answers are complete for this revision.

Questions to answer in the PR description

1. What behavior does this PR change, and why is the existing behavior insufficient?
2. What did you validate, and what remains unverified? Include commands, hardware/configuration and results.
3. What behavior would each changed op test fail on before this PR, and why is that coverage not already provided by an existing canonical test?
4. For every changed op_tests/ file, what behavior does it cover and why is it not redundant? Write None if no op test changed.

Edit the PR Self-Review section in the PR description. Include the current revision SHA; this comment refreshes automatically.

Missing or invalid answers

• FAIL — missing ### Reviewed revision in the PR description
• FAIL — missing ### Behavior and intent in the PR description
• FAIL — missing ### Validation and unverified scope in the PR description
• FAIL — missing ### PR-specific question in the PR description
• FAIL — missing ### Changed op tests in the PR description
• FAIL — Reviewed revision must contain at least the first 7 characters of the current head SHA
• FAIL — Behavior and intent has no substantive answer
• FAIL — Validation and unverified scope has no substantive answer
• FAIL — PR-specific question has no substantive answer
• FAIL — changed op test not explained in Changed op tests: op_tests/op_benchmarks/triton/bench_chunk_kda.py
• FAIL — changed op test not explained in Changed op tests: op_tests/op_benchmarks/triton/bench_flash_kda.py
• FAIL — changed op test not explained in Changed op tests: op_tests/op_benchmarks/triton/bench_fused_kda_decode.py
• FAIL — changed op test not explained in Changed op tests: op_tests/op_benchmarks/triton/bench_kda.py
• FAIL — Changed op tests must explain why each changed op test is necessary and non-redundant

_{Shadow mode reports what would block but keeps the status successful. Set AITER_PR_SELF_REVIEW_MODE=enforce only after reviewing real false positives.}

[Boss2002n]

The KDA benchmarks now rotate their inputs in a CUDA graph, which is good, but something must change before merge.

Must fix before merge

1. --state_dtype bf16 has no effect in bench_chunk_kda.py (bench_chunk_kda.py:90, :146): the pool makes the state without a dtype, so the Gluon and walk lines time an fp32 state but count bf16 bytes. Fix: pass state_dtype to their generators, and keep fp32 for the triton line.

Nice to fix

• Build beta and its in_proj once on the fixed side: the pool counts all of in_proj, so 4 of 12 default shapes rotate 132-243 MB, not 256 MB (bench_chunk_kda.py:108-125).
• Add the four fused-only fields to the pool only with --fused; an fla that does not take them fails with TypeError (bench_kda.py:208-212).
• Delete the dead _chunk_kda_fn (bench_flash_kda.py:171-199).
• Write the pool code in the _p_gen and functools.partial form of the benchmarks in #7121 (bench_kda.py:156-212).
• Correct the "cache-warm" header (bench_kda.py:352-355) and the two false docstrings (bench_fused_kda_decode.py:175-177, bench_flash_kda.py:9-13).
• Cut the eight long comments to one short sentence each, and delete the repeated one at :494-495 (bench_fused_kda_decode.py:268-276).

Merge order

Merge #7121 first. With #7225, the later PR must add conv1d_layout to _generate_kda_pool_inputs, or bench_kda.py fails with TypeError; with #6912, the later PR must keep an eager --varlen path.

[rahulbatra85]

The KDA benchmarks now rotate their inputs in a CUDA graph, which is good, but something must change before merge.

Must fix before merge

1. --state_dtype bf16 has no effect in bench_chunk_kda.py (bench_chunk_kda.py:90, :146): the pool makes the state without a dtype, so the Gluon and walk lines time an fp32 state but count bf16 bytes. Fix: pass state_dtype to their generators, and keep fp32 for the triton line.

Nice to fix

• Build beta and its in_proj once on the fixed side: the pool counts all of in_proj, so 4 of 12 default shapes rotate 132-243 MB, not 256 MB (bench_chunk_kda.py:108-125).
• Add the four fused-only fields to the pool only with --fused; an fla that does not take them fails with TypeError (bench_kda.py:208-212).
• Delete the dead _chunk_kda_fn (bench_flash_kda.py:171-199).
• Write the pool code in the _p_gen and functools.partial form of the benchmarks in [Triton/Gluon] : Change RotatingBuffersBench API and update any existing benchmarks #7121 (bench_kda.py:156-212).
• Correct the "cache-warm" header (bench_kda.py:352-355) and the two false docstrings (bench_fused_kda_decode.py:175-177, bench_flash_kda.py:9-13).
• Cut the eight long comments to one short sentence each, and delete the repeated one at :494-495 (bench_fused_kda_decode.py:268-276).

Merge order

Merge #7121 first. With #7225, the later PR must add conv1d_layout to _generate_kda_pool_inputs, or bench_kda.py fails with TypeError; with #6912, the later PR must keep an eager --varlen path.

1. Fixed