[shay-li77]

Motivation

aiter's MXFP4 paged MQA-logits op (aiter.ops.opus.pa_mqa_logits_mxfp4) serves the DeepSeek-V4
sparse-attention indexer on gfx1250 and gfx950. Its API was built around one kernel family:

• the caller names a kernel variant per buffer pool (pa_mqa_logits_mxfp4_plan_buffers(..., variant="qlen1_kv64")), so every new instance is a caller change;
• every call needs a per-forward plan (tile cut + schedule) and a row_to_batch map;
• only 64 index heads (V4) are covered. DeepSeek-V4.1's indexer (32 heads, full layers on
  PAGEs of 8-row pages, candidate layers with one table row per query row) has no OPUS path;
• gfx950 has one-row kernels only, so MTP > 1 decode reads a sequence's KV once per row.

This PR replaces that API with one interface for V4 and V4.1 on both arches, where the caller
passes static bounds and the op picks the kernel instance and the schedule mode itself.

Technical Details

API

from aiter.ops.opus.pa_mqa_logits_mxfp4 import (
    mqa_logits_mxfp4, mqa_logits_mxfp4_config, mqa_logits_mxfp4_schedule)

cfg   = mqa_logits_mxfp4_config(dev, heads=64, page_size=64, max_rows=..., max_seqs=...,
                                max_qlen=4, max_seq_len=...)   # a bucket's or a step's shapes
width = cfg.block_table_width                                   # min block_tables.shape[1]
buf   = cfg.schedule_buffers()                                  # optional, caller-held
sched = mqa_logits_mxfp4_schedule(cfg, query_start_loc, row_ends, buffers=buf)  # optional, per forward
out   = mqa_logits_mxfp4(q, q_scale, kv_cache, kv_scale, weights, block_tables,
                         query_start_loc, row_ends, cfg=cfg, schedule=sched, out=out)

• One argument list. Rows are grouped into sequences by query_start_loc. Sequence b
  reads its keys through block_tables[b], an entry naming pages_per_block pages, and row r
  sees the prefix [0, row_ends[r]). There is no row_to_batch and no row_starts. A band of
  rows is the sliced row tensors plus query_start_loc.clamp(start, start + count) - start.
  out is [rows, >= max_seq_len], or flat with out_offsets.
• No instance names. The config reads static bounds only, picks the instance and its
  schedule-free mode, and sizes the grid from the instance's measured occupancy.
• The config is a value. It holds no device memory and compares by its bounds and pick, and
  the 1024 most recent are cached, so a caller may ask
  for one a step with that step's shapes: prefill's varying rows, sequences and widths do not
  accumulate configs. A schedule or schedule buffers built from one serve any equal one.
• Optional schedule. With a schedule the launch runs Table. Without one it runs Free
  (decode) or Prefill. No path reads device data on the host, so every path can be captured in
  a CUDA graph.
• Layouts are main's: natural 3-D on gfx1250, MFMA-permuted 4-D on gfx950. The op checks the
  ndim.

schedule mode	how work is assigned
Table	a 32-byte record per CTA, built on device once a forward and reused by every layer
Free	every CTA derives its share from query_start_loc / row_ends
Prefill	every row group is cut into a static number of parts

Free cuts a decode batch of up to 64 sequences with the same number of rows each, whatever
their lengths, in one pass over the lane-per-sequence bounds; other batches take the general cut.

Instances and selection

The C++ registry exports each instance's attributes. Python's per-arch rule picks by attributes,
and the launch dispatches by registry id. Names read
h<heads>[g]_r<rows a wave>w<waves>s<stages><T|L>[n][x4]:

• g: the instance gathers 8-row entries;
• T: the key scales travel with the tile; L: they are loaded per lane;
• n: the KV is loaded nontemporal (decode: no other CTA reads a walk again);
• x4: a CTA's four waves are schedule slots of their own.

The picks below are the config's own, and each was the best instance measured for its scenario.

arch	model	instances	decode pick, MTP 1 / 2 / 4 / 6 / 8	prefill pick
gfx1250	V4.1 full layers	h32_r1w1s2L, h32_r1w2s2T, h32_r3w1s2T, h32_r2w2s2T, h32_r3w2s2T, h32_r4w2s2T	h32_r1w1s2L / h32_r1w2s2T / h32_r2w2s2T / h32_r3w2s2T / h32_r4w2s2T (MTP 3: h32_r3w1s2T), Free	h32_r4w2s2T, Prefill
gfx1250	V4.1 candidate layers	h32g_r1w1s4L, h32g_r1w1s3L	h32g_r1w1s4L, Prefill	h32g_r1w1s3L (from 16384 rows), Prefill
gfx1250	V4	h64_r1w1s2T, h64_r1w2s2T, h64_r1w4s2T, h64_r1w6s2T	h64_r1w1s2T / h64_r1w2s2T / h64_r1w4s2T / h64_r1w6s2T / h64_r1w4s2T, Table	h64_r1w4s2T, Prefill
gfx950	V4.1 full layers	h32_r{1,2,3,6,8}w1s3Ln, h32_r4w1s2Ln, each also x4; h32_r4w1s3L	h32_r1w1s3Ln / h32_r2w1s3Ln / h32_r4w1s2Ln / h32_r6w1s3Ln / h32_r8w1s3Ln (MTP 3: h32_r3w1s3Ln; x4 from max_rows x max_seq_len >= 12M), Free	h32_r4w1s3L, Prefill
gfx950	V4.1 candidate layers	h32g_r1w1s2L	h32g_r1w1s2L, Prefill	h32g_r1w1s2L, Prefill
gfx950	V4	h64_r1w1s2L (main's qlen1_kv64), h64_r1w2s4T, h64_r1w4s4T (new: a sequence's rows share each KV tile through LDS)	h64_r1w1s2L / h64_r1w2s4T / h64_r1w4s4T / h64_r1w2s4T / h64_r1w4s4T, Table	h64_r1w1s2L, Prefill

The caller only chooses whether to pass a schedule. Each configuration was ranked by kernel
time, with a schedule build charged over the layers that share it: V4's 30 indexer layers, and
V4.1's three ratio-2 full layers or four candidate layers.

• V4 decode should pass one. On gfx1250, Free reaches 0.88-0.90 of the best.
• V4.1 should not. A 9-34 us build does not pay off over three layers, and Free reaches
  0.98-1.00 of the best.
• Prefill should not either. Table and Prefill tie on kernel time, so the build (up to 47 us
  on gfx1250 and 104 us on gfx950 at 16384 rows) decides it.

A tier above a CTA's rows runs two or three CTAs a sequence. Not registered on gfx950:

• six- and eight-row qshare CTAs: two or three smaller CTAs were faster at every MTP;
• main's qlen1_kv256: it was never the best choice.

V4.1 on gfx950 reads V4's 64-row layout: an 8-row page is an aligned eighth of a 64-row
group, and a full layer's table entry names a run of groups.

• Full layers run the rows kernel (_gfx950_rows.cuh): a wave scores up to eight rows of one
  sequence off each KV tile, with MFMA 32x32x64, a row's 32 heads one m-tile. Decode loads the
  KV nontemporal; prefill keeps it cached for the sequence's next rows.
• Decode in Free: one- and two-row waves get at least 32 row-tiles a slot (and no fewer slots
  than SIMDs), and large launches use the x4 CTAs, whose four slots are cut once.
• Candidate layers gather 8-row pages in the one-row kernel. A CTA's walk grows with the
  launch's size, and a row's parts are a multiple of the 8 XCDs, so the rows of a request, which
  share most of their pages, meet in one L2.
• ATOM writes gfx950's FP4 plane in FlyDSL row-group's 8-row page layout; scoring it with this op
  needs the plane in the 64-row layout.

Sources: csrc/kernels/opus_mqa_logits/pa_mqa_logits_mxfp4/, laid out as opus_gemm /
opus_moe are:

• .h: the interface;
• .cu: the registry, the Table builder and the launch;
• _pybind.cu;
• _common.cuh: shared device helpers;
• a set per arch: _<arch>.cu (traits, instances, registry), _<arch>_common.cuh (CTA assignment,
  tile math) and one .cuh a kernel (_gfx1250_{h32,h64}, _gfx950_{kv,qshare,rows}).

Caller contract

What the device cannot check, stated in the header and the module docstring:

1. row_ends does not decrease along a sequence's rows;
2. row_ends[r] <= max_seq_len, within what block_tables[b] covers;
3. query_start_loc[-1] <= rows; rows past it (graph padding) are neither read nor written;
4. on gfx1250 a page's offset fits 32 bits (its plane row); gfx950 addresses each page on its own,
   so the kv_cache is not bounded at 4 GiB, except that a candidate layer reads within its first
   2^32 keys (256 GiB, which the launch checks).

Breaking changes

• Removed: pa_mqa_logits_mxfp4, pa_mqa_logits_mxfp4_plan,
  pa_mqa_logits_mxfp4_plan_buffers, pa_mqa_logits_mxfp4_block_table_width,
  pa_mqa_logits_mxfp4_variants, MqaLogitsVariant, MqaLogitsBuffers, MqaLogitsPlan. A stale
  caller fails at import.

• ATOM (9df8e84) imports the removed names at atom/model_ops/indexer/rows.py:11,
  atom/model_ops/indexer/score.py:27 and atom/model_ops/attentions/deepseek_v4_attn.py:555,
  and moves to the new API in the same upgrade. TODO: link the ATOM PR, and add ci:atom here
  so aiter CI runs ATOM's tests against this branch.

  ATOM today	after
  plan_buffers(variant=...) a pool	a cfg a bucket (max_qlen), with cfg.schedule_buffers()
  plan(cu_seqlens_q, local_ends, row_to_batch=...)	mqa_logits_mxfp4_schedule(cfg, query_start_loc, row_ends)
  a prefill chunk's host cumsum + row_to_batch slice	query_start_loc.clamp(start, start + count) - start
  launch with plan	launch with query_start_loc, row_ends, cfg, schedule

• JIT: the module keeps main's name (module_pa_mqa_logits_mxfp4_opus). Rebuild a build
  left over from main with AITER_REBUILD=1.

Test Plan

Op test (rewritten for the new API), on each arch:

python3 op_tests/test_pa_mqa_logits_mxfp4_opus.py

It runs every instance the arch registers in every mode it compiles, against a torch reference
(gfx950's reference rounds after each 32-element block, as its MFMA does):

• decode: V4 at bs 1 / 16 / 64 / 128 x MTP 1 / 4 / 6, 64K context; prefill: 16384 rows at
  4096 / 16384 rows a sequence (V4.1 full-layer and candidate shapes as well, on both arches);
• semantics:
  • rows: every row against the reference, every mode bit-identical to its instance's Table;
  • edges: nine sequences, each carrying one source of NaN / Inf (E8M0 0xFF key or query
    scales, with and without a zero q block; NaN / Inf weights; overflowing products), and
    nothing written past a window;
  • bands: a forward scored in bands of rows through clamped query_start_loc;
  • padding: graph-padding rows neither read nor written;
  • graph: capture and replay with inputs refreshed in place;
  • errors: the other arch's layout, a head count other than the config's, a narrow block
    table or out, a schedule built from other tensors, rows over max_rows, int64 row_ends,
    a page geometry no instance reads -- each rejected with a message.

Performance against what ATOM runs today.

• gfx1250, V4.1: FlyDSL's flydsl_fp8_paged_mqa_logits ([FlyDSL] [gfx1250] Paged FP8 MQA logits with 8-token KV pages #6863), which ATOM runs for V4.1's
  index scoring on gfx1250 since [gfx1250] DeepSeek-V4.1: run on gfx1250 (MI455X) with FlyDSL kernels ATOM#2617 (atom/model_ops/indexer/score.py). It is called
  as ATOM calls it:

  • V4.1's default FP8 index plane, preshuffled in 8-row pages (Preshuffle=True, KVBlockSize=8);
  • every row its own item (next_n = 1) on its own table row, for full and candidate layers
    alike.

  ATOM's FP4 index plane (--index-cache-dtype fp4) is gfx950-only, so on gfx1250 V4.1 FP4
  scoring exists only through this PR. An FP8 key is 132 bytes against FP4's 68, and an FP8 row
  reads its own window where our rows of a sequence share theirs. Part of that gap is the format
  and the sharing, not the kernel.

• gfx1250, V4: main's OPUS op (the API this PR removes), as ATOM calls it today: qlen1_kv64 at
  MTP 1, qlen4_kv64 otherwise and for prefill (ROCm/ATOM 9df8e84,
  atom/model_ops/indexer/rows.py), its plan built once a forward.

• gfx950: FlyDSL row-group ([HIP] [FlyDSL] feat(flydsl): row-group FP4 paged MQA logits for 8- and 64-row pages #6145), which ATOM main uses for decode (feat(v41): mono decode, FP4 index plane, shared mono framework ATOM#2479) and prefill
  (refactor: one graph mechanism per region, settled DP step shapes, cheaper V4.1 decode prep ATOM#2688), called as ATOM calls it (make_fp4_mqa_plan for the call's shape). For V4.1
  it is ATOM's FP4 scorer (--index-cache-dtype fp4, flydsl_pa_mqa_logits_fp4 on 8-row pages),
  timed on its own 8-row page layout of the same plane. ATOM's default FP8 plane, scored by
  Triton's deepgemm_fp8_paged_mqa_logits, is not compared.

Method (gfx1250; harness in ROCm/opus-ops opus_logits/gfx1250/tools/mqa_v2/, run_pr.sh +
merge_min.py + report_pr.py):

• Timing: device time of the kernels a call launches, from the torch profiler, cold (a 2 GiB read
  before every call), the median of 50 calls (20 for prefill). Each arm is the minimum over every
  pass that timed it: 2-3 full passes of 2 reps each (the second rep in reverse order). The card
  is shared, and another job only ever slows a call down. The schedule build is timed separately
  (it runs once a forward, not once a layer).
• main's op runs in a second process right after ours, on the same inputs from the same seeds.
  The two builds share a module name, so they cannot be loaded in one process.
• Checks: every OPUS arm, main's included, is compared with an f64 reference on 16 sampled rows
  (to 1e-5 of the largest logit), and every mode must equal its instance's first bit for bit. The
  FP8 arm reads random FP8 data and is timed only. Its call was checked separately: on 16-row
  pages it equals Triton's deepgemm_fp8_paged_mqa_logits (relative error 2e-7), and on 8-row
  pages it writes exactly each row's window.
• Shapes (head dim 128; context lengths are raw tokens):
  • V4.1 full layers: 32 heads, ratio 2 (layers 2 / 8 / 14: the row at position p sees
    (p + 1) // 2 keys), 8-row pages, a 256-token PAGE a table entry. Decode at MTP 1 / 2 / 4 /
    6 / 8 on V4's contexts below, at most 16M keys a call (22 fixed and 13 ragged shapes a MTP);
    prefill on V4's fresh and chunked shapes.
  • V4.1 candidate layers: each row has its own table of the 2048 8-row blocks kept for it (at
    most 16384 keys). Decode at bs 1-256 x MTP 1 / 6, and two 16384-row prefills.
  • V4: 64 heads, 64-row pages, CSA ratio 4 (the row at position p sees (p + 1) // 4 keys).
    • Decode, MTP 1 / 2 / 4 / 6 / 8 (each sequence scores its last MTP tokens): bs 1 / 4 / 16 /
      64 / 128 / 256 x 8K / 32K / 128K / 512K (23 shapes), and ragged batches with each context
      drawn uniformly from [L / 8, L], bs 4 / 16 / 64 / 128 / 256 x L = 32K / 128K / 512K (14
      shapes). A call reads at most 16M KV keys, which leaves out bs 256 at 512K.
    • Prefill, 16384 rows a call: fresh (ctx = rows a sequence) and chunked.
    • ATOM's long-prompt bands: one request's 1024-8192 consecutive rows at 64K-1M context, the
      band sizes ATOM's workspace rule gives.

Method (gfx950; harness in ROCm/opus-ops opus_logits/gfx950/tools/mqa_v2/, run_pr.sh +
report_pr.py): device time of the kernels a call launches, from the torch profiler, cold (a
1 GiB read before every call evicts L2 and the MALL), the minimum over 3 reps of the median of 50
calls (20 for prefill), every arm on one GPU in one process. The schedule build is timed
separately (it runs once a forward, not once a layer). Before timing, every arm's output is
compared with h64_r1w1s2L Table: OPUS arms bit-exact, row-group within 1e-5 of the row's
largest logit. Shapes: gfx1250's V4 shapes above.

Method (gfx950, V4.1): the same cold, kernel-only timing, with every call timed on its own. The
MI355X was shared with another job during these runs, and another job only ever slows a call
down, so the two arms alternate call by call, 30-60 calls an arm (40 for prefill), and each arm
is its minimum over 2-6 passes on two GPUs. Ours is the config's own pick for the call's shapes,
without a schedule; both arms get the batch's longest window as the logits width. Before timing,
sweep_v41.py compares 16 sampled rows of every arm with the reference (to 1e-5 of the largest
logit) and every mode of an instance bit for bit. Shapes, on ATOM's 128-token PAGE (8 / 16 pages
a table entry at ratio 2 / 1):

• full layers, ratio 2 and 1: decode at MTP 1 / 2 / 3 / 4 / 6 / 8, bs 1 / 4 / 16 / 32 / 64 / 128
  / 256 x 32K / 128K tokens, with fixed contexts and ragged ones (each context drawn from
  [L / 8, L], each sequence 1-MTP rows): 168 + 168 shapes. Prefill: 16384 rows, fresh at 512 /
  1024 / 4096 / 16384 rows a sequence, chunked at 32K-256K, and ATOM's bands of one request's
  1024-16384 rows at 32K-1M: 27 shapes.
• candidate layers: decode at bs 1 / 4 / 16 / 48 / 64 / 128 x MTP 1 / 6 x 32K / 128K (24 shapes);
  prefill on 11 of the full-layer prefill shapes.
• ATOM's width: the full layers' logits width at max_model_len / ratio (V4.1's 1M tokens), the
  contexts ragged at 4K-32K or 16K-128K, bs 1-64 at MTP 1 / 6 and bs 128 / 256 at MTP 1, and a
  forward's eight calls (three ratio-2, one ratio-1, four candidate).

Environment

	gfx1250	gfx950
GPU	gfx1250, 256 CUs	MI355X (gfx950)
image	rocm/fw-bringup-mlperf@sha256:e7fcdb088de972009dcd02778eadd5552374a30a3edfd26a3b35218cb7df116d (gfx1250-atom-20261012)	rocm/atom-dev@sha256:066841903594e8d691b1c528123dfd2e17eacd293acfde947d4fb02afffd57db
ROCm	10.1.0 (AMD clang 23.0.0git d177931e)	7.2.4 (AMD clang 22.0.0git, roc-7.2.4 26084)
PyTorch	2.11.0+rocm10.1.0a20260811	2.10.0+rocm7.2.4
FlyDSL	0.3.3 (the image's), for the V4.1 baseline	0.4.0 (aiter's pin)
aiter	d7bbac8bb; main 778d74d57 for the baseline	d7bbac8bb (V4), 791fc8874 (V4.1)

Test Result

Correctness

check	gfx1250	gfx950
decode / prefill tables, every instance x mode	all err = 0: V4.1 24 decode, 4 prefill shapes x 18 arms; candidate 12, 2 x 4; V4 12, 2 x 12	all err = 0: V4.1 24 decode, 4 prefill shapes x 39 arms; candidate 12, 2 x 3; V4 12, 2 x 9
semantics: rows / edges / bands / padding / graph / errors	V4.1 180 / 18 / 18 / 18 / 18 / 8, candidate 16 / 4 / 4 / 4 / 4 / 8, V4 204 / 12 / 12 / 12 / 12 / 9, all PASS	V4.1 390 / 39 / 39 / 39 / 39 / 9, candidate 12 / 3 / 3 / 3 / 3 / 9, V4 171 / 9 / 9 / 9 / 9 / 9, all PASS
performance sweeps, 16 sampled rows a shape	401 shapes (V4.1, candidate and V4), plus main's op on the 204 V4 ones: every OPUS arm within 1e-5 of the f64 reference, every mode of an instance bit-identical	V4: every OPUS arm bit-exact with h64_r1w1s2L Table, row-group within 1e-5. V4.1: 336 full-layer decode shapes (every decode instance in Free) and 27 prefill shapes, 24 candidate decode and 6 prefill shapes: every OPUS arm within 1e-5 of the reference, every mode of an instance bit-identical

Performance

Each scenario runs its pick (Instances and selection) against what ATOM runs today (Test Plan).
x = baseline us / ours: above 1, ours is faster. Timing is cold and kernel-only, and decode
cells are geomeans at MTP 1 / 2 / 4 / 6 / 8. The per-shape tables are in ROCm/opus-ops
opus_logits/common/docs/PR6952_PERF_DETAILS.md.

arch	model	baseline	decode, fixed contexts	decode, ragged contexts	prefill	faster on (range)
gfx1250	V4.1 full layers	FlyDSL FP8 (flydsl_fp8_paged_mqa_logits)	1.47 / 2.00 / 2.70 / 3.00 / 3.15x	2.04 / 3.06 / 4.52 / 4.59 / 5.17x	fresh 4.21x, chunked 6.90x	decode 172 / 175 (0.89-10.39x), prefill 8 / 8 (3.38-7.04x)
gfx1250	V4.1 candidate layers	FlyDSL FP8	1.87x (MTP 1 and 6)	-	4.05x	decode 12 / 12 (1.27-2.78x), prefill 2 / 2 (3.87-4.24x)
gfx1250	V4	main's OPUS op	1.07 / 1.51 / 1.22 / 1.43 / 1.28x	1.06 / 1.66 / 1.25 / 1.53 / 1.34x	fresh 1.32x, chunked 1.44x, bands 1.46x	decode 183 / 185 (0.98-2.41x), prefill 19 / 19 (1.29-1.69x)
gfx950	V4	FlyDSL row-group	1.16 / 1.31 / 1.17 / 1.11 / 1.15x	1.25 / 1.46 / 1.28 / 1.20 / 1.18x	fresh 1.42x, chunked 1.01x, bands 0.98x	decode 156 / 185 (0.80-2.21x), prefill 8 / 19 (0.96-1.56x)
gfx950	V4.1 full layers	FlyDSL row-group (FP4)	1.06 / 1.06 / 1.04 / 1.04 / 1.02x	1.15 / 1.13 / 1.19 / 1.17 / 1.13x	fresh 1.22-2.29x (16K-token prompts 1.22-1.30x), chunked 1.12x, bands 1.10x	decode 242 / 336 (0.88-1.76x), prefill 27 / 27 (1.06-2.29x)
gfx950	V4.1 candidate layers	FlyDSL row-group (FP4)	0.98 / 1.03x (MTP 1 / 6)	-	fresh 2.30-6.55x (16K-token prompts 2.30x), chunked 1.37x, bands 1.29x	decode 13 / 24 (0.91-1.10x), prefill 11 / 11 (1.22-6.55x)

• gfx1250, V4.1. ATOM's FP4 index plane is gfx950-only, so this PR is V4.1's FP4 scorer on
  gfx1250.
  • The gain grows with MTP: a sequence's rows share each KV tile, while an FP8 row reads its
    own window.
  • Part of every gap is the format: an FP8 key is 132 bytes, an FP4 key 68.
  • Ours is behind only on a single sequence at 8K, by up to 11% at about 4 us.
• gfx1250, V4.
  • MTP 2 / 6 gain the most (1.43-1.66x): main runs both on its four-row CTA, which is half
    empty there.
  • MTP 1 / 4 / 8 use main's CTA shapes, so their gain is the kernel's.
  • Our schedule build costs about what main's plan does (8-25 against 7-22 us).
  • Long prefill runs at 6300-7600 TFLOPS, against main's 4300-5600.
• gfx950, V4.
  • Row-group leads at MTP 6 on identical contexts from bs 16-64 up, by up to 1.25x. Row-group
    reads a 6-row sequence's KV twice; h64_r1w2s4T reads it three times.
  • On one long request's prefill bands, row-group is 1-4% ahead (both run at 3000-3260
    TFLOPS). h64_r1w1s2L needs fewer cycles, but it moves more data out of L2 and runs at a
    lower clock under the power limit.
• gfx950, V4.1. FP4 against FP4: the format ATOM selects with --index-cache-dtype fp4.
  • Full-layer decode is ahead at every MTP, more on ragged contexts. It trails on 94 of 336
    shapes, by at most 12%.
  • At ATOM's logits width (max_model_len / ratio, 1M tokens unless lowered) the gap widens:
    ratio-2 calls 1.62x, ratio-1 1.47x, candidate 0.98x; a forward's eight calls 1.31x
    (1.02-2.17x). Our cut follows each row's window, whatever the width.
  • Long full-layer prefill runs at 2850-3300 TFLOPS, against row-group's 2600-3030.
  • Candidate decode is at parity.
  • Candidate prefill gains most on short fresh prompts. Row-group sizes its grid by the 16384-key
    candidate table (16 workgroups a row) whatever a row's window, so its time hardly falls with
    the prompt: 481 us at 512-token prompts, 1354 us at 16K. Ours follows the keys: 97 / 589 us.

Submission Checklist

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

🤖 Generated with Claude Code

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
ci:triton-355	Run the full Triton test suite on MI35X, not only the tests the change affects
ci:windows	Windows JIT op smoke test on the gfx1201 ROCm runner
ci:multi-gpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: DeepSeek-V4-Pro, MiniMax-M3-MXFP4
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 6952 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.

[Copilot]

🟡 Changes recommended

Missing dense-output width validation permits out-of-bounds writes, and previously supported gfx950 functionality has no replacement.

2 open issues

• Preserve gfx950 V4 OPUS support during API migration · New
• Reject undersized dense outputs before kernel launch · New

What changed in this PR

This PR replaces the legacy OPUS MXFP4 MQA-logits API with a registry-based interface for DeepSeek-V4 and V4.1 on gfx1250.

Changes:

• Adds configuration-driven kernel selection and Table, Free, and Prefill scheduling.
• Introduces specialized 64-head and 32-head kernels, including candidate-page gathering and flat outputs.
• Updates Python bindings and JIT configuration while removing the legacy implementation.

File	Description
csrc/​kernels/​opus_mqa_logits/​pa_mqa_logits_mxfp4/​pa_mqa_logits_mxfp4_opus.h	Removes the legacy API and architecture traits.
csrc/​kernels/​opus_mqa_logits/​pa_mqa_logits_mxfp4/​pa_mqa_logits_mxfp4_kernels.cu	Removes the legacy dual-architecture launcher.
csrc/​kernels/​opus_mqa_logits/​pa_mqa_logits_mxfp4/​pa_mqa_logits_mxfp4_gfx1250.cuh	Removes the previous gfx1250 kernel.
csrc/​kernels/​opus_mqa_logits/​mqa_logits_mxfp4/​mqa_logits_mxfp4.h	Defines the unified API, arguments, and registry.
csrc/​kernels/​opus_mqa_logits/​mqa_logits_mxfp4/​mqa_logits_mxfp4_sched.cuh	Adapts Table scheduling to sequence-based row groups.
csrc/​kernels/​opus_mqa_logits/​mqa_logits_mxfp4/​mqa_logits_mxfp4_pybind.cu	Connects the replacement Python bindings.
csrc/​kernels/​opus_mqa_logits/​mqa_logits_mxfp4/​mqa_logits_mxfp4_launcher.cu	Adds registry discovery, validation, and launch dispatch.
csrc/​kernels/​opus_mqa_logits/​mqa_logits_mxfp4/​mqa_logits_mxfp4_gfx1250.cu	Registers gfx1250 kernel instances and occupancy queries.
csrc/​kernels/​opus_mqa_logits/​mqa_logits_mxfp4/​mqa_logits_mxfp4_gfx1250_v41.cuh	Implements V4.1 full-layer and candidate-layer kernels.
csrc/​kernels/​opus_mqa_logits/​mqa_logits_mxfp4/​mqa_logits_mxfp4_gfx1250_traits.h	Defines shared layouts and kernel-specific traits.
csrc/​kernels/​opus_mqa_logits/​mqa_logits_mxfp4/​mqa_logits_mxfp4_gfx1250_h64.cuh	Implements the replacement V4 64-head kernel.
csrc/​kernels/​opus_mqa_logits/​mqa_logits_mxfp4/​mqa_logits_mxfp4_gfx1250_common.cuh	Adds shared reduction, NaN, and priority helpers.
csrc/​kernels/​opus_mqa_logits/​mqa_logits_mxfp4/​mqa_logits_mxfp4_gfx1250_assign.cuh	Assigns kernel work across the three scheduling modes.
csrc/​include/​rocm_ops.hpp	Replaces legacy binding declarations.
aiter/​ops/​opus/​pa_mqa_logits_mxfp4.py	Removes the legacy Python interface.
aiter/​ops/​opus/​mqa_logits_mxfp4.py	Adds configuration, scheduling, and unified launch APIs.
aiter/​jit/​optCompilerConfig.json	Updates source files and per-source compiler flags.

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

The retained findings identify incorrect large-cache addressing and missing dense-output width validation that can permit out-of-bounds writes.

4 open issues

• Use 64-bit page offsets for qshare and scale copies · New
• Restore host-side dense output width validation · New
• Reject undersized dense outputs before kernel launch
• Preserve gfx950 V4 OPUS support during API migration

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

The substantial GPU pipeline, scheduling, and NaN-handling changes need hardware correctness and performance validation plus final human review.

4 open issues

• Use 64-bit page offsets for qshare and scale copies
• Reject undersized dense outputs before kernel launch
• Preserve gfx950 V4 OPUS support during API migration
• Align advertised page range with launcher shift limit · New

1 resolved since last review

• Restore host-side dense output width validation

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

Human assessment is needed because a critical masked-store offset issue can overwrite valid tail logits and cause cross-chunk write races in both kernel variants.

No open issues 🎉

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

Human assessment is needed because duplicate-lane stores can corrupt tail logits and race other CTAs for accepted large row lengths.

No open issues 🎉

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

The architecture-specific pipeline rewrite requires GPU correctness and graph-replay validation on gfx950 and gfx1250, but the supplied results omit gfx950 V4.1 coverage for the newly selected 32-head kernels.

No open issues 🎉

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

A critical lane-bias wraparound allows unused head-reduction partners to overwrite valid output columns at supported large widths, requiring store guards in both kernel variants, including the gathered path.

1 open issue

• Mask unused head-reduction lanes to prevent output overwrite · New

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

Critical kernel bugs leave rows unwritten when batches exceed the cached prefix and allow supposedly masked stores to overwrite valid wide-row outputs and race their correct writers.

2 open issues

• Disable the free-mode shortcut when the batch exceeds the cached prefix · New
• Predicate stores for redundant lanes to prevent wide-row overwrites · New

1 resolved since last review

• Mask unused head-reduction lanes to prevent output overwrite

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

The architecture-specific kernel rewrite changes synchronization and compiler-sensitive NaN handling, requiring correctness and latency validation on gfx950 and gfx1250 GPUs unavailable in this review environment.

No open issues 🎉

2 resolved since last review

• Predicate stores for redundant lanes to prevent wide-row overwrites
• Disable the free-mode shortcut when the batch exceeds the cached prefix

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

The GPU pipeline and scheduling rewrite requires head-matched correctness and performance validation on both gfx950 and gfx1250, which was unavailable in this review.

No open issues 🎉

_{Review effort: Balanced}