推論系統
TileMix Mixes FP16 and INT8 Within the Attention Matrix, Reaching 2.22× FlashAttention Throughput for 4K Prefill on A100
Rather than removing token connections, TileMix selects either the FP16 or INT8 Tensor Core path for each attention-score tile within a single fused kernel. It reaches 31.80K token/s on a 4K LLaMA 3.2 3B benchmark, though the results are currently limited to A100, static routing, and the prefill stage.
TileMix refines mixed-precision execution from the level of an entire tensor or kernel invocation down to hardware-aligned tile groups within the attention-score matrix. A routing bit determines whether each valid tile group uses FP16 or computes QK scores with INT8 Tensor Cores and INT32 accumulation. After rescaling, both paths jointly update the same FP16 online-softmax maximum, normalization factor, and output accumulator. TileMix therefore preserves the full dense attention connectivity—unlike sparse attention, which directly removes some token interactions—and requires no model retraining.
To prevent routing decisions from slowing the FlashAttention-style inner loop, the system packs decisions for each query tile row into a 64-bit bitmask, enabling constant-time lookup through shifts and masks. For longer sequences, one bit can control multiple adjacent key tiles, so routing metadata grows only with the number of KV heads and query tile rows. The public Triton implementation also supports grouped-query attention, unpadded variable-length batches, an INT8 KV cache interface, and integration scripts for LLaMA and Qwen models.
On a single NVIDIA A100 40GB with a batch size of 8, the authors included quantization, rescaling, routing, data movement, and scheduling in end-to-end prefill time. With a 4K input for LLaMA 3.2 3B-Instruct, the SpTrans configuration—which routes 75% of tile groups to INT8—reached 31.80K token/s. Under the same wrapper, FlashAttention achieved 14.33K token/s, making TileMix 2.22× faster, while the all-INT8 path reached 29.80K token/s. Results on LongEval and LV-Eval, which includes both Chinese- and English-language data, show that quality depends on both the spatial placement of FP16 tiles and the proportion covered by INT8. Mixed configurations can often recover accuracy lost by the all-INT8 approach, but high INT8 coverage does not match FP16 on every dataset.
The practical value of this work lies in providing a tunable accuracy–throughput interface, rather than claiming a universally lossless speedup. Key questions for future evaluation include whether adaptive routing justifies its additional cost, as well as performance with FP8, Hopper and Blackwell GPUs, the decode stage, and larger models. The current evidence covers only the FP16/INT8 path on A100. The library also has only two commits and no official release, so independent reproduction and integration testing are still needed.