[FlyDSL] pa_decode: MXFP4 K/V in the shared PA tile kernel (gfx950, query_length 1) by ganyi1996ppo · Pull Request #7056 · ROCm/aiter

ganyi1996ppo · 2026-10-28T02:17:42Z

Ready for review once the pending items under "Not validated" are settled.

Summary

MXFP4 (e2m1 + e8m0 scale per 32) K/V decode on gfx950 is now a mode of the existing fp8/bf16 compile_pa_decode_tile (kv_dtype="mxfp4"), not a separate kernel. This replaces the dedicated kernel this PR carried before (pa_decode_mxfp4_kernel.py, deleted). pa_decode takes the MXFP4 cache directly; the separate pa_decode_mxfp4 entry was removed (review asked to merge it into pa_decode), and ATOM#2770 calls pa_decode.

Size against origin/main, 9 files: source +857/-51, tests +716/-50 (the previous head 13a5814 was +1637/-0, three new files). The additions land inside the original fp8 files: pa_decode_kernel.py, pa_decode/{context,traits,op_lds,op_gemm,op_epilog,pipeline}.py (+579/-47 with the kernel file), pa_decode.py (+278/-4), and the MXFP4 cases in the existing op_tests/test_flydsl_pa_decode.py.

Input: key_cache/value_cache are uint8 data views [pages, kv_heads, 128, 64], key_scale/value_scale uint8 e8m0 views [pages, kv_heads, 128, 4] of the same pages (shared page/head strides, multiples of 16 B). query is float8_e4m3fn (used as stored) or bf16 (quantized to e4m3 per head row in the kernel); output is bf16. work_plan=None (MXFP4 only) runs one task per sequence and writes the row directly; a plan with max_partitions == 1 is served the same way.

Supported: head_dim 128, block size 128, query_length 1, up to 16 query heads per KV head, scale bytes other than 0/255 in valid rows.
Refused with ValueError/NotImplementedError/TypeError: block size 16/64, fp16 output, other head dims, more than 16 query heads per KV head, query_length > 1, sliding window, sinks, plan-free output with strides other than the query's, a task over 4096 pages (pass a smaller max_context_length), plan-free max_partitions other than 1, a query or output spanning more than 2**31 - 1 bytes. Also not offered, unlike the previous head: tile_depth, exact bf16 Q, fp16 output, plan-free multi-partition. ATOM#2770 and vLLM use none of these.

Dependencies

Pages come from the MXFP4 KV write kernel in #7057; the index-score kernels are #7060 (decode) and #7061 (prefill). No shared kernel files; each PR adds its own lazy-export lines. ATOM integration: ROCm/ATOM#2770 (draft).

PR Self-Review

Reviewed revision

f2a0579

Behavior and intent

Before: the fp8/bf16 PA could not read an MXFP4 cache, and this PR added a second 660-line kernel plus wrapper and test file for it. After: one tile kernel with an MXFP4 mode. Each wave owns 64 tokens of the 256-token tile and all 128 dims, so softmax, P (registers) and O are wave-private, the tile loop has no barrier, and the four waves merge once after it. QK is one scaled MFMA (fp4 K x e4m3 Q); V is widened to bf16 through a wave-private LDS tile and read back transposed for the bf16 PV MFMA. The fp8/bf16 paths are not changed. Kernel name carries the dtype, direct-output and staged-page values.

Validation and unverified scope

MI355X (gfx950), image rocm/atom-dev:nightly_202610261558 (FlyDSL 0.4.0), FlyDSL cache off, 2026-10-29/30. Aiter tree at this head's kernel sources over origin/main cf831b2 unless noted.

python3 -m pytest op_tests/test_flydsl_pa_decode.py on gfx950 (one idle card, aiter tree = this head over a 0dd2716 container tree; FlyDSL 0.4.1): 265 passed, 16 skipped (16 MXFP4 unowned-row geometries filtered by design), 4 min 20 s cold. Not run on gfx942 or gfx1250.
Deliberate breaks are caught: one output element off by 0.05 (bf16, fp8, mxfp4), a NaN, a wrong fp8 value_scale, rolled MXFP4 K block scales, unowned NaN rows with mask_unowned_kv=False, a corrupted plan record, a refusal that raises nothing or the wrong error. With the K scale byte selection shifted 15 of the 30 MXFP4 cases failed (earlier run).
fp8/bf16 ISA: the 85 original specializations dump byte-identical to origin/main (combined sha256 107a1be7ca6f, base and this head).
Accuracy against an fp64 reference on the same e4m3 Q, bf16 output, rel-L2: 'mid' data 2.24e-3 (ctx 2048), 2.26e-3 (2000), 2.15e-3 (16384); random data 0.9e-3 to 1.1e-3. The same 'mid' number as the dedicated kernel.
ATOM end to end (ganyi/m31-mono-fp4-pr f328ca296, TP4, MXFP4 band, real acceptance, one run, aiter tree combo2 = origin/main 2db75fc + [HIP] fused_qknorm_idxrqknorm: write MXFP4 index-K and main KV caches, emit an fp8 main Q #7057/[FlyDSL] MiniMax-M3 index score: MXFP4 index-cache decode scorer #7060/[FlyDSL] MiniMax-M3 index score: MXFP4 index-cache prefill scorer #7061 + this PA): GSM8K 1319, 5-shot, c32: flexible 0.9439 +- 0.0063, strict 0.8954 +- 0.0084 (previous 0.9416/0.8757 on the dedicated PA).
vLLM TP4 on the first version of this replacement (b2febdaa15, same kernel sources, the then-standalone wrapper): GSM8K 1319 flexible 95.30 +- 0.58, strict 89.76 +- 0.83 (reference 95.38/89.76); fixed-B decode p50 at about 130k context B=16 9.48 ms (reference 9.50). One run each.

Not validated: the final dispatch form (pa_decode taking the MXFP4 cache) has not been run in vLLM; the vLLM numbers above are from the earlier standalone wrapper. The ATOM throughput A/B (c16 AgentX, same session, dedicated vs this PA) is not reported here; it is pending and does not gate this push. gfx942 and block size 16/64 are refused, not tested. CI has not run these tests: the PR is a draft.

Microbench, same card, CUDA graph, ctx 2048, group 16, 3 interleaved runs, median us per call, rows 8/32/64/128/256/512/4096:

path	dedicated (`13a5814`)	this head
planned	8.59 / 9.37 / 11.83 / 14.90 / 20.95 / 33.46 / 240.93	8.00 / 9.17 / 12.16 / 15.21 / 20.92 / 32.95 / 239.03
plan-free (max_partitions 1)	18.56 / 19.02 / 19.22 / 19.61 / 20.48 / 29.40 / 214.25	15.36 / 15.76 / 15.95 / 16.44 / 17.77 / 29.24 / 214.87

At 64 and 128 rows planned is 2.8% and 2.1% slower than the dedicated kernel; everything else is equal or faster. The existing fp8 flydsl PA is unchanged (origin/main vs this head, 3 runs: 7.71/9.27/11.47/18.70/28.77/52.86/412.94 vs 7.72/9.28/11.48/18.63/28.82/52.85/413.92). Fold resources: 132 VGPR, 75,728 B LDS, no spills.

Bandwidth (TB/s)

Same card, CUDA graph, head bc2b541. Cell = FP8 flydsl / MXFP4 planned, both with a bf16 query. Query is [rows, 16, 128], one query token per row, 16 query heads, 1 KV head, page 128 (GQA 16); rows = num_q.

num_q \ ctx	512	1024	2048	8192
8	0.112 / 0.056	0.222 / 0.111	0.451 / 0.219	0.968 / 0.616
128	1.573 / 0.807	2.629 / 1.270	3.426 / 2.122	5.192 / 4.170
512	3.474 / 2.884	4.448 / 3.660	5.200 / 4.636	5.965 / 5.659

TB/s is the nominal KV payload per call (rows x ctx x the KV bytes the cache layout holds for that dtype, MXFP4 block scales included) divided by the median per-call rocprofv3 --kernel-trace GPU-timestamped dispatch time of tile kernel plus reducer. It is not an HBM counter measurement, and the nominal KV bytes differ per arm (MXFP4 holds about half of FP8's), so equal TB/s does not mean equal time. I draw no speedup from this table.

Notes and limits:

Both arms share one plan from the old benchmark policy get_recommended_splits(rows, 1), which caps max partitions at 8/4/1 for rows 8/128/512. This is not the ATOM default plan; the numbers are for this benchmark's plan, not a deployment.
Short shapes are latency and reducer dominated (rows 8 and 128 at ctx 512/1024 take about 10 to 14 us, of which the reducer is about 4 us), so their TB/s is not a bandwidth statement. At ctx 512/1024 the plan has fewer tiles per sequence than its cap, so columns are not a pure ctx scaling.
Rows 512: MXFP4 runs the direct-output kernel with no reducer; FP8 still runs tile plus a one-partition reducer (about 4.3 to 4.6 us of its time). That reducer is part of the FP8 vs MXFP4 gap there. Elsewhere both arms launch tile plus reducer.
ctx 8192 is a generic kernel shape, not the MiniMax M3.1 sparse-attention deployment shape.
One card, one container, one sweep, gfx950; single run per cell, 10 CUDA-graph replays, no repeat-round claim. Median over all calls of the cell (320 to 7680). Not cold-cache, no lower-bound claim.
Raw rocprofv3 traces and per-cell audits are on the author's machine only, not attached here: /shared/amdgpu/home/gyu_qle/ganyi/task-pa-replace/r7056-bw-runs/ctx2d_1103_000249/ (main_audit.csv, summary.md). Scripts: .../r7056-bw-runs/scripts/prof_cmp.py, prof_audit.py.

PR-specific question

The MXFP4 cases are added to the existing test_flydsl_pa_decode.py, whose test_pa_decode already compares against an fp32 reference over logical K/V. The cases that MXFP4 adds and the existing ones cannot express: the e2m1/e8m0 page packing and dequantization reference, unequal 32-element scale blocks (a block/scale mix-up), NaN/Inf and e8m0 255 garbage in rows no sequence owns, plan-free versus one-partition-plan bitwise equality, empty rows, strided output, and the refusals listed above. None of these is covered by the fp8/bf16 cases.

Changed op tests

op_tests/test_flydsl_pa_decode.py keeps the existing pytest structure and its three test functions (test_pa_decode_plan, test_pa_decode, test_pa_decode_unowned_kv_rows); the MXFP4 cases are new parametrized cases and helpers in it, and the file is run with pytest. .github is unchanged; op_tests/test_flydsl_pa_decode_mxfp4.py stays deleted.

MXFP4 operand validation

Review found the MXFP4 branch of pa_decode skipped checks the generic path makes. Confirmed on GPU against the previous head: a non-contiguous block_tables, a windowed work_plan (with one partition the plan is ignored, with more it is used) and a query/output with a strided head_dim axis gave wrong output or an illegal memory access; a context_lengths/block_tables batch different from the query was not refused; key_scale=None raised an AttributeError; an output or user scratch (exp_sums, max_logits, temporary_output, used with more than one partition) on another device reached the kernel as a device pointer. Now refused, with the generic path's checks and messages: batch mismatch, non-contiguous or off-device block_tables/context_lengths, non-unit head_dim stride, non-tensor scales, output/scratch on another device, work_plan.sliding_window != 0. Eleven regression cases in the refused-input table (they fail on c57f106, pass here). pa_decode.py is otherwise unchanged and the generic path is not touched.

🤖 Generated with Claude Code

github-actions · 2026-10-28T02:18:08Z

🏷️ 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 7056 --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.

…length 1) pa_decode_mxfp4 reads K and V as e2m1 data + e8m0 scales (one per 32 elements, along head_dim), each as two uint8 views with their own page/head strides, so a caller may use two tensors or one single-region page [block_size x 64 B data][block_size x 4 B scale]. It shares pa_decode's work plan and partition reducer and leaves pa_decode itself untouched. Per 16-token tile and wave: QK is one mixed fp4(A) x fp8(B) v_mfma_scale_f32_16x16x128_f8f6f4 with the K block's e8m0 byte as scaleA and Q used as stored (query_scale, a [1] fp32, multiplies the logits); a bf16 Q takes a reference path that widens K with v_cvt_scalef32_pk_bf16_fp4. V is widened to bf16 (exact) into a wave-private LDS tile and read back transposed with ds_read_b64_tr_b16 as the PV operand; P stays bf16. Invalid tokens get a zero V scale and a -inf score, so tail garbage (e8m0 0/255 included) cannot leak. The fp8 B operand of the 16x16x128 MFMA is not contiguous per lane: lane (g, s) holds Q head s elements 16g..16g+15 in bytes 0..15 and 64+16g.. in bytes 16..31 (measured on gfx950; A lane group g holds K elements 32g..32g+31). tile_depth defaults to 2: with independent blocks (2048 tokens/row, 16 q heads) it is 0.2-2.2 us faster than 4 for 1-512 rows and equal at 4096. Arithmetic is unchanged (a test pins default == depth 4 bitwise). Tests: rel-L2 ~1e-3 vs an fp32 reference over the dequantized K/V (the bf16 output rounding floor), ctx 1..2048, block_size 16/64/128, groups 1..16, several KV heads, plan tasks of 1..N tiles, garbage tails, query_scale. Co-Authored-By: Claude <noreply@anthropic.com>

… one task With a plan of max_partitions == 1 each sequence is a single task, so the kernel normalizes and stores the output row itself and the reduce kernel, its scratch and the pmax/psum/pout partials disappear (direct_output, auto on for that plan). Bitwise equal to split + reduce (tested for rows 1/8/64/512, bf16/f16, strided output, graph replay). A context-0 row has no task and is left unwritten on this path; the reduce path zero-fills it. Co-Authored-By: Claude <noreply@anthropic.com>

Tasks follow from the block index (CTA slot = sequence * P + part, P tiles-per-task chosen on the host from max_partitions, default the planner's 2 * CU task budget), so no plan_pa_decode launch and no work_info read: the context length, block-table entries and Q are loaded in parallel. The kernel writes the reduce info itself and zero-fills context-0 rows on both the reduce and the direct (P == 1) path. Bitwise equal to the planned path wherever the two splits coincide; tested for ragged contexts (1..5000, tasks of several tiles), block sizes 16/64/128, empty rows, the default-split switch at 255/256/257 rows, graph replay. Co-Authored-By: Claude <noreply@anthropic.com>

…ew checks Generated ISA is unchanged (82 specializations, identical hash). Co-Authored-By: Claude <noreply@anthropic.com>

Module-level gfx950 skip, checkAllclose + TFLOPS in the benchmark, and run the pytest checks from main() so `python3 op_tests/...` (what CI runs) tests it. Co-Authored-By: Claude <noreply@anthropic.com>

…s from the CU count Co-Authored-By: Claude <noreply@anthropic.com>

…equences The test no longer queries the GPU at import. Co-Authored-By: Claude <noreply@anthropic.com>

…stored) Co-Authored-By: Claude <noreply@anthropic.com>

…ved) launch_pa_decode_ps_reduce now takes the bound as max_partitions; same value as before: min(parts, capacity, tiles). Co-Authored-By: Claude <noreply@anthropic.com>

Removes the bf16-query reference path, the tile_depth and direct_output arguments, multi-partition plan-free and the fp8 comparison in the test benchmark. Generated ISA and outputs are unchanged for what remains. Co-Authored-By: Claude <noreply@anthropic.com>

… pytest Plain check functions with asserts, a case table and a pandas summary, try/except for the rejection cases (expect_error), one PA_DECODE name for the implementation under test. Same 39 correctness cases and the same timing sweep. Co-Authored-By: Claude <noreply@anthropic.com>

aiter-bot · 2026-10-30T02:15:05Z

NOT READY FOR HUMAN REVIEW

PR Self-Review

Revision 8e06e6fe328b · mode shadow

Waiting for the PR author. The readiness report will appear after the self-review answers are complete for this revision.

Questions to answer in the PR description

What behavior does this PR change, and why is the existing behavior insufficient?
What did you validate, and what remains unverified? Include commands, hardware/configuration and results.
What behavior would each changed op test fail on before this PR, and why is that coverage not already provided by an existing canonical test?
For every changed op_tests/ file, what behavior does it cover and why is it not redundant? Write None if no op test changed.

Edit the PR Self-Review section in the PR description. Include the current revision SHA; this comment refreshes automatically.

Missing or invalid answers

FAIL — Reviewed revision must contain at least the first 7 characters of the current head SHA

_{Shadow mode reports what would block but keeps the status successful. Set AITER_PR_SELF_REVIEW_MODE=enforce only after reviewing real false positives.}

ganyi1996ppo · 2026-10-30T02:15:08Z

Pushed e90aed1: the op test is rewritten in the aiter script style, with no pytest. Plain test_* check functions with asserts, a case table, a pandas summary printed before the timing table, try/except (expect_error) instead of pytest.raises, and one module-level name PA_DECODE for the implementation under test. Same 39 correctness cases as before (6 ctx_sweep, 2 garbage_tail, 4 heads, 4 task_shapes, 8 one_task_per_row, 3 plan_free_vs_reference, 4 plan_free_bitwise, 8 single cases) and the same timing sweep; a mismatch raises and the script exits non-zero. No kernel or wrapper code changed in this commit. Ran on an MI355 card with a fresh FlyDSL cache: python3 op_tests/test_flydsl_pa_decode_mxfp4.py exits 0 from the repo root and from /tmp.

…ode_mxfp4 directly Replace the custom rel_l2/assert_rel with assertAllclose (rtol=atol=1e-2, at most 1% of elements outside; the measured worst case is 0.1%), drop the PA_DECODE alias, and make the timing section header say what is timed. Same 39 cases and timing sweep. Co-Authored-By: Claude <noreply@anthropic.com>

ganyi1996ppo · 2026-10-30T02:38:21Z

Current head: 13a5814. Read this one: SHAs quoted in my older replies on this PR belong to history that a rebase rewrote.

New at this head, test file only (no kernel or wrapper change): the custom rel_l2/assert_rel are replaced by assertAllclose (rtol=atol=1e-2, at most 1% of the elements outside; the worst of the 39 cases is 0.098%); the PA_DECODE alias is gone and pa_decode_mxfp4 is called directly; the timing section header now says what is timed (MXFP4 only). The per-row slice checks are dropped because the whole-tensor elementwise check covers every row at the same tolerance.

Ran on an MI355 card from a fresh tree with FlyDSL cache off: python3 op_tests/test_flydsl_pa_decode_mxfp4.py exits 0 from the repo root (26 s wall-clock) and from /tmp (14 s), 39 cases then the timing table. With the reference scaled by 1.05 (not committed) it exits 1. CI has not run this file (draft). The PR description now has the PR Self-Review sections, and the "0.3-1.2 us slower" statement is corrected to the table (0.29-0.91 us vs fp8 page 16).

…e dedicated kernel Fold MXFP4 K/V (e2m1 data + e8m0 scale per 32) into the fp8/bf16 compile_pa_decode_tile (kv_dtype="mxfp4"); pa_decode dispatches a uint8 key cache to it and pa_decode_mxfp4 stays as a thin entry. Delete pa_decode_mxfp4_kernel.py; the MXFP4 cases move into test_flydsl_pa_decode.py. Co-Authored-By: Claude <noreply@anthropic.com>

ganyi1996ppo · 2026-10-30T04:32:08Z

Pushed 356de44: the dedicated MXFP4 kernel is replaced by an MXFP4 mode of the existing fp8/bf16 compile_pa_decode_tile, entered through pa_decode. pa_decode_mxfp4 stays as a thin entry with the same signature, so ATOM#2770 and vLLM are unchanged. Why: the fold reuses the fp8 pipeline (plan, reducer, page-id staging) and needs 761 source lines inside original files instead of a 989-line second kernel and wrapper; the tests move into test_flydsl_pa_decode.py (+345/-40) and the separate op test is gone.

Narrower than the previous head: block size 128, head_dim 128, bf16 output, no sliding window, query_length 1 only. fp8/bf16 ISA is byte-identical to main, the existing file passes 226 tests, ATOM GSM8K flexible 0.9439 / strict 0.8954. The final dispatch form has not been run in vLLM and the ATOM throughput A/B is pending. Numbers and scope are in the PR body.

Co-Authored-By: Claude <noreply@anthropic.com>

Co-Authored-By: Claude Code <noreply@anthropic.com>

Resolve pa_decode/pipeline.py: keep main's tile()/peeled mask_tail loop and dispatch the MXFP4 tile from tile(). Co-Authored-By: Claude Code <noreply@anthropic.com>

Copilot

🟡 Changes recommended

The promised MXFP4 compatibility entry point is missing, stored e4m3 queries are rejected, and unchecked query address spans can cause offset wraparound and incorrect loads.

4 open issues

_{Review effort: Balanced}

ganyi1996ppo · 2026-11-03T08:32:40Z

+    if query.dtype != torch.bfloat16:
+        raise TypeError(f"MXFP4 pa_decode needs a bf16 query, got {query.dtype}")


Fixed in f73e934, which is newer than the commit this was reviewed on (74fc9d5). _pa_decode_mxfp4 accepts torch.float8_e4m3fn or bf16, passes query_dtype="fp8" or "bf16" from the tensor dtype, and the query loader copies stored e4m3 bytes with unit scale (output stays bf16). Both dispatch modes have e4m3 cases in test_pa_decode: mxfp4-g16-e4m3-q, mxfp4-g8-e4m3-q (planned) and mxfp4-plan-free-e4m3-q. They pass in the full run on gfx950 at c47551a (265 passed, 16 skipped). ATOM#2770 sparse_attn_mxfp4.py requires torch.float8_e4m3fn and passes it straight to pa_decode. Only e4m3fn is accepted, not fnuz: MXFP4 is gfx950-only, where e4m3 is OCP.

Re-add the stored e4m3 query path removed in 491b929 (accepted by ROCm/ATOM#2770's sparse_pa_mxfp4), selected by a Q_IN_DTYPE trait rather than a bool. The in-kernel BF16 query to FP8 quantization and the BF16 MXFP4 path are unchanged. Co-Authored-By: Claude Code <noreply@anthropic.com>

Copilot

🔵 Needs a closer look

Accepted Q/O layouts can overflow the kernel’s signed 32-bit offset arithmetic, causing output stores to overwrite other rows or access memory outside the allocation.

4 open issues

_{Review effort: Balanced}

The kernel adds query row/head offsets in i32, so a query whose last head_dim chunk ends past 2 GiB wrapped to an earlier row, and a direct output store did the same. Refuse such query and output views before launch; two refused-input cases in test_pa_decode cover them. Co-Authored-By: Claude <noreply@anthropic.com>

Copilot

🔵 Needs a closer look

Human assessment is needed because fixed probability scaling can discard a finite attention contribution even when cache scales are valid and values fit BF16.

2 open issues

2 resolved since last review

_{Review effort: Balanced}

Serve MiniMax-M3.1 with `--kv-cache-dtype mxfp4` and `indexer_kv_dtype="mxfp4"` on gfx950, TP4 (one KV head per rank): - the fused QK-norm / RoPE kernel writes MXFP4 K, V and index-K pages and hands back the main q as e4m3 (new `fused_qknorm_idxrqknorm_mxfp4` wrapper); - the index scores come from AITER's FlyDSL MXFP4 scorers, the forced init/local blocks are pinned with AITER's sentinels, and the top-k is the one the NVFP4 indexer uses; - the attend turns each query token's top-k into a one-block-per-page table and runs AITER's `pa_decode_mxfp4` on it, for decode and prefill. The sparse block-table builders gain a `pages_per_block` argument (1 for a page-128 cache). Needs an AITER with the MXFP4 FlyDSL ops (ROCm/aiter#7056, vllm-project#7057, vllm-project#7060, vllm-project#7061). Mono, indexer CP, other TP sizes and the draft layers are not covered. Signed-off-by: ganyi <ygan@amd.com> Co-Authored-By: Claude <noreply@anthropic.com>

Resolve conflicts in pa_decode_kernel.py and test_flydsl_pa_decode.py: keep the MXFP4 stage_pages/direct_output support next to main's gfx1250 merge_partitions/cluster_merge, and keep both test skips and case keys. Co-Authored-By: Claude Code <noreply@anthropic.com>

Copilot

🟡 Changes recommended

Critical issues remain because post-accumulation softmax scaling can produce FP32 overflow, while 32-bit block-table offsets can wrap and read incorrect page IDs.

4 open issues

_{Review effort: Balanced}

+        scale_b = fx.full(4, self.ctx.scale_qk * q_mantissa, fx.Float32)
+        scores = [frag_Ss[a] * scale_b for a in range_constexpr(t.NCHUNK)]


+            raise ValueError(f"{name} must be contiguous on the query device ({dev})")
+    if num_seqs == 0:
+        return
+    max_blocks_per_seq = block_tables.shape[1]


valarLip · 2026-11-04T01:09:18Z

@valar-cop review

valar-cop

some changes needed..

ci green.

valar-cop · 2026-11-04T01:17:38Z

+        for a in range_constexpr(t.NCHUNK):
+            self.kv.mxfp4_stage_v_bf16(v_raw, a, tile_valid)
+        scale_b = fx.full(4, self.ctx.scale_qk * q_mantissa, fx.Float32)
+        scores = [frag_Ss[a] * scale_b for a in range_constexpr(t.NCHUNK)]


scale applied after mfma can saturate the acc, byte 249 -> inf. fold it into the q e8m0 byte? otherwise fix the docstring and add a byte-249 case with fp64 ref

valar-cop · 2026-11-04T01:17:38Z

+        ("context_lengths", context_lengths),
+    ):
+        if t.device != dev or not t.is_contiguous():
+            raise ValueError(f"{name} must be contiguous on the query device ({dev})")


block_tables not bounded here, int32 index wraps past 2**31. same bound as query/output, or widen to i64

valar-cop · 2026-11-04T01:17:38Z


 from .. import dpp_utils
 from ..kernels_common import create_llvm_ptr
+from ..moe_2stage_a16wmix.utils import _e8m0_byte_to_f32


why import from moe utils? move _e8m0_byte_to_f32 into mxfp4_gemm_common, it's already the home for e8m0 helpers

Resolve conflicts with the gfx1250 in-kernel partition merges (#7079), MXFP4 K/V (#7056) and the D64 BF16 wave-pair pipeline (#7384): - Gfx1250Traits keeps context_partition_size alongside merge_partitions and cluster_merge; the per-wave FP8 and BF16 pipelines from main size their token ranges by the plan's partition, not KV_COMPUTE_BLOCK. - Batched FP8 softmax now covers two or three query tiles; four tiles take main's per-wave pipeline, which is 7-20% faster on gfx1250 (D128/page128/MTP4/GQA16, CUDA graphs). BF16 four-tile plain-V caches likewise take main's wave-pair pipeline (20-45% faster than the packed pipeline). - PADecodePlan keeps arrival_counters and context_partition_size, passed by keyword; gfx1250-strict supports the same partition sizes as gfx1250. - pa_decode bounds in-kernel merges and the reducer by the plan's partition size; MXFP4 accepts only 256-token partitions. - Tests take main's new merge, split-group and MXFP4 cases; the plan test stays removed as on main. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

zhuyuhua-v mentioned this pull request Oct 29, 2026

[MiniMax-M3.1] FP4 index cache and NVFP4 sparse KV cache ROCm/ATOM#2744

Open

ganyi1996ppo and others added 9 commits October 28, 2026 21:20

[FlyDSL] pa_decode_mxfp4: reuse shared constants and ceildiv, fold vi…

fdeb3b2

…ew checks Generated ISA is unchanged (82 specializations, identical hash). Co-Authored-By: Claude <noreply@anthropic.com>

[Test] pa_decode_mxfp4: follow op_test conventions

5f172f3

Module-level gfx950 skip, checkAllclose + TFLOPS in the benchmark, and run the pytest checks from main() so `python3 op_tests/...` (what CI runs) tests it. Co-Authored-By: Claude <noreply@anthropic.com>

[FlyDSL] Export pa_decode_mxfp4 lazily; derive default-split test row…

1ebcd98

…s from the CU count Co-Authored-By: Claude <noreply@anthropic.com>

[FlyDSL] pa_decode_mxfp4: check output device, return early on zero s…

e013556

…equences The test no longer queries the GPU at import. Co-Authored-By: Claude <noreply@anthropic.com>

[FlyDSL] pa_decode_mxfp4: drop query_scale (the fp8 query is used as …

090b0a0

…stored) Co-Authored-By: Claude <noreply@anthropic.com>

[FlyDSL] pa_decode_mxfp4: follow main's reducer (partition_bound remo…

765c6cc

…ved) launch_pa_decode_ps_reduce now takes the bound as max_partitions; same value as before: min(parts, capacity, tiles). Co-Authored-By: Claude <noreply@anthropic.com>

ganyi1996ppo force-pushed the ganyi/flydsl-pa-decode-mxfp4 branch from 8efda80 to 765c6cc Compare October 29, 2026 02:28

ganyi1996ppo mentioned this pull request Oct 29, 2026

[MiniMax-M3.1] Mono decode over the MXFP4 KV and index caches ROCm/ATOM#2812

Draft

r and others added 2 commits October 29, 2026 23:11

Merge remote-tracking branch 'origin/main' into ganyi/pa-replace-push2

1f87f82

ganyi1996ppo changed the title ~~[FlyDSL] Add an MXFP4 KV-cache paged-attention decode (gfx950, query_length 1)~~ [FlyDSL] pa_decode: MXFP4 K/V in the shared PA tile kernel (gfx950, query_length 1) Oct 30, 2026

[Test] pa_decode: drop an unused unpacked name (ruff RUF059)

c57f106

Co-Authored-By: Claude <noreply@anthropic.com>

ganyi1996ppo and others added 2 commits November 3, 2026 01:22

flydsl pa_decode: clean up FP8 input-query handling

491b929

Co-Authored-By: Claude Code <noreply@anthropic.com>

Merge origin/main into ganyi/flydsl-pa-decode-mxfp4

74fc9d5

Resolve pa_decode/pipeline.py: keep main's tile()/peeled mask_tail loop and dispatch the MXFP4 tile from tile(). Co-Authored-By: Claude Code <noreply@anthropic.com>

ganyi1996ppo dismissed fsx950223’s stale review via 74fc9d5 November 3, 2026 07:29

Copilot AI balanced review requested due to automatic review settings November 3, 2026 07:29

Copilot started reviewing on behalf of ganyi1996ppo November 3, 2026 07:29

View session

Copilot AI reviewed Nov 3, 2026

View reviewed changes

Copilot AI balanced review requested due to automatic review settings November 3, 2026 08:01

Copilot started reviewing on behalf of ganyi1996ppo November 3, 2026 08:02

View session

Copilot AI reviewed Nov 3, 2026

View reviewed changes

Copilot AI balanced review requested due to automatic review settings November 3, 2026 08:31

Copilot started reviewing on behalf of ganyi1996ppo November 3, 2026 08:31

View session

Copilot AI reviewed Nov 3, 2026

View reviewed changes

fsx950223 mentioned this pull request Nov 3, 2026

[FlyDSL] [Misc] pa_decode: key the tile kernel on FlyDSL dtypes and drop the batch-first plan grid #7388

Merged

4 of 10 tasks

Copilot AI balanced review requested due to automatic review settings November 3, 2026 11:06

Copilot started reviewing on behalf of ganyi1996ppo November 3, 2026 11:07

View session

Copilot AI reviewed Nov 3, 2026

View reviewed changes

valar-cop reviewed Nov 4, 2026

View reviewed changes

valarLip approved these changes Nov 4, 2026

View reviewed changes

valarLip merged commit 5469555 into main Nov 4, 2026
62 checks passed

valarLip deleted the ganyi/flydsl-pa-decode-mxfp4 branch November 4, 2026 01:37

ganyi1996ppo mentioned this pull request Nov 4, 2026

[FlyDSL] Move _e8m0_byte_to_f32 into mxfp4_gemm_common #7424

Merged

3 of 4 tasks

fsx950223 mentioned this pull request Nov 4, 2026

[FlyDSL] Optimize gfx1250 PA partitions and packed attention #7081

Open

XiaobingSuper mentioned this pull request Nov 10, 2026

fix(attention): accept optional explicit FlyDSL plans ROCm/ATOM#3052

Merged

lirui927 mentioned this pull request Nov 10, 2026

fix(attention): drop FlyDSL aiter API probes ROCm/ATOM#3162

Closed

		if query.dtype != torch.bfloat16:
		raise TypeError(f"MXFP4 pa_decode needs a bf16 query, got {query.dtype}")

		scale_b = fx.full(4, self.ctx.scale_qk * q_mantissa, fx.Float32)
		scores = [frag_Ss[a] * scale_b for a in range_constexpr(t.NCHUNK)]

Conversation

ganyi1996ppo commented Oct 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Dependencies

PR Self-Review

Reviewed revision

Behavior and intent

Validation and unverified scope

Bandwidth (TB/s)

PR-specific question

Changed op tests

MXFP4 operand validation

Uh oh!

github-actions Bot commented Oct 28, 2026

🏷️ CI Guide

Uh oh!

aiter-bot commented Oct 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

NOT READY FOR HUMAN REVIEW

PR Self-Review

Questions to answer in the PR description

Missing or invalid answers

Uh oh!

ganyi1996ppo commented Oct 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ganyi1996ppo commented Oct 30, 2026

Uh oh!

ganyi1996ppo commented Oct 30, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

ganyi1996ppo Nov 3, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🔵 Needs a closer look

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🔵 Needs a closer look

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

valarLip commented Nov 4, 2026

Uh oh!

valar-cop left a comment

Choose a reason for hiding this comment

Uh oh!

valar-cop Nov 4, 2026

Choose a reason for hiding this comment

Uh oh!

valar-cop Nov 4, 2026

Choose a reason for hiding this comment

Uh oh!

valar-cop Nov 4, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

ganyi1996ppo commented Oct 28, 2026 •

edited

Loading

aiter-bot commented Oct 30, 2026 •

edited

Loading

ganyi1996ppo commented Oct 30, 2026 •

edited

Loading