Source-linked AI summary
FlowTT: Exploiting Computation Flow Reuse in Irregular Tensor-Train Embedding
Jongmin Seok, Chae Eun Rhee
TL;DR
Large embedding tables are compressed by TT decomposition, but TT gather still repeats partially shared computation and materializes intermediates off chip. FlowTT exposes prefix-shared irregular flows through grouped, fused, on-chip execution with dynamic scheduling. Across Meta-240, Meta-480, and Meta-788, it achieves the lowest latency, including up to 42.2% lower inference and 49.2% lower training latency than EcoRec at batch size 32,768.
Problem
TT-based embedding lookup does not fully reuse shared computation across input indices and repeatedly materializes intermediate results between sequential TT-core contractions.
Method
FlowTT combines prefix-based index grouping, fused on-chip TT execution, and persistent-thread scheduling with work stealing and L2 checkpointing.
Results
FlowTT consistently achieves the lowest latency on Meta-240, Meta-480, and Meta-788, with up to 42.2% lower inference latency and 49.2% lower training latency than EcoRec at batch size 32,768.
Takeaways & Limitations
The results indicate that exposing prefix-shared computation is central to efficient TT-based embedding execution, beyond compression alone.
Takeaways & Limitations
Existing on-chip pipeline techniques assume static, uniform dense data flow, whereas FlowTT addresses input-dependent irregular paths.
Abstract
from arXiv · showhide
Tensor-Train (TT) decomposition effectively compresses large embedding tables in recommendation models, but TT-based embedding lookup remains inefficient because partially shared computation flows across input indices are not fully reused and intermediate results are repeatedly materialized off-chip between sequential TT-core contractions. We present FlowTT, a flow-aware GPU execution framework that reformulates TT gather as a set of prefix-shared irregular computation flows. FlowTT combines flow-aligned prefix-based index grouping, a fused TT-embedding execution path with on-chip intermediate retention, and persistent-thread scheduling with chunk-based work stealing and L2 checkpointing to preserve reuse under skewed workloads. By co-designing task formation, data buffering, and scheduling with the structure of TT gather, FlowTT reduces redundant TT-core operations, global-memory traffic, and load imbalance. On Meta's synthetic recommendation benchmarks (Meta-240, Meta-480, and Meta-788), FlowTT consistently achieves the lowest latency compared to existing methods. At batch size 32,768, it reduces latency by up to 42.2% in inference and 49.2% in training relative to EcoRec, while also achieving the lowest inference peak memory usage. These results show that exposing prefix-shared computation is key to efficient TT-based embedding execution.
1 INTRODUCTION
TT decomposition compresses recommendation-model embeddings but leaves TT gather inefficient because shared computation and intermediate results are not fully reused. FlowTT addresses this with flow-aware grouping, on-chip execution, and dynamic scheduling.
- 1 INTRODUCTION: Embedding lookup dominates recommendation-model training and inference because large tables create high bandwidth demand and frequent off-chip movement.Distributed execution reduces some pressure but introduces communication overhead and imbalanced memory access.
- 1 INTRODUCTION: TT decomposition compresses embedding tables with low-dimensional TT cores, but lookup still requires sequential matrix operations across cores.These repeated operations create structural inefficiencies beyond ordinary embedding lookup.
- 1 INTRODUCTION: FlowTT reframes TT gather as partially shared computation flows and reorganizes execution around those flows.The framework groups common operations, retains reusable data on chip, and schedules irregular flow units dynamically.
- 1 INTRODUCTION: FlowTT combines a flow-aware execution model, flow-aligned on-chip pipelines, and dynamic scheduling for irregular workloads.Its stated goals are reducing redundant computation and off-chip traffic while alleviating workload imbalance.
2 RELATED WORK
Prior TT-embedding systems improve operator organization, communication, locality, or pipeline execution, but FlowTT targets finer-grained reuse of prefix-shared intermediates in irregular TT gather. The key challenge is coordinating prefix-aware task formation, on-chip buffering, and dynamic load balancing.
- 2.2 Optimization of Tensor-Train Embedding: TT-Rec reconstructs sparse embeddings from TT cores within TT-EMB modules before combining them with dense features for CTR prediction.EL-Rec and EcoRec improve practicality through operator reordering, contraction patterns, micro-batching, and pipeline scheduling, while FlowTT targets finer execution granularity.
- 2.3 On-Chip Pipeline and Scheduling Techniques: On-chip fusion techniques reduce intermediate DRAM spill but generally assume static, uniform dense pipelines, unlike input-dependent TT gather paths.Persistent threads and task stealing address irregular scheduling, but generic mechanisms do not by themselves specify prefix-intermediate reuse.
- 2.4 Problem Definition: TT gather combines shared prefix computation with input-specific suffix computation, forming an irregular flow rather than independent GEMM chains or duplicate lookups.Inputs sharing early mixed-radix digits can reuse intermediate contractions while differing at later stages.
- 2.4 Problem Definition: FlowTT promotes prefix-shared intermediates to GPU block-level execution units, unlike prior methods organized around operators, TT pairs, micro-batches, or tables.Shared-memory buffering lets inputs with a common prefix reuse the same intermediate within a block.
- 2.4 Problem Definition: FlowTT must co-design prefix-based task formation, shared-memory-centered buffering, and dynamic scheduling because irregular prefix frequencies create both reuse opportunities and load imbalance.Simple GEMM fusion or locality improvements do not determine which input-dependent intermediates to reuse or resolve skewed workloads.
3 PROPOSED METHOD
FlowTT reorganizes TT gather around prefix-shared computation flows, combining fused on-chip execution, prefix-based grouping, and dynamic scheduling to reuse intermediates while balancing irregular work.
- 3.3 Flow-Aligned Prefix-Based Index Grouping: FlowTT forms prefix groups from shared high-order TT-core indices and assigns them as execution units that remove redundant TT-core operations.Grouping uses sorting, mixed-radix conversion, and run-length encoding.
- 3.2 Flow-Aware Fused TT Embedding Execution: Fused forward and backward kernels combine sequential TT-core contractions and aggregation so shared intermediates are computed once and reused.The forward path computes C = A·B, then E = C·D, followed by pooling; the backward path reverses this structure.
- 3.2 Flow-Aware Fused TT Embedding Execution: FlowTT retains reusable intermediates in shared memory and uses L2 for inter-block reuse, avoiding repeated global-memory materialization and reloads.Ping-pong buffering overlaps TT-slice loading with computation, while batched shared-memory aggregation reduces atomic-operation frequency.
- 3.4 Dynamic Load Balancing via Persistent Threads and Work Stealing: Persistent threads first process assigned prefix groups, then steal unfinished chunks from other groups to reduce idle time while preserving prefix continuity.Chunk-level stealing addresses the wide variation in prefix-group lengths caused by long-tail recommendation workloads.
4 EVALUATION
FlowTT is evaluated against representative TT-based embedding implementations on Meta benchmarks, with results covering latency, memory, end-to-end training cost, and ablations of scheduling and prefix reuse.
- Latency and Memory Comparison under Varying Batch Sizes: FlowTT consistently achieves the lowest inference and training latency across Meta-240, Meta-480, and Meta-788 as batch size increases.The performance gap widens in larger-batch regimes, especially for training.
- End-to-End Latency Breakdown Including Index Grouping: Backward execution remains the dominant training cost, while index-grouping overhead becomes proportionally smaller at larger batch sizes.The breakdown includes index sorting, unique extraction, mixed-radix conversion, prefix extraction, and run-length encoding.
- Ablation of Dynamic Scheduling and Inter-Block Reuse: Dynamic scheduling reduces normalized latency from 1.00 to 0.87, and L2 checkpointing reduces it further to 0.85 on Meta-788 at batch size 32,768.The ablation compares static prefix mapping, persistent threads with work stealing, and persistent threads with work stealing plus L2 checkpointing.
- Ablation of Prefix-Based Reuse over Exact Deduplication: Prefix-Based Reuse consistently achieves lower latency than Unique Only by reusing computation among inputs with common prefixes.Unique Only removes identical indices, whereas Prefix-Based Reuse additionally groups shared higher-order TT-core paths.
5 CONCLUSION
FlowTT reframes TT gather as prefix-shared irregular computation and maps reusable states onto GPU execution and memory mechanisms. Across Meta benchmarks, it achieves the lowest latency, with substantial gains over EcoRec and the lowest inference peak memory usage.
- 5 CONCLUSION: FlowTT redefines TT gather as prefix-shared irregular computation and combines prefix grouping, fused on-chip execution, work stealing, and L2 checkpointing.The design targets redundant TT-core computation, global-memory traffic, and load imbalance.
- 5 CONCLUSION: 42.2% lower inference latency and 49.2% lower training latency than EcoRec are achieved at batch size 32,768 across the evaluated Meta benchmarks.FlowTT also attains the lowest inference peak memory usage and the best latency across evaluated settings.
- 5 CONCLUSION: The results indicate that exposing shared computation flows is a central performance opportunity for TT embedding beyond compression alone.FlowTT preserves reuse while mapping irregular execution onto the GPU memory hierarchy.