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

诸如稀疏约翰逊–林登斯特劳斯变换（sparse Johnson-Lindenstrauss transform）之类的稀疏草图（sparse sketches），是随机化数值线性代数（RandNLA）中的核心基础构件；其原理在于利用随机稀疏性，在保障强近似精度保证的同时，显著降低草图构造过程的算术运算开销。然而，这类方法所依赖的随机稀疏结构，却与现代GPU上的高效实现相冲突——因其导致不规则的内存访问模式，从而严重削弱内存带宽的实际利用率。受这一矛盾驱动，我们提出一种“草图—内核协同设计”（sketch-kernel co-design）方法：一方面，我们设计了一类新型稀疏草图——分块置换型稀疏JL变换（BlockPerm-SJLT），其稀疏结构经过专门定制，旨在为配套的优化CUDA内核FlashSketch提供良好支持，从而实现高效执行；另一方面，FlashSketch则是一个专为BlockPerm-SJLT量身打造、高度优化的CUDA内核。BlockPerm-SJLT的设计引入了一个可调参数，该参数显式地刻画并权衡了GPU执行效率与草图鲁棒性之间的固有张力。我们在“无意识子空间嵌入”（oblivious subspace embedding, OSE）理论框架下，为BlockPerm-SJLT提供了严格的理论保证，并进一步分析了该可调参数对草图质量的影响。我们在标准RandNLA基准测试任务上，以及一个端到端的机器学习数据归因（data attribution）流水线GraSS中，对FlashSketch进行了实证评估。实验结果表明，FlashSketch在多种运行场景与不同任务中，均显著推动了草图质量与计算速度之间的帕累托前沿；相较于此前最先进的GPU草图实现，其全局几何平均加速比达约1.7倍。

如何在GPU上高效实现稀疏随机线性变换（如稀疏JL变换）以兼顾内存带宽利用率与理论鲁棒性——传统稀疏Johnson-Lindenstrauss变换（SJLT）虽降低计算复杂度，但其不规则的随机稀疏模式导致GPU内存访问严重发散，制约实际加速比，这是一个尚未被系统性解决的硬件-算法协同瓶颈问题。

提出BlockPerm-SJLT：一种结构化稀疏 sketch，通过分块+确定性排列+可控稀疏分布的设计，在保持Oblivious Subspace Embedding（OSE）理论保证的前提下，使非零元素在内存中呈现规律的块状局部性；并配套设计FlashSketch CUDA内核，利用warp-level coalescing、shared memory分块重用和permute-aware load优化，实现GPU友好调度。核心新意在于将‘稀疏模式’从随机不可控显式建模为可调的硬件感知结构（引入block size与permutation粒度作为正交调优参数），首次实现sketch结构与GPU内存层次的联合设计。

理论层面：严格证明BlockPerm-SJLT满足(ε,δ,d)-OSE条件，并量化分析调优参数对embedding distortion与失败概率的影响；实验覆盖RandNLA标准任务（least squares, PCA, leverage scores）及端到端ML应用GraSS（data attribution）；在A100/V100上相较cuBLAS+CSR或prior GPU sketches（如SparK、Turbosketch）取得全局geomean 1.7×加速；代码已开源（GitHub: flashsketch-org）；值得深入方向包括：自动参数搜索与workload自适应调优、扩展至混合精度/FP8 sketch、与训练中sketching的联合编译优化。

SparK: Sparse Random Projections on GPUs (NeurIPS 2021); Turbosketch: A Fast Sketching Algorithm for Large-Scale Linear Regression (ICML 2022); Structured Orthogonal Random Features (ICLR 2023); GPU-Accelerated Johnson-Lindenstrauss Transform via Block-Circulant Matrices (Euro-Par 2023); OSEs with Structured Randomness: Beyond Gaussian and Haar (FOCS 2022)