You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
gfx950 (#5332) and gfx1250 (#5656) shipped the same operator as two parallel implementations: two
JIT modules, two schedule builders, two Python surfaces and two op tests. Every fix had to land
twice. Two of them had already diverged into the same kind of hole: a compile-time stride sitting
behind a check that looked only at the rank. This PR folds both into one module, one schedule, one
API and one test. The device bodies keep their math.
Technical Details
One JIT module.module_pa_mqa_logits_mxfp4_opus carries both device bodies. The sources live
in csrc/kernels/opus_mqa_logits/pa_mqa_logits_mxfp4/, which leaves room for an fp8 sibling:
file
contents
pa_mqa_logits_mxfp4_opus.h
public ABI, kargs, CTA record, both arches' traits
pa_mqa_logits_mxfp4_sched.cuh
the shared, arch-agnostic schedule builder
pa_mqa_logits_mxfp4_gfx950.cuh
gfx950 body: MFMA 32x32x64, wave64
pa_mqa_logits_mxfp4_gfx1250.cuh
gfx1250 body: WMMA 32x16x128, wave32, TDM
pa_mqa_logits_mxfp4_kernels.cu
shape checks, launcher, runtime-arch dispatch
Each body compiles only for its own target (__gfx950__ / __gfx1250__) and compiles to an empty
stub on every other pass, so one TU links both launch symbols. The fat binary is built per --offload-arch, the idiom opus_gemm already uses. Namespaces: opus_logits holds the shared
schedule. opus_logits::gfx950 and opus_logits::gfx1250 each hold their arch's pa_mqa_logits_mxfp4_traits and pa_mqa_logits_mxfp4_kernel. The four per-arch source files and
the separate module_pa_mqa_logits_mxfp4_gfx1250_opus are deleted.
One schedule. gfx1250's two-phase build_tiles + build_sched is now the only builder, and
gfx950's per-row builder is retired. gfx950 runs at q_per_block == 1. There a tile is one row,
and the tile-union window is that row's own, so the shared builder emits exactly the per-row
record the MFMA kernel already reads.
The tile cut is skipped at q_per_block == 1. At that setting the cut is the identity (tile t is row t), so build_tiles is not launched. That is one kernel fewer per forward on
every gfx950 call and on gfx1250's MTP=1 decode instance, and plan_buffers no longer allocates a cu_tiles buffer that nothing reads.
One runtime-arch dispatch.pa_mqa_logits_mxfp4_fwd_sched reads the arch once per process,
then dispatches on (q_per_block, block_k). Any combination it does not recognise raises; none
falls back to a default. Mixed-arch nodes (one process driving both gfx950 and gfx1250) are out of
scope, as with opus_get_arch_info().
Host-visible bounds, on both arches.
The plan's row count (num_rows) goes down to C++. There it is checked against q, weights
and out, and against the lengths of local_ends, local_starts and row_to_batch.
block_tables must be at least as wide as the KV tiles need. A tile is 64 or 256 tokens, and a
CTA reads every page of its last tile. pa_mqa_logits_mxfp4_block_table_width returns the
required width.
Both kernels drop any schedule record with row_id >= num_rows. This test is uniform across
the CTA, so a mismatched cta_info loses rows instead of reading or writing out of bounds.
Caller contract (module docstring and header):
Within a tile, the window is non-decreasing in row index.
A write stays within the row's window.
A row that belongs to no live sequence has an empty window, and only such a row may have a
negative row_to_batch.
At q_per_block > 1, cu_seq_q[batch] <= total_q, because the tile cut reads cu_seq_q on
device. With row_to_batch=None, a tile uses its first row's block_tables row, so the rows
of a sequence must share it.
At q_per_block == 1, cu_seq_q does not trim rows: every row of total_q is scheduled.
Condition 3 is what keeps pad rows inert. ATOM already meets it: its decode pad rows carry row_to_batch = -1 and local_ends = 0.
One Python surface.aiter/ops/opus/pa_mqa_logits_mxfp4.py exposes plan_buffers / plan / pa_mqa_logits_mxfp4, parameterized by a per-arch MqaLogitsVariant(name, q_per_block, block_k, cta_resident, arch):
arch
instance
q_per_block
block_k
cta_resident
gfx1250
qlen4_kv64 (default)
4
64
768
gfx1250
qlen1_kv64
1
64
3072
gfx950
qlen1_kv64 (default)
1
64
1024
gfx950
qlen1_kv256
1
256
1024
The gfx950 cta_resident is #5332's SCHED_CTA_TARGET = 1024, used for both widths.
One op test, op_tests/test_pa_mqa_logits_mxfp4_opus.py, runs on both arches. It builds
natural layouts on gfx1250 and MFMA-permuted layouts on gfx950.
-mllvm -enable-post-misched=1 applies to gfx950 only. It is load-bearing there: aiter's base
JIT flags set it to 0, and without =1 the tuned MFMA body is rescheduled. The gfx1250 body keeps
the base =0, as its own module on main did. The flag is a per-arch expression in optCompilerConfig.json, the form module_deepgemm_opus already uses. The gfx1250 A/B that
decided this is below. The expression is true whenever the build targets gfx950 (the running GPU or GPU_ARCHS), and
compiler flags are per TU. So a single build for both gfx950;gfx1250 would still give the
gfx1250 body =1; aiter's release default (gfx942;gfx950) does not build gfx1250.
kargs loses six fields that no kernel read (144 → 112 B). The offsetof pins are re-derived.
API changes
Retired (gfx950, [HIP] [OPUS] [JIT] [GFX950]OPUS PA MQA Logits MXFP4 #5332):aiter.ops.opus.pa_mqa_logits_mxfp4_sched, _build_sched, _sched_slots and _sched_buffer_ints. Before this PR they were also reachable as from aiter import .... They have no callers: gfx950 users, ATOM included, still run FlyDSL. They
are replaced by the shared plan_buffers / plan / pa_mqa_logits_mxfp4. The cta_target / block_k knobs become MqaLogitsVariant, chosen by name.
Kept (gfx1250, [HIP] [OPUS] [JIT] [GFX1250]OPUS PA MQA Logits MXFP4 #5656): all nine public names in aiter.ops.opus.pa_mqa_logits_mxfp4, and
the signatures of pa_mqa_logits_mxfp4, _plan_buffers and _variants. The only removal is plan(block_k=...) together with MqaLogitsPlan.block_k, which was gfx950-only and on gfx1250
raised for anything but None. The additions are backward-compatible: trailing MqaLogitsVariant.arch and MqaLogitsPlan.num_rows, and a device= keyword on pa_mqa_logits_mxfp4_block_table_width.
Test Plan
Run the unified op test on both arches. It covers:
corner cases on every compiled instance;
ATOM-shaped padded decode (MTP 1 and 2, with and without padding);
host-side raises on row-count mismatches;
the kernel row guard, driven through a raw launch;
NaN propagation from the E8M0 scales, including a zero-NaN control;
the four perf tables.
On gfx950, cross-check every in-window cell against FlyDSL, with a negative control.
Compare the kernel ISA and resources against the standalone modules on main.
Compare kernel time against main, same card, alternating runs.
A/B -mllvm -enable-post-misched=1 on gfx1250, including VGPR and occupancy.
Test Result
gfx1250 (ROCm 10.1.0, HIP 7.16.26315, AMD clang 23.0.0git d177931e), head dc13ff647, module rebuilt clean
Re-measured at the final head. The previous table in this PR was taken at 79463da, before the
review round; its values are in brackets for comparison.
check
result
op test, both instances
corner 38/38, 19 per instance (now includes ATOM decode MTP 1/2, host raises, raw row guard, qlen4 cu_seq_q over-claim), NaN E8M0 8/8, err = 0 on all 27 perf rows [was: 28/28, 6/6, 50 rows]
kernel resources, qlen4_kv64 / qlen1_kv64
269 VGPR / occupancy 3 / no spill / LDS 9728 on both, kernarg 112 B [was: 269 / 3 / 0 / 9728 on both]
plan-stage launches
qlen1_kv64: 2 → 1 (small), 3 → 2 (split); qlen4_kv64 unchanged (structural, from the code)
plan-stage cost
9.1–9.2 µs saved per forward on MTP=1 decode, 12.3–19.8 µs on prefill [was: 8.7–9.1 / 12–20 µs]
22 shared shapes, 16 alternating rounds: geomean head/main 0.9979, per-shape min ratio 0.960–1.020. Causal 1.000, fresh 1.001, chunked 0.994, decode 0.997 [was: geomean 0.9982 over 6 vs 16 runs]
post-misched on gfx1250: off. With =1 the flag took effect (it is appended after the base =0, and the last one wins) as a pure post-RA reschedule. Resources were identical on and off:
269 VGPR / occupancy 3 / no spill / LDS 9728 on both instances. The reordered code had fewer
instructions (qlen4_kv64 1623 vs 1641: more VOPD, fewer s_delay_alu and nops).
Timing: 16 paired alternating runs, done twice, min-of-16 per shape.
Prefill causal bs1–8 and fresh: 1.0–3.1 % faster with the flag on, 14–16/16 pairs.
Chunked bs4 (kvlen 25000): 3.6–5.5 % slower with the flag on (0.9–1.9 % by paired median),
11–12/16 pairs, in both runs and on both instances. The 79463da data already had it (median
1160.8 vs 1127.0 µs, 12/16 pairs), but there the min ratio (0.998) hid it.
Decode: no direction.
A hand-scheduled body should not carry a flag that regresses a production shape, so dc13ff647
gates the flag on gfx950. The gfx1250 ISA is now line-for-line the flag-off build, with the same
flags as main's gfx1250 module. The cost is the prefill gain above: a third 16-pair run of the
final head against the flag-on build measures prefill 0.9–2.7 % slower, chunked bs4 no longer
slower, and a geomean of 1.0095 over 27 shapes.
gfx950 (ROCm 7.2.4), head dc13ff647, module rebuilt clean
gfx950 now runs through the unified module, schedule and API. Every call runs at q_per_block == 1. qlen1_kv64 is the production instance. qlen1_kv256 is compiled, and it is
tested for both correctness and perf.
The median is +0.14 % and no shape moves by 5 % or more. Fresh bs1/bs2 and the ragged
m1s128L16 / m4s32L16 / m4s32L4 shapes come out +1.6 to +3.4 % in both rounds. Nothing in the
review round moved them, and the long prefill (chunked) is about 1 % faster.
…hedule/API/test
Collapse the two parallel implementations into one: a single JIT module with both
device bodies arch-guarded by version macros, one shared per-tile schedule
(gfx950 runs it at q_per_block=1), one runtime-arch fwd dispatch, one Python
surface, and one test covering both arches. The gfx950 MFMA body is unchanged.
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 5761 --add-label <label>
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 opt this PR out.
github-actionsBot
changed the title
[HIP] [OPUS] Unify gfx950+gfx1250 MXFP4 MQA-logits
[HIP] [OPUS] [JIT] Unify gfx950+gfx1250 MXFP4 MQA-logits
Sep 22, 2026
…onvention
gfx950's two instances are q_per_block=1, so rename mfma_kv64/mfma_kv256 to
qlen1_kv64/qlen1_kv256. gfx950's qlen1_kv64 shares the name with gfx1250's (same
q_per_block/block_k); the per-arch variant tables keep them distinct.
Use runtime GPU architecture for test buffer layout
op_tests/test_pa_mqa_logits_mxfp4_opus.py:203
This test selects the buffer layout from get_gfx(), which is the build-time/GPU_ARCHS value, while the dispatcher selects the kernel and validates layouts from the tensor's runtime device. With a multi-arch GPU_ARCHS setting (or a non-current device), _is_permuted() can build natural inputs for gfx950 or permuted inputs for gfx1250, making the test fail or exercise the wrong ABI. Use the same runtime-device arch probe as the operation for both the gate and layout construction.
[verified] On gfx950 every compiled variant has q_per_block == 1, so _compute_tiles only needs an identity row boundary, but this path still launches _build_tiles_raw before _build_sched_raw. The deleted gfx950 implementation built the schedule directly from local_ends; this adds a GPU launch and a batch-prefix scan to every forward's metadata setup, which is especially visible for short decode where the layer kernel is small. Author must bypass or fold the q-per-block-1 tile cut, or provide a measured justification for the added per-forward launch.
Use runtime architecture detection for variant selection
op_tests/test_pa_mqa_logits_mxfp4_opus.py:203
[verified] _is_permuted() uses get_gfx(), but that helper is build-time/env-driven (GPU_ARCHS) rather than the runtime device probe. With a multi-arch setting such as GPU_ARCHS=gfx942;gfx950 on a gfx942 runner, it can select the gfx950 layout/variants for the actual gfx942 device, causing the test to build the wrong ABI or fail before exercising the target. Author must use the runtime arch probe consistently for the gate, layout builder, and variant lookup.
Use runtime GPU architecture for permutation layout selection
op_tests/test_pa_mqa_logits_mxfp4_opus.py:203
[verified] _is_permuted() uses get_gfx(), which is a build-time probe and deliberately selects the last entry for GPU_ARCHS values such as gfx950;gfx1250; this test is running on the live GPU and must choose the layout for that device. With that common multi-arch setting on a gfx950 machine, the test builds natural gfx1250 layouts and the wrapper rejects them or would exercise the wrong path. Author must use the runtime/device-specific arch probe for both the layout choice and the arch gate.
The reason will be displayed to describe this comment to others. Learn more.
Reviewed at head of lxx/dev/pa_mqa_logits_mxfp4_gfx950_unify against merge-base 25ae2bfa3e, in isolated worktrees.
The motivation for unifying is well argued, and the surface this PR exposes is clean — I checked that the three defs in rocm_ops.hpp match the header declarations and the Python @compile_ops stubs positionally, and that no stale reference to the deleted module, the deleted pybind macro, or the removed Python names survives anywhere in the tree.
My concerns are that the unification stops at one file rather than one module, that it adds host-side cost to a path that runs 61 times per forward, and that several guards only survived the merge in one direction. Taking those in that order.
Structure
The output is a 2061-line header — 2.5x the repo's 800-line soft ceiling — whose two arch bodies share no code. Measured layout: shared ABI and schedule at L1-655 (~30%), gfx950 traits L657-791, gfx1250 traits L792-1109, the gfx950 wave64 body L1111-1656, the gfx1250 wave32/TDM body L1658-2059. So 68% is arch-specific, sitting in disjoint #if islands that each re-declare their own pin_sgpr (:1156 vs :1692) and permlane_head_reduce (:1164 vs :1706). The file now also carries three spellings of a workgroup barrier: qshare::cta_sync, sched_detail::sync_block, and opus.hpp's own.
In kernels.cu, check_shapes_mfma (L37-121) and check_shapes_wmma (L190-296) are ~70 near-identical lines differing only in the trait constant names for the same quantity, and launch_sched_mfma / launch_sched_wmma repeat the same eleven kargs assignments and the same launch tail. The block_tables width gap below is the drift this duplication has already produced.
Some of what the merge carried over is dead: opus_logits::sched_slots and max_tiles_for have no C++ caller (Python re-implements both), ceil_div_i has no caller, logits_fp4_qshare_traits_4q is unused, and four kargs fields (split_kv, num_batches, ptr_cu_seq_q, ptr_row_to_batch) are pinned at fixed offsets by a 22-line offsetof block without being read.
The naming also inverts: the module-generic names — opus_mqa_logits_fp4_traits, pa_mqa_logits_mxfp4_kernel — are the gfx950-only ones, while the gfx1250 twins carry the qualifier.
A three-file split (_opus.h for ABI and schedule, _mfma_body.h, _wmma_body.h) plus one templated check_shapes and launch_sched keeps every benefit of the merge and narrows the JIT rebuild blast radius — as it stands, touching the gfx950 body forces a rebuild of the 400 lines of hand-scheduled gfx1250 TDM code.
Performance
const std::string arch = get_gpu_arch(); now runs on everyfwd_sched call. get_gpu_arch (aiter_hip_common.h:537) does no caching: a hipGetDeviceCount, a hipGetDevice, a hipGetDeviceProperties, then a std::string construct plus find plus substr. The outer HipDeviceGuard also duplicates one the inner launchers already take, and HipDeviceGuard (:643-648) has no prev == target fast path. Per-layer host cost goes from three HIP device-management calls to six, plus a full hipDeviceProp_t fill, which takes a runtime lock.
Neither base TU called get_gpu_arch at all — the probe is new with the unification, and both arches pay it. The module's own docstring budgets this op at 61 launches per DeepSeek-V4 forward against ~16 us of host dispatch, so I think this more than cancels the one build_tiles launch the PR advertises saving. It's also invisible under cudagraph capture, so a graphed A/B won't surface it. The repo idiom is one directory over — fmha_fwd_bf16_opus_launch.h:280 does static const std::string arch_id = get_gpu_arch(); — though here it wants a per-device cache, since the guard keys on q.device_id.
The gfx950 grid policy changes twice, with no measurement behind either. _sched_slots now rounds the grid up to a multiple of cta_resident: base was max(num_rows, 1024) with no round-up, HEAD is ceil(max(n,r)/r)*r. At total_q=1025 that launches 2048 CTAs instead of 1025 — 1023 surplus CTAs, each dispatching and issuing a 32-byte s_load_dwordx8, 61 times per forward, with cta_info allocated for the padded count too. Separately qlen1_kv256 ships cta_resident=512, halving base's SCHED_CTA_TARGET=1024 split aim, so settle_and_emit gives each CTA roughly twice the chunks on the long-window prefill regime that instance exists to serve.
Both values are self-labelled "UNTUNED PLACEHOLDERS" at py:158-161, while the class docstring two screens up (:99-100) says cta_resident is "TUNED per instance". And because decode_shapes names only qlen1_kv64/None and _resolve_variant falls back to the arch default, qlen1_kv256 appears in no perf table on either arch — a large load-balance regression wouldn't show up in any number this PR reports.
One compile-side item: folding the gfx1250 sources into module_pa_mqa_logits_mxfp4_opus silently applies -mllvm -enable-post-misched=1 to the hand-scheduled WMMA body. The deleted module_pa_mqa_logits_mxfp4_gfx1250_opus entry carried exactly ['-ffast-math', '-fno-finite-math-only'], and -mllvm flags are per-TU. That body is explicitly hand-scheduled — __builtin_amdgcn_sched_barrier(0) around every store (:1981-1983, :2036), the two-accumulator ping-pong (:1831-1836), a TDM pipeline sized by TENSORCNT_KEEP (:1072), pin_sgpr inline asm. Post-RA rescheduling can reorder the TDM issue/wait pairing or raise pressure past the occupancy the hardcoded cta_resident=768/3072 were swept for. The PR does report a gfx1250 A/B for the flag, but not gfx1250 VGPR/occupancy, which is what cta_resident depends on. The same JSON already gates this exact flag on get_gfx()=='gfx1250' around line 354, so a per-arch expression is available.
Cleanliness
general_emit's new even split does const int nc = (n + safe - 1) / safe; then const int q = n / nc, r = n % nc; — a device-side divide by zero whenever n == 0 — and that nc shadows the outer nc at :416 with a different value. The inline comment asserts ">= 1: a zero-chunk tile emits no slot", but that property isn't local: it holds only because the upper_bound(s_excl, j) - 1 search at :431-436 resolves a run of equal prefix values to the last index, so a zero-nc tile can only be selected if it's the block's trailing tile, which j < block_total already excludes. Switch that search to lower_bound, hoist it, or make the scan inclusive, and the assumption breaks — and integer divide by zero on AMD is an undefined result rather than a fault, so the table silently gets garbage chunk_start/chunk_count. A n > 0 ? ... : 1 guard and a rename for the shadowing local would close it.
The cross-arch buffer guard if buffers.variant not in _variants_for(arch) rests entirely on frozen-dataclass value equality, and the only field separating gfx950's qlen1_kv64 from gfx1250's is cta_resident — 1024 vs 3072. MqaLogitsVariant (:89) has no arch field, and the comment at :158-161 asks someone to measure the gfx950 numbers. The moment that tuning lands on 3072 the two records compare equal, and both this guard and _as_variant (:494-500) stop rejecting cross-arch buffers; the only backstop left is _check_layout's ndim sniff, which the module itself documents cannot catch an array reshaped to the right ndim. A pure perf edit would silently disable a correctness guard. An arch field, or keying the guard on arch, avoids that.
tile_kv_tiles_cached reads cta_info[t].chunk_count back through a const ... * __restrict__ pointer that emit_fast wrote through a different__restrict__ pointer in the same kernel. Both helpers inline into mqa_logits_build_sched (:545), LLVM emits !noalias/!alias.scope per inlined call site, and the two restrict-qualified parameters land in disjoint scopes — formally telling the optimizer the load in settle_safe cannot see the stores in emit_fast. Only the s_barrier inside block_reduce_stats separates them. Base's settle_safe recomputed the count from local_starts/local_ends and never read the table back. If the load is hoisted or forwarded, ctas_for() sums garbage and general_emit emits a schedule that either overruns num_ctas (tiles dropped, their rows keeping the caller's -inf) or under-splits. Dropping __restrict__ on the read-back overload, or carrying the counts through LDS, removes the hazard.
pa_mqa_logits_mxfp4_block_table_width now calls _device_arch(torch.cuda.current_device() ...) unconditionally — including when the caller passed explicit MqaLogitsVariant objects that need no lookup. Base's docstring said verbatim "Pure arithmetic -- no GPU, no JIT build -- so it is safe to call at startup"; HEAD downgrades that to "Past that arch probe this does no GPU work". A serving stack sizing block_tables at config time, on a CPU-only rank, under CUDA_VISIBLE_DEVICES='', or inside a fork-based dataloader worker, now either raises or initializes a CUDA context in a process that was deliberately keeping it lazy. The arch is only needed on the variant-is-None and by-name branches.
_is_permuted() (:198) and _variants() (:400) are @functools.cache on zero-argument functions that probe the current device, so the first device touched permanently fixes the scale layout and the variant list for the process.
Two places where the gfx950 arm still drifts from the shared ABI the unification is meant to establish. Its kernel template defaults SCHED = mqa_logits_sched::Prefill (:1114 and :1299), an arm whose kargs fields the MFMA launcher never assigns, while gfx1250 defaults to Table and static_asserts it (:1716). kernels.cu:176 passes Table explicitly so nothing is broken today, but any new call site omitting the argument instantiates the Prefill arm and reads ptr_row_to_batch / ptr_local_starts / num_rows, all left zero-initialized by launch_sched_mfma (:154-170) — a NULL deref, not an error. Same shape of trap for the kargs invariant at header:117 ("ptr_local_ends may not [be NULL], being the store's upper bound"): it's false on the gfx950 arm, so the next gfx950 edit that adopts gfx1250's per-row store mask faults with the header telling the author it can't. Defaulting SCHED to Table (or dropping the parameter) and either assigning the pointers or scoping the invariant per arm fixes both.
build_sched validates q_per_block >= 1 and block_k >= 1 but dropped base gfx950's AITER_CHECK(cta_target >= 1), and the Python mirror isn't equivalent: C++ sched_target/settle_and_emit treat resident <= 0 as "do not split" and accept it silently, while _sched_slots (py:299) computes -(-n // r) * r, which raises ZeroDivisionError at r == 0 and returns a negative grid at r < 0 that _cta_info then hands to torch.empty as a first dimension.
Correctness
Three device-array length guards did not survive the merge. Base had AITER_CHECK(local_ends.numel() >= num_rows) plus the same for local_starts and row_to_batch (base kernels.cu:190, 217, 225); HEAD has none of the three, only dtype and contiguity. At the identity cut (q_per_block == 1, i.e. every gfx950 launch) num_tilesis the row count, and emit_fast reads local_ends[t] and row_to_batch[r0] for every t < num_tiles. So pa_mqa_logits_mxfp4_plan(cu_seq_q, local_ends=<1024 entries>, total_q=4096) gives num_tiles=4096 and reads local_ends[1024..4095] past the allocation; a short row_to_batch gives an arbitrary batch_id, hence the wrong block_tables row and garbage page ids. The compensating Python check at pa_mqa_logits_mxfp4.py:762-773 compares num_tiles against q.shape[0], never against local_ends.numel(), and it runs at launch — after the builder has already executed.
The gfx950 arm has neither of the two bounds its gfx1250 twin gained. check_shapes_wmma computes bt_cols = ceil(max_seq_len/KV_TILE_SIZE) * PAGES_PER_TILE and raises below it (:249-254); check_shapes_mfma has only AITER_CHECK(block_tables.dim() == 2, ...) at :51 and then takes block_tables.size(1) as an unvalidated row stride at :164, while the kernel indexes g_bt at (chunk_start + tile_idx) * pages_per_tile (header:1462) with nothing else bounding the column. This PR is what makes that reachable, since it's the first to expose the gfx950 qlen1_kv256 instance (4 pages per tile) through the shared variant API: a caller sizing block_tables by ceil(max_seq_len/64) page rounding and running qlen1_kv256 reads up to three int32 past each row. Three Python docstrings now assert the opposite for both arches — :58 ("The C++ header states the resulting bound", where the paragraph that stated it was deleted with the gfx1250 header), :520 ("the C++ launcher raises below it"), and :481 ("They all take ... the same block_tables width"), which gfx950's 64-vs-256 pair contradicts 40 lines later at :522.
Likewise the Table arm's if(rec.chunk_count <= 0) return; has no num_rows bound where gfx1250 added || rec.row_id >= num_rows, and kargs.num_rows is write-only on the MFMA arm (set at kernels.cu:164, never read). pa_mqa_logits_mxfp4_fwd_sched is exported through PA_MQA_LOGITS_MXFP4_PYBIND and reachable via _fwd_sched_raw, so a schedule built for total_q=4096 launched against q.shape[0]==1024 gives row_id up to 4095 indexing q_base = p_q + row_id*Q_ROW_BYTES (4 KB/row) and out + row_id*stride_out — reads and writes ~12 MB past both allocations. That bound belongs in the kernel next to the gfx1250 one rather than in one Python wrapper.
Finally, gfx1250's pre-existing qlen1_kv64 instance quietly stops honouring cu_seq_q. Base ran _gfx1250_compute_tiles for both gfx1250 variants, and mqa_logits_build_tiles writes cu_tiles[t] = cu_seq_q[batch] for every t past the real tile count, so padding tiles got r0 == r1 → tile_union_end() == 0 → chunk_count = 0 → a surplus slot. HEAD takes num_tiles = _max_tiles_for(n, batch, v) = total_q and gives tile t the real row t. A cudagraph caller that pins q/local_ends/out at capacity 256 but has cu_seq_q=[0,...,100] this forward: rows 100..255 carry stale non-zero local_ends from a previous replay, HEAD schedules them, reads q[100..255] and block_tables[row_to_batch[r]] (row 255 against a 100-row table), and writes into out[100..255] that base left at -inf. The new docstring documents the contract change ("the row domain is total_q, not cu_seq_q[batch]"), but nothing checks it and the PR body doesn't mention the gfx1250 behaviour change.
Tests and exports
The deleted op_tests/test_pa_mqa_logits_opus.py cross-checked gfx950 against an independent flydsl_prefill/flydsl_decode implementation, and nothing replaces it — while the surviving file's docstring asserts "There is no second implementation to cross-check against on this target", a gfx1250-only claim now stated in a file that runs on gfx950, where aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4*.py still exist. More importantly, the HEAD reference and the kernel now consume the same _scale_to_opus permutation (:206), so a coordinated MFMA scale-permutation error — the exact silent failure the module docstring says "is caught HERE and nowhere else" — passes.
check_nan_scale('prefill', 1, 4, [0,1,2,3], kv_row=2) is labelled "the control that says the NaN is not smeared", but row 3's window [0,3) contains column 2, so want.sum() == 1, not 0, and the le == 0 row asserts nothing; its precondition is a bare assert that -O strips. And check_row_id_bound is skipped entirely on gfx950 (:686), so the PR's new gfx950-only host row bound (py:762-773) is untested.
aiter/ops/opus/__init__.py drops pa_mqa_logits_mxfp4_sched / _build_sched / _sched_slots / _sched_buffer_ints from both imports and __all__ and adds nothing back, so the op has no package-level or top-level aiter.* export at all. In-tree grep is clean and ATOM already uses the submodule path, but an out-of-tree gfx950 caller breaks with an error pointing at no successor, and the retired cta_target / block_k knobs have no supported replacement.
Summary
I'd like to see the header split before this lands — one file with 68% arch-specific content in disjoint #if islands and two 70-line copies of the same shape check is the structure the merge was supposed to remove, not produce. On performance, the uncached get_gpu_arch on a 61-launches-per-forward path likely costs more than the build_tiles launch it saves, and the two gfx950 grid changes should carry a measurement, ideally on an instance that some perf table actually exercises. And the four guards that only survived in one direction — the three length checks, the block_tables width bound, the row_id bound, and gfx1250's cu_seq_q contract — are the ones I'd treat as blocking, since each one turns a host-side raise into an unbounded device read.
…0 row guard, trimmed ABI
Launcher and ABI
- One check_shapes/launch_sched template for both arches (traits share
KV_PAGE_BYTES / KVS_PAGE_BYTES / PAGES_PER_TILE / READS_ROW_WINDOWS).
- The plan's row count (num_rows) reaches build_sched and fwd_sched (private
ABI); every per-row array is checked against it at any q_per_block, and
fwd_sched checks it against q / weights / out on both arches. Replaces the
num_tiles-based length checks, which used the wrong unit above
q_per_block == 1, and the gfx950-only Python row bound.
- Arch cached per device (SynchronizedCache); a single device guard per launch.
- build_sched requires cta_resident >= 1.
- Drop the six kargs fields no kernel reads (144 -> 112 B).
Kernels
- gfx950: CTA-uniform `row_id >= num_rows` guard, as gfx1250 has; SCHED
defaults to Table and is static_asserted; the unused Prefill/Decode arms are
removed; traits and kernel renamed *_mfma_*.
Schedule builder
- Skip the row_to_batch load on empty tiles (read one past the end on surplus
tiles); guard general_emit's divide; drop __restrict__ on the cta_info
read-back; __syncthreads() instead of a hand-written copy; remove dead helpers.
Tables are byte-identical for qlen1_kv64.
Python
- MqaLogitsVariant gains a trailing, defaulted `arch`; cross-arch guards compare
it instead of relying on value equality. block_table_width does no device
probe for explicit variant objects. MqaLogitsPlan gains a defaulted num_rows.
- qlen1_kv256 cta_resident back to 1024.
- Document the empty-window contract for rows outside any live sequence, and
that `out` is only written inside each row's window.
Tests
- FlyDSL cross-check on gfx950, a real NaN control, host-raise and raw row-guard
probes on both arches, ATOM-shaped padded decode, kv256 perf rows.
Comments condensed throughout; no code change from that part (ISA identical).
Validated on gfx950: corner 36/36, NaN 8/8, err 0 on all perf tables; MFMA ISA
unchanged apart from the kargs and row-guard prologue; perf vs pre-PR main within
+/-3% on #5332 and #5656 shapes.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
- fwd_sched caches the arch in a function-local static (the
opus_get_arch_info() model) instead of a per-device SynchronizedCache.
Mixed-arch nodes are not supported; the comment says so.
- Ruff F821: check_row_count_raises rebinds `inp` instead of `del`, since
the lambdas above close over it.
- Black: reformat pa_mqa_logits_mxfp4.py and the op test (shorten two
comments that Black would otherwise wrap awkwardly).
gfx950: op test passes (corner 36/36, NaN 8/8, err 0); ISA of all 9
kernels byte-identical to the previous build.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
…mqa_logits_mxfp4/
Pure move, no content change: the parent directory leaves room for an
fp8 sibling. Only optCompilerConfig.json's two srcs referenced the path;
the files include each other by bare name.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Each arch's traits and device body now live in opus_logits::gfx950 /
opus_logits::gfx1250 under the same generic names,
pa_mqa_logits_mxfp4_traits and pa_mqa_logits_mxfp4_kernel. The shared
schedule stays in opus_logits. This replaces the gfx950 body sitting
directly in opus_logits (opus_mqa_logits_fp4_mfma_traits,
pa_mqa_logits_mxfp4_mfma_kernel) and gfx1250's opus_logits::qshare
(opus_mqa_logits_fp4_qshare_traits,
mqa_logits_mxfp4_32x16x128_qshare_kernel). The header's opening comment
states the layout. Kernel symbol names change in profiles.
Also drop the gfx950 body's `if (tile_count <= 0) return;`, which can
never fire: the record check above already returns on chunk_count <= 0.
gfx950: op test passes (corner 36/36, NaN 8/8, err 0). The ISA is
unchanged apart from kernel symbol names and LLVM block-label comments.
Perf A/B vs 6f2fb13 (3 alternating rounds, 56 shapes): median +0.07%
(kv64) and +0.17% (kv256), and every out fingerprint is identical.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Reviewed at head of lxx/dev/pa_mqa_logits_mxfp4_gfx950_unify against merge-base 25ae2bfa3e, in isolated worktrees.
The motivation for unifying is well argued, and the surface this PR exposes is clean — I checked that the three defs in rocm_ops.hpp match the header declarations and the Python @compile_ops stubs positionally, and that no stale reference to the deleted module, the deleted pybind macro, or the removed Python names survives anywhere in the tree.
My concerns are that the unification stops at one file rather than one module, that it adds host-side cost to a path that runs 61 times per forward, and that several guards only survived the merge in one direction. Taking those in that order.
Structure
The output is a 2061-line header — 2.5x the repo's 800-line soft ceiling — whose two arch bodies share no code. Measured layout: shared ABI and schedule at L1-655 (~30%), gfx950 traits L657-791, gfx1250 traits L792-1109, the gfx950 wave64 body L1111-1656, the gfx1250 wave32/TDM body L1658-2059. So 68% is arch-specific, sitting in disjoint #if islands that each re-declare their own pin_sgpr (:1156 vs :1692) and permlane_head_reduce (:1164 vs :1706). The file now also carries three spellings of a workgroup barrier: qshare::cta_sync, sched_detail::sync_block, and opus.hpp's own.
In kernels.cu, check_shapes_mfma (L37-121) and check_shapes_wmma (L190-296) are ~70 near-identical lines differing only in the trait constant names for the same quantity, and launch_sched_mfma / launch_sched_wmma repeat the same eleven kargs assignments and the same launch tail. The block_tables width gap below is the drift this duplication has already produced.
Some of what the merge carried over is dead: opus_logits::sched_slots and max_tiles_for have no C++ caller (Python re-implements both), ceil_div_i has no caller, logits_fp4_qshare_traits_4q is unused, and four kargs fields (split_kv, num_batches, ptr_cu_seq_q, ptr_row_to_batch) are pinned at fixed offsets by a 22-line offsetof block without being read.
The naming also inverts: the module-generic names — opus_mqa_logits_fp4_traits, pa_mqa_logits_mxfp4_kernel — are the gfx950-only ones, while the gfx1250 twins carry the qualifier.
A three-file split (_opus.h for ABI and schedule, _mfma_body.h, _wmma_body.h) plus one templated check_shapes and launch_sched keeps every benefit of the merge and narrows the JIT rebuild blast radius — as it stands, touching the gfx950 body forces a rebuild of the 400 lines of hand-scheduled gfx1250 TDM code.
Performance
const std::string arch = get_gpu_arch(); now runs on everyfwd_sched call. get_gpu_arch (aiter_hip_common.h:537) does no caching: a hipGetDeviceCount, a hipGetDevice, a hipGetDeviceProperties, then a std::string construct plus find plus substr. The outer HipDeviceGuard also duplicates one the inner launchers already take, and HipDeviceGuard (:643-648) has no prev == target fast path. Per-layer host cost goes from three HIP device-management calls to six, plus a full hipDeviceProp_t fill, which takes a runtime lock.
Neither base TU called get_gpu_arch at all — the probe is new with the unification, and both arches pay it. The module's own docstring budgets this op at 61 launches per DeepSeek-V4 forward against ~16 us of host dispatch, so I think this more than cancels the one build_tiles launch the PR advertises saving. It's also invisible under cudagraph capture, so a graphed A/B won't surface it. The repo idiom is one directory over — fmha_fwd_bf16_opus_launch.h:280 does static const std::string arch_id = get_gpu_arch(); — though here it wants a per-device cache, since the guard keys on q.device_id.
The gfx950 grid policy changes twice, with no measurement behind either. _sched_slots now rounds the grid up to a multiple of cta_resident: base was max(num_rows, 1024) with no round-up, HEAD is ceil(max(n,r)/r)*r. At total_q=1025 that launches 2048 CTAs instead of 1025 — 1023 surplus CTAs, each dispatching and issuing a 32-byte s_load_dwordx8, 61 times per forward, with cta_info allocated for the padded count too. Separately qlen1_kv256 ships cta_resident=512, halving base's SCHED_CTA_TARGET=1024 split aim, so settle_and_emit gives each CTA roughly twice the chunks on the long-window prefill regime that instance exists to serve.
Both values are self-labelled "UNTUNED PLACEHOLDERS" at py:158-161, while the class docstring two screens up (:99-100) says cta_resident is "TUNED per instance". And because decode_shapes names only qlen1_kv64/None and _resolve_variant falls back to the arch default, qlen1_kv256 appears in no perf table on either arch — a large load-balance regression wouldn't show up in any number this PR reports.
One compile-side item: folding the gfx1250 sources into module_pa_mqa_logits_mxfp4_opus silently applies -mllvm -enable-post-misched=1 to the hand-scheduled WMMA body. The deleted module_pa_mqa_logits_mxfp4_gfx1250_opus entry carried exactly ['-ffast-math', '-fno-finite-math-only'], and -mllvm flags are per-TU. That body is explicitly hand-scheduled — __builtin_amdgcn_sched_barrier(0) around every store (:1981-1983, :2036), the two-accumulator ping-pong (:1831-1836), a TDM pipeline sized by TENSORCNT_KEEP (:1072), pin_sgpr inline asm. Post-RA rescheduling can reorder the TDM issue/wait pairing or raise pressure past the occupancy the hardcoded cta_resident=768/3072 were swept for. The PR does report a gfx1250 A/B for the flag, but not gfx1250 VGPR/occupancy, which is what cta_resident depends on. The same JSON already gates this exact flag on get_gfx()=='gfx1250' around line 354, so a per-arch expression is available.
Cleanliness
general_emit's new even split does const int nc = (n + safe - 1) / safe; then const int q = n / nc, r = n % nc; — a device-side divide by zero whenever n == 0 — and that nc shadows the outer nc at :416 with a different value. The inline comment asserts ">= 1: a zero-chunk tile emits no slot", but that property isn't local: it holds only because the upper_bound(s_excl, j) - 1 search at :431-436 resolves a run of equal prefix values to the last index, so a zero-nc tile can only be selected if it's the block's trailing tile, which j < block_total already excludes. Switch that search to lower_bound, hoist it, or make the scan inclusive, and the assumption breaks — and integer divide by zero on AMD is an undefined result rather than a fault, so the table silently gets garbage chunk_start/chunk_count. A n > 0 ? ... : 1 guard and a rename for the shadowing local would close it.
The cross-arch buffer guard if buffers.variant not in _variants_for(arch) rests entirely on frozen-dataclass value equality, and the only field separating gfx950's qlen1_kv64 from gfx1250's is cta_resident — 1024 vs 3072. MqaLogitsVariant (:89) has no arch field, and the comment at :158-161 asks someone to measure the gfx950 numbers. The moment that tuning lands on 3072 the two records compare equal, and both this guard and _as_variant (:494-500) stop rejecting cross-arch buffers; the only backstop left is _check_layout's ndim sniff, which the module itself documents cannot catch an array reshaped to the right ndim. A pure perf edit would silently disable a correctness guard. An arch field, or keying the guard on arch, avoids that.
tile_kv_tiles_cached reads cta_info[t].chunk_count back through a const ... * __restrict__ pointer that emit_fast wrote through a different__restrict__ pointer in the same kernel. Both helpers inline into mqa_logits_build_sched (:545), LLVM emits !noalias/!alias.scope per inlined call site, and the two restrict-qualified parameters land in disjoint scopes — formally telling the optimizer the load in settle_safe cannot see the stores in emit_fast. Only the s_barrier inside block_reduce_stats separates them. Base's settle_safe recomputed the count from local_starts/local_ends and never read the table back. If the load is hoisted or forwarded, ctas_for() sums garbage and general_emit emits a schedule that either overruns num_ctas (tiles dropped, their rows keeping the caller's -inf) or under-splits. Dropping __restrict__ on the read-back overload, or carrying the counts through LDS, removes the hazard.
pa_mqa_logits_mxfp4_block_table_width now calls _device_arch(torch.cuda.current_device() ...) unconditionally — including when the caller passed explicit MqaLogitsVariant objects that need no lookup. Base's docstring said verbatim "Pure arithmetic -- no GPU, no JIT build -- so it is safe to call at startup"; HEAD downgrades that to "Past that arch probe this does no GPU work". A serving stack sizing block_tables at config time, on a CPU-only rank, under CUDA_VISIBLE_DEVICES='', or inside a fork-based dataloader worker, now either raises or initializes a CUDA context in a process that was deliberately keeping it lazy. The arch is only needed on the variant-is-None and by-name branches.
_is_permuted() (:198) and _variants() (:400) are @functools.cache on zero-argument functions that probe the current device, so the first device touched permanently fixes the scale layout and the variant list for the process.
Two places where the gfx950 arm still drifts from the shared ABI the unification is meant to establish. Its kernel template defaults SCHED = mqa_logits_sched::Prefill (:1114 and :1299), an arm whose kargs fields the MFMA launcher never assigns, while gfx1250 defaults to Table and static_asserts it (:1716). kernels.cu:176 passes Table explicitly so nothing is broken today, but any new call site omitting the argument instantiates the Prefill arm and reads ptr_row_to_batch / ptr_local_starts / num_rows, all left zero-initialized by launch_sched_mfma (:154-170) — a NULL deref, not an error. Same shape of trap for the kargs invariant at header:117 ("ptr_local_ends may not [be NULL], being the store's upper bound"): it's false on the gfx950 arm, so the next gfx950 edit that adopts gfx1250's per-row store mask faults with the header telling the author it can't. Defaulting SCHED to Table (or dropping the parameter) and either assigning the pointers or scoping the invariant per arm fixes both.
build_sched validates q_per_block >= 1 and block_k >= 1 but dropped base gfx950's AITER_CHECK(cta_target >= 1), and the Python mirror isn't equivalent: C++ sched_target/settle_and_emit treat resident <= 0 as "do not split" and accept it silently, while _sched_slots (py:299) computes -(-n // r) * r, which raises ZeroDivisionError at r == 0 and returns a negative grid at r < 0 that _cta_info then hands to torch.empty as a first dimension.
Correctness
Three device-array length guards did not survive the merge. Base had AITER_CHECK(local_ends.numel() >= num_rows) plus the same for local_starts and row_to_batch (base kernels.cu:190, 217, 225); HEAD has none of the three, only dtype and contiguity. At the identity cut (q_per_block == 1, i.e. every gfx950 launch) num_tilesis the row count, and emit_fast reads local_ends[t] and row_to_batch[r0] for every t < num_tiles. So pa_mqa_logits_mxfp4_plan(cu_seq_q, local_ends=<1024 entries>, total_q=4096) gives num_tiles=4096 and reads local_ends[1024..4095] past the allocation; a short row_to_batch gives an arbitrary batch_id, hence the wrong block_tables row and garbage page ids. The compensating Python check at pa_mqa_logits_mxfp4.py:762-773 compares num_tiles against q.shape[0], never against local_ends.numel(), and it runs at launch — after the builder has already executed.
The gfx950 arm has neither of the two bounds its gfx1250 twin gained. check_shapes_wmma computes bt_cols = ceil(max_seq_len/KV_TILE_SIZE) * PAGES_PER_TILE and raises below it (:249-254); check_shapes_mfma has only AITER_CHECK(block_tables.dim() == 2, ...) at :51 and then takes block_tables.size(1) as an unvalidated row stride at :164, while the kernel indexes g_bt at (chunk_start + tile_idx) * pages_per_tile (header:1462) with nothing else bounding the column. This PR is what makes that reachable, since it's the first to expose the gfx950 qlen1_kv256 instance (4 pages per tile) through the shared variant API: a caller sizing block_tables by ceil(max_seq_len/64) page rounding and running qlen1_kv256 reads up to three int32 past each row. Three Python docstrings now assert the opposite for both arches — :58 ("The C++ header states the resulting bound", where the paragraph that stated it was deleted with the gfx1250 header), :520 ("the C++ launcher raises below it"), and :481 ("They all take ... the same block_tables width"), which gfx950's 64-vs-256 pair contradicts 40 lines later at :522.
Likewise the Table arm's if(rec.chunk_count <= 0) return; has no num_rows bound where gfx1250 added || rec.row_id >= num_rows, and kargs.num_rows is write-only on the MFMA arm (set at kernels.cu:164, never read). pa_mqa_logits_mxfp4_fwd_sched is exported through PA_MQA_LOGITS_MXFP4_PYBIND and reachable via _fwd_sched_raw, so a schedule built for total_q=4096 launched against q.shape[0]==1024 gives row_id up to 4095 indexing q_base = p_q + row_id*Q_ROW_BYTES (4 KB/row) and out + row_id*stride_out — reads and writes ~12 MB past both allocations. That bound belongs in the kernel next to the gfx1250 one rather than in one Python wrapper.
Finally, gfx1250's pre-existing qlen1_kv64 instance quietly stops honouring cu_seq_q. Base ran _gfx1250_compute_tiles for both gfx1250 variants, and mqa_logits_build_tiles writes cu_tiles[t] = cu_seq_q[batch] for every t past the real tile count, so padding tiles got r0 == r1 → tile_union_end() == 0 → chunk_count = 0 → a surplus slot. HEAD takes num_tiles = _max_tiles_for(n, batch, v) = total_q and gives tile t the real row t. A cudagraph caller that pins q/local_ends/out at capacity 256 but has cu_seq_q=[0,...,100] this forward: rows 100..255 carry stale non-zero local_ends from a previous replay, HEAD schedules them, reads q[100..255] and block_tables[row_to_batch[r]] (row 255 against a 100-row table), and writes into out[100..255] that base left at -inf. The new docstring documents the contract change ("the row domain is total_q, not cu_seq_q[batch]"), but nothing checks it and the PR body doesn't mention the gfx1250 behaviour change.
Tests and exports
The deleted op_tests/test_pa_mqa_logits_opus.py cross-checked gfx950 against an independent flydsl_prefill/flydsl_decode implementation, and nothing replaces it — while the surviving file's docstring asserts "There is no second implementation to cross-check against on this target", a gfx1250-only claim now stated in a file that runs on gfx950, where aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4*.py still exist. More importantly, the HEAD reference and the kernel now consume the same _scale_to_opus permutation (:206), so a coordinated MFMA scale-permutation error — the exact silent failure the module docstring says "is caught HERE and nowhere else" — passes.
check_nan_scale('prefill', 1, 4, [0,1,2,3], kv_row=2) is labelled "the control that says the NaN is not smeared", but row 3's window [0,3) contains column 2, so want.sum() == 1, not 0, and the le == 0 row asserts nothing; its precondition is a bare assert that -O strips. And check_row_id_bound is skipped entirely on gfx950 (:686), so the PR's new gfx950-only host row bound (py:762-773) is untested.
aiter/ops/opus/__init__.py drops pa_mqa_logits_mxfp4_sched / _build_sched / _sched_slots / _sched_buffer_ints from both imports and __all__ and adds nothing back, so the op has no package-level or top-level aiter.* export at all. In-tree grep is clean and ATOM already uses the submodule path, but an out-of-tree gfx950 caller breaks with an error pointing at no successor, and the retired cta_target / block_k knobs have no supported replacement.
Summary
I'd like to see the header split before this lands — one file with 68% arch-specific content in disjoint #if islands and two 70-line copies of the same shape check is the structure the merge was supposed to remove, not produce. On performance, the uncached get_gpu_arch on a 61-launches-per-forward path likely costs more than the build_tiles launch it saves, and the two gfx950 grid changes should carry a measurement, ideally on an instance that some perf table actually exercises. And the four guards that only survived in one direction — the three length checks, the block_tables width bound, the row_id bound, and gfx1250's cu_seq_q contract — are the ones I'd treat as blocking, since each one turns a host-side raise into an unbounded device read.
Thanks for the thorough review. Every item is addressed below: fixed, or answered with the reason
for no change. Commits: de2fe0c39 (header split), 8e03a1a62 (length and width guards), 3bcbfccf7 (the rest of the review), 6f2fb13d0 (arch probe, lint), 530fdfdfb (directory move), d552ccfde (per-arch namespaces), dc13ff647 (post-misched on gfx950 only), a758f0c81
(contract condition 4, documentation only).
Validated:
gfx950, clean rebuild at dc13ff647: corner 36/36, NaN 8/8, err = 0 on every perf table,
FlyDSL cross-check mismatch 0. The ISA of all 9 kernels is identical to the d552ccfde build,
because dc13ff647 leaves gfx950's compiler flags unchanged.
gfx1250 (ROCm 10.1.0, clang 23), clean rebuild at dc13ff647: corner 38/38, NaN 8/8, err = 0
on all 27 perf rows. d552ccfde gave the same.
Blocking
1. Three dropped length guards. Fixed. The plan's row count (num_rows, which is total_q)
now reaches build_sched. build_sched checks local_ends, local_starts and row_to_batch
against it at every q_per_block. At the identity cut it also checks num_tiles == num_rows.
8e03a1a62 first compared the lengths against num_tiles. That is the wrong unit above q_per_block == 1, where num_tiles is an upper bound counted in tiles: it could reject valid
inputs and miss real overruns. 3bcbfccf7 replaces it.
num_rows travels in the private build_sched ABI and in a trailing, defaulted MqaLogitsPlan.num_rows field. No user-facing signature changes.
Test: check_row_count_raises runs on every instance of both arches. A 10-entry local_ends with total_q = 14 raises from build_sched.
2. gfx950 block_tables width. Fixed in 8e03a1a62. It is part of the one check_shapes
template shared by both arches: bt_cols = ceil(max_seq_len / KV_TILE_SIZE) * PAGES_PER_TILE. The
contradicting docstrings are corrected. pa_mqa_logits_mxfp4_block_table_width is the sizing
helper.
3. Table arm row_id bound. Fixed, on the host and in the kernel, as you suggested.
Host: C++ fwd_sched raises on num_rows > q.size(0) on both arches, and also checks weights
and out. The gfx950-only Python check is gone.
Kernel: gfx950's Table arm now returns on rec.chunk_count <= 0 || rec.row_id >= num_rows, the
CTA-uniform clause gfx1250 already had. A raw fwd_sched caller with a mismatched cta_info
now drops rows instead of reading past q and writing past out.
Test: check_raw_row_guard drives exactly that case through the private launcher, on both arches.
It passes on gfx1250 on both instances, qlen4_kv64 included. The kernel change costs one
prologue compare. It is not measurable (qlen1_kv64 median −0.46 % over 27 shapes).
4. gfx1250 qlen1_kv64 and cu_seq_q. Documented as an explicit contract; not enforced in the
builder. You are right that at q_per_block == 1 the schedule no longer trims rows by cu_seq_q:
the identity cut schedules every row of total_q. We keep that behaviour, and the requirement it
rests on is now part of the contract (header, and condition 3 of the module docstring): a row that
belongs to no live sequence must carry an empty window (local_ends <= local_starts), and only
such a row may have a negative row_to_batch. An empty window yields chunk_count = 0, so the
row reads neither q nor block_tables and writes nothing to out.
This is what the shipping caller (ATOM's DeepSeek-V4 path) already does:
Decode (CUDAGraph). Pad rows come from one liveness mask that is refreshed every forward,
outside the graph. On a pad row, row_to_batch is -1 and local_ends is 0, and both come
from that same mask. cu_seq_q is never used to mark padding: it is per sequence, and its pad
tail stays flat at the real token count. total_q equals q.shape[0], which equals the local_ends length: the padded token count.
Prefill (eager). Each chunk gets its own plan, with a chunk-local cu_seq_q from which empty
sequences are dropped, so cu_seq_q[batch] equals the chunk's row count. Pad rows are left out of
every plan, and their top-k is filled with -1.
The consumer of the logits reads only [start, end) of each row. So out can be a torch.empty,
and pad rows are never read.
The scenario in the review (stale, non-empty local_ends on pad rows, with padding marked only by cu_seq_q) describes a caller that breaks this condition. Everything that remains visible to the
host is now checked on both arches (items 1 and 3), and both kernels drop row_id >= num_rows.
We did not add a device-side clamp to cu_seq_q[batch] in the builder, because no caller needs it
and the contract is now explicit. The same holds for the grouped cut at q_per_block > 1, where cu_seq_q is read on device. Condition 4 (a758f0c81) requires cu_seq_q[batch] <= total_q
there, because the builder reads every tiled row's local_ends / row_to_batch. ATOM's flat pad
tail meets it. The test gained ATOM-shaped coverage: check_atom_decode pads
the batch with a flat cu_seq_q tail and pad rows with row_to_batch = -1 and local_ends = 0,
at MTP 1 and 2, padded and unpadded.
The same audit also found and fixed a real bug: row_to_batch[r0] was read one past the end on
empty tiles, which ATOM's unpadded MTP decode hits. It is covered by the atom decode mtp=2 5/5
case. On gfx1250, check_atom_decode (MTP 1 and 2, padded and unpadded) and check_row_id_bound
(device cu_seq_q over-claiming 14 rows against a 10-row q) pass on both instances.
Header split. Done (de2fe0c39). The sources now live in csrc/kernels/opus_mqa_logits/pa_mqa_logits_mxfp4/:
file
contents
lines
_opus.h
ABI, kargs, traits
606
_sched.cuh
shared builder
450
_gfx950.cuh
gfx950 body
472
_gfx1250.cuh
gfx1250 body
355
_kernels.cu
launcher
457
Structure
Three barrier spellings. Down to one custom spelling. The builder's sched_detail::sync_block
was HIP's __syncthreads() written out by hand (release fence, s_barrier, acquire fence), so
it is now __syncthreads(). The builder does need those fences: it hands values through LDS and
through global cta_info, which it reads back. gfx1250's cta_sync stays different on purpose:
the hot loop waits by hand (TENSORCNT_KEEP) and brackets the barrier with sched_barrier,
because a fence there would add conservative waits to every phase. opus.hpp's barrier is no
longer used.
Duplicated check_shapes / launch. Merged into one check_shapes<Traits> and one launch_sched<Traits>. The traits share constant names, and if constexpr picks the kernel.
Dead code. Removed:sched_slots, max_tiles_for, ceil_div_i, logits_fp4_qshare_traits_4q, and gfx950's Prefill/Decode arms. Also removed are all six kargs
fields that no kernel reads: your four, plus max_seq_len and kv_block_size. kargs goes from
144 to 112 B, and the offsetof pins are re-derived. On gfx950 the MFMA kernels issue one s_load fewer, and VGPR/SGPR counts are unchanged.
Inverted naming. Fixed symmetrically. The arch qualifier now lives in the namespace: opus_logits::gfx950 and opus_logits::gfx1250 each hold pa_mqa_logits_mxfp4_traits and pa_mqa_logits_mxfp4_kernel, and the shared schedule stays in opus_logits. This also separates
the two pin_sgpr / permlane_head_reduce definitions. Kernel symbol names change in profiles.
JIT rebuild blast radius. Kept as one module. This is the in-tree structure for multi-arch
OPUS ops (module_deepgemm_opus covers gfx942, gfx950 and gfx1250 the same way), and a rebuild
costs about 9 s. Narrowing it would take a second TU or module. Per-arch flags do not need that
split: the module config takes a per-arch expression, which is how post-misched is now handled
(see Performance).
Performance
Uncached get_gpu_arch, duplicated guard. Fixed.fwd_sched probes the arch once per
process. This is exactly the pattern at fmha_fwd_bf16_opus_launch.h:280 that you pointed to
(static const std::string arch_id = get_gpu_arch();), and the same model as opus_get_arch_info(). Like that launcher, it is not keyed per device, and mixed-arch nodes are
out of scope; the code comment says so. The probe runs under HipDeviceGuard(q.device_id), so the
first call reads q's own device. The launcher's inner HipDeviceGuard is gone, so fwd_sched
holds the only one.
gfx950 grid round-up. Kept, measured. Base's max(n, cta_resident) is never faster, and it
is about 10× slower at 1536 rows (122.7 vs 12.8 µs, kv64). The extra CTAs are the spare slots
that let long windows split.
The measurement also exposed a cliff that base had too. When the row count is an exact multiple
of cta_resident, both policies give grid == rows, so nothing can split: 1024 rows with 16
long rows takes 121.7 µs. (n // r + 1) * r for n >= r brings that to 8.65 µs. That fix goes
in a separate PR, because it changes the grid on gfx1250 as well.
qlen1_kv256cta_resident halved. Back to 1024, [HIP] [OPUS] [JIT] [GFX950]OPUS PA MQA Logits MXFP4 #5332's value for both widths. kv256 is now
in the decode table and in chunked prefill. For the record, 512 would be 11–13 % faster on MTP=4
kv256 decode under the shared builder. That can be revisited with a per-instance sweep in the
grid PR.
-mllvm -enable-post-misched=1 on the WMMA body. Fixed as you suggested: gfx950 only
(dc13ff647). The flag is now a per-arch expression in optCompilerConfig.json, in the form
line ~354 uses. On gfx950 it stays load-bearing: the base JIT sets =0, and without =1 the
tuned MFMA body is rescheduled. The gfx1250 body goes back to the base =0, the flags its own
module had on main. Its ISA is now line-for-line the flag-off build. One limit: flags are per TU,
so a single build targeting both gfx950;gfx1250 would still give the gfx1250 body =1. The module_deepgemm_opus form has the same limit, and aiter's release default (gfx942;gfx950) does
not build gfx1250.
Your occupancy concern did not materialise. With the flag on it was a pure post-RA reschedule,
and the resources were identical on and off: VGPR 269 / 269, occupancy 3 / 3, spill 0 / 0, LDS
9728 / 9728, on both instances. Timing is what decided it. We ran 16 paired alternating runs
twice, min-of-16 per shape:
prefill (causal bs1–8, fresh) was 1.0–3.1 % faster with the flag on, 14–16/16 pairs;
chunked bs4 (kvlen 25000) was 3.6–5.5 % slower (0.9–1.9 % by paired median), 11–12/16
pairs, in both runs and on both instances;
decode showed no direction.
The earlier run at 79463da reported that no shape regressed. That was wrong: its data already
had the chunked bs4 regression (median 1160.8 vs 1127.0 µs, 12/16 pairs), but its min ratio of
0.998 hid it. A hand-scheduled body should not carry a flag that regresses a production shape,
so gfx1250 gives up the prefill gain.
Cleanliness
general_emit divide-by-zero and shadowing. Guarded and renamed. The invariant does hold:
the tile that owns slot j has at least one chunk. But it rests on the form of the search, as you
say, so n > 0 is now guarded explicitly, and the shadowing local is renamed to nchunks.
Cross-arch guard rests on value equality. Fixed.MqaLogitsVariant gains a trailing arch: str | None = None, which the module's table fills in. Existing positional and keyword
construction is unaffected. _as_variant refuses an object that names another arch, and plan()
refuses buffers sized for another arch. A hand-built arch=None object is matched only against
the running device's own instances.
__restrict__ aliasing. Dropped from tile_kv_tiles_cached and settle_safe. We believe the
described miscompile cannot happen: scoped noalias from two inlined callees does not make one
callee's load noalias with the other's store, and a fenced barrier sits between them. The
qualifier was still a promise the code did not keep.
block_table_width touches the GPU. Fixed. Given MqaLogitsVariant objects and no device,
it is now pure arithmetic against their own arch's table, and a mix of arches raises. Names or None still need the device, which they did before too. The CPU-only-rank scenario is blocked
earlier than this function anyway: import aiter already initializes the GPU stack, and on a
rank with no GPU it fails in the Triton arch_info import.
functools.cache freezes device probes. Fixed._is_permuted() and _variants() now ask
the current device on every call, and the underlying probe is cached per device.
gfx950 SCHED default. Fixed. The default is now Table, and a static_assert enforces it,
as on gfx1250. The dead arms are deleted, and the kargs comment scopes ptr_local_ends to
gfx1250.
Dropped cta_target >= 1. Restored as build_sched checking cta_resident >= 1. On the
Python side, an import-time check of the instance table rejects any value below 1.
Tests and exports
Deleted cross-check not replaced. Replaced. On gfx950, the op test now checks every
in-window cell against FlyDSL. FlyDSL is fed its own preshuffled layout, built from the natural
E8M0 rather than from _scale_to_opus, so a shared permutation bug cannot hide. A negative
control misroutes the scales by 16 rows and gives mismatch 1.0.
Windows whose start is not a multiple of 4 are skipped and shown as NaN, because FlyDSL at num_warps=1 drops leading cells there. The docstring is corrected.
NaN control asserts nothing. Fixed. The existing case is relabelled: rows 1 and 2 are the
per-row edge. A real control is added: no window holds the poisoned row, so 0 NaN is expected.
The bare assert is now a ValueError.
gfx950 row bound untested. Now tested.check_row_count_raises runs on every instance of
both arches: a 14-row plan over a 10-row q raises from C++. check_row_id_bound now keeps total_q honest and lets only the device cu_seq_q over-claim, which is the path only the kernel
guard can see.
Missing exports. Retired on purpose. Please confirm you are OK with this.[HIP] [OPUS] [JIT] [GFX950]OPUS PA MQA Logits MXFP4 #5332's gfx950 entry
points (pa_mqa_logits_mxfp4_sched / _build_sched / _sched_slots / _sched_buffer_ints) have
no callers: gfx950 users, ATOM included, still run FlyDSL. So this PR retires them before anyone
adopts them, in favour of the one API both arches share (plan_buffers / plan / pa_mqa_logits_mxfp4). cta_target / block_k become the per-arch MqaLogitsVariant, chosen by
name. The shared API lives at aiter.ops.opus.pa_mqa_logits_mxfp4, the path ATOM already
imports. If you would rather keep deprecated forwarding wrappers for a release, I'll add them.
Perf vs main (gfx950, 8253efc40 vs d552ccfde)
qlen1_kv64, min-of-15 over 2 alternating rounds on one idle card:
* [Config] Add gfx942 a8w8 blockscale GEMM tunings for Qwen3/Qwen3.5/GLM/DSV4 shapes (#5839)
MI325X (gfx942, 304 CUs): 286 (M, N, K) rows over 28 weight shapes that no
existing model_configs table covers, tuned on main with
gemm_a8w8_blockscale_tune.py --libtype ck --splitK. Keys already present in
model_configs and rows slower than the heuristic default are left out.
The rows go into each model's existing model_configs table; Qwen3-14B and
Qwen3.5-122B-A10B get new tables. A shape two models share is kept in one
table only, since the config merge rejects duplicate keys across files.
Co-authored-by: Cursor <cursoragent@cursor.com>
(cherry picked from commit 8260de6)
* [HIP] [OPUS] [FlyDSL] gfx950 MXFP8 e8m0 GEMM + DeepSeek-V4/V4.1 tuned configs (#5896)
* feat(flydsl): gfx950 MXFP8 e8m0 GEMM on FlyDSL + DeepSeek-V4/V4.1 configs
- New FlyDSL gfx950 MXFP8 batched GEMM (bmm_a8w8_mxscale_gfx950): scaled
MFMA 16x16x128, async LDS DMA pipeline, 32x32 / 128x128 e8m0 blocks,
split-K with same-XCD last-arrival reduction, B direct to registers, XCD
tile order, non-temporal B, per-stage or preloaded scale panels, K % 64
tail, column-major (blockscale) x_scale. Scale rows that are not whole
dwords (e.g. 128-wide blocks at K = 384 / 768) are copied a byte per LDS
slot with exact buffer bounds. check_bmm_config is the single legality
source.
- Arch-neutral front door flydsl.batched_gemm_a8w8 dispatching by (arch,
w_scale block, x_scale layout); gfx950 module with kernelName parsing,
untuned-shape heuristic and a compiled-launcher cache (host 71 -> 21 us).
- gemm_a8w8_blockscale_bpreshuffle on gfx950 routes e8m0 x_scale + w_scale
to it through the tuned table AITER_CONFIG_GEMM_A8W8_BLOCKSCALE_MXSCALE_
BPRESHUFFLE; FP32 scales keep CK / asm. gemm_a8w8_blockscale with
isBpreshuffled forwards native group32 operands there. The V4 wo_a
batched GEMM uses the same kernel.
- AOT precompiles kernelId=bmm rows (both x_scale layouts for 128x128).
- Tuned gfx950 configs: DeepSeek-V4-Pro 128x128 linears (980 rows, TP1/4/8,
1.30x vs CK/asm), DeepSeek-V4.1 group32 (1224 rows, 1.22x vs Triton),
V4 wo_a batched (680 rows).
- Also carries the other local changes in the tree: inverse_rope_group_quant,
opus policy / bmm tune, gfx1250 bmm wrapper, tensor_shim helpers.
* refactor(flydsl): tidy gfx950 MXFP8 bmm kernel per FlyDSL cleanup guide
- Use fx.copy instead of fx.copy_atom_call for the single-atom copies.
- Build the async-LDS DMA destination from the LDS pointer with
fx.add_offset + fx.to_llvm_ptr instead of ptrtoint/inttoptr with a
hand-picked address space. The raw buffer_load_async_lds stays: the
BufferLoadAsyncLDS atom has no 1-byte size and no cache-policy operand.
- Factor the compile-and-run + leaked ir.Context recovery out of
tensor_shim._run_compiled into _compile_and_run, and use it for the
bmm wrapper's per-config compile cache too.
Generated ISA is byte-identical on 9 reference configs.
(cherry picked from commit 40d524b)
* [Triton/Gluon] [ASM] [HIP] Mha v4: adds bf16 sparse, LSE support, KV varlen, fixes, etc (#5798)
Motivation: adds log-sum-exp output to MHA v4 gfx950 kernels (only dense variants now), so it can run under ring / context parallelism, which merges per-rank partials via LSE. Also extends block-sparse to the BF16 recipes, canonicalizes the MXFP4 rows, and improves K/V quantization accuracy.
# Kernels:
- New BF16 and BF16FP8 sorted-sparse kernels.
- Dense LSE epilogue on all ten hd128 recipes, gated at runtime on s_lse.
- f8f6, f6f4 and mxfp4 sparse rows moved to FP6-P V, matching their dense siblings.
- Dense mxfp4 claims the canonical FP6-P V order; the duplicate f4f4 row is disabled.
# Host:
- bugfix: per-channel V amax clamped so empty heads cannot quantize to NaN.
- Optional lse output on mha_v4 / mha_v4_packed, plumbed through asm_mha_v4_fwd.cu.
+ ABI unchanged: ptr_lse / s_lse / s_lse_Hs were already reserved in the kernarg.
- _LSE_CAPABLE_QV gates the supported format pairs; sorted-sparse still raises.
- K mean is subtracted before quantizing (K-smoothing), fused into the MX quantizer kernels.
- Dense MHA v4 accepts per-batch key lengths (ragged seqlen_k).
# Minor:
- Retired MXFP4 Q/K + FP8 V and the deprecated mha_v4_mxfp8 alias.
- bench_sage.py: improve input distributions, diffusion-calibrated default, BF16 sparse modes.
- Split the block-sparse cases into op_tests/test_mha_v4_sparse.py.
(cherry picked from commit 105615e)
* [tuner] Fail a task as soon as its worker process exits (#5841)
A GPU memory fault aborts the worker process, so its task never returns a
result. mp_tuner only noticed at the task timeout (1800s by default) or
never without one. Record which worker started each task and, when that
process is gone, fail the task and restart the pool right away, like the
existing accelerator-error path. Fixes#5840.
Co-authored-by: Cursor <cursoragent@cursor.com>
(cherry picked from commit bcb56d9)
* [CI] Update Aiter artifact downloads to v8.0.1 (#5908)
(cherry picked from commit c03689f)
* [MLA v4 nm] Test fix _run_one_point reading packed BF16 as FP32 partials (#5905)
* [MLA v4 nm] Fix _run_one_point reading packed BF16 as FP32 partials
Since #4311 the v4 nm dispatcher derives out_16_nosplit from num_kv_splits
and ignores the caller's value, so a single-split launch writes packed BF16
into the logits buffer. _run_one_point still read logits_buf[:, 0] as FP32
for num_kv_splits == 1, which compared every other BF16 element (plus the
never-written tail of the buffer) against the reference and printed spurious
`fp8_dequant_ref vs asm` checkAllclose failures in the script-mode sweep.
Read output_buf, which holds the final result for every split count.
Co-authored-by: Cursor <cursoragent@cursor.com>
* [MLA v4 nm] Gate accuracy checks so a failed! fails the test
checkAllclose only raises on a catastrophic delta; otherwise it logs
`failed!` and returns the mismatch fraction. Four of the six accuracy
checks in test_mla_v4_nm.py dropped that return value, which is how the
packed-BF16 readback bug fixed in the previous commit passed both pytest
and the script-mode CI run.
Route them through _gated_allclose, which asserts the mismatch fraction
against the same tol_err_ratio checkAllclose uses for `failed!`. The
script-mode sweep keeps going past a failing shape, prints a summary, and
exits non-zero so aiter_test.sh reports the file as failed.
Co-authored-by: Cursor <cursoragent@cursor.com>
---------
Co-authored-by: Cursor <cursoragent@cursor.com>
(cherry picked from commit 049fae4)
* [AOT] Inline the FlyDSL FP8 FMHA head shapes, drop the config CSVs (#5901)
Follow-up to #5796. The AOT job list for the gfx950 FlyDSL FP8 flash
attention was driven by a header-only aiter/configs/fmha_fp8_aot.csv
merged with aiter/configs/model_configs/*_fmha_fp8_aot.csv. A single row
does not justify the CSV plumbing, so the head shapes now live in a
DEFAULT_SHAPES table in the module, the way mega_moe.py already does it,
and the comments introduced by #5796 are trimmed.
- fmha_fp8.py: DEFAULT_SHAPES replaces parse_csv()/DEFAULT_CSVS;
default_jobs() replaces the CSV walk; --csv is gone (--shape still
overrides). The cu_num column is dropped with the CSVs: the kernel is
gfx950-only and AOT_ARCH is fixed, so every non-gfx950 row was warned
about and skipped anyway.
- common.py: FMHA_FP8 returns default_jobs() next to MEGA_MOE instead of
going through collect_aot_jobs().
- jit/core.py: drop AITER_CONFIG_FMHA_FP8_AOT and its config-file
property, now unused.
- README.md: document the table instead of the CSVs.
_variant_space()/jobs_for_shape() are unchanged, so coverage is
unchanged: --list still emits the same 92 kernel names for Kimi-K3 TP8
(12:12:192:128 varlen_cross) as the CSV-driven version, and setup.py
source builds still compile them via run_aot().
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
(cherry picked from commit 71a31be)
* [FlyDSL] gfx942 fp8_mqa_logits: let _auto_variant choose rows_per_block (#4963)
* [FlyDSL] gfx942 fp8_mqa_logits: let _auto_variant choose rows_per_block
_auto_variant returned f"mfma_r2_w{wpb}", so rows_per_block was pinned at 2 and
seven of the nine registered variants -- including every member of the r4 family
-- could never be selected.
r4 amortizes each KV tile load over twice as many query rows and is faster from
seq_len 8 upward. The gap is widest exactly where it costs most: vLLM chunks
indexer prefill to fit VLLM_SPARSE_INDEXER_MAX_LOGITS_MB (512 MB), which caps
seq_len at 1024 when seq_len_kv is 131072, so the existing seq_len >= 2048 branch
cannot fire at long context and every such call took mfma_r2_w4 -- the median of
the nine by speed, with the best 1.65x faster.
Measured on MI325X (gfx942), seq_len_kv 131072, best variant vs the r2 pick:
seq_len 1 r2 25.2 us r4 79.1 us r4 3.1x worse
seq_len 4 r2 24.3 us r4 31.5 us r4 1.3x worse
seq_len 8 r2 34.0 us r4 32.9 us r4 1.03x better
seq_len 16 r2 56.5 us r4 47.5 us r4 1.19x better
seq_len 1024 r2 2564.3 us r4 1520.3 us r4 1.69x better
Below seq_len 8 the host padding of seq_len up to a multiple of RPB dominates --
at seq_len 1 an r4 kernel computes 4 rows to obtain 1 -- so r2 is kept there and
behaviour for those shapes is unchanged.
Logits are bitwise identical across all variants at every shape tested, so this
is purely a blocking/occupancy change.
End to end on 8x MI325X, TP8, GLM-5.2-FP8, 131072 in / 1024 out, concurrency 8:
median TPOT improves 7.22% and output throughput 6.60%. This kernel is 16.8% of
GPU time at that point.
Signed-off-by: Jin Tao <jin.tao@amd.com>
* [FlyDSL] gfx942 fp8_mqa_logits: pick RPB on element count, and keep it a divisor
Refines the previous commit's rule after a 2-D sweep. That rule keyed RPB off
seq_len alone with a crossover measured only at seq_len_kv=131072; sweeping the
other contexts shows the crossover is not a seq_len threshold at all, and that a
second effect was being read as one.
RPB tracks the logits element count. Over seq_len 1..8192 x seq_len_kv
1024..262144 on MI325X, the boundaries land on the same element count at every
context: RPB=1 wins below 2**19 elements (27/27 shapes), RPB=2 at 2**19 (6/6),
RPB=4 from 2**21 up (38/38), with 2**20 a transition band split 3/3. Keying off
seq_len instead put the previous rule on the wrong side at low context: at
seq_len 16, seq_len_kv 1024 it chose RPB=4 and ran 1.26x slower than RPB=1.
RPB must also divide seq_len. When it does not, the launcher pads with four
torch.cat calls; that is a flat ~44 us of host-side overhead, independent of
seq_len_kv, and it is the whole of the "small seq_len" penalty the previous
commit attributed to wasted rows. At seq_len 1, seq_len_kv 131072: RPB=1 23.1 us,
RPB=2 67.8 us, of which the four cats are 44.1 us and pre-padding by hand
recovers all of it (21.9 us). So the penalty is not proportional to the padding
-- 1 wasted row of 2 costs the same as 3 of 4 -- and it applies to every odd
seq_len, which the old rule sent to RPB=2 unconditionally.
Stepping down to a divisor is only right while the kernel is cheap relative to
that fixed cost, so it is gated to the same 2**21 elements: at seq_len 1025,
seq_len_kv 131072 the dividing RPB=1 takes 3880 us against 2601 us for a padded
RPB=2.
Measured on MI325X, no FLYDSL_FP8_MQA_LOGITS_VARIANT set, old pick vs new:
seq_len seq_len_kv old new speedup
1024 131072 r2_w4 2598.7 r4_w4 1623.5 1.60x
1025 131072 r2_w4 2609.3 r4_w4 1594.0 1.64x
512 131072 r2_w4 1190.1 r4_w4 795.3 1.50x
700 50000 r2_w4 661.6 r4_w4 440.3 1.50x
333 12000 r2_w4 118.4 r4_w4 97.1 1.22x
16 131072 r2_w4 58.3 r4_w4 47.4 1.23x
3 131072 r2_w4 66.4 r1_w4 28.1 2.36x
1 131072 r2_w4 69.4 r1_w4 23.2 2.99x
1 1024 r2_w4 66.2 r1_w4 22.2 2.98x
No shape measured regresses; the smallest gain is 1.04x. Against the best of the
nine variants at each shape, pooled over held-out data (non-power-of-two shapes,
a fine seq_len sweep, and head counts 16 and 64), the geometric mean cost falls
from 1.45x to 1.03x and the worst case from 3.17x to 1.41x.
Logits are bitwise identical across all nine variants at all 180 shapes swept
(1620 timings), so this remains purely a blocking/occupancy change.
WPB is deliberately left alone. It is worth a few percent at most here, and
unlike RPB its optimum moves with the head count -- at 64 heads the current
WPB rule costs 1.61x worst case where a fixed WPB=4 costs 1.07x -- so it needs
its own sweep rather than a change fitted to one head count.
Signed-off-by: Jin Tao <jin.tao@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
* [FlyDSL] gfx942 fp8_mqa_logits: trim _auto_variant comments per review
Move RPB element-count thresholds into _auto_variant and shorten the
docstring; logic unchanged.
Signed-off-by: Jin Tao <jin.tao@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
* [FlyDSL] gfx942 fp8_mqa_logits: scope the _auto_variant step-down note
The docstring stated the divisor step-down as an unconditional rule, but it
only applies in the middle band: above the top threshold RPB stays 4 and the
padding is accepted. Say which band it applies to, and why the top band is
exempt. No logic change.
Co-authored-by: Cursor <cursoragent@cursor.com>
* [FlyDSL] gfx942 fp8_mqa_logits: unit-test the variant selector
The shape sweep in test_flydsl_fp8_mqa_logits.py never reaches r4 -- its
largest default shape is 1024 x 1560, below the 2**21 threshold -- and nothing
asserted _auto_variant directly, so a regression could re-pin RPB to 2 with
every correctness test still green.
Cover the RPB bands at both thresholds, that the edges track
seq_len * seq_len_kv rather than seq_len alone, the middle-band step-down on
odd seq_len (and that it stops above the top threshold), the production
1024 x 131072 shape, the unchanged WPB rule, and _resolve_variant precedence
so the auto path is confirmed to be the default. Pure shape arithmetic, so no
kernel launch.
Co-authored-by: Cursor <cursoragent@cursor.com>
* [FlyDSL] gfx942 fp8_mqa_logits: put selector tests on the CI path
CI only discovers op_tests/test_*.py, so the r1/r2/r4 selector coverage
never ran. Move it there, run pytest from __main__, and add two r4
shapes to the GPU sweep so auto-selected padding is launched.
Signed-off-by: Jin Tao <jin.tao@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
* [FlyDSL] gfx942 fp8_mqa_logits: fold the selector pins into the op test
aiter op tests are plain scripts, not pytest, so drop
test_flydsl_fp8_mqa_logits_variant.py and pin the gfx942 auto-selected
variant in test_flydsl_fp8_mqa_logits.py instead: a host-only
verify_auto_variant table at both RPB band edges, either side of the
odd-seq_len step-down, the long-context prefill shape and the WPB switch.
It runs in the default verify scenario on gfx942, and its failures count
toward the same exit code as the kernel sweep. The selector import sits in
the file's existing ImportError guard.
Signed-off-by: Jin Tao <jin.tao@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
---------
Signed-off-by: Jin Tao <jin.tao@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: jin.tao@amd.com <jin.tao@amd.com@tus1-p15-g43.tus.tensorwave.lan>
Co-authored-by: Felix Li <felix.li@amd.com>
(cherry picked from commit bbaccdf)
* [Triton] [MHA] Add a tuned gfx1101 config, split small_head/default (#4493)
gfx1101 (RDNA3, e.g. RX 7800 XT) ships no MHA config, so `_get_config` in
`_triton_kernels/attention/mha.py` finds no
`configs/gfx1101/triton/attention/mha/DEFAULT.json` and every call to
`aiter.ops.triton.attention.mha.flash_attn_func` fails before a kernel runs.
Of the architectures in RDNA_ARCHS, only gfx1151 ships one today.
Nine of the eleven entries are taken verbatim from the gfx1151 donor
(RDNA3.5, added in #3423, tuned in #3560), which is the nearest tuned
architecture. Two forward entries are tuned on gfx1101 instead of inherited,
and they differ from each other only in `num_warps` and `num_stages`:
fwd/default BLOCK_M 128, num_warps 8, num_stages 3 (donor: 64 / 4 / 2)
fwd/small_head BLOCK_M 128, num_warps 4, num_stages 1
The split uses the `small_head` bucket added in #4414, which is opt-in per
architecture by the mere presence of the key, so this stays a data-only
change. It is needed because one `fwd/default` cannot serve both halves on
this card: measured against the donor, `M128 w4 s1` is 0.935x on head_dim 64
but 1.269x on head_dim 128, while `M128 w8 s3` is 0.914x on head_dim 128 but
1.062x on head_dim 64.
Measured on Windows native ROCm, triton 3.8.0, fp16, 10 independent repeats
of 20 iterations after 5 warmups, `torch.cuda.synchronize()` per iteration;
a result counts only when the [min, max] intervals across repeats are
disjoint. Ratios are against the gfx1151 donor entry, i.e. against what a
donor-inherited config would do.
head_dim 128, `default` Flux joint 0.914x / 0.891x / 0.892x on three
independent torch+ROCm stacks (2.11/7.15,
2.11/10.1, 2.15/10.1), all pinned to the same
triton; llama3-8B 0.882x, mixtral-7B 0.880x,
kimik25-tp4 0.849x at seqlen 16384
head_dim <= 64, `small_head` SDXL self-attn 0.935x; deepseek-V3 0.658x and
glm47fp8-tp4 0.857x at seqlen 16384
The LLM shapes come from `op_tests/op_benchmarks/triton/utils/model_configs.json`
(prefill, causal, GQA, batch 1); the sequence lengths are not in that file and
are chosen here. Not covered: batch > 1, varlen/thd, sliding window, decode.
Note that the `small_head` comment in `_get_config` does not describe gfx1101.
It states that 16 < d <= 64 suffers a num_stages=1 pipelining pathology which
num_stages=3 cures. On this card the ordering is the opposite -- on the tuned
M128/N32/w4 tile, num_stages 1/2/3 measured 2.300 / 2.399 / 2.492 ms. The
bucket is still the right mechanism here, for a different reason: the two
head_dim ranges want a different `num_warps`, not a different `num_stages`.
Signed-off-by: Martin Domanský <ragua@email.cz>
Co-authored-by: Mehmet Cagri <mehmet.kaymak@amd.com>
(cherry picked from commit f1f95ce)
* [Triton/Gluon] Consolidate tuning harnesses (#5874)
* [Triton/Gluon] Consolidate tuning harnesses
* [Triton/Gluon] Simplify tuning harness
(cherry picked from commit af3514a)
* [Bugfix][Gluon][MLA] Follow-up: fix stale comments + test hardening (#5860)
Addresses review feedback on #5648 (kept separate to not disturb the approved PR):
- mla_gluon.py: the >2GB global_load calls now carry a bounds mask + other=0.0,
so the old "No mask needed / in-bounds" comments above them were stale and
gave the opposite (unsafe) guidance. Replace with an accurate one-liner.
- test_mla.py: allocate the output with device=q_nope.device instead of relying
on the module default device; call torch.cuda.empty_cache() before the
mem_get_info() free-memory gate to avoid allocator-state flakiness; and assert
exact parity (atol=rtol=0) between the >2GB and <2GB paths, which read
identical KV and must be bit-identical (a loose tolerance could hide a
regression).
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
(cherry picked from commit a966245)
* [Triton/Gluon] [Config] gemm_a16w16: gfx950 tuned per-shape configs (#5830)
* [Triton] gfx950: tuned defaults for gemm_a16w16, gmm, MHA fwd and PA decode
Tuned on MI355X (gfx950) and validated for correctness and no regressions on
broader shape sets than tuned. All changes are gated to gfx950 configs/arch.
- gemm_a16w16: per-shape configs for N,K = 2048/2048, 4096/4096, 8192/8192,
10240/8192, 57344/8192, 8192/28672 (non-persistent and persistent). Only the
tuned M bucket differs from DEFAULT.json. 1.07-2.5x.
- gmm: new "large_kn" config (8 warps), selected for K, N >= 4096 with
>= 256 rows per group. 1.14-1.26x there; smaller problems keep "default".
- MHA fwd (Triton): new "mid_head" config (128x128, 2 stages) for bf16/fp16
with 64 < head_dim <= 128. 1.03-1.16x; d<=64 and d>128 unchanged.
- PA decode (Triton): use v2 whenever there is more than one partition for
bf16/fp16 KV. 1.9-5.3x for the batch/context sizes that previously hit v1.
Unit tests: test_gemm_a16w16 360 passed, test_gmm 48 passed,
test_pa_decode 816 passed, test_mha 2060 passed.
* [Triton] pa_decode: sort imports (ruff I001)
* Move gmm, MHA fwd and PA decode changes to their own PRs
Per review, each kernel gets its own PR; this PR keeps only the gemm_a16w16
tuned configs. The removed changes are on branches gfx950-tuned-gmm,
gfx950-tuned-mha-fwd and gfx950-pa-decode-v2.
* [Triton] gemm_a16w16_persistent: gfx950 tuned config for N=K=2048
(cherry picked from commit a651db0)
* replacing the fp32 mfma with two bf16 mfma. fn is split inside the kernel into hi = bf16(fn) and lo = bf16(fn -hi). a simple bf16 downcast for fn inside the kernel drops accuracy more and performs worse for large M. also change the concatenation of the four streams from column major to stream major (k = stream * TILE_K + column). retuned the gfx950 configs. (#5885)
(cherry picked from commit f10cd2a)
* perf(fused-moe): add tuned DSV4.1 TP4 configs (#5919)
(cherry picked from commit 977ae79)
* [Config] Add Qwen3.8-27B TP1 a8w8 blockscale GEMM tunings for gfx942 (#5585)
* [Config] Add Qwen3.8-27B TP1 a8w8 blockscale GEMM tunings for gfx942
#3324 covered this model family at TP=2/4/8 only, so the five widths
Qwen3.8-27B drives at TP1 have no tuned entries for gfx942/cu_num=304
and every call falls back to the default -- 275 "not found tuned
config" messages per profile run, now 4. This adds 641 rows in one
file, 249 decode (M <= 512) and 392 prefill (M 907-65536), from a shape
list extracted out of real vLLM server logs rather than generated as an
M ladder.
Decode, against AITER's own default, measured in a real vLLM serving
run -- Qwen3.8-27B-FP8 at 128K context, max-concurrency 1, on one
MI325X, byte-identical images, the config bind-mount the only variable,
three interleaved repeats per arm on an idle node:
decode step 15.444 ms -> 12.449 ms 1.24x
Per-shape decode GEMM vs the default: down 2.33x, out_proj 1.51x,
in_proj 1.12x, qkv 1.00x, gate_up 1.00x. No decode shape regresses.
Decode-step spread across repeats was 0.037 ms and 0.012 ms.
The 392 prefill rows ship and are tuned, but their end-to-end effect is
currently unmeasurable on this workload: gemm_a8w8_blockscale silently
stores only M*N mod 2^31 output elements once M*N reaches 2^31, and at
a 65,536-token prefill chunk the 34,816-wide gate_up projection crosses
that limit. The truncation depends only on M*N and not on which kernel
the config selects -- AITER's default truncates identically -- so it is
neither introduced nor worsened by this change.
Tuned with gemm_a8w8_blockscale_tune.py --libtype ck --splitK, decode
rows then re-timed interleaved because the tuner's min() over short
samples is noisy on near-ties; only rows beating the default by >=3%
were kept, following #5421. test_csv_validation (15 tests) and
test_config_shape_collision (17) both pass, there are zero duplicate
(gfx, cu_num, M, N, K) keys against the 12 other sources in the runtime
merge set on main, and max relative Frobenius error is 4.11e-3 against
bf16's 2^-8 = 3.9e-3 rounding floor.
Scope: only the decode rows have in-situ evidence, and MI300X shares
the gfx942/304 key but is untested here. Two rows pin AITER's default
kernel (qkv and gate_up at M=1) after the trace measured the tuned
picks slower -- pinned, not deleted, because getPaddedM(1,N,K,0) == 16.
Six decode near-ties ship no row. gate_up ships 72 of 80 prefill rows;
the 8 missing have M*N > INT32_MAX and fault during tuning, but inherit
the tuned M=8192 row at runtime via the same padding collapse.
Signed-off-by: Pham Binh <phamhuuthanh.binh@amd.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
* [Config] Thin Qwen3.8-27B TP1 GEMM rows to the padded M ladder
Lookup already rounds M through getPaddedM, so a row per traced M is redundant.
Keep the power-of-two M list used by the other gfx942 blockscale tables, and
omit shapes that main already ships.
---------
Signed-off-by: Pham Binh <phamhuuthanh.binh@amd.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: akii96 <aakif.nawaz@amd.com>
(cherry picked from commit 330f127)
* [HIP] [OPUS] [JIT] Unify gfx950 + gfx1250 MXFP4 paged MQA-logits into one module, schedule and API (#5761)
Fold the gfx950 (#5332) and gfx1250 (#5656) OPUS MXFP4 paged MQA-logits ops into one
implementation. The device math of both bodies is unchanged.
- One JIT module, module_pa_mqa_logits_mxfp4_opus. Sources move to
csrc/kernels/opus_mqa_logits/pa_mqa_logits_mxfp4/:
- _opus.h: ABI, kargs, traits;
- _sched.cuh: shared builder;
- _gfx950.cuh / _gfx1250.cuh: arch bodies, each an empty stub on the other arch;
- _kernels.cu: launcher.
Namespaces are opus_logits::gfx950 / ::gfx1250, with generic names (pa_mqa_logits_mxfp4_traits
and _kernel). The per-arch sources and module_pa_mqa_logits_mxfp4_gfx1250_opus are deleted.
- One schedule: gfx1250's build_tiles + build_sched serve both arches. build_tiles is skipped at
q_per_block == 1, where the cut is the identity, which saves one launch per forward on every
gfx950 call and on gfx1250's MTP=1 decode.
- One runtime-arch dispatch: fwd_sched probes the arch once per process, then dispatches on
(q_per_block, block_k). An unmatched config raises.
- Host-visible bounds on both arches:
- num_rows is checked against q, weights and out, and against the lengths of local_ends,
local_starts and row_to_batch;
- block_tables width is checked in KV tiles;
- both kernels drop records with row_id >= num_rows.
The caller contract (four conditions) is documented in the module docstring and the header.
- One Python API (aiter.ops.opus.pa_mqa_logits_mxfp4: plan_buffers / plan / pa_mqa_logits_mxfp4),
with per-arch MqaLogitsVariant instances. gfx1250: qlen4_kv64 / qlen1_kv64. gfx950: qlen1_kv64 /
qlen1_kv256, cta_resident 1024.
- Retired: #5332's gfx950 entry points (pa_mqa_logits_mxfp4_sched / _build_sched /
_sched_slots / _sched_buffer_ints).
- Kept: #5656's public gfx1250 API, except the gfx950-only plan(block_k=).
- -mllvm -enable-post-misched=1 is applied on gfx950 only. It is load-bearing there; on gfx1250
it slowed chunked prefill.
- kargs drops six unread fields (144 -> 112 B).
- One op test for both arches, op_tests/test_pa_mqa_logits_mxfp4_opus.py. It adds ATOM-shaped
padded decode, host-raise and raw row-guard cases, a real NaN control, and on gfx950 a FlyDSL
cross-check.
(cherry picked from commit 9b3885f)
* [FlyDSL] feat(mega_moe/gfx1250): bind mori tokoff-ext allocator on the mori dispatch path (#5810)
* feat(mega_moe/gfx1250): bind mori tokoff-ext allocator on the mori dispatch path
mori's op layer builds the tokoff-ext slot allocator in
EpDispatchCombineOpHip.__init__; driving EpDispatchPlan directly bypasses
that and leaves EpArgs.tokOffPeers null, keeping dispatch on the
serializing cco-window atomic. Build it here from mori's TokOffExt, gated
exactly like mori (default on; MORI_EP_TOKOFF_EXT=0 opts out), pass its
peer-pointer array to plan.launch, and free it in close().
Needs mori with a public TokOffExt (dispatch_combine_v2.hip_backend).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* style(mega_moe/gfx1250): black-format the tok_off_peers ternary
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---------
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: HaonanWang98 <hwang@amd.com>
(cherry picked from commit 475cf0f)
* [Triton] Do not write segment softmax state at NUM_SEGMENTS_PER_SEQ == 1 (#5935)
mla_decode_fwd splits each sequence's KV into NUM_SEGMENTS_PER_SEQ
segments and lets a reduce kernel merge the per-segment softmax max and
expsum. At one segment there is nothing to merge, so the host skips the
reduce kernel and hands both scratch pointers the output buffer itself:
else:
segm_output = out
segm_max = out # dummy ptr
segm_expsum = out # dummy ptr
The decode kernels store M and L unconditionally, so at one segment they
write the softmax state straight over the attention output. Nothing
faults and nothing warns; the result is simply wrong.
Guard both stores. NUM_SEGMENTS_PER_SEQ is a constexpr, so the branch
costs nothing when it is greater than one.
Only the gfx1250 Gluon kernel reaches this today: select_3d_config ends
its gfx12 branch at max(1, ...), while the other branch floors at
MIN_SEGMENTS >= 8. The segment count falls as batch x heads grows, so on
gfx1250 it reaches one at large batch -- a DeepSeek-R1 TP2 serving run at
--max-running-requests 256 sits in that range throughout, and its decode
is silently corrupted. The plain Triton kernel shares both the unguarded
store and the host-side aliasing, so guard it as well, before a future
tuning change makes it reachable.
test_mla_decode_fwd stays green either way, which is why this went
unnoticed. Its grid does reach one segment for the larger head count,
but the corruption lands only at out.flat[token * num_query_heads +
head] -- one element in kv_lora_rank, measured at 0.04% to 0.12% of the
output -- while the assertion allows a tol_err_ratio of 0.01. Catching
it needs a case with no error budget.
Checked on gfx1250. DeepSeek-R1 TP2 at page size 64, gsm8k over 2000
questions, scores 0.945 with the fix. A standalone sweep over batch 1 to
256, bf16 and fp8 e4m3 caches, sequence lengths that are not page
multiples and pages scattered through the pool gives a worst relative
error of 0.0033 for decode and prefill alike. Compared against a torch
reference with no error budget, a one-segment batch differs at a 0.08%
ratio before this change -- softmax maxima around 100 sitting where
attention outputs near 0.02 belong -- and matches to the last element
after it. The plain Triton kernel, forced on by disabling the gfx12
branch, matches the same reference at batch 1 to 256, confirming the
guard leaves the multi-segment reduce path intact. test_mla.py is
unchanged: its 128 pre-existing failures are all shuffled_kv_cache=True
on the pipelined kernel, identical before and after.
Signed-off-by: Lin, Soga <soga.lin@amd.com>
(cherry picked from commit 2b6ff3d)
* [Triton/Gluon] Enable unified-attention skip-mask for gfx950 hd256 FP8 prefill (#5598)
Long FP8 prefill at head size 256 matches `D_GEQ_256` (`BLOCK_M=16`, no `SPLIT_UNMASKED_LOOP`) because D outranks Q in the lookup, so it never reaches the hd128 skip key. Add the same `Q_GEQ_256` family under `D_GEQ_256`, with `SW` and `SHUF` siblings left skip-off, so the existing kernel fast path fires only for non-windowed, non-shuffled prefill.
(cherry picked from commit e289613)
* [Triton/Gluon] [Kimi-K3][ROCm] Add merged MoE front (#5321)
* feat(kimi-k3): add minimal large-M MoE front
Signed-off-by: jiacao-amd <jiahui.cao@amd.com>
* feat(kimi-k3): enable merged front for M=7 decode
Signed-off-by: jiacao-amd <jiahui.cao@amd.com>
* feat(kimi-k3): tune merged front for M=14 decode
Signed-off-by: jiacao-amd <jiahui.cao@amd.com>
* perf(kimi-k3): tune merged front decode bucket
Signed-off-by: jiacao-amd <jiahui.cao@amd.com>
* test(kimi-k3): validate the merged front across the MTP decode bucket
test_decode_full_front_matches_reference stopped at m=16, so the shapes an MTP
server actually decodes at were never checked against the torch.mm reference:
with num_speculative_tokens=3 a pure-decode step submits num_seqs * 4 tokens,
which is 32 and up for any server past four concurrent sequences.
Extend the parametrization to 32, 48, 64, 80, 96, 112, 128 and 192 -- the shapes
the companion vLLM change enables, and for which kimik3_bf16_tuned_gemm.csv
already ships tuned solutions at N=6016, K=7168, gfx950, cu_num=256. No kernel
or config change is needed; this only closes the correctness gap left by the
old parametrization.
gfx950: 9 passed -> 17 passed.
Signed-off-by: jiacao-amd <jiahui.cao@amd.com>
* style(kimi-k3): satisfy the pinned ruff on this PR's own files
`ruff==0.16.0`, the version .github/workflows/pre-checks.yaml pins, reports two
errors in files this PR introduces: RUF022 on the `__all__` list in
kimi_k3_moe_front.py and I001 on the import block in its test. Both are
autofixes and neither changes behaviour -- `__all__` ordering only affects
`import *`, and every symbol here is imported by name.
The two EXE001 findings that remain in the repo are in unrelated flydsl files
and predate this PR.
Signed-off-by: jiacao-amd <jiahui.cao@amd.com>
* perf(kimi-k3): retune the M=16 merged-front FP32 row
The decode shape is 4 * concurrency, and getPaddedM(gl=0) rounds every M <= 256
up to 16, so this row serves both the C1 (M=4) and C4 (M=16) decode buckets --
about 40% of iterations in the 8k/1k replay.
It was tuned against a single weight buffer. The merged front's weight is
6016x7168 BF16 = 86.2 MB and MI355X has 256 MB of LLC, so that regime serves
most of the GEMM from cache; the shipped 16.005 us is faster than anything the
shape can reach when the weight is actually streamed. In the real model 92
distinct MoE layers stream 7.9 GB per decode step and nothing is resident.
Re-searched all 681 hipblaslt solutions with the weight chained across 604 MB
of distinct buffers (past LLC) and timed inside a CUDA graph, which is how vLLM
runs decode:
solidx 443935 (shipped) 22.364 us 3857 GB/s
solidx 443486 (this) 19.279 us 4474 GB/s -13.8%
443486 is not a new solution -- the table already carries it at M=7. The right
kernel was present and keyed to the wrong M.
Max relative error vs a torch.mm FP32 reference is 4.9e-07, so this is a
dispatch change only. The us/tflops/bw columns are the streamed measurement and
are therefore not comparable to the cache-hot numbers in neighbouring rows.
* perf(kimi-k3): size the front-GEMM config cache to the whitelist
_KIMI_K3_MERGED_FRONT_TOKEN_COUNTS admits 28 distinct token counts but the
config cache held 16, so the M values a mixed decode/prefill server cycles
through evicted each other.
Minor: the tuned table is cached upstream, so a miss costs a few get_padded_m()
extension calls rather than a CSV parse, and in the graphed decode path it is
paid at capture instead of replay. This just makes the cache cover the set it
was sized for.
* refactor(kimi-k3): separate Triton epilogue layers
Move the launchable epilogue kernel under _triton_kernels/moe, keep only the public wrapper in ops/triton/moe, and add a config-aware kernel repr and CUDA Graph benchmark. Remove Kimi-specific weight packing and hipBLASLt orchestration from the Triton module.
Signed-off-by: jiacao-amd <jiahui.cao@amd.com>
* refactor(moe): generalize the SiTU epilogue
Rename the op to describe the SiTU epilogue it actually performs rather than implying that it projects the incoming activation. Make all branch widths caller-provided, mask partial tiles for arbitrary shapes, and reuse the shared Triton tanh helper.
Update the numerical and CUDA Graph coverage plus the benchmark for the generic API.
Signed-off-by: jiacao-amd <jiahui.cao@amd.com>
* test(triton): use Triton arch helper for SiTU epilogue
Signed-off-by: jiacao-amd <jiahui.cao@amd.com>
* fix(triton): avoid branch-local mask type collisions
Signed-off-by: jiacao-amd <jiahui.cao@amd.com>
---------
Signed-off-by: jiacao-amd <jiahui.cao@amd.com>
Co-authored-by: Jiahui Cao <jiacao@crs-m2m-cpu-spur-014.us-east2-a.compute.internal>
(cherry picked from commit aa84815)
* [Triton/Gluon] Add SonicMoE pure-Triton grouped GEMM MoE (#5725)
* [Triton] Add SonicMoE pure-Triton grouped GEMM MoE
Add a pure-Triton, expert-major grouped GEMM MoE with full autograd
support (SonicMoE), covering top-k and general routing, GLU/elementwise
activations, and blockwise FP8 scales for the grouped GEMM. Includes
gfx942/gfx950 tuned configs, unit tests, and a benchmark script.
Public entry points: moe_TC_softmax_topk_layer, moe_general_routing_inputs,
moe_pre_routed_inputs (aiter.ops.triton.sonicmoe).
This is the first of two PRs splitting the SonicMoE contribution; a
follow-up PR stacks the hipBLASLt/multistream grouped GEMM backend on
top of this pure-Triton implementation.
* [Triton] Drop unused SonicMoE stream_id dels and restating comments
Keep the positional stream argument for caller compatibility as
_stream_id, and remove comments that only restated the next GEMM call.
* [Triton] Move SonicMoE autograd API out of _triton_kernels
Keep @triton.jit kernels under _triton_kernels/moe/sonicmoe and put
the autograd wrappers plus public entry points in aiter.ops.triton.sonicmoe.
* [Triton] Use AMD copyright headers on SonicMoE kernels
Replace third-party author banners and drop external source-link
comments so new files match aiter's SPDX header.
* Address SonicMoE review feedback
* [Triton] Move SonicMoE routing kernels into moe_routing
* [Triton] Collapse SonicMoE host wrappers into one module
Keep the public API and tests in a single file instead of a set of sibling wrappers.
(cherry picked from commit 2e62094)
* [Triton/Gluon] Move the KDA_DECODE configs into the nested config layout (#5941)
(cherry picked from commit 92192fd)
* [Triton/Gluon] Move _triton_kernels/gated_delta_rule/ to _triton_kernels/gated_delta_net/ (#5943)
(cherry picked from commit cf89ffd)
* [Triton/Gluon] Move _gluon_kernels/gfx1250/norm/ to _gluon_kernels/gfx1250/normalization/ (#5944)
(cherry picked from commit 9ef0d08)
* [CI] Select impacted Triton and Gluon unit tests (#5878)
Select impacted Triton and Gluon unit tests
(cherry picked from commit 32b1cb0)
* [Triton/Gluon] [Config] Drop no-op kpack from RDNA GEMM configs (#5917)
* [Triton] [Config] Drop no-op kpack from RDNA GEMM configs
* [Triton] [Config] Drop redundant matrix_instr_nonkdim from selected RDNA GEMM configs
* [Triton] [Config] Extend RDNA matrix_instr_nonkdim cleanup
* [Triton] [Config] Retain matrix_instr_nonkdim for shared GEMM consumers
(cherry picked from commit 9ef3e41)
* mxfp8 gemm cga update, update (#5055)
Co-authored-by: Satya Nikhil Kodukula <nikhil.kodukula@gmail.com>
(cherry picked from commit 693701c)
* stop fused_bmm_rope_kv_cache from using batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant configs (#5877)
(cherry picked from commit a2650cd)
* [Triton/Gluon] Add Triton-based Conv3D kernels (#5952)
* conv3d implementation on Triton
* Remove unnecessary convolution compatibility aliases
* Update the README file.
* Add Conv2D weight-pack cache cleanup and updated the tests.
* Simplify the Conv2D Winograd test guard
* Move Conv3D device validation out of shape helper
* Add diagnostics to convolution test assertions
* Removing benchmark related tests from unit test
* Add context to Conv3D numerical assertions
* Deduplicate Conv2D and Conv3D blocked-layout kernels
* Simplified benchmark and fixed formatting issue on test file
* Split Conv2D and Conv3D cache-clear tests into their owning suites
* Deduplicate Conv2D and Conv3D Winograd transforms
* Deduplicate the Conv2D and Conv3D Winograd filter transform
* Consolidate Conv2D Winograd launch paths
* Deduplicate Conv2D and Conv3D prepack helpers and cache wrappers
* Black formatting
* Replace convolution test helper prints with logging
* Use safe configs for Conv3D on CDNA
(cherry picked from commit c8325e0)
---------
Signed-off-by: Jin Tao <jin.tao@amd.com>
Signed-off-by: Martin Domanský <ragua@email.cz>
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
Signed-off-by: Pham Binh <phamhuuthanh.binh@amd.com>
Signed-off-by: Lin, Soga <soga.lin@amd.com>
Signed-off-by: jiacao-amd <jiahui.cao@amd.com>
Co-authored-by: siliangchen-amd <SiLiang.Chen@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Lingpeng Jin <103567126+valarLip@users.noreply.github.com>
Co-authored-by: Jesús Carabaño <jcaraban@users.noreply.github.com>
Co-authored-by: Leo <drleonid@amd.com>
Co-authored-by: liyjiang <liying.jiang@amd.com>
Co-authored-by: gbyu-amd <Guanbao.Yu@amd.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: Jin Tao <jintao12@amd.com>
Co-authored-by: jin.tao@amd.com <jin.tao@amd.com@tus1-p15-g43.tus.tensorwave.lan>
Co-authored-by: Felix Li <felix.li@amd.com>
Co-authored-by: Martin Domanský <8312516+Ragua1@users.noreply.github.com>
Co-authored-by: Mehmet Cagri <mehmet.kaymak@amd.com>
Co-authored-by: Satya Nikhil Kodukula <nikhil.kodukula@gmail.com>
Co-authored-by: Rohan Potdar <rohanpotdar138@gmail.com>
Co-authored-by: Nimit Patel <61071220+NimitPtl@users.noreply.github.com>
Co-authored-by: Muhammad Ahmed <mm.ahmed2202@gmail.com>
Co-authored-by: yifehuan-amd <yifehuan@amd.com>
Co-authored-by: Pham Binh <phamhuuthanh.binh@amd.com>
Co-authored-by: akii96 <aakif.nawaz@amd.com>
Co-authored-by: li,xiangxiang <xiangxli@amd.com>
Co-authored-by: jhchouuu <jiahzhou@amd.com>
Co-authored-by: HaonanWang98 <hwang@amd.com>
Co-authored-by: sogalin_codegen <39478626+sogalin@users.noreply.github.com>
Co-authored-by: vorapolsiloai <115975949+vorapolsiloai@users.noreply.github.com>
Co-authored-by: jiacao-amd <jiahui.cao@amd.com>
Co-authored-by: Jiahui Cao <jiacao@crs-m2m-cpu-spur-014.us-east2-a.compute.internal>
Co-authored-by: WuLei-AMD <leiwu@amd.com>
Co-authored-by: Hyunjune Kim <132782704+hyjuunn@users.noreply.github.com>
Co-authored-by: Shao-Chun Lee <Shao-Chun.Lee@amd.com>
Co-authored-by: Saeid Rostami <123997133+saeid-rostami@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
github-actions
Bot
changed the title
valarLip
left a comment
Motivation
gfx950 (#5332) and gfx1250 (#5656) shipped the same operator as two parallel implementations: two
JIT modules, two schedule builders, two Python surfaces and two op tests. Every fix had to land
twice. Two of them had already diverged into the same kind of hole: a compile-time stride sitting
behind a check that looked only at the rank. This PR folds both into one module, one schedule, one
API and one test. The device bodies keep their math.
Technical Details
One JIT module.
module_pa_mqa_logits_mxfp4_opuscarries both device bodies. The sources livein
csrc/kernels/opus_mqa_logits/pa_mqa_logits_mxfp4/, which leaves room for an fp8 sibling:pa_mqa_logits_mxfp4_opus.hpa_mqa_logits_mxfp4_sched.cuhpa_mqa_logits_mxfp4_gfx950.cuhpa_mqa_logits_mxfp4_gfx1250.cuhpa_mqa_logits_mxfp4_kernels.cuEach body compiles only for its own target (
__gfx950__/__gfx1250__) and compiles to an emptystub on every other pass, so one TU links both launch symbols. The fat binary is built per
--offload-arch, the idiomopus_gemmalready uses. Namespaces:opus_logitsholds the sharedschedule.
opus_logits::gfx950andopus_logits::gfx1250each hold their arch'spa_mqa_logits_mxfp4_traitsandpa_mqa_logits_mxfp4_kernel. The four per-arch source files andthe separate
module_pa_mqa_logits_mxfp4_gfx1250_opusare deleted.One schedule. gfx1250's two-phase
build_tiles+build_schedis now the only builder, andgfx950's per-row builder is retired. gfx950 runs at
q_per_block == 1. There a tile is one row,and the tile-union window is that row's own, so the shared builder emits exactly the per-row
record the MFMA kernel already reads.
The tile cut is skipped at
q_per_block == 1. At that setting the cut is the identity (tiletis rowt), sobuild_tilesis not launched. That is one kernel fewer per forward onevery gfx950 call and on gfx1250's MTP=1 decode instance, and
plan_buffersno longer allocates acu_tilesbuffer that nothing reads.One runtime-arch dispatch.
pa_mqa_logits_mxfp4_fwd_schedreads the arch once per process,then dispatches on
(q_per_block, block_k). Any combination it does not recognise raises; nonefalls back to a default. Mixed-arch nodes (one process driving both gfx950 and gfx1250) are out of
scope, as with
opus_get_arch_info().Host-visible bounds, on both arches.
num_rows) goes down to C++. There it is checked againstq,weightsand
out, and against the lengths oflocal_ends,local_startsandrow_to_batch.block_tablesmust be at least as wide as the KV tiles need. A tile is 64 or 256 tokens, and aCTA reads every page of its last tile.
pa_mqa_logits_mxfp4_block_table_widthreturns therequired width.
row_id >= num_rows. This test is uniform acrossthe CTA, so a mismatched
cta_infoloses rows instead of reading or writing out of bounds.Caller contract (module docstring and header):
negative
row_to_batch.q_per_block > 1,cu_seq_q[batch] <= total_q, because the tile cut readscu_seq_qondevice. With
row_to_batch=None, a tile uses its first row'sblock_tablesrow, so the rowsof a sequence must share it.
At
q_per_block == 1,cu_seq_qdoes not trim rows: every row oftotal_qis scheduled.Condition 3 is what keeps pad rows inert. ATOM already meets it: its decode pad rows carry
row_to_batch = -1andlocal_ends = 0.One Python surface.
aiter/ops/opus/pa_mqa_logits_mxfp4.pyexposesplan_buffers/plan/pa_mqa_logits_mxfp4, parameterized by a per-archMqaLogitsVariant(name, q_per_block, block_k, cta_resident, arch):qlen4_kv64(default)qlen1_kv64qlen1_kv64(default)qlen1_kv256The gfx950
cta_residentis #5332'sSCHED_CTA_TARGET = 1024, used for both widths.One op test,
op_tests/test_pa_mqa_logits_mxfp4_opus.py, runs on both arches. It buildsnatural layouts on gfx1250 and MFMA-permuted layouts on gfx950.
-mllvm -enable-post-misched=1applies to gfx950 only. It is load-bearing there: aiter's baseJIT flags set it to 0, and without
=1the tuned MFMA body is rescheduled. The gfx1250 body keepsthe base
=0, as its own module on main did. The flag is a per-arch expression inoptCompilerConfig.json, the formmodule_deepgemm_opusalready uses. The gfx1250 A/B thatdecided this is below. The expression is true whenever the build targets gfx950 (the running GPU or
GPU_ARCHS), andcompiler flags are per TU. So a single build for both
gfx950;gfx1250would still give thegfx1250 body
=1; aiter's release default (gfx942;gfx950) does not build gfx1250.kargs loses six fields that no kernel read (144 → 112 B). The
offsetofpins are re-derived.API changes
aiter.ops.opus.pa_mqa_logits_mxfp4_sched,_build_sched,_sched_slotsand_sched_buffer_ints. Before this PR they were also reachable asfrom aiter import .... They have no callers: gfx950 users, ATOM included, still run FlyDSL. Theyare replaced by the shared
plan_buffers/plan/pa_mqa_logits_mxfp4. Thecta_target/block_kknobs becomeMqaLogitsVariant, chosen by name.aiter.ops.opus.pa_mqa_logits_mxfp4, andthe signatures of
pa_mqa_logits_mxfp4,_plan_buffersand_variants. The only removal isplan(block_k=...)together withMqaLogitsPlan.block_k, which was gfx950-only and on gfx1250raised for anything but
None. The additions are backward-compatible: trailingMqaLogitsVariant.archandMqaLogitsPlan.num_rows, and adevice=keyword onpa_mqa_logits_mxfp4_block_table_width.Test Plan
-mllvm -enable-post-misched=1on gfx1250, including VGPR and occupancy.Test Result
gfx1250 (ROCm 10.1.0, HIP 7.16.26315, AMD clang 23.0.0git
d177931e), headdc13ff647, module rebuilt cleanRe-measured at the final head. The previous table in this PR was taken at 79463da, before the
review round; its values are in brackets for comparison.
cu_seq_qover-claim), NaN E8M0 8/8,err = 0on all 27 perf rows [was: 28/28, 6/6, 50 rows]qlen4_kv64/qlen1_kv64qlen1_kv64: 2 → 1 (small), 3 → 2 (split);qlen4_kv64unchanged (structural, from the code)post-misched on gfx1250: off. With
=1the flag took effect (it is appended after the base=0, and the last one wins) as a pure post-RA reschedule. Resources were identical on and off:269 VGPR / occupancy 3 / no spill / LDS 9728 on both instances. The reordered code had fewer
instructions (
qlen4_kv641623 vs 1641: more VOPD, fewers_delay_aluand nops).Timing: 16 paired alternating runs, done twice, min-of-16 per shape.
11–12/16 pairs, in both runs and on both instances. The 79463da data already had it (median
1160.8 vs 1127.0 µs, 12/16 pairs), but there the min ratio (0.998) hid it.
A hand-scheduled body should not carry a flag that regresses a production shape, so
dc13ff647gates the flag on gfx950. The gfx1250 ISA is now line-for-line the flag-off build, with the same
flags as main's gfx1250 module. The cost is the prefill gain above: a third 16-pair run of the
final head against the flag-on build measures prefill 0.9–2.7 % slower, chunked bs4 no longer
slower, and a geomean of 1.0095 over 27 shapes.
gfx950 (ROCm 7.2.4), head
dc13ff647, module rebuilt cleangfx950 now runs through the unified module, schedule and API. Every call runs at
q_per_block == 1.qlen1_kv64is the production instance.qlen1_kv256is compiled, and it istested for both correctness and perf.
Correctness. Op test, both instances:
err = 0on every rowKernel and schedule vs #5332.
the prologue: one scalar load fewer after the kargs trim, plus the new
row_idguard.qlen1_kv64's schedule table is byte-identical to [HIP] [OPUS] [JIT] [GFX950]OPUS PA MQA Logits MXFP4 #5332's. This was verified at the unificationcommit, and it is unchanged since: 36/36 grid cells match from 8e03a1a to 3bcbfcc.
build_tilesis skipped atq_per_block == 1.Performance vs main.
8253efc40(#5332 standalone) against this head,qlen1_kv64,min-of-15 over 2 alternating rounds, same idle card, µs:
The median is +0.14 % and no shape moves by 5 % or more. Fresh bs1/bs2 and the ragged
m1s128L16 / m4s32L16 / m4s32L4 shapes come out +1.6 to +3.4 % in both rounds. Nothing in the
review round moved them, and the long prefill (chunked) is about 1 % faster.
Submission Checklist