Chefang/pa decode opus

[fangche123]

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

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

[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; main branch always runs both MI35X and MI300X
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: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
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 4794 --add-label <label>

[Copilot]

Pull request overview

Adds an OPUS-based paged-attention decode implementation targeting gfx950, wires it into the Python API via JIT/pybind, and introduces correctness/perf coverage for the new path.

Changes:

• Add gfx950 OPUS paged-attention decode kernel (HIP/C++) plus pybind/JIT module plumbing.
• Expose aiter.pa_decode_opus(...) Python wrapper and integrate it into existing paged-attention perf testing.
• Add a standalone numerical test and a perf comparison benchmark vs the existing ASM decode path; pin vcs_versioning in setup.py to avoid stale transitive deps.

Reviewed changes

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

Show a summary per file

File	Description
setup.py	Pin vcs_versioning>=2.0.0 in setup_requires to avoid setuptools_scm import failures with stale 1.x installs.
op_tests/test_pa.py	Add perf-test hook for pa_decode_opus and improve ASM timing reporting for bf16.
op_tests/test_pa_decode_opus.py	New pytest numerical correctness test comparing against a torch reference.
op_tests/bench_pa_decode_opus_vs_asm.py	New benchmark script comparing OPUS decode vs ASM decode (with optional correctness verification).
csrc/pybind/pa_decode_opus_pybind.cu	New pybind module definition for the OPUS decode entrypoint.
csrc/py_itfs_cu/pa_decode_opus_kernels.cu	Host-side launcher/validation + per-stream split-KV scratch allocation/management.
csrc/include/rocm_ops.hpp	Add PA_DECODE_OPUS_PYBIND macro wiring for pybind.
csrc/include/pa_decode_opus.h	New single-header kernel implementation and public C++ API declaration.
aiter/ops/pa_decode_opus.py	New Python API wrapper around the JIT-compiled module.
aiter/jit/optCompilerConfig.json	Register module_pa_decode_opus sources and HIP compile flags.
aiter/init.py	Export pa_decode_opus from the package top-level import surface.

---

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