FlashLib: высокоскоростная библиотека классического машинного обучения для современных GPU

Репозиторий проекта: https://github.com/FlashML-org/flashlib Сайт проекта и технический блог: https://flashml-org.github.io/ Имя пакета в PyPI: flashlib Лицензия: Apache-2.0 Рекомендуемый раздел: 11-auxiliary-tools

FlashLib — библиотека операторов классического машинного обучения для GPU с открытым исходным кодом от FlashML-org. Это не VLA, не фреймворк для управления роботом и не модель мира. Скорее это быстрая низкоуровневая библиотека, которая переписывает KMeans, KNN, PCA, SVD, UMAP, t-SNE, HDBSCAN, ANN-поиск, регрессию, классификацию, GEMM и другие классические операторы машинного обучения и линейной алгебры под современные GPU.

В учебнике по воплощённому ИИ она включена в раздел «Вспомогательные инструменты». Причина простая: в проектах воплощённого ИИ часто приходится работать с большими embedding, признаками точечных облаков, визуальными признаками, признаками траекторий, базами для поиска, результатами кластеризации и визуализацией после понижения размерности. Часто узкое место конвейера — не большая языковая модель, а окружающие её операции: поиск ближайших соседей, кластеризация, PCA/SVD, KNN, t-SNE или HDBSCAN. Если заменить эти операции CPU/scikit-learn или обычным GPU-путём на более эффективные GPU-примитивы, подготовка данных, анализ признаков и онлайн-поиск станут легче.

1. Какую проблему решает

У классических алгоритмов машинного обучения есть типичная проблема: сами алгоритмы известны, но многие реализации рассчитаны на пакетную обработку офлайн. Например:

• KMeans многократно вычисляет расстояния от каждой точки до всех центров кластеров;
• KNN / ANN выполняют top-k поиск по большой базе embedding;
• PCA / SVD строят матрицу Грама или ковариационную матрицу, а затем выполняют разложение на собственные значения;
• t-SNE / UMAP / HDBSCAN часто используются для визуализации и неконтролируемого анализа, но на больших объёмах данных работают медленно;
• GEMM, QR, eigh, polar и другие примитивы линейной алгебры — ключевое узкое место многих высокоуровневых алгоритмов.

Если эти операторы реализовать «в лоб», возникают огромные промежуточные матрицы, частые обращения к HBM, конкуренция за атомарные операции и шаблоны доступа к памяти, которые плохо подходят для TensorCore и архитектуры Hopper. Цель FlashLib — переписать классические операторы в виде fused / tiled реализаций, учитывающих особенности аппаратуры.

На главной странице проекта FlashLib описан как инструмент, который превращает классические операторы машинного обучения из офлайн-утилит в онлайн-примитивы, пригодные для agentic AI, поиска, проверки, сжатия и циклов оркестрации. Для воплощённого ИИ этот подход тоже полезен. Во многих звеньях робототехнической системы нужны не трансформеры сами по себе, а:

• поиск похожих демонстраций по визуальным признакам;
• кластеризация embedding траекторий;
• сжатие высокоразмерных состояний или визуальных признаков с помощью PCA / SVD;
• поиск ближайших соседей для точечных облаков и объектов сцены;
• визуализация представлений политики с помощью UMAP / t-SNE;
• очистка данных, поиск аномальных траекторий и группировка типов отказов с помощью KMeans / HDBSCAN.

FlashLib особенно подходит для ускорения именно таких операторов, которые находятся рядом с моделью, но критически влияют на общую скорость.

2. Какие операторы поддерживаются

В официальном README указано, что текущий релиз охватывает 18 high-level примитивов. Они делятся на несколько групп:

В API каждый примитив обычно доступен в двух формах:

• функция верхнего уровня flash_*, которую удобно вызывать прямо в конвейере PyTorch;
• класс в стиле sklearn, например KMeans, PCA, HDBSCAN, IVFFlat, IVFPQ, CAGRA, с методами fit(...) и kneighbors(...).

3. Базовая установка

Самый простой способ установки:

pip install flashlib

Установка из исходников:

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

Тестировать лучше в окружении с PyTorch и доступной CUDA. FlashLib работает с тензорами GPU, типичный вход — torch.Tensor(device="cuda"). Если NVIDIA GPU нет, можно читать исходный код и использовать интерфейс оценки стоимости flashlib.info, но основной выигрыш в скорости получается именно на GPU.

4. Минимальный пример KMeans

Пример написан в стиле официального README. Он создаёт на GPU 1 миллион 128-мерных векторов и разбивает их на 1024 кластера с помощью flash_kmeans.

import 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)

Этот smoke test доказывает только три вещи:

1. flashlib корректно импортируется;
2. тензоры CUDA PyTorch попадают в примитивы FlashLib;
3. прямой проход KMeans завершается и возвращает метки, центроиды и число итераций.

Он не доказывает, что на вашей реальной задаче будет быстрее. Для сравнения скорости нужно зафиксировать размер входных данных, тип данных, модель GPU, число прогревочных запусков и базовый вариант, например cuML, FAISS, scikit-learn или собственную реализацию на PyTorch.

5. Поиск ближайших соседей: IVF-Flat, IVF-PQ и CAGRA

FlashLib также предоставляет примитивы приблизительного поиска ближайших соседей (ANN) на GPU. В официальном README выделены три варианта:

• IVFFlat: сначала кластеризация и построение инвертированного индекса, затем поиск по нескольким спискам, заданным параметром nprobe;
• IVFPQ: IVF с product quantization, которая позволяет хранить векторы в меньшем объёме памяти;
• CAGRA: строит граф близости и выполняет быстрый жадный обход (fused greedy traversal) для запросов с высокой пропускной способностью.

Пример IVF-Flat:

import 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)

Главный параметр полноты результатов (recall) — nprobe. Чем больше nprobe, тем больше списков просматривается: полнота обычно выше, а скорость ниже. Этот параметр часто встречается в RAG, поиске похожих демонстраций и поиске по embedding траекторий.

Пример IVF-PQ:

from 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)

Если исходные векторы 128-мерные и хранятся в fp32, каждый занимает около 512 байт. PQ-код при m=16 занимает около 16 байт, то есть теоретическое сжатие примерно в 32 раза. Этот вариант подходит, когда база векторов очень большая и главное ограничение — видеопамять, но нужно учитывать потолок полноты и ошибку квантования.

Пример CAGRA:

from flashlib import CAGRA

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

Главные параметры скорости и полноты у CAGRA — itopk_size и graph_degree. При онлайн-запросах с небольшим батчем графовые методы могут подходить лучше IVF. При большом батче и очень высокой полноте выгоднее IVF-Flat с полным сканированием через матричное умножение. В учебнике не стоит утверждать, что какой-то индекс абсолютно оптимален: выбирайте его по размеру батча, требованиям к полноте, размеру базы и бюджету видеопамяти.

6. flashlib.info: оценка стоимости без GPU

Интересная особенность FlashLib — flashlib.info. В официальном README сказано, что он может на чистом CPU оценить время выполнения примитива, число операций с плавающей точкой (FLOPs) и объём обращений к HBM. Оценка занимает около 5 микросекунд и не импортирует torch, triton или cutlass.

Пример:

import flashlib.info as info

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

print(est.summary_line())

Этот интерфейс полезен в agentic workflow. Например, агент автоматической обработки данных может перед выбором между «сначала PCA или сначала KMeans» оценить стоимость каждого шага и только потом отправлять задачу на GPU. Для обработки данных воплощённого ИИ это тоже работает: если нужно обработать несколько миллионов визуальных embedding или embedding траекторий, сначала можно оценить порядок величины через info.estimate(...), чтобы не запустить задачу, которая не уложится в видеопамять или временной бюджет.

7. Связь с воплощённым ИИ

FlashLib не является библиотекой для робототехники, но в проектах воплощённого ИИ у неё есть несколько прямых применений.

7.1 Кластеризация больших наборов демонстраций и траекторий

После сбора данных DROID, BridgeData, RoboTwin, LIBERO или собственных данных с реального робота демонстрации часто нужно разбить на группы по задаче, сцене, типу движения или типу отказа. Можно сначала извлечь embedding эпизода с помощью VLM или энкодера политики, а затем применить KMeans или HDBSCAN.

эпизоды робота -> embedding эпизодов -> flash_kmeans / flash_hdbscan -> метки кластеров -> очистка данных / анализ отказов

7.2 Поиск похожих демонстраций

Во многих задачах дообучения VLA или подбора контекста для промпта нужно найти, какие исторические демонстрации больше всего похожи на текущую задачу. По сути это поиск ближайших соседей. На небольшом объёме подойдёт flash_knn. На большом — IVFFlat, IVFPQ или CAGRA.

embedding задачи / визуальный embedding -> ANN-индекс -> top-k демонстраций -> промпт / тёплый старт адаптера / выборка для replay

7.3 Понижение размерности и визуализация представлений

При обучении VLA, моделей мира или навигационных политик часто нужно проверить, разделяется ли пространство embedding по задачам, сценам и стадиям действия. Для такого анализа обычно используют PCA, t-SNE и UMAP. FlashLib может ускорить подготовку данных для визуализации, особенно при большом числе примеров.

7.4 Обработка облаков точек и 3D-признаков

Реконструкция 3D-сцен, сегментация облаков точек, локальное сопоставление геометрии и фильтрация предложений захвата часто требуют KNN, PCA или кластеризации. Операторы KNN, PCA и KMeans из FlashLib можно использовать как ускоряющие компоненты при прототипировании. Но заменить ими существующий 3D-конвейер можно только с учётом формата данных, требований к точности и организации батчей.

8. Идея проекта: не просто «sklearn на GPU»

На главной странице FlashLib принципы проекта сведены к четырём пунктам. Ниже они изложены в более инженерной форме.

Первый — математически эквивалентная переформулировка. Наивная реализация многих операторов создаёт большие промежуточные матрицы. FlashLib по возможности переписывает формулы в удобном для GPU виде. Например, назначение точек в KMeans не материализует матрицу расстояний N x K целиком, а хранит локальные минимумы в регистрах.

Второй — аппаратно-ориентированные ядра. Оптимальная реализация одного и того же оператора зависит от формы данных, модели GPU и размера батча. FlashLib использует Triton и CuteDSL для нескольких вариантов ядер и выбирает подходящий по эвристике в зависимости от нагрузки.

Третий — диспетчеризация по допустимой погрешности (tolerance-driven dispatch). В одних научных расчётах нужна строгая точность, в задачах кластеризации или поиска embedding иногда допустима небольшая ошибка ради скорости. FlashLib предоставляет параметр tol, который задаёт бюджет точности: диспетчер выбирает более быстрый путь, если остаточная ошибка остаётся в допустимых пределах.

Четвёртый — предсказуемая стоимость через API. flashlib.info показывает время выполнения, FLOPs, объём обращений к HBM и тип ограничения (bound regime), поэтому конвейер или LLM-агент может оценить стоимость до запуска.

Именно поэтому FlashLib стоит включить в учебник, а не просто назвать «более быстрым KMeans»: он превращает классический оператор машинного обучения в базовый примитив, который можно комбинировать, заранее оценивать по стоимости и планировать автоматически.

9. На что обратить внимание

FlashLib пока относительно новая библиотека, поэтому в учебнике рекомендуется использовать её осторожно.

1. Сначала проверьте численное совпадение. Передайте один и тот же набор входных данных в FlashLib и в базовую реализацию и сравните метки, расстояния, индексы, объяснённую дисперсию, embedding или результаты классификации. Для кластеризации, t-SNE и UMAP, где результат не единственен, сравнивайте целевые метрики и итоговый эффект на последующих задачах, а не поэлементное совпадение.
2. Различайте точные и приблизительные методы. flash_knn, IVFFlat, IVFPQ и CAGRA имеют разную семантику. Полноту ANN-индекса нужно оценивать явно.
3. Фиксируйте оборудование и версии перед сравнением скорости. Официальные бенчмарки в основном проводились на H200, CUDA 13.0, PyTorch 2.11, Triton 3.6 и cuML 25.10. При переходе на A100, 4090, L40S или другие версии CUDA и PyTorch выводы о скорости нужно проверять заново.
4. Учитывайте прогрев JIT. При первом вызове путей Triton и CuteDSL возникают накладные расходы на компиляцию. В бенчмарках первый вызов нужно исключать или отдельно указывать холодный старт.
5. Не заменяйте этой библиотекой все остальные. Для нескольких тысяч примеров CPU-версия sklearn или PyTorch может быть достаточно быстрой. FlashLib лучше подходит для больших тензоров, которые постоянно находятся в памяти GPU.

10. Рекомендуемый smoke test

Ниже минимальная проверка, которую удобно записать в журнал экспериментов.

python - <<'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

Ожидаемый результат:

• cuda available равно True;
• форма labels — (100000,);
• форма centroids — (256, 64);
• выходные тензоры находятся на устройстве CUDA;
• n_iter не превышает max_iters.

Если проверка не проходит, в первую очередь проверьте:

• установлен ли flashlib в текущем окружении Python;
• видит ли PyTorch CUDA;
• действительно ли входной тензор находится на GPU;
• совместимы ли версии CUDA, PyTorch и Triton с вашей платформой;
• не запускаете ли вы примитив GPU на машине без NVIDIA GPU.

11. Место в учебнике

FlashLib стоит рассматривать вместе с flash-attn: быстрая установка, Быстрая загрузка с Hugging Face без лишних веток файлов, Библиотека Xget и Быстрая загрузка с ModelScope. Это не основная линия алгоритмов воплощённого ИИ, но она заметно влияет на эффективность экспериментов.

При дальнейшем развитии раздела можно добавить три небольших практических эксперимента:

• кластеризация embedding эпизодов робота с помощью flash_kmeans;
• поиск похожих демонстраций с помощью IVFFlat или CAGRA;
• визуализация скрытых состояний VLA или embedding траекторий с помощью flash_pca и flash_tsne.

Ссылки

• GitHub FlashLib: https://github.com/FlashML-org/flashlib
• Сайт проекта FlashLib и технический блог: https://flashml-org.github.io/
• Страница установки в PyPI: https://pypi.org/project/flashlib/