优化稀疏检索模型的内存瓶颈,实现更快更省显存的推理。
Sparton: Fast and Memory-Efficient Triton Kernel for Learned Sparse Retrieval
- 将矩阵乘法、激活函数和归约操作融合为单个核函数
- 减少峰值内存使用超90%,最高提速4.8倍
- 适用于大规模词表的多语言检索模型
当前先进的学习型稀疏检索(LSR)模型(如Splade)通常采用语言模型(LM)头将隐藏状态映射为词汇锚定的对数概率矩阵。该中间矩阵通过逐元素操作(ReLU、Log1P)及序列维度上的最大池化转换为稀疏词汇表示。然而,由于词表规模大(|V|可达3万至25万以上),该矩阵带来显著内存瓶颈,限制模型扩展。此外,各操作间的数据传输也拖慢性能。本文提出Sparton,一种针对LSR中LM头设计的快速、内存高效的Triton内核。Sparton将分块矩阵乘法、ReLU、Log1P与最大值归约融合为单一GPU内核,通过在原始对数概率块上进行早期在线归约,避免完整矩阵驻留内存。实验表明,仅在内核层面即实现最高4.8倍加速和一个数量级的内存降低。集成至Splade(|V| ~ 30k)后,支持33%更大的批大小,训练提速14%且无效果损失;在多语言骨干网络(|V| ~ 250k)上,批大小扩大26倍,训练速度提升2.5倍。
原文摘要 · Abstract (English)
State-of-the-art Learned Sparse Retrieval (LSR) models, such as Splade, typically employ a Language Modeling (LM) head to project latent hidden states into a lexically-anchored logit matrix. This intermediate matrix is subsequently transformed into a sparse lexical representation through element-wise operations (ReLU, Log1P) and max-pooling over the sequence dimension. Despite its effectiveness, the LM head creates a massive memory bottleneck due to the sheer size of the vocabulary (V), which can range from 30,000 to over 250,000 tokens in recent models. Materializing this matrix creates a significant memory bottleneck, limiting model scaling. The resulting I/O overhead between operators further throttles throughput and runtime performance. In this paper, we propose Sparton, a fast memory-efficient Triton kernel tailored for the LM head in LSR models. Sparton utilizes a fused approach that integrates the tiled matrix multiplication, ReLU, Log1P, and max-reduction into a single GPU kernel. By performing an early online reduction directly on raw logit tiles, Sparton avoids materializing the full logit matrix in memory. Our experiments demonstrate that the Sparton kernel, in isolation, achieves up to a 4.8x speedup and an order-of-magnitude reduction in peak memory usage compared to PyTorch baselines. Integrated into Splade (|V| ~ 30k), Sparton enables a 33% larger batch size and 14% faster training with no effectiveness loss. On a multilingual backbone (|V| ~ 250k), these gains jump to a 26x larger batch size and 2.5x faster training.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。