cs.LGSep 24, 2026

DanLing NestedTensor: Composable Multi-Ragged Tensors for Deep Learning

Authors: Zhiyuan Chen

Organizations: DanLing Team

Abstract

Variable-size inputs are common in deep learning, but dense batching allocates a shared envelope and spends computation on padding. The cost multiplies across varying axes: an explicit pair state allocates BNmax⁡2BN_{\max}^2 positions instead of ∑iNi2\sum_i N_i^2. Packing removes that waste, but composing packed operations still requires the logical axes and sample boundaries a flat buffer no longer exposes. We present DanLing NestedTensor, a PyTorch tensor abstraction that makes multi-ragged structure a property of the tensor itself. Packed values carry tensor-backed partitions and logical dimension order, so broadcasting creates ragged axes, feature transformations retain them, and reductions consume them. The same representation carries through autograd and both eager and compiled execution. On an A100, the geometric-mean speedup over same-mode padding is 2.74×\times eager and 3.39×\times compiled across four BERT scales, and 1.97×\times eager across four FCN backbones. A four-block Pairformer-style workload runs 2.40-4.32×\times faster than a padded reference using native PyTorch kernels across square length regimes in eager execution, with peak allocation falling from 38.08 to 5.41 GiB on its high-variation batch. The tensor interface lets model code built from its supported operators compose efficient variable-size computation without managing offsets at any call site. Code will be released publicly upon publication.

Figures & tables

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Explore similar work

CardsList
  1. BARTER: Compact Dynamic Sparse Training with Standard Dense Kernels

    Oct 12, 2026Yueyuan Sui, Payal Mohapatra, Qi Zhu +1Dynamic Sparse TrainingNeural Network Pruning

  2. RLX: A Unified Multi-Backend Tensor Compiler and Distributed Runtime in Rust

    Sep 29, 2026Eugene Hauptmann, Nataliya KosmynaEfficient Neural Network Inference

  3. DynamiQ: Accelerating Gradient Synchronization using Compressed Multi-hop All-reduce

    Feb 9, 2026Wenchen Han, Shay Vargaftik, Michael Mitzenmacher +1Communication-Efficient Distributed TrainingHigh-Performance Computing