FlashAttention-T：面向张量化的注意力机制

[FlashAttention-T: Towards Tensorized Attention](https://dl.acm.org/doi/10.1145/3774934.3786425) FlashAttention-T, introduced by Tri Dao and colleagues (ACM DOI: 10.1145/3774934.3786425), represents a pivotal advancement in the FlashAttention lineage, pushing Transformer attention mechanisms toward full hardware-native tensorization. Standard scaled dot-product attention in Transformers incurs O(N²) time and memory complexity for sequence length N, bottlenecking long-context models (e.g., LLMs with millions of tokens) due to excessive HBM (High Bandwidth Memory) traffic. Prior FlashAttention iterations addressed this via IO-aware tiling: partitioning Q, K, V into SRAM-fit blocks, online softmax computation to avoid materializing the full N² attention matrix, and gradient recomputation during backpropagation, reducing memory to O(N) and yielding 2-4x speedups with kernel fusion. The key innovation of FlashAttention-T lies in fully tensorizing the attention kernel to exploit modern GPU Tensor Cores (e.g., NVIDIA Hopper/Blackwell WMMA instructions), transforming all operations—including nonlinear softmax—into high-throughput Tensor Matrix Multiply-Accumulate (TMA) primitives. Core technical contributions include: 1) **Tensorized Softmax**: Conventional per-element exp/logsumexp is replaced by block-wise approximations via tensor GEMMs, using lookup tables for exp or low-order Taylor expansions, with causal masking via tensor broadcasting for exactness. This enables single TMA calls for softmax(QK^T / √d). 2) **Multi-Head Tensorization**: Multi-head attention (MHA) is elevated to higher-order tensors (e.g., [batch, heads, seq_len/block_len]), processing all heads in one TMA invocation, slashing kernel launch overheads by 10x compared to batched serial execution. 3) **Low-Precision TMA Fusion**: Native FP8/INT4 support with block-floating-point scaling preserves accuracy while leveraging Tensor Core peaks (e.g., 1 PFLOPS FP16 on H100). Backward pass employs 'tensorized recomputation,' recomputing only tiled tensors on-the-fly.