Facility Location Submodular Optimization on GPU for Token Subset Selection

Token subset selection for vision-language model and in-context learning is often modeled as facility location maximization: choose $k$ of $n$ source tokens that best cover $m$ target tokens under a similarity matrix. Its objective function is monotone submodular, so the greedy is near-optimal. However, every existing solver is slow. CPU libraries submodlib and apricot run the lazy greedy as a serial chain of row reads. Moreover, PyTorch loops recompute all marginal gains on the GPU at every iteration, moving the whole similarity matrix through HBM $k$ times. To improve the execution speed, we present Flash Greedy, an IO-aware persistent CUDA kernel that runs the entire selection in one launch. The similarity matrix is tiled over the streaming multiprocessors (SM) and placed in SRAM, L2 cache or HBM by its size. Every SM keeps a private copy of the coverage vector and evaluates exactly one candidate token per iteration. In addition, a replicated selection loop admits several tokens per grid barrier. A bool position vector is used to certify which candidate gains are still exact, thereby selecting the multiple tokens. Most importantly, our method returns exactly the naive greedy solution, and we give its IO complexity per placement regime. In the evaluation, we use three similarity matrices from MMTok, SCOPE, and Selective Annotation; whose token features are extracted from image, video, and text. Flash Greedy returns the same subset as the submodlib, apricot and original PyTorch implementation. However, Flash Greedy is $9\times$ faster than the PyTorch loops, and $26\times$ faster than the best CPU library.