Single-cell GPU Acceleration with Seurat

Normalize and scale on the GPU

pbmc <- RunGGML(pbmc, op = "normalize")   # -> assay "data" layer
pbmc <- RunGGML(pbmc, op = "scale")       # -> assay "scale.data" layer

These behave like NormalizeData() and ScaleData() — the transformed matrices now live in the standard assay layers:

dim(LayerData(pbmc, layer = "data"))
dim(LayerData(pbmc, layer = "scale.data"))

The GPU "normalize" result matches Seurat’s NormalizeData() to floating-point tolerance:

gpu_data <- as.matrix(LayerData(pbmc, layer = "data"))
ref_data <- as.matrix(LayerData(
  NormalizeData(pbmc, verbose = FALSE), layer = "data"))
max(abs(gpu_data - ref_data))

PCA (embed)

pbmc <- RunGGML(pbmc, op = "embed", n_components = 20, reduction_name = "ggml")
Embeddings(pbmc, "ggml")[1:3, 1:4]

The reduction is an ordinary DimReduc, so every downstream step — Seurat’s own or ggmlR’s — uses it directly.

UMAP (umap)

op = "umap" lays the cells out in 2-D on the GPU. Both heavy phases are custom Vulkan compute shaders: a tiled f32 pairwise-distance kernel that builds the kNN graph (it sidesteps mul_mat, whose f16 accumulation would reorder nearest neighbours and corrupt the graph), and an SGD layout kernel that runs one dispatch per epoch with lock-free Hogwild updates. With a kd-tree kNN search and a sparse fuzzy graph in between, the whole UMAP runs an order of magnitude faster than a naive reference while matching its layout to float32 precision.

pbmc <- RunGGML(pbmc, op = "umap", reduction = "ggml", dims = 1:20,
                reduction_name = "umap")
Embeddings(pbmc, "umap")[1:3, ]

Neighbour graphs (neighbors)

op = "neighbors" is the GPU equivalent of FindNeighbors(): it builds the binary kNN graph and the shared-nearest-neighbour (SNN) graph whose weights are the Jaccard overlap of each pair’s neighbourhoods. The graphs land in @graphs under Seurat’s naming convention, so FindClusters() consumes them unchanged.

pbmc <- RunGGML(pbmc, op = "neighbors", reduction = "ggml", dims = 1:20)
Graphs(pbmc)                                   # <assay>_nn and <assay>_snn

pbmc <- FindClusters(pbmc, graph.name = paste0(DefaultAssay(pbmc), "_snn"),
                     verbose = FALSE)
table(pbmc$seurat_clusters)

DimPlot(pbmc, reduction = "umap", group.by = "seurat_clusters", label = TRUE)

The SNN weights are numerically identical to FindNeighbors() when both use the same exact kNN (Seurat’s default Annoy is approximate, so an exact match needs nn.method = "rann").

Speed-up

A representative run — 11 % subsample (23 279 cells), 2 000 variable features, 50 PCs, --gpu-knn — on an RDNA-class GPU:

step              cpu (s)    gpu (s)   speedup
largest_gene         8.92       0.28    32.22x
normalize            2.84       1.61     1.76x
scale                1.84       1.94     0.95x
embed (PCA)         15.92       3.03     5.26x
neighbors            4.43       1.45     3.05x
umap                15.32       5.34     2.87x
TOTAL (gpu ops)     49.27      13.66     3.61x

The pattern is worth reading rather than just the bottom line:

• PCA is the biggest matrix-multiply win (5.3×) — the covariance multiply is exactly what the GPU is for. Its own breakdown shows where the time goes: centre 0.47s matmul_cov 1.27s eigen 0.06s matmul_proj 0.98s — the eigensolve (CPU) is negligible, so keeping it off the GPU costs nothing.
• scale comes out ~1× (or slightly slower): it is a memory-bound elementwise pass with nothing to accelerate, so the host↔︎VRAM copy eats the gain. That is why op = "scale" defaults to the CPU even under Vulkan (scale_backend = "vulkan" forces the GPU path).
• umap here runs its kNN on the FNN kd-tree and its SGD layout on the CPU by default — the 2.9× comes from the GPU distance/graph phase. The SGD default is deliberate (see the UMAP section); sgd_backend = "vulkan" opts into the GPU shader.
• largest_gene looks like a 32× GPU win but is CPU on both sides — the speed-up is the sparse @x argmax replacing stock Seurat’s qlcMatrix path, not the GPU.

Agreement

The script densifies a fixed random column sample and compares the two arms. Every accelerated step matches Seurat to float32-vs-float64 noise, and the downstream clustering/embedding preserve the same structure:

normalize   max abs err  9.26e-07   (2000 features x 5000 sampled cells)
largest_gene top-gene agree  1.0000   percent max abs err  0.00e+00
scale       max abs err  2.38e-06
PCA         min |cor| over PC1-10  1.0000
clusters    ARI  0.9364   (11 GPU vs 11 CPU communities)
UMAP        within/total SS  gpu 0.048  cpu 0.060  (lower = tighter)

The comparisons account for the maths, not just the bytes: PCs are eigenvectors so their sign is arbitrary (compare |cor|), clusterings are label-permuted (the adjusted Rand index is invariant), and UMAP is stochastic and initialised differently in each arm (so instead of coordinates, measure how tightly each clustering separates in the embedding — the GPU layout is actually a hair tighter here, 0.048 vs 0.060).