Skip to content

[HIP] [OPUS] [JIT] [GFX950]OPUS PA MQA Logits MXFP4 #5332

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
junhaha666 merged 13 commits into from
Sep 21, 2026
Merged

[HIP] [OPUS] [JIT] [GFX950]OPUS PA MQA Logits MXFP4 #5332

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
b7fb5bf
add opus pa mqa logits mxfp4
shay-li77 Sep 7, 2026
1135f24
support kv compress for decode
shay-li77 Sep 8, 2026
efe944c
fix ruff and black findings in op_tests
shay-li77 Sep 8, 2026
f9a06bd
fold store-alignment cases into the mqa logits op test
shay-li77 Sep 8, 2026
3deb560
validate launcher inputs and guard the device kernel on gfx950
shay-li77 Sep 9, 2026
fd4fbb0
rework the mqa logits op test to the aiter op-test standard
shay-li77 Sep 9, 2026
a97b5eb
propagate a NaN e8m0 scale through the mqa logits relu
Sep 11, 2026
8351413
add the per-row schedule for prefill and decode
shay-li77 Sep 17, 2026
127d7e9
drop the schedule-free launches and the window builder
shay-li77 Sep 17, 2026
88c79f6
merge upstream main
shay-li77 Sep 17, 2026
37a6726
guard the schedule builder on gfx950 too
shay-li77 Sep 17, 2026
5f26fe1
validate the weights row width in the launcher
shay-li77 Sep 17, 2026
220706b
score every row of a corner case, not eight of them
shay-li77 Sep 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions aiter/jit/optCompilerConfig.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1394,6 +1394,22 @@
"f'{CK_DIR}/example/ck_tile/50_sparse_attn/generate.py -d fwd_vsa --receipt 600 --output_dir {{}}'"
]
},
"module_pa_mqa_logits_mxfp4_opus": {
"srcs": [
"f'{AITER_CSRC_DIR}/py_itfs_cu/pa_mqa_logits_mxfp4_kernels.cu'",
"f'{AITER_CSRC_DIR}/pybind/pa_mqa_logits_mxfp4_pybind.cu'"
],
"flags_extra_cc": [],
"flags_extra_hip": [
"'-ffast-math'",
"'-fno-finite-math-only'",
"'-mllvm -enable-post-misched=1'"
],
"extra_ldflags": "None",
"extra_include": [],
"verbose": "False",
"blob_gen_cmd": "''"
},
"module_mha_fwd": {
"srcs": [
"f'{AITER_CSRC_DIR}/kernels/mha_common.cu'",
Expand Down
33 changes: 33 additions & 0 deletions aiter/ops/opus/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,17 @@ def _stub(*_args, **_kwargs):
opus_gemm_workspace_release_all,
)

# gfx950-only, unlike the gemm entries above: _SUPPORTED is wider than this
# kernel, so the wrappers re-check at call time rather than at import.
from .pa_mqa_logits_opus import (
compute_prefill_windows,
pa_mqa_logits_mxfp4_decode,
pa_mqa_logits_mxfp4_fwd_decode,
pa_mqa_logits_mxfp4_fwd_prefill,
Comment thread
junhaha666 marked this conversation as resolved.
Resolved
pa_mqa_logits_mxfp4_prefill,
pa_mqa_logits_mxfp4_prefill_windows,
)

def opus_gemm_a8w8_blockscale_bpreshuffle_tune(*args, **kwargs):
from .gemm_op_a8w8 import (
opus_gemm_a8w8_blockscale_bpreshuffle_tune as _impl,
Expand All @@ -69,14 +80,36 @@ def opus_gemm_a8w8_blockscale_bpreshuffle_tune(*args, **kwargs):
opus_gemm_workspace_release_all = _make_unsupported_arch_stub(
"opus_gemm_workspace_release_all"
)
compute_prefill_windows = _make_unsupported_arch_stub("compute_prefill_windows")
pa_mqa_logits_mxfp4_decode = _make_unsupported_arch_stub(
"pa_mqa_logits_mxfp4_decode"
)
pa_mqa_logits_mxfp4_fwd_decode = _make_unsupported_arch_stub(
"pa_mqa_logits_mxfp4_fwd_decode"
)
pa_mqa_logits_mxfp4_fwd_prefill = _make_unsupported_arch_stub(
"pa_mqa_logits_mxfp4_fwd_prefill"
)
pa_mqa_logits_mxfp4_prefill = _make_unsupported_arch_stub(
"pa_mqa_logits_mxfp4_prefill"
)
pa_mqa_logits_mxfp4_prefill_windows = _make_unsupported_arch_stub(
"pa_mqa_logits_mxfp4_prefill_windows"
)


__all__ = [
"bmm_a8w8_mxscale_opus",
"compute_prefill_windows",
"gemm_a16w16_opus",
"opus_gemm_a8w8_blockscale_bpreshuffle_tune",
"opus_gemm_a16w16_tune",
"opus_gemm_workspace_init",
"opus_gemm_workspace_release",
"opus_gemm_workspace_release_all",
"pa_mqa_logits_mxfp4_decode",
"pa_mqa_logits_mxfp4_fwd_decode",
"pa_mqa_logits_mxfp4_fwd_prefill",
"pa_mqa_logits_mxfp4_prefill",
"pa_mqa_logits_mxfp4_prefill_windows",
]
Loading
Loading