[FlyDSL] Fused online Hadamard rotation + MXFP4 quantization (flydsl_rot_quant)

[jiangyon-amd]

[FlyDSL] Fused online Hadamard rotation + MXFP4 quantization (flydsl_rot_quant)

What

Adds aiter.ops.flydsl.flydsl_rot_quant(x, RS, shuffle_scales=) — a single gfx950
kernel that takes a bf16 activation, applies a block-diagonal Hadamard rotation,
quantizes to MXFP4, and writes the e8m0 scales, optionally already in the CK-gemm
swizzle that gemm_a4w4 consumes.

from aiter.ops.flydsl import flydsl_rot_quant

xq, xs = flydsl_rot_quant(x, RS=64, shuffle_scales=True)
y = gemm_a4w4(xq, shuffle_weight(wq, layout=(16, 16)), xs,
              shuffle_scale(ws), dtype=torch.bfloat16)[:M]

Why

Rotated MXFP4 checkpoints (QuaRot / SpinQuant style, as produced by Quark) fold the
offline rotation into the weights. The matching online rotation of the activation
then has to run at every projection, every layer, every token. Written with today's
aiter ops that is three passes over HBM on a purely memory-bound op:

step	today	this PR
rotate	(no aiter op at RS=64 — a torch blockwise Hadamard matmul)	fused
pack MXFP4 + e8m0	dynamic_mxfp4_quant	fused
scale swizzle for gemm_a4w4	shuffle_scale	fused

Relation to existing ops — this is a gap, not a duplicate

• rotate_activation_fp4quant_inplace (csrc/kernels/dsv4_rotate_quant.cu) rotates,
  fp4-quantizes and dequantizes back to bf16 in place, at a fixed dim of 128 or
  256. That is the simulation/QAT path; it does not produce a gemm input.
• dynamic_mxfp4_quant produces packed fp4 + e8m0, but does not rotate.
• shuffle_scale re-lays-out the scales in a separate pass.

Nothing currently produces rotated, packed fp4 + e8m0 scales in one launch, and
nothing emits the CK scale swizzle from inside the quantizing kernel.

Performance

MI350X (gfx950), K=4096, RS=64, CUDA-graph timed, 50 launches per replay.
Baseline is the three-pass path above (torch blockwise Hadamard +
dynamic_mxfp4_quant + shuffle_scale).

M	3-pass (µs)	fused (µs)	speedup	fused GB/s
1	12.46	2.80	4.46×	4
32	12.79	3.09	4.13×	107
256	15.31	3.26	4.69×	814
1024	17.92	4.49	3.99×	2365
4096	42.60	11.61	3.67×	3657
16384	117.05	36.72	3.19×	4626

At small M both sides are launch-bound (three launches vs one); the win there is
kernel count, not bandwidth. At large M the kernel reaches ~4.5 TB/s of moved bytes
(read 2 B + write 0.53 B per element), which is where a rotate+quant+swizzle fusion
should land on this part.

Arithmetic intensity: 2.53 bytes moved per element for log2(RS) f32 add/sub
(6 per element at RS=64) plus one v_cvt_scalef32_pk_fp4_f32 per pair — ~2.4
FLOP/byte against a ~575 FLOP/byte machine balance. Firmly memory-bound, which is
exactly why the three passes are worth collapsing into one.

Scope of the claim: the numbers above are the claim being made. This op is one of
several per-projection steps in a decode iteration, so the end-to-end serving effect
is model-, shape- and gemm-tuning-dependent and is a small fraction of the
microbenchmark ratio — in our own measurements it is a low single-digit percentage of
total throughput, positive but not the 3–4× above.

Implementation

Work = M * (K//COLS) independent thread-blocks; one thread owns a COLS-wide chunk
holding COLS//RS independent RS-wide Hadamard blocks, so the FWHT needs no
cross-lane communication — it is scalar-unrolled in registers (Sylvester,
(a+b, a-b) pairing). At RS=COLS=64 that is 64 f32 registers per thread. Loads are
128-bit (vec8 bf16); fp4 stores are 128-bit (4× packed i32).

Written against the post-#4501 stable FlyDSL interface: buffer_ops / vector from
aiter.ops.flydsl.kernels, tensor_shim._run_compiled / _to_raw for the launch
path, and quant_utils.emit_mx_e8m0_scale for the block scale. No new dependency floor:
verified on flydsl 0.3.0 (the version requirements.txt pins) and on 0.2.4
(the floor aiter/ops/flydsl/__init__.py enforces), identical results on both.

Two details worth calling out for review:

1. The 1/sqrt(RS) normalization is folded into the e8m0 exponent when
   log2(RS)/2 is an integer (RS=64), replacing a per-element f32 multiply with an
   integer subtract on the shared group exponent. RS=32 and RS=128 take the
   multiply path.
2. shuffle_scales=True scatters each e8m0 byte directly to its
   shuffle_scale destination, computed by bit-slicing the global thread id. This
   removes the separate swizzle pass and its padding-tile zeroing: the padded cells
   are torch.empty and never written, exactly as shuffle_scale itself leaves them,
   because the gemm slices its output back to [:M].

Scale convention

No bespoke convention here: the e8m0 block scale is built by the shared
quant_utils.emit_mx_e8m0_scale IR helper in aiter's default
MxScaleRoundMode::RoundUp (ceil_pow2(amax / 6)), and the torch reference in the
test uses its CPU mirror fp4_utils.f32_to_mx_e8m0_scale with the same mode. The
fold_k exponent fold in (1) is exact on top of it because RoundUp is
exponent-linear: e8m0(amax · 2⁻ᵏ) = e8m0(amax) − k.

Testing

op_tests/test_flydsl_rot_quant.py. All three checks matter independently, because
all three failure modes are silent — wrong numbers, no exception:

1. Bit-exactness — packed fp4 and e8m0 scales compared byte-for-byte against a
   pure-torch reference (_ref_rot_quant, in the test file — no external dependency).
   Byte equality, not allclose: allclose would hide a systematic off-by-one-binade
   scale bug. Covers RS ∈ {32, 64, 128} × K ∈ {2048, 4096} × M ∈ {1, 32, 256, 1024,
   4096, 16384}, plus an outlier pass that drives elements into the saturating
   [6, 8) band where an encoder that clamps differently from the hardware would
   diverge. All 39 configurations are exact.
2. Fused swizzle parity — the in-kernel scatter vs shuffle_scale() of the
   natural-layout output, compared only on the live cells (a deterministic 0/1 mask
   pushed through the same reshape/permute; shuffle_scale pads with torch.empty,
   so the mask cannot be recovered from its output).
3. End-to-end through gemm_a4w4 — the consumer the swizzle exists for. A
   quantizer/gemm layout mismatch produces garbage with no error, so unit-level parity
   is not sufficient evidence that the op works. Gated at N ≥ 512: aiter's CK asm
   gemm_a4w4 is numerically wrong below that (rel err ~0.94 at N=64, exact at
   N ≥ 512), so it is not a valid reference there. That is a property of the gemm,
   not of this op.

Plus test_rot_quant_rejects_bad_input: 7 guards (fp16/fp32 input, non-contiguous,
3D, K % RS != 0, unsupported RS, bad AMAX) must raise rather than compute. The
kernel hardcodes bf16 buffer loads and linear element offsets, so every one of these
would otherwise be read as contiguous bf16 and yield plausible garbage.

Registered in .github/scripts/split_tests.sh (FILE_TIMES=150; ~9 s warm, the
rest is cold FlyDSL JIT for the 27 kernel variants the matrix compiles). Peak GPU
memory 812 MiB — below the level that would warrant a MEMORY_WEIGHT_FLOOR entry.

Scope / limits

• gfx950 only — the quantizer is v_cvt_scalef32_pk_fp4_f32. Raises on other
  architectures.
• Input must be bf16, 2D, contiguous; RS ∈ {32, 64, 128}; K % RS == 0.
• _pick_block / _pick_svec defaults are swept on MI350X. They are overridable
  per call; other parts should re-sweep.
• No new dependency: flydsl is already a hard requirement of amd-aiter, and the
  export sits inside the existing is_flydsl_available() guard in
  aiter/ops/flydsl/__init__.py.

Files

file	change
aiter/ops/flydsl/kernels/rot_quant.py	new, 523 lines — the kernel
aiter/ops/flydsl/__init__.py	+2 — import and __all__ entry
op_tests/test_flydsl_rot_quant.py	new, 267 lines — the tests
.github/scripts/split_tests.sh	+1 — CI scheduling weight

[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:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
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: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
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 4549 --add-label <label>