FlashSketch: Sketch-Kernel Co-Design for Fast Sparse Sketching on GPUs

TL;DR

FlashSketch introduces a co-designed sparse sketching method and optimized GPU kernel that significantly improves the speed of randomized linear algebra tasks while maintaining sketch quality, enabling faster computations on modern GPUs.

Contribution

The paper presents BlockPerm-SJLT, a new sparse sketch family, and FlashSketch, an optimized CUDA kernel, achieving efficient GPU implementation with tunable robustness-speed trade-offs.

Findings

Achieves ~1.7x speedup over prior GPU sketches.

Provides theoretical guarantees under the OSE framework.

Balances sketching robustness and GPU efficiency effectively.

Abstract

Sparse sketches such as the sparse Johnson-Lindenstrauss transform are a core primitive in randomized numerical linear algebra because they leverage random sparsity to reduce the arithmetic cost of sketching, while still offering strong approximation guarantees. Their random sparsity, however, is at odds with efficient implementations on modern GPUs, since it leads to irregular memory access patterns that degrade memory bandwidth utilization. Motivated by this tension, we pursue a sketch-kernel co-design approach: we design a new family of sparse sketches, BlockPerm-SJLT, whose sparsity structure is chosen to enable FlashSketch, a corresponding optimized CUDA kernel that implements these sketches efficiently. The design of BlockPerm-SJLT introduces a tunable parameter that explicitly trades off the tension between GPU-efficiency and sketching robustness. We provide theoretical guarantees for BlockPerm-SJLT under the oblivious subspace embedding (OSE) framework, and also analyze the effect of the tunable parameter on sketching quality. We empirically evaluate FlashSketch on standard RandNLA benchmarks, as well as an end-to-end ML data attribution pipeline called GraSS. FlashSketch pushes the Pareto frontier of sketching quality versus speed, across a range of regimes and tasks, and achieves a global geomean speedup of roughly 1.7x over the prior state-of-the-art GPU sketches.