Diagonal-Tiled Mixed-Precision Attention for Efficient Low-Bit MXFP Inference

Yifu Ding; Xinhao Zhang; Jinyang Guo

cs.LGApr 5, 2026

Diagonal-Tiled Mixed-Precision Attention for Efficient Low-Bit MXFP Inference

Authors: Yifu Ding, Xinhao Zhang, Jinyang Guo

Organizations: State Key Laboratory of Complex & Critical Software Environment (SKLCCSE), Beihang University · School of Computer Science and Engineering, Beihang University · Beijing Jiaotong University · School of Artificial Intelligence, Beihang University

Abstract

Transformer-based large language models (LLMs) have demonstrated remarkable performance across a wide range of real-world tasks, but their inference cost remains prohibitively high due to the quadratic complexity of attention and the memory bandwidth limitations of high-precision operations. In this work, we present a low-bit mixed-precision attention kernel using the microscaling floating-point (MXFP) data format, utilizing the computing capability on next-generation GPU architectures. Our Diagonal-Tiled Mixed-Precision Attention (DMA) incorporates two kinds of low-bit computation at the tiling-level, and is a delicate fused kernel implemented using Triton, exploiting hardware-level parallelism and memory efficiency to enable fast and efficient inference without compromising model performance. Extensive empirical evaluations on NVIDIA B200 GPUs show that our kernel maintains generation quality with negligible degradation, and meanwhile achieves significant speedup by kernel fusion. We release our code at https://github.com/yifu-ding/MP-Sparse-Attn.

Figures & tables

Name	Block	Element		Scale
	Size	Format	Bits	Format	Bits
MXFP8	32	FP8 (E4M3 / E5M2)	8	E8M0	8
MXFP4	32	FP4 (E2M1)	4	E8M0	8
NVFP4	16	FP4 (E2M1)	4	E4M3	8

Table 1: Comparison of representative MXFP data formats, including their block sizes, element formats, and shared scale formats.

Figure 1: Visualization of quantization error of MXFP4 and NVFP4 format for query, key and attention score.

Format	Cos Sim ( $\uparrow$ )	PSNR ( $\uparrow$ )	L1 ( $\downarrow$ )	RMSE ( $\downarrow$ )
MXFP8	0.988	71.70	0.246	0.003
MXFP4	0.714	60.82	0.924	0.009
NVFP4	0.982	69.37	0.309	0.003
NVFP4+	0.983	69.63	0.312	0.003
Ours	0.988	71.70	0.248	0.003

Table 2: Quantization error of attention score using different data formats. “+” means combining with tokenwise quantization.

Figure 2: Overview workflow of our Diagonal-Tiled Mixed-Precision Attention. It first applies fused mixed-precision quantization to produce low-bit and high-bit $Q$ and $K$ representations, and then performs diagonal-tiled mixed-precision attention, using higher precision near the diagonal and lower precision elsewhere to balance accuracy and efficiency.

	LLaMA3.1-8B		LLaMA3.2-3B
Task	Native	Ours	Native	Ours
2wikimqa	39.15	32.54	30.75	36.33
dureader	28.40	32.93	28.11	33.04
gov_report	33.87	34.89	32.26	33.04
hotpotqa	49.89	50.49	46.18	48.59
lcc	49.46	56.96	42.98	45.54

Table 3: Comparison of attention implementations on LLaMA3.2-3B and LLaMA3.1-8B. “Native” refers to SDPA implementation supported by PyTorch.

Format	MP Size	Attn (ms)	Quant (ms)	Total (ms)
MXFP4	–	12.491	0.242	12.980
NVFP4	–	12.941	0.204	13.404
MXFP8	–	16.480	0.044	16.771
Ours	128	7.110	0.382	7.776
Ours	256	15.056	0.382	15.720

Table 4: Latency breakdown of different block-scale types and configurations. “MP Size” denotes the mixed-precision block size used for the higher-bit diagonal and sink blocks.

Diag.	Sink	Bit ${}_{\text{high}}$ (%)	Cos Sim $\uparrow$	Rel. L1 $\downarrow$	RMSE $\downarrow$	PSNR $\uparrow$
-	-	0.0	0.778	0.620	0.065	43.715
-	-	100.0	0.819	0.547	0.059	44.568
0	128	1.15	0.781	0.780	0.072	42.817
128	0	1.15	0.782	0.644	0.066	43.635
128	128	2.30	0.822	0.539	0.059	44.657
512	512	9.22	0.826	0.542	0.058	44.731

Table 5: Similarity metrics under different token numbers for diagonal and sink windows.

Encode	Pack	Scale Cvt.	MP	$L=2\text{k}$ ( $\mu$ s)	$L=8\text{k}$ ( $\mu$ s)
✗	✗	✗	✗	7262.41	22628.96
✓	✗	✗	✗	802.90	1113.77
✓	✓	✗	✗	740.64	942.67
✓	✓	✓	✗	179.97	299.69
✓	✓	✓	✓	97.87	282.46

Table 6: Ablation study of kernel fusion components to analyze their effect on throughput (TOPS) and latency. Encode : encoding FP16 to FP4/FP8 MX format. Pack : packing two FP4 values to one UINT8. Scale Cvt. : converting microscaling scalar to E8M0 format. MP : fusing quantization for mixed bitwidth to a single kernel.

Operator	Time	Time (%)
Not fused	18.320 ms	–
- Element encoding	17.742 ms	100.0%
MinOps	2.110 ms	11.89%
ArgMinOps	2.054 ms	11.58%
Direct_Copy	1.256 ms	7.08%
CompareEq	890.18 us	5.02%

Table 7: Latency breakdown of the non-fused MX encoding pipeline and our fused implementation.

Granu.	Latency $\downarrow$	Cos Sim $\uparrow$	Rel. L1 $\downarrow$	RMSE $\downarrow$	PSNR $\uparrow$
Per-Tensor	6.276 ms	0.732	0.560	0.067	43.479
Per-Block	6.366 ms	0.736	0.558	0.067	43.531
Per-Token	7.131 ms	0.822	0.539	0.059	44.657

Table 8: Performance comparison under different quantization granularities. Latency are reported in milliseconds.

LLM Inference Acceleration

LLM inference acceleration develops algorithms and systems that reduce language-model latency, memory use, and serving cost while preserving output quality.

See all 1,084 papers

TileMix: Tile-Centric Mixed-Precision Attention for LLM Inference Acceleration
Aug 18, 2026Hanzhi Zhang, Qiao Zhang, Qinglei Cao +4LLM Inference Acceleration Long-Context Language Model Inference
LLaVi Lab, Computer Science & Engineering, University of North Texas · Computer Science, Saint Louis University · Data Science, University of North Texas
LampAttention: Look-Ahead Mixed-Precision FlashAttention for Dedicated Accelerators
Sep 30, 2026Stanislav Budzinskiy, Marian Gloser, Tolunay Yilmaz +5AI Accelerator Inference LLM Inference Acceleration
Faculty of Mathematics University of Vienna, Austria · Huawei Heisenberg Research Center, Munich, Germany · Huawei Technologies Co. Ltd
Accelerating GPU Inference of Large Language Models with Moderately Unstructured Sparse Weight Matrices
Jun 13, 2026Tao Lu, Haoyu Wang, Zonghui Wang +3LLM Pruning GPU Kernel Optimization
National University of Singapore · Zhejiang University

Aug 18, 2026J/K move · Enter open · S save

TileMix: Tile-Centric Mixed-Precision Attention for LLM Inference Acceleration

Hanzhi Zhang, Qiao Zhang, Qinglei Cao, Heng Fan+3

LLaVi Lab, Computer Science & Engineering, University of North Texas · Computer Science, Saint Louis University · Data Science, University of North Texas

LLM Inference Acceleration Long-Context Language Model Inference Transformer Attention Mixed-Precision Quantization

Open paper AskPDF arXiv

Long-context prefill in large language models (LLMs) incurs substantial computation and memory traffic because dense self-attention computes quadratic query-key scores. Existing methods either use a uniform low-precision path or select token interactions, leaving spatial precision routing over hardware-aligned score tiles outside fused dense attention. We introduce TileMix, a tile-centric precision-routing kernel that makes numerical precision an executable spatial decision over score-tile groups within fused dense attention. TileMix partitions the attention matrix into hardware-aligned score tiles, packs routing decisions into compact bitmasks, and dispatches each tile group through FP16 or INT8 score computation while both paths update a shared online-softmax state. Scalable precision grouping lets each routing bit govern multiple adjacent key tiles, preserving hardware-aligned compute tiles and compact metadata at long contexts. By routing all legal tile groups, TileMix preserves dense token connectivity, requires no training, and supports grouped-query attention, variable-length batches, and INT8 key/value caches. Across LongEval, LV-Eval, and A100 prefill benchmarks on LLaMA, Qwen, and Vicuna, TileMix recovers long-context quality lost under uniform INT8 and improves prefill throughput over FP16, yielding a controllable accuracy-efficiency frontier across model families. The implementation is available at https://github.com/HanzhiZhang-Ulrica/TileMix.

Abstract

LLM Inference Acceleration

Explore similar work

TileMix: Tile-Centric Mixed-Precision Attention for LLM Inference Acceleration

LampAttention: Look-Ahead Mixed-Precision FlashAttention for Dedicated Accelerators

Accelerating GPU Inference of Large Language Models with Moderately Unstructured Sparse Weight Matrices