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

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.7 \times$ over the prior state-of-the-art GPU sketches.

Doing math on computers is usually split into two subproblems. On the one hand, mathematicians design theoretically elegant and fast algorithms purely in the abstract, with the idea that a certain notion of fast would translate to actual efficiency on real computers. The second aspect is to take an abstract mathematical algorithm and actually create concrete instructions to implement this algorithm on real computers. In the context of modern computers, specifically GPUs (Graphics Processing Units), a divide exists between what is mathematically fast and what ends up being actually fast when implemented. We investigate one specific mathematical task called "sketching" that we want to do fast on GPUs. Sketching is a mechanism that greatly reduces the size of large data into something much more manageable, at the cost of some accuracy. It compresses data by approximation. The utility of this is that some downstream tasks that we would want to do on this data could now be much easier and faster, and still accurate since they can tolerate approximations. An inherent aspect of sketching is that there is some "randomness" in this procedure. This is what makes the approximation mathematically good. But this randomness also makes the real concrete instructions to run this procedure on GPUs very slow. The reason for this is that the sketching procedure was developed independently in the abstract, without explicitly taking into account the special characteristics of GPUs. To resolve this tension, we develop FlashSketch. The core principle is to design a new sketching procedure that still maintains good properties from a mathematical perspective, while also being amenable to a much faster and more elegant implementation on modern GPUs. This is a "co-design" approach to solving the overarching problem of doing this specific math problem on modern computers. It resolves, in this setting, a fundamental mismatch between mathematical efficiency and hardware efficiency. The result is that FlashSketch pushes the "Pareto" frontier, which means that it achieves the same level of approximation quality at a much faster speed, roughly 1.7x faster than previous approaches which were not co-designed with the GPU in mind. We also mathematically show that the approximation it produces remains strong.