Searcharxiv⌕ Search

arXiv subjects

Junchao Yang

Publications and source records attributed to Junchao Yang.

2 recordsLinked to original sources

SparseLift: Automatically Optimizing Opaque Sparse Paths in Industrial-Scale Recommendation Model Training

The sparse path, comprising sparse operations and collective communication, is a critical component of large-scale recommendation models. However, these operations are often represented as opaque nodes in the training computation graph, hindering whole-graph optimization. We present SparseLift, a low-overhead compiler for the sparse path of large-scale recommendation training. It exposes sparse computation and communication to whole-graph analysis to expand the optimization space and improve training efficiency. By analyzing model structure and runtime information, SparseLift reconstructs opaque nodes as analyzable subgraphs, recovering the resource semantics of sparse computation and communication within the graph. Building on the recovered semantics, we design incremental cost models and a constraint-aware planner to automatically select and organize subgraph optimizations into an execution plan that is reused during subsequent steady-state training. We evaluate SparseLift on an industrial video-platform recommendation task with 100B sparse parameters and the public DLRM benchmark. SparseLift improves steady-state throughput for ordinary industrial training steps by 16.5\% over the native baseline. On DLRM, it achieves $1.60\times$ and $1.26\times$ model-execution speedups over DLRM-Baseline and TorchRec, respectively.

cs.LG↗

PolyGraph: Plan-Buffered CUDA Graphs for Adaptive Speculative Decoding

Adaptive speculative decoding changes the active batch size B and draft length d from one decode step to the next, whereas a CUDA Graph replays a fixed launch configuration. The resulting verification width M = B d forces serving systems to pre-capture many shapes, capture new ones on the serving path, or pad every step to the largest draft. We present PolyGraph, a plan-buffer mechanism for custom Triton kernels that avoids this choice. PolyGraph captures one maximum-grid graph and writes the active shape to GPU memory before each replay. Kernels skip tiles beyond that shape, while graph topology, allocations, layouts, pointers, and launch grids remain fixed; the update transfers only 28 bytes. On one NVIDIA H800, a trace-driven prototype with Qwen2.5-3B dimensions replaces 64 shape-specific graphs with one, matches the warmed baseline in observed step latency (17.7 vs. 17.8 ms P50), and reduces measured pre-serving graph-build time from 3,915 to 93 ms. Fixed-draft controls show that the benefit comes from avoiding padded payload: the P99 difference from MaxPad is negligible at zero padding waste but reaches an 18.7% reduction at 100% waste. Our evidence establishes a bounded work-shape mechanism for custom kernels, not a production serving system, and is limited to synthetic traces on one GPU.

cs.LG↗