SURGE让超大规模文本嵌入在低内存下实现高速流式处理
SURGE: SuperBatch Unified Resource-efficient GPU Encoding for Heterogeneous Partitioned Data
- 通过分析建模实现动态批处理，避免固定大小批的内存瓶颈
- 1000万文本处理速度达26,413条/秒，内存仅需2.6GB（原需32.7GB）
- 适合超大规模文本嵌入场景，尤其在资源受限节点上表现优异
我们提出SURGE，一个部署于生产环境的流式GPU编码系统，用于对超过8亿条文本、跨4万个逻辑分区的数据生成嵌入。生产级嵌入管道面临逻辑分区与高效利用GPU之间的矛盾：独立处理每个分区导致$P$次进程间通信（IPC），其开销限制了计算轻量模型的吞吐。本工作贡献为理论分析：(i) 构建成本模型（定理1），在三个编码器上预测吞吐误差小于2%，参数范围覆盖15倍；(ii) 提出内存安全边界（引理3），支持流式双阈值策略，峰值内存降至$O(B_{\min} + n_{\max})$而非$O(N)$；(iii) 建立$ϕ$/CV决策框架，界定该模式适用范围。简单批量处理需$O(N)$峰值内存（1000万文本时32.7 GB，超过6000万即不可行），输出延迟高且无容错。SURGE以$O(B_{\min} + n_{\max})$内存（2.6 GB）达成相同吞吐，首次输出时间快68倍，支持以SuperBatch粒度崩溃恢复。在1000万文本、4块NVIDIA L4 GPU上，达到26,413条/秒，内存消耗仅为固定批处理的1/12.6。验证基于bge-base（109M参数，$d$=768，误差1.3%）及$σ$∈{1.0, 1.72, 2.5}的对数正态分布，速度提升稳定在±3%内。相比分区批处理基线（PB-PBP-LB），SURGE仍保持7%吞吐优势和2.5倍更快的首次输出时间。配套工程优化——零拷贝Arrow序列化（提速22-25×）、异步I/O流水线（最高获益93%）——虽显著提升性能，但非本文核心贡献。
原文摘要 · Abstract (English)
We present SURGE, a streaming GPU encoding system deployed in production to generate embeddings for over 800 million texts across 40,000 logical partitions. Production embedding pipelines face a tension between logical data partitioning and efficient GPU utilization: processing each partition independently incurs $P$ inter-process communication (IPC) calls whose overhead limits throughput for compute-light models. Our contributions are analytical: (i) a cost model (Theorem 1) predicting throughput within 2% across three encoders spanning a 15$\times$ parameter range; (ii) a memory-safety bound (Lemma 3) enabling a streaming two-threshold policy with peak memory $O(B_{\min} + n_{\max})$ rather than $O(N)$; and (iii) a $ϕ$/CV decision framework characterizing when the pattern applies beyond our workload. The naive fix of batching at fixed size requires $O(N)$ peak memory (32.7 GB at 10M texts; infeasible beyond ~60M on 192 GB nodes), produces no output until all encoding completes, and offers no fault tolerance. SURGE achieves the same throughput with $O(B_{\min} + n_{\max})$ bounded memory (2.6 GB), 68$\times$ faster time-to-first-output, and crash recovery at SuperBatch granularity. On 10M texts with 4 NVIDIA L4 GPUs, SURGE delivers 26,413 texts/s -- matching fixed-batch throughput while using 12.6$\times$ less memory. We validate on bge-base (109M, $d$=768, error 1.3%) and across log-normal $σ$ in {1.0, 1.72, 2.5} (speedup invariant within $\pm$3%), and compare against a partition-batched baseline (PB-PBP-LB), against which SURGE retains a 7% throughput edge and 2.5$\times$ faster TTFO. Complementary engineering -- zero-copy Arrow serialization (22-25$\times$ speedup) and async I/O pipelining (up to 93% benefit) -- realizes the design but is not the contribution.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。