logo Mr.Robota 具身智能学习手册
/
中文
Chapter 11: Other auxiliary tools

FlashLib: A high-speed classic machine learning operator library for modern GPUs

Project repository: flashlib Project homepage and technical blog: 11-其他辅助工具 PyPI package name: https://github.com/FlashML-org/flashlib Open source license: Apache-2.0 Suitable location: https://flashml-org.github.io/

FlashLib is an open-source GPU classical machine learning operator library developed by FlashML-org. It is neither VLA, nor a robot control framework, nor a world model; rather, it resembles a high-speed underlying toolkit that reworks classic machine learning and linear algebra operators such as KMeans, KNN, PCA, SVD, UMAP, t-SNE, HDBSCAN, ANN retrieval, regression, classification, and GEMM into versions more suitable for modern GPUs.

In the embodied AI tutorial, it is suitable to be included under “Other auxiliary tools”. The reason is straightforward: in embodied AI projects, large-scale embeddings, point cloud features, visual features, trajectory features, retrieval databases, clustering results, and dimensionality reduction visualizations are often processed. Often, what actually stalls the pipeline is not necessarily large model inference, but the surrounding retrieval, clustering, PCA/SVD, KNN, t-SNE, or HDBSCAN operations. If these operators are replaced with more efficient GPU primitives from CPU/scikit-learn/regular GPU paths, data preprocessing, feature analysis, and online retrieval processes will become much smoother.

1. What problem does it solve

Traditional classic machine learning operators have a problem: the algorithms themselves are common, but many implementations are still designed for the offline batch processing era. For example:

  • KMeans repeatedly calculates the distance from points to all cluster centers;
  • KNN / ANN perform top-k on large-scale embedding databases;
  • PCA / SVD construct the Gram matrix or covariance matrix, and then perform feature decomposition;
  • t-SNE / UMAP / HDBSCAN are commonly used for visualization and unsupervised analysis, but they are slow with large data volumes;
  • Linear algebra primitives such as GEMM, QR, eigh, and polar are the core bottlenecks of many upper-level algorithms.

If these operators are written directly using the naive formula, they tend to generate large intermediate matrices, require frequent read and write operations on HBM, cause atomic contention, or result in memory access patterns that are not suitable for TensorCore / Hopper architectures. The goal of FlashLib is to reorganize these classic operators into fused/tiled/hardware-aware implementations that are more compatible with modern GPUs.

The positioning outlined on the project homepage is: transforming classical ML operators from offline tools into online primitives that can be integrated into agentic AI, retrieval, verification, compression, and orchestration loops. This perspective also applies to embodied AI. In robot systems, many intermediate steps are not the Transformer itself, but rather:

  • Retrieve similar demos from visual features;
  • Cluster trajectory embeddings;
  • Compress high-dimensional states or visual features using PCA / SVD;
  • Perform nearest neighbor search on point clouds or scene objects;
  • Visualize policy representations using UMAP / t-SNE;
  • Perform data cleaning, anomaly trajectory detection, and failure pattern grouping using KMeans/HDBSCAN.

FlashLib is designed to accelerate precisely these “key but peripheral to the model” operators.

2. Which operators are covered

The official README states that the current release covers 18 high-level primitives, which are mainly categorized into several types:

Category Operator
Clustering flash_kmeans, flash_dbscan, flash_hdbscan, flash_spectral_clustering
Nearest neighbors flash_knn, flash_ivf_flat, flash_ivf_pq, flash_cagra
Decomposition flash_pca, flash_truncated_svd
Manifold learning flash_umap, flash_tsne
Regression flash_linear_regression, flash_ridge, flash_logistic_regression
Classification flash_multinomial_nb, flash_random_forest
Preprocessing flash_standard_scaler
Linear algebra cov_gemm, gram_gemm, ab_gemm, eigh, polar, msign, cholqr2, split_basis, Multi-precision GEMM variants

In API design, each primitive usually provides two forms simultaneously:

  • Top-level flash_* function, suitable for direct invocation in the PyTorch pipeline;
  • sklearn-style class, such as KMeans, PCA, HDBSCAN, IVFFlat, IVFPQ, CAGRA, suitable for use in fit(...) / kneighbors(...) style.

3. Basic Installation

The simplest installation method is:

bashpip install flashlib

Install from source code:

bashgit clone https://github.com/FlashML-org/flashlib.git
cd flashlib
pip install -e .

It is recommended to test in a PyTorch environment that supports CUDA. FlashLib is designed for GPU tensors, and the typical input is torch.Tensor(device="cuda"). If the machine does not have a NVIDIA GPU, you can still read the source code and use the flashlib.info cost estimation interface, but the main acceleration benefit comes from the GPU.

4. Example of Minimum KMeans

The following example directly follows the usage style in the official README. It generates 1 million 128-dimensional vectors on the GPU and aggregates them into 1024 classes using flash_kmeans.

pythonimport torch
from flashlib import flash_kmeans

x = torch.randn(1_000_000, 128, device="cuda", dtype=torch.float32)
labels, centroids, n_iter = flash_kmeans(
    x,
    n_clusters=1024,
    max_iters=20,
)

print(labels.shape)
print(centroids.shape)
print(n_iter)

This smoke test proves only three things:

  1. flashlib can be correctly imported;
  2. PyTorch CUDA tensors can be passed to FlashLib primitives;
  3. The forward pipeline of KMeans can complete and return labels / centroids / iteration count.

It does not prove that a certain real task is definitely faster. To compare speeds, it is necessary to fix the input size, dtype, GPU model, number of warmup iterations, and baseline, such as in cuML, FAISS, scikit-learn, or the self-written version of PyTorch.

5. ANN Search: IVF-Flat, IVF-PQ, and CAGRA

FlashLib also provides GPU approximate nearest neighbor primitives. The official README highlights three types:

  • IVFFlat: First cluster and build an inverted index table, then query several lists according to nprobe;
  • IVFPQ: Add product quantization on top of IVF, using less memory to store vectors;
  • CAGRA: Build a proximity graph and use fused greedy traversal for high-throughput queries.

IVF-Flat Example:

pythonimport torch
from flashlib import IVFFlat

db = torch.randn(1_000_000, 128, device="cuda")
queries = torch.randn(10_000, 128, device="cuda")

index = IVFFlat(nlist=1024, nprobe=16).fit(db)
distances, indices = index.kneighbors(queries, n_neighbors=10)

nprobe is the main recall knob. Increasing nprobe will scan more inverted lists, usually resulting in higher recall but slower speed. This parameter is commonly used in RAG retrieval, similar demo retrieval, and trajectory embedding retrieval.

IVF-PQ Example:

pythonfrom flashlib import IVFPQ

index = IVFPQ(nlist=1024, m=16, nprobe=16).fit(db)
distances, indices = index.kneighbors(queries, n_neighbors=10)
print(index.compression_ratio)

If the original vector is 128-dimensional fp32, each vector occupies approximately 512 bytes; the PQ code of m=16 is about 16 bytes. Theoretically, compression can reach approximately 32 times. It is suitable for scenarios where the vector library is large and memory is the main bottleneck, but recall ceiling and quantization errors must be accepted.

CAGRA Example:

pythonfrom flashlib import CAGRA

index = CAGRA(graph_degree=32, itopk_size=64).fit(db)
distances, indices = index.kneighbors(queries, n_neighbors=10)

itopk_size and graph_degree are the main speed/recall adjustment parameters of CAGRA. For online small batch queries, graph search-based methods may be more suitable than IVF; when dealing with large batches and high recall, the matrix multiplication fine-scan of IVF-Flat may be more advantageous. In the tutorial, do not claim that a specific index is the absolute best; instead, choose based on batch size, recall requirements, database size, and GPU memory budget.

6. flashlib.info: Cost estimation without running the GPU

One interesting feature of FlashLib is flashlib.info. The official README states that it can predict the runtime, FLOPs, and HBM bytes of primitives on a pure CPU, taking about 5 microseconds, without importing torch, triton, or cutlass.

Example:

pythonimport flashlib.info as info

est = info.estimate(
    "kmeans",
    shape=(100_000, 64),
    params={"K": 256, "max_iters": 20},
    device="H200",
)

print(est.summary_line())

This interface is very useful for agentic workflows. For example, an automated data processing agent can estimate the cost of each step before deciding whether to perform PCA or KMeans first, and then decide whether to assign the task to the GPU. The same applies to embodied data processing: if dealing with millions of visual or trajectory embeddings, one can use info.estimate(...) to estimate a magnitude to avoid discovering insufficient memory or time budget after starting the process.

FlashLib is not a robot-specific library, but it has several direct uses in embodied AI projects.

7.1 Large-scale demo / trajectory clustering

After collecting data from DROID, BridgeData, RoboTwin, LIBERO, or self-built real robots, it is often necessary to group the teaching based on tasks, scenarios, action modes, or failure types. First, use VLM/policy encoder to extract episode embeddings, and then use KMeans or HDBSCAN for clustering.

textrobot episodes -> episode embeddings -> flash_kmeans / flash_hdbscan -> cluster labels -> 数据清洗 / 失败模式分析

Many VLA fine-tuning or prompt retrieval processes look for “which historical demos are most similar to the current task”. This is essentially nearest neighbor search. For small-scale cases, it can be directly flash_knn; for large-scale cases, IVFFlat, IVFPQ, or CAGRA can be used.

texttask embedding / visual embedding -> ANN index -> top-k demos -> prompt / adapter warm start / replay sampling

7.3 Representation Dimensionality Reduction and Visualization

When training VLA, world model, or navigation policy, it is often necessary to check whether the embedding space is separated by task, scenario, and action stage. PCA, t-SNE, and UMAP are commonly used analysis tools. FlashLib can make these visualization preprocessing steps faster, especially when the number of samples is large.

7.4 Point Cloud and 3D Feature Processing

3D scene reconstruction, point cloud segmentation, local geometric matching, and grasp proposal filtering often require KNN, PCA, or clustering. The KNN/PCA/KMeans in FlashLib can serve as prototype acceleration components, but whether it can replace the existing 3D pipeline depends on the data format, accuracy requirements, and batch organization method.

8. Design concept: Not simply “moving sklearn to GPU”

The FlashLib project homepage summarizes the design principles into four categories. Here, it is translated in a more engineering terms.

First, mathematical equivalent rewriting. Many simple representations of operators result in large intermediate matrices. FlashLib tries to convert formulas into GPU-friendly forms. For example, in KMeans assignments, the N x K distance matrix is not explicitly materialized, but the local minima are stored in registers.

Second, hardware-aware kernel. The optimal implementation of the same operator varies under different shapes, GPUs, and batch sizes. FlashLib uses Triton and CuteDSL to build multiple kernel variants, and makes heuristic dispatch based on the workload.

Third, tolerance-driven dispatch. Some scientific calculations require strict precision, while some embedding clustering or retrieval tasks can accept small errors in exchange for speed. FlashLib exposes a precision budget like tol, allowing the dispatcher to choose a faster path while satisfying the residual constraint.

Fourth, cost-predictable API. flashlib.info exposes runtime, FLOPs, HBM bytes, and bound regime, allowing the upper-level pipeline or LLM agent to estimate costs before execution.

This is also why it is more worthy of inclusion than just a “slightly faster KMeans”: it transforms the classical ML operator into a composable, budget-friendly, and schedulable basic primitive that can be managed by an agent.

9. What to note when using it

FlashLib is still relatively new, so it is recommended to use it cautiously in the tutorials.

  1. First, align the numerical values. Provide the same set of inputs to both FlashLib and baseline, and compare labels, distances, indices, explained variance, embeddings, or classification results. For non-unique outputs such as clustering and t-SNE/UMAP, compare the target metrics and downstream performance, rather than ensuring perfect consistency across all elements.
  2. Distinguish between exact and approximate. flash_knn, IVFFlat, IVFPQ, and CAGRA have different meanings. The recall of ANN index needs to be evaluated explicitly.
  3. Report performance after fixing hardware and version. Official benchmarks are mainly compared on H200, CUDA 13.0, PyTorch 2.11, Triton 3.6, and cuML 25.10. If using A100, 4090, L40S, or different CUDA/PyTorch versions, the performance results must be remeasured.
  4. Be aware of JIT warmup. The first call to the Triton/CuteDSL path incurs compilation overhead; the benchmark should exclude the first call or report cold start separately.
  5. Do not replace all libraries with this method. If there are only a few thousand data points, CPU sklearn or PyTorch may already be fast enough; FlashLib is more suitable for large-scale GPU-resident tensors.

Here is the minimum inspection process suitable for writing into experimental records.

bashpython - <<'PY'
import torch
from flashlib import flash_kmeans

print("torch", torch.__version__)
print("cuda available", torch.cuda.is_available())

x = torch.randn(100_000, 64, device="cuda", dtype=torch.float32)
labels, centroids, n_iter = flash_kmeans(x, n_clusters=256, max_iters=5)

print("labels", tuple(labels.shape), labels.dtype, labels.device)
print("centroids", tuple(centroids.shape), centroids.dtype, centroids.device)
print("n_iter", n_iter)
PY

Expected phenomena:

  • cuda available is equal to True;
  • labels has a shape of (100000,);
  • centroids has a shape of (256, 64);
  • The output tensor is on the CUDA device;
  • n_iter does not exceed max_iters.

If this fails, prioritize checking:

  • Whether flashlib is installed in the current Python environment;
  • Whether PyTorch can detect CUDA;
  • Whether the input tensor is actually on the GPU;
  • Whether the CUDA / PyTorch / Triton versions are compatible with the current platform;
  • Whether GPU primitives are executed incorrectly on machines without a NVIDIA GPU.

11. Location in this tutorial library

FlashLib should be placed together with flash-attn快速安装.md, hf快速下载同时避免下载不需要的其他分支文件.md, Xget库.md, and modelscope快速下载.md. It is not a core embodied AI algorithm, but it can significantly improve experimental efficiency.

If expansion is needed later, three types of small experimental measurements can be added:

  • Use flash_kmeans to cluster robot trajectory embeddings;
  • Use IVFFlat / CAGRA for similarity demo retrieval;
  • Use flash_pca / flash_tsne to visualize VLA hidden states or trajectory embeddings.