[penglu12-quark]

Summary

End-to-end speedup for 2-bit (MatMulNBits bits=2) phi4-14B on Strix Halo (gfx1151):

• Against 2-bit on main: decode is 2.7╬ô├ç├┤2.9x faster and time to first token 2.0╬ô├ç├┤3.5x shorter. This uses a last-token lm_head / fp16 KV re-export of the model. On the unmodified export, commits 2+3 alone give 1.5╬ô├ç├┤2.3x decode and 1.4╬ô├ç├┤2.3x time to first token.
• Against 4-bit phi4: 2-bit now decodes 1.37╬ô├ç├┤1.65x faster at every prompt length. Time to first token matches 4-bit at 128 tokens, beats it at 16k, and is 6% behind at 2048.

The PR has four commits, each independently reviewable:

1. perf(kernels): bits=2 MatMulNBits. fp32 activations and output are handled natively, a uniform fp16 zero point becomes a scalar, and GEMV autotune uses cold timing.
2. perf(gqa): fp32 KV cache prefill. Attention runs on the fp16 flash prefill kernel instead of the decomposed fp32 path, which is quadratic in prompt length.
3. fix(hip-fusion): tensor.dim uses. Shape-only tensor.dim reads no longer block the existing Q/DQ fusions (QMatMul for lm_head, SwiGLU, layout requantize). Requantize also passes through hip.gather, for embedding lookups.
4. feat(hip-fusion): new fused ops. hip.qswiglu (quantized SwiGLU chain) and hip.fake_quantize (Q╬ô├Ñ├åDQ round trip), each exactly equal to the unfused chain.

Related issue or design

None.

Why

Profiling the Quark 2-bit phi4 export showed it running slower than 4-bit on main, for reasons unrelated to 2-bit math:

• Fusions never fired. convert-onnx-to-hip sizes each DPS init from tensor.dim of the producer, and those reads counted as uses. As a result, every token dequantized the full 2 GB lm_head and embedding table to fp32 and ran the quantized SwiGLU unfused (commit 3).
• fp32 KV cache attention fell back to fp32 GEMMs. At 16k tokens this made prefill attention dominate time to first token (commit 2).
• Per-layer quantized activation chains. Q(gate) ╬ô├Ñ├å qsigmoid ╬ô├Ñ├å 2Γö£├╣ qmul ╬ô├Ñ├å Q(up) ╬ô├Ñ├å DQ costs about 10 launches and about 1.9 GB of traffic per layer at 2048 tokens, and each QDQ round trip costs two launches plus an integer tensor (commit 4).
• 2-bit matmul overheads. fp32 graphs paid a cast-A / cast-C round trip around every 2-bit matmul (about 400 extra launches per decode token). The fp16 zero points, all 1.5, cost as many bytes per token as the scales. The u2 GEMV autotuner timed cache-hot replays (about 1 TB/s readings) and kept glitch winners (commit 1).

What

• Commit 1, kernels:
  • u2 GEMV and WMMA are templated on the A/C element types. GEMV reads fp32 A and writes fp32 C directly.
  • The WMMA kernels re-read each A tile once per N block, so the u2 WMMA path casts A to fp16 once and writes fp32 C from the epilogue. Reading fp32 A directly regressed prefill by 20% at 2048 tokens, because A no longer fits in the last-level cache.
  • New hip_matmul_nbits_ex(..., float uniform_zp). The runtime detects an all-equal fp16 zero-point tensor once per buffer (hip_matmul_nbits_fp16_uniform_value, cached in MatmulNbitsState), and the kernels then never read it.
  • u2 GEMV autotune uses GemvWeightRotation + timeConfigRotated, the median of cold bursts.
  • Documented in docs/design/custom_kernel_design.md.
• Commit 2, GQA: split, RoPE and the cache append stay fp32, so the cache contents and decode are unchanged. Roped Q and cache rows [0, total_seq) are cast into fp16 scratch, run through the flash prefill kernel, and the output is cast back. HIPDNN_EP_GQA_DISABLE_FP32_CACHE_FLASH_PREFILL=1 restores the old path.
• Commit 3, fusion fix: tensor::DimOp canonicalization runs inside hip-fusion-transform's pattern set, folding shape reads onto the producer's own init. hip.gather is added to the layout ops that requantize can pass through.
• Commit 4, new ops: a second greedy phase in hip-fusion-transform (qdq_chain_fusion.cpp). It runs after the first so it cannot preempt a first-phase fusion that consumes the DQ.
  • Each fusion declines unless every consumer reads with the scale and zero point its producer wrote.
  • The work covers the ops, bufferization, HIP-to-LLVM lowering, real and mock runtime wrappers, and kernels.
  • Documented in docs/supported-operations.md and docs/pipeline_pass_menu.md.

Performance

Hardware and method:

• Hardware: Strix Halo APU (gfx1151, wave32), Windows 11, Release build, HIP_ARCHITECTURES=gfx1151.
• Harness: onnxruntime-genai model_benchmark -l <prompt> -g 128 -r 3 -w 1, runs strictly serial, no other GPU or CPU load.
• Settings: no HIPDNN_EP_PERF/DEBUG.
• Prompt length: the long prompt is 16256 tokens, because 16384 + 128 exceeds phi4's 16384-token context.

Model: phi4-14B Quark QAD 2-bit export, MatMulNBits bits=2 with block 64, fp16 zero points (all 1.5), fp32 activations, UINT16 activation QDQ. The cumulative table uses the same export re-saved with a last-token lm_head slice (as the 4-bit baseline export already does) and an fp16 KV cache. Those re-exports are model-side and not part of this PR.

Same 2-bit model, main vs this PR (commits 2+3), time to first token in seconds / decode tokens/s:

Prompt	main	This PR (commits 2+3)	Speedup
128	0.286 s / 14.70	0.200 s / 34.23	1.43x / 2.33x
2048	5.035 s / 13.75	3.256 s / 28.54	1.55x / 2.08x
16256	66.28 s / 9.28	28.24 s / 13.83	2.35x / 1.49x

Cumulative, pruned lm_head + fp16 KV export of the same model (time to first token / decode tokens/s / prefill tokens/s):

Prompt	Commits 2+3	+ commit 4 (fused Q/DQ)	+ commit 1 (u2 kernels)	Commit 4 + commit 1 gain
128	0.163 / 34.02 / 787	0.151 / 36.48 / 850	0.141 / 42.03 / 906	TTFT 1.16x, decode 1.24x
2048	2.581 / 31.86 / 793	2.260 / 32.88 / 906	2.175 / 38.44 / 942	TTFT 1.19x, decode 1.21x
16256	22.27 / 21.81 / 730	19.61 / 23.34 / 829	19.00 / 24.74 / 856	TTFT 1.17x, decode 1.13x

Overall, 2-bit main vs 2-bit with this PR (with the last-token lm_head and fp16 KV exports): time to first token is 2.0x / 2.3x / 3.5x shorter and decode is 2.9x / 2.8x / 2.7x faster at 128 / 2048 / 16256 tokens.

Against 4-bit phi4 (phi4-14B-int4-rtn-block-32_directml-pruned_lm_head):

Prompt	4-bit, main	2-bit, this PR	2-bit vs 4-bit
128	0.143 s / 25.44	0.141 s / 42.03	TTFT 1.01x, decode 1.65x
2048	2.041 s / 24.11	2.175 s / 38.44	TTFT 0.94x, decode 1.59x
16256	19.77 s / 18.07	19.00 s / 24.74	TTFT 1.04x, decode 1.37x

4-bit regression check (same int4 model):

• main: 0.143 / 2.041 / 19.77 s, decode 25.44 / 24.11 / 18.07.
• This PR, commits 2+3: 0.141 / 1.919 / 16.59 s, decode 25.92 / 24.42 / 18.30.
• This PR, commits 2╬ô├ç├┤4, at 2048: 1.872 s / 24.52.
• Commit 1: only changes bits=2 dispatch. The shared fp16 GEMM template's new A/C type parameters default to fp16, so the 4-bit instantiations are unchanged.

Mixed precision {2, 4}-bit: Qwen3.8-27B study

These kernels let one graph mix bits=4 and bits=2 MatMulNBits node by node, with no per-config code in hip-ep. To measure the accuracy/latency trade-off, I ran a study on Qwen3.8-27B (int4 k_quant g128; 48 Gated DeltaNet + 16 GQA layers) on gfx1151. The study scripts are outside this PR; only the kernels it exercises are here.

Method:

• 2-bit weights: derived from the int4 checkpoint (dequantize, then re-quantize to 2-bit asymmetric, group 64, 2.28 bpw, with a clip search weighted by per-channel E[x^2] from WikiText-2 train activations). int4 is therefore the accuracy ceiling.
• Units: 256 = 64 layers x 4 roles (attn_in, attn_out, mlp gate/up, mlp_down). lm_head stays int4.
• Sensitivity (sens*): for each unit alone at 2-bit, measured dNLL on 8k calibration tokens, then a greedy knapsack on dNLL/params up to the 2-bit budget.
• Baselines at the same 47% budget: proxy47 ranks units by weight reconstruction error; role47 puts only mlp gate/up at 2-bit.
• Accuracy: WikiText-2 test perplexity (full set, 2048-token windows), plus GSM8K chain-of-thought 8-shot greedy on the first 250 questions (lm-eval gsm8k_cot prompt; "keep" = accuracy / int4 accuracy).
• Latency: model_benchmark -l 128/2048/16384 --use_random_tokens -g 128 -r 3 -w 1 (OGA 0.17 / ORT 1.30), text-only prompts, HIPDNN_EP_STRICT=1 and CPU fallback disabled. Run serially on this branch rebased onto ddb207c7.

Config	2-bit share	Linear weights (GB)	WikiText-2 PPL	GSM8K strict (keep)	Decode tok/s @128 / 2048 / 16384	Decode speedup @2048	TTFT @2048 / 16384 (s)
int4	0%	13.31	7.217	0.888 (1.000)	14.57 / 14.53 / 13.65	1.000	3.87 / 32.66
sens25	26%	11.85	7.400	0.864 (0.973)	15.78 / 15.60 / 14.58	1.074	4.00 / 33.29
sens47	47%	10.64	8.090	0.756 (0.851)	17.15 / 16.90 / 15.67	1.163	4.03 / 33.90
proxy47	47%	10.63	8.590	0.768 (0.865)	17.06 / 16.93 / 15.67	1.165	3.92 / 33.20
role47	47%	10.64	9.056	0.664 (0.748)	17.18 / 16.97 / 15.67	1.168	4.10 / 34.36
sens75	75%	9.03	10.418	0.424 (0.477)	19.30 / 18.98 / 17.47	1.306	4.15 / 34.64
all 2-bit	100%	7.60	26.48	0.028 (0.032)	21.92 / 21.64 / 19.68	1.489	4.19 / 34.90

• Decode speed depends only on the 2-bit share: decode is weight-bandwidth bound, and mixing bit widths in one graph costs nothing measurable.
• Prefill is compute bound, so 2-bit gives no gain there (TTFT +1-8%).
• Allocation matters: at the same 47% budget, sensitivity-based allocation keeps much more accuracy than role-based (GSM8K strict 0.756 vs 0.664, McNemar p = 0.006). It matches the weight-error proxy on GSM8K strict (0.756 vs 0.768, p = 0.80) with lower perplexity (8.09 vs 8.59).
• Operating points: sens25 costs 2.7% of GSM8K strict for 7% faster decode, and sens47 costs 15% for 16%.
• Same-day main (ddb207c7), decode tok/s @128 / 2048 / 16384: int4 14.67 / 14.49 / 13.59 (this PR within +-1%), sens47 16.37 / 15.88 / 15.17 (this PR +3-6%), all 2-bit 20.95 / 20.32 / 18.34 (this PR +5-7%).
• Correctness on the rebased build: greedy tokens for sens47 and all 2-bit match their pre-rebase reference.

Limitations:

• 2-bit weights come from int4, not from bf16.
• GSM8K uses a 250-question subset.
• The published Qwen3.8 benchmark adds about 1.9k image tokens to the prompt, which is not reproduced here.
• Accuracy was measured before the rebase. The rebased build reproduces the sens47 and all 2-bit token references.

Test plan

• LIT suite (rebased tree): ctest -R MorphizenMLIRLitTests. 675 discovered: 666 passed, 9 unsupported, 0 failed. New tests: qdq_chain.mlir, test_qswiglu.mlir, test_fake_quantize_model.mlir, plus new cases in qmatmul.mlir, qdq_round_trip.mlir and swiglu.mlir.
• 2-bit kernel unit tests (rebased tree): hipdnn-kernel-ut-matmul-nbits-xfp32-wu2-yfp32 and -xfp16-wu2-yfp16, 72/72 cases each.
• Numeric, strict mode (HIPDNN_EP_STRICT=1, --backend ort_ep):
  • test_qdq_chain.py: 5/5 passed.
  • test_gqa.py: 8 passed, 18 failed. Plain main at ddb207c7 fails exactly the same 18 tests in repeated runs (shared-buffer KV limitation); all 4 new fp32-cache prefill cases pass. Those cases compare only the present-cache rows [0, total) that the kernel writes: with separate past/present buffers the tail is not written and holds whatever the allocator left.
  • test_topk.py: 13/13 passed (main's tests; TopK itself is unchanged by this PR).
• Full-model accuracy (2-bit phi4, strict, teacher-forced perplexity on a fixed paragraph):
  • Perplexity: ONNX Runtime CPU 9.401. Before commit 1: 9.363, next-token accuracy 0.506. With commit 1: 9.375, next-token accuracy 0.513.
  • Logit relative L2 vs CPU: 1.016e-1 before commit 1, 1.010e-1 with it.
  • Single decode step logits vs CPU: relative L2 1.1e-2, same top-1 token.
  • Commit 4 alone gives bit-identical logits to the unfused build.
• pre-commit: pre-commit run --from-ref origin/main --to-ref HEAD passes (lintrunner, license headers, whitespace).

Notes for reviewers

• Scope: 36 files / about 2.8k lines, kept as one PR at the submitter's request. The four commits are separate concerns, so the PR can be split along commit boundaries if reviewers prefer. Only the final tree was built and tested, not each commit individually.
• Rebase onto main: benchmarks were taken on the pre-rebase base (9d38afb9). After rebasing onto 702a9481, the only conflict was adopting #639's Hip_DpsOp_AutoReifyInfer for the two new ops. The branch is now rebased onto ddb207c7. Two TopK commits briefly added here were dropped, because #1264 landed the same largest/sorted conversion fix and a broader block-parallel TopK. LIT, unit and numeric tests above were rerun on the ddb207c7 rebase.
• Same-day check on ddb207c7 (Qwen3.8-27B int4 k_quant, gfx1151, model_benchmark -l 128/2048/16384 --use_random_tokens -g 128 -r 3 -w 1, text-only): int4 decode is unchanged vs main (14.57 / 14.53 / 13.65 vs 14.67 / 14.49 / 13.59 tokens/s). With 47% of linear weights re-quantized to 2-bit g64 it is +3-6% (17.15 / 16.90 / 15.67 vs 16.37 / 15.88 / 15.17), and with all linear weights at 2-bit +5-7% (21.92 / 21.64 / 19.68 vs 20.95 / 20.32 / 18.34). Time to first token for the 2-bit variants is 1-3% longer than main, which is close to run-to-run variation.
• Wheel packaging: the local wheel target fails because this machine's ROCm SDK has no tensilelite-host.dll. This is environment-only; all compile and link targets build.
• Follow-ups:
  • u2 tuning caches are keyed by exact M, so each new prompt length retunes once (about 15 s); bucketing M would remove that.
  • No unit test yet covers the fp16 uniform zero-point path; it was verified through the full model.
  • tools/kernel_resources_*.json were not regenerated for the renamed u2 instantiations.

Checklist

• The change is focused, or links a design/series explaining its scope. (Four independent commits; see Notes.)
• Relevant tests were added or updated and the results are documented.
• User-facing or design documentation was updated when needed.
• Substantial AI assistance is disclosed, and I reviewed and understand the result.

AI assistance: the implementation, profiling, benchmarks and this description were produced with Cursor (Claude Opus 5.5) and validated with the LIT, kernel unit, numeric and full-model accuracy runs listed above. Commits carry Co-Authored-By / Made-with trailers.

[github-actions]

Thanks for opening a PR!

This project follows LLVM's incremental-development and AI-tool-use
guidance. See CONTRIBUTING.md
for the project workflow.

Before requesting review, please check that:

1. The change is focused. Substantial work links the relevant issue
   or design discussion.
2. The PR documents relevant test results and updates affected
   documentation.
3. If AI tools provided substantial assistance, the description
   explains what was assisted and how it was validated, and commit
   trailers identify the tool. The contributor has reviewed and
   understands the result.

Reviewers are assigned through
CODEOWNERS where ownership
is configured.

[github-actions]

Kernel Unit Test Results

Op	Leaf	Arch	Cases	Failed	Suite (ms)	Verdict
matmul_nbits	Xfp16_Wu2_Yfp16	gfx1151	12	0	6775.3	PASS
matmul_nbits	Xfp16_Wu3_Yfp16	gfx1151	12	0	6784.5	PASS
matmul_nbits	Xfp16_Wu4_Yfp16	gfx1151	12	0	4354	PASS
matmul_nbits	Xfp16_Wi8_Yfp16	gfx1151	12	0	5772.3	PASS
matmul_nbits	Xfp32_Wu2_Yfp32	gfx1151	12	0	6927	PASS
matmul_nbits	Xfp32_Wu3_Yfp32	gfx1151	12	0	6801.5	PASS
matmul_nbits	Xfp32_Wu4_Yfp32	gfx1151	12	0	4213.5	PASS
matmul_nbits	Xfp32_Wi8_Yfp32	gfx1151	12	0	5924.7	PASS
gqa/decode	Xfp16_Wfp16_Yfp16	gfx1151	17	0	3347.4	PASS
gqa/decode	Xfp16_Wi8_Yfp16	gfx1151	13	0	2460.5	PASS
gqa/prefill	Xfp16_Wfp16_Yfp16	gfx1151	26	0	5106.1	PASS
gqa/prefill	Xfp16_Wi8_Yfp16	gfx1151	9	0	3733.1	PASS
qmoe_ragged_mm	Xfp16_Wu4_Yfp16	gfx1151	40	0	1637.6	PASS
gather_block_quantized	Xi64_Wu4_Yfp16	gfx1151	32	0	386.9	PASS
gather_block_quantized	Xi64_Wu8_Yfp16	gfx1151	32	0	399	PASS
gather_block_quantized	Xi64_Wfp8_Yfp16	gfx1151	24	0	972.9	PASS

Leaves: 16 (0 failed) | Cases: 289 (0 failed)

Mode: lookup | Coverage: 1 | Run: 8120 - Commit: 2bda8fa

[github-actions]

L2 Accuracy Results (EP vs CPU)

Model	Combined L2	Total Elems	Skipped NaN/Inf
conv_test_hybrid	1.12256E-07	64	0
GroupQueryAttention_seq256	59.2561	2621440	0
MatMulNBits_o_seq128	259.905	368640	0
QMoE_seq128	24.2583	368640	0

Threshold: 0.01 | Run: 8120 - Commit: 2bda8fa

[github-actions]

Binary sizes

Baseline: run 40919992229 - ddb207c (dll-sizes)

File	This (MB)	Main (MB)	Delta (MB)
hipgpu.dll	59.45	59.37	+0.07
custom_kernels_gfx1170.dll	5.79	5.60	+0.19
custom_kernels_gfx11-generic.dll	6.04	5.85	+0.19

MorphiZen EP Performance Results

Model	QPS	Session (s)	1st Infer (ms)	CPU%	Mem (MB)
full_model_seq128	8.93	7.10	245	3	1160
GroupQueryAttention_seq128	11049.60	2.58182	0	5	168
matmul_down_seq128	740.03	3.35	117	3	269

EPContext Export Performance

Model	QPS	Session (s)	1st Infer (ms)	CPU%	Mem (MB)
full_model_seq128	8.97	47.07	245	3	15506

EPContext Import Performance

Model	QPS	Session (s)	1st Infer (ms)	CPU%	Mem (MB)
full_model_seq128	8.97	10.04	253	3	15545

OGA Benchmark Results

Model	Warmup	Reps	Prompt Len	Gen Tokens	TTFT (ms)	TPS	Peak Mem (GB)	GPU Mem (GB)
gpt-oss-20b-webgpu-int4-rtn-block-32	1	5	128	128	91.7	82.5	1.25	12.35
Llama-3.1-8B-awq-g128-int4-asym-fp16-onnx-dml	1	5	128	128	118.4	49.6	1.15	5.41

OGA Wheel Smoke (Python benchmark_e2e.py)

Model	TTFT (ms)	TPS
Llama-3.1-8B-awq-g128-int4-asym-fp16-onnx-dml	115	49.8

VLM Benchmark Results

Model	TTFT (ms)	TPS	Output
gemma-4-E4B-it-fp16-int4-kquant-prune-logits	3763.0	8.0	This is a high-quality, outdoor photograph featuring a beautiful, fluffy white dog, likely a West Highland White Terrier or a si
Qwen3.8-27B-fp16-ve-fp16-int4-k_quant-gs128-text-dml-onnx	4855.0	14.6	The user wants a detailed description of the image. Let me carefully observe all elements before writing. Subject: A white do

TTFT / TPS vs main

Baseline: run 40919992229 - ddb207c | Gate: TTFT +20% / TPS -20%

Suite	Model	TTFT (ms)	Main	Delta	TPS	Main	Delta
oga	gpt-oss-20b-webgpu-int4-rtn-block-32	91.7	94.3	-2.7%	82.5	80.6	+2.4%
oga	Llama-3.1-8B-awq-g128-int4-asym-fp16-onnx-dml	118.4	119.1	-0.5%	49.6	49.2	+0.8%
vlm	gemma-4-E4B-it-fp16-int4-kquant-prune-logits	3763.0	3948.3	-4.7%	8.0	7.5	+7.6%
vlm	Qwen3.8-27B-fp16-ve-fp16-int4-k_quant-gs128-text-dml-onnx	4855.0	4929.0	-1.5%	14.6	14.6	+0.1%

Run: 8120 - Commit: 2bda8fa