Search papers, labs, and topics across Lattice.
This paper introduces the Direct-P method for optimizing FP4 tensor cores in attention mechanisms, addressing the limitations of softmax conversion and on-chip dependencies that hinder performance. By mapping scores directly to FP4 probabilities, the method achieves a throughput improvement of up to 2.13 times compared to bfloat16 on NVIDIA GB200 hardware. Additionally, the causal path enhances training efficiency by reconstructing probabilities from quantized queries and keys, leading to a 1.14 times acceleration in single-GPU updates for 8-billion-parameter models.
Achieving over 2x throughput improvements in attention mechanisms could redefine efficiency benchmarks for large-scale model training.
Blackwell's 4-bit floating-point (FP4) tensor cores do not automatically make attention faster because softmax conversion and on-chip dependencies dominate once its matrix products shrink. We address this with \emph{Direct-P} for noncausal inference and a causal path that passes the forward quantization directly into backward. Direct-P maps scores directly to FP4 probabilities and reaches up to 2.13$\times$ the bfloat16 (BF16) forward throughput on an NVIDIA GB200. The causal path reconstructs probabilities from saved quantized queries and keys and uses 8-bit floating-point (FP8) gradient operands, accelerating a complete single-GPU 8-billion-parameter update by up to 1.14$\times$. Matched distributed training retains FP8 probabilities and values; every tested MXFP4 probability/value training trajectory diverges.
