跳到正文
r/MachineLearning· /u/Code_MasterCody·· 4 小时前AI 评分49

TKintegers K5 在 NVIDIA L4 上对比 TorchAO FP8:prefill 吞吐 3.21×、decode 吞吐 3.10×

TKintegers K5: 3.21× prefill and 3.10× decode throughput vs TorchAO FP8 on NVIDIA L4 [R] [P]

AI 导读

TKintegers 作者发布的 K5 平衡三进制权重表示在 NVIDIA L4 上跑出高于 TorchAO FP8 的吞吐:Qwen2.5-0.5B-Instruct 上 prefill 3,402.85 tokens/s 对 1,061.58(3.21×),decode 26.75 tokens/s 对 8.63(3.10×)。

正文

I built TKintegers, a balanced-ternary weight representation, converter, and GPU inference implementation.

The result that stood out to me: the NVIDIA L4 has native FP8 tensor-core support, but my K5 implementation achieved higher throughput than the TorchAO FP8 baseline—even though the GPU has no native K5 instructions.

On Qwen2.5-0.5B-Instruct, using a fixed matched workload:

- Prefill: K5 achieved 3,402.85 tokens/s versus FP8’s 1,061.58 — 3.21× throughput.

- Decode: K5 achieved 26.75 tokens/s versus FP8’s 8.63 — 3.10× throughput.

The baseline was standard TorchAO dynamic E4M3 FP8, running in eager PyTorch. Decode measured fixed-token model forwards with a carried KV cache.

What is K5?

TK stands for Ternary Kluster. K5 represents each weight code with five balanced ternary digits, giving 243 possible values. In this layout, each code occupies one byte. Five trits does not mean five binary bits.

My kernels expand bounded tiles for computation and use FP16 tensor-core operands with FP32 accumulation. Some tensors remain in higher precision.

For me, this is why the runtime matters alongside the representation: K5 can outperform this FP8 baseline on existing hardware without requiring a processor built specifically for K5.

These measurements apply to this model, workload, and implementation; they don’t establish a win over every FP8 runtime. The FP16 control was faster than both.

The repository includes the converter, native kernels, benchmark scripts, and recorded evidence:

https://github.com/CodeMasterCody3D/tkintegers

https://youtu.be/JFDBuwMoL0s?is=ucya9errGEh4xJ2L

I’d welcome technical feedback on the kernel design and benchmark methodology, especially independent replication.

submitted by /u/Code_MasterCody
[link] [留言]

来源:r/MachineLearning · reddit.com