Skip to content

[HIP] [OPUS] [JIT] Unify gfx950+gfx1250 MXFP4 MQA-logits #5761

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
shay-li77 merged 13 commits into from
Sep 29, 2026
Merged

[HIP] [OPUS] [JIT] Unify gfx950+gfx1250 MXFP4 MQA-logits #5761

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
2c2542e
[HIP] [OPUS] Unify gfx950+gfx1250 MXFP4 MQA-logits into one module/sc…
shay-li77 Sep 22, 2026
4325dd6
[OPUS] Name gfx950 MQA-logits instances by the shared qlen{Q}_kv{K} c…
shay-li77 Sep 22, 2026
82cc2ad
style: black-format the unified mqa-logits test
shay-li77 Sep 22, 2026
79463da
[OPUS] Skip the tile cut at q_per_block == 1, where it is the identity
shay-li77 Sep 23, 2026
d140bbc
Merge branch 'main' into lxx/dev/pa_mqa_logits_gfx950_unify
shay-li77 Sep 23, 2026
de2fe0c
[OPUS] Split pa_mqa_logits header (sched/gfx950/gfx1250) under csrc/k…
shay-li77 Sep 24, 2026
8e03a1a
[OPUS] Bound local_ends/local_starts/row_to_batch and block_tables wi…
shay-li77 Sep 24, 2026
3bcbfcc
[OPUS] Address PR #5761 review: one row count, shared launcher, gfx95…
shay-li77 Sep 25, 2026
6f2fb13
[OPUS] pa_mqa_logits: probe arch once per process; fix Black/Ruff
shay-li77 Sep 25, 2026
530fdfd
[OPUS] Move pa_mqa_logits sources to csrc/kernels/opus_mqa_logits/pa_…
shay-li77 Sep 25, 2026
d552ccf
[OPUS] pa_mqa_logits_mxfp4: per-arch namespaces with generic names
shay-li77 Sep 25, 2026
dc13ff6
[OPUS] pa_mqa_logits_mxfp4: apply -enable-post-misched=1 on gfx950 only.
shay-li77 Sep 27, 2026
a758f0c
[OPUS] pa_mqa_logits_mxfp4: document cu_seq_q <= total_q and row_to_b…
shay-li77 Sep 27, 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
21 changes: 3 additions & 18 deletions aiter/jit/optCompilerConfig.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -619,21 +619,6 @@
"verbose": "False",
"blob_gen_cmd": "''"
},
"module_pa_mqa_logits_mxfp4_gfx1250_opus": {
"srcs": [
"f'{AITER_CSRC_DIR}/py_itfs_cu/pa_mqa_logits_mxfp4_gfx1250_kernels.cu'",
"f'{AITER_CSRC_DIR}/pybind/pa_mqa_logits_mxfp4_gfx1250_pybind.cu'"
],
"flags_extra_cc": [],
"flags_extra_hip": [
"'-ffast-math'",
"'-fno-finite-math-only'"
],
"extra_ldflags": "None",
"extra_include": [],
"verbose": "False",
"blob_gen_cmd": "''"
},
"module_moe_sorting_opus": {
"srcs": [
"f'{AITER_CSRC_DIR}/py_itfs_cu/moe_sorting_opus_kernels.cu'",
Expand Down Expand Up @@ -1526,14 +1511,14 @@
},
"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'"
"f'{AITER_CSRC_DIR}/kernels/opus_mqa_logits/pa_mqa_logits_mxfp4/pa_mqa_logits_mxfp4_kernels.cu'",
"f'{AITER_CSRC_DIR}/kernels/opus_mqa_logits/pa_mqa_logits_mxfp4/pa_mqa_logits_mxfp4_pybind.cu'"
],
"flags_extra_cc": [],
"flags_extra_hip": [
"'-ffast-math'",
"'-fno-finite-math-only'",
"'-mllvm -enable-post-misched=1'"
"('-mllvm -enable-post-misched=1' if (get_gfx() == 'gfx950' or 'gfx950' in os.environ.get('GPU_ARCHS', '')) else '')"
Comment thread
shay-li77 marked this conversation as resolved.
],
"extra_ldflags": "None",
"extra_include": [],
Expand Down
14 changes: 0 additions & 14 deletions aiter/ops/opus/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,16 +9,6 @@

from .dispatch import _opus_dispatch

Comment thread
shay-li77 marked this conversation as resolved.
# gfx950-only, unlike the GEMM entries below, and nothing here gates on arch at import:
# the wrappers re-check at call time so a non-gfx950 caller gets a named error instead of
# a JIT failure.
from .pa_mqa_logits_opus import (
pa_mqa_logits_mxfp4_build_sched,
pa_mqa_logits_mxfp4_sched,
pa_mqa_logits_mxfp4_sched_buffer_ints,
pa_mqa_logits_mxfp4_sched_slots,
)


def opus_gemm(
XQ: Tensor,
Expand Down Expand Up @@ -112,8 +102,4 @@ def gemm_a16w16_opus(
"gemm_a16w16_opus",
"opus_bmm",
"opus_gemm",
"pa_mqa_logits_mxfp4_build_sched",
"pa_mqa_logits_mxfp4_sched",
"pa_mqa_logits_mxfp4_sched_buffer_ints",
"pa_mqa_logits_mxfp4_sched_slots",
]
782 changes: 345 additions & 437 deletions aiter/ops/opus/pa_mqa_logits_mxfp4.py
View file
Open in desktop

Large diffs are not rendered by default.

265 changes: 0 additions & 265 deletions aiter/ops/opus/pa_mqa_logits_opus.py
View file
Open in desktop

This file was deleted.

Loading
Loading