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.

翻译：稀疏素描（如稀疏Johnson-Lindenstrauss变换）是随机数值线性代数中的核心原语，因其利用随机稀疏性来降低素描的算术成本，同时仍能提供强近似保证。然而，其随机稀疏性与现代GPU上的高效实现存在矛盾，因为它会导致不规则的内存访问模式，从而降低内存带宽利用率。受此矛盾启发，我们采用素描-内核协同设计方法：设计了一个新的稀疏素描家族BlockPerm-SJLT，其稀疏结构经过专门选择，以支持FlashSketch——一个高效实现这些素描的优化CUDA内核。BlockPerm-SJLT的设计引入了一个可调参数，该参数明确权衡了GPU效率与素描鲁棒性之间的张力。我们在无意识子空间嵌入（OSE）框架下为BlockPerm-SJLT提供了理论保证，并分析了可调参数对素描质量的影响。我们在标准RandNLA基准测试以及一个名为GraSS的端到端机器学习数据归因流程上对FlashSketch进行了实证评估。FlashSketch在一系列场景和任务中，推动了素描质量与速度的帕累托前沿，并相比先前最先进的GPU素描实现了约1.7倍的全局几何平均加速。