Search papers, labs, and topics across Lattice.
This paper introduces Windowed Batch Matrix Multiplication (WBMM), a novel approach that enhances the efficiency of large kernel depthwise convolutions by partitioning inputs into contiguous windows and utilizing a compact relative position bias table for regular memory access. Unlike traditional methods that degrade in performance with larger kernels, WBMM shows improved throughput as window sizes increase, achieving significant speedups and a larger per-layer receptive field. Experimental results demonstrate that WBMM outperforms standard depthwise convolution baselines while maintaining or exceeding accuracy on several benchmark datasets, leading to substantial training speedups across various hardware platforms.
WBMM flips the script on convolution efficiency, showing that larger windows can actually boost throughput rather than hinder it.
Large kernel depthwise convolutions achieve strong performance but suffer from significant degradation as kernel size grows due to irregular memory access from gather-based computation; while Large Kernel Acceleration (LKA) helps on small feature maps, it becomes counterproductive on large feature maps, even slower than non-accelerated implementations. We propose Windowed Batch Matrix Multiplication (WBMM), which partitions input into contiguous windows and indexes a compact relative position bias table to construct weight matrices, enabling regular memory access via batched matrix multiplication. This yields a unique property: WBMM's throughput improves with larger windows, opposite to depthwise convolutions that degrade with larger kernels. Operator-level benchmarks show WBMM with 14x14 windows outperforms 5x5 depthwise convolution baselines in speed while providing a 7.8x larger per-layer receptive field. Combined with inter-block cross-window communication and hierarchical window reparameterization, WBMM achieves comparable or higher accuracy on ImageNet-1K, COCO, and ADE20K with 1.31-1.88x training speedup, and demonstrates consistent advantages across GPU, CPU, and edge devices without requiring specialized acceleration kernels. Our code is available at http://github.com/wansong-s/WBMM
