ruBLAS Руководство пользователя

English | 简体中文 | 日本語 | Deutsch | Русский

Вычислительные библиотеки · Среда выполнения API · 中文

1. Обзор и возможности

ruBLAS предоставляет операции линейной алгебры и интерфейсы тензоров устройств.

Пакет Cargo — rublas. Выберите необходимые функции и отключите ненужные значения по умолчанию для путей общего назначения. См. Cargo.toml.

2. Умножение матриц, векторы и INT4

Умножение тензорной матрицы

rublas::tensor_matmul::matmul принимает lhs, rhs, необязательный выход out, MatmulStrategy и выходной DType. Возвращает тензор устройства или MatmulSetupError. Передайте существующий выход через out либо не задавайте его для выделения нового.

Стратегии включают Ruda, CmmaResidueFirst, Naive и Autotune при включённом tensor-matmul-autotune. Без этого feature по умолчанию используется Ruda, с ним — Autotune. Для квантованных входов Naive может после сбоя исходного пути деквантовать данные и выполнить вычисление; это не нативное вычисление INT4.

CmmaResidueFirst явно выбирает путь Tensor Core, который сначала обрабатывает частичный фрагмент K32, не требуя материализованной дополненной матрицы. lhs и rhs должны иметь совпадающие неквантованные BF16 или F16 dtype и совместимые размеры матрицы. Устройство должно поддерживать матричные инструкции пути. Эта функция возвращает вывод F32:

use ruda_core::tensor::DType;
use ruda_kernel::{dsl::Runtime, tensor::RudaTensor};
use rublas::{
    kernel_ir::definition::MatmulSetupError,
    tensor_matmul::{MatmulStrategy, matmul},
};

fn residue_matmul<R: Runtime>(
    lhs: RudaTensor<R>,
    rhs: RudaTensor<R>,
) -> Result<RudaTensor<R>, MatmulSetupError> {
    matmul(lhs, rhs, None, MatmulStrategy::CmmaResidueFirst, DType::F32)
}

Входы [M, K] и [K, N] создают [M, N]. Неверная настройка или отсутствие возможностей устройства возвращают MatmulSetupError вместо переключения на Naive.

См. точку входа умножения матрицы.

Векторное произведение

rublas::tensor_vector::cross(lhs, rhs, dim) требует длину 3 выбранного измерения и возвращает тензор устройства. Вычисление вдоль измерения, отличного от последнего, включает перестановку осей и преобразование в непрерывное размещение. См. cross.rs.

AWQ INT4

rublas::tensor_int4::AwqGemm::new использует qweight, qzeros, scales, дополнительные bias и group_size для построения упакованных весов. forward принимает входные данные F16 и возвращает выходные данные F16 или Int4Error.

Для входного измерения K, выходного измерения N и размера группы G:

K, N и G должны быть ненулевыми; K должен делиться на G и N на 8. Конечная размерность входных данных — K, заменяемая на N на выходе. Все операнды должны использовать общее устройство. Порядок упакованных битов должен соответствовать ядру AWQ; произвольный файл INT4 не может использоваться непосредственно как qweight.

См. Интерфейс INT4 для объекта и его проверок.

3. Групповое умножение матриц

rublas::tensor_grouped::grouped_matmul_nt<R: Runtime> принимает входные данные, веса и row_experts как значения RudaTensor<R>. Он возвращает Result<RudaTensor<R>, GroupedMatmulError>.

Строка m выбирает матрицу весов, индексированную row_experts[m], и вычисляет скалярное произведение между входной строкой и каждой строкой этой матрицы. Последние два весовых измерения участвуют в транспонированной форме, не требуя материализованного транспонирования.

4. Семантика группового умножения

• K, N и E должны быть положительными; М может быть нулевым.
• Индексы экспертов вне допустимого диапазона обозначают заполнение; соответствующие выходные строки заполняются нулями.
• Ядро накапливает данные в FP32, а затем преобразует их во входные данные dtype.
• Точка входа при необходимости делает ввод и веса непрерывными, что может привести к копированию данных.
• Количество соответствующих элементов должно соответствовать индексу U32.
• Неверные аргументы возвращают GroupedMatmulError. Обработка ошибок асинхронного выполнения во время обратного чтения или синхронизации.

Источник: сгруппированный интерфейс и ядро.

5. Интеграция с другими библиотеками

ruDNN MoE использует групповое умножение для экспертных прогнозов. Квантованные веса используют отдельный модуль tensor_int4; обычное групповое умножение с плавающей запятой не является экспертным вычислением INT4.

grouped_matmul_nt использует скалярное накопление. Измеряйте стратегии умножения матриц для своего dtype, формы и бэкенда.

6. Сегментированное умножение матриц экспертов

При включённом tensor-grouped функция rublas::tensor_grouped::grouped_matmul_nt_segmented(input, weights, row_experts, offsets, strategy) добавляет к существующему групповому интерфейсу сегменты экспертов на устройстве. Input имеет форму [M, K], weights — [E, N, K], row_experts — U32 [M], offsets — U32 [E + 1]; результат — [M, N] с dtype входа. Операнды должны быть неквантованными и использовать одно устройство и одну очередь исполнения.

Это unsafe Rust API: offsets должен быть неубывающим эксклюзивным префиксом, начинающимся с 0 и заканчивающимся M, а каждая строка в [offsets[e], offsets[e + 1]) должна принадлежать эксперту e. row_experts должен описывать то же неизменяемое распределение. Проверки форм не проверяют эти значения на устройстве.

GroupedStrategy::Scalar использует существующее скалярное ядро. TensorCore явно требует совпадающих входов F16/BF16, поддержки кооперативных матричных операций 16×16×16, plane из 32 линий, достаточной общей памяти и допустимой сетки запуска; неподдерживаемая конфигурация возвращает GroupedMatmulError. Ядро накапливает в FP32 и обрабатывает неполные тайлы. Auto выбирает этот путь при наличии поддержки, иначе — скалярное ядро. Ошибки компиляции, запуска или численных вычислений не являются условиями перехода на запасной путь.

Для безопасного входа MoE с внутренним построением offsets используйте rudnn::moe::SwiGluExperts::forward_dispatched_with_strategy. Существующий forward_dispatched по умолчанию остаётся скалярным.

7. Сегментированный обратный проход

grouped_matmul_nt_backward_segmented(input, weights, grad_output, row_experts, offsets) возвращает GroupedBackward { dinput, dweights }. Формы: input [M, K], weights [E, N, K], grad_output [M, N]; dtype F32/F16/BF16, устройство и очередь должны совпадать. dinput сохраняет dtype входа; dweights имеет FP32 и форму [E, N, K]. Пустые сегменты экспертов получают нулевые градиенты весов. Идентификаторы строк и offsets соблюдают те же U32-форматы и неизменяемые префиксные инварианты, что и прямой проход. Это unsafe API без чтения метаданных устройства на хост для проверки.

Стандартная обёртка выбирает GroupedStrategy::Scalar. grouped_matmul_nt_backward_segmented_with_strategy(..., strategy) независимо от прямого прохода выбирает Scalar, Auto или TensorCore. Кооперативный путь использует тайлы 16×16×16, FP32-накопление и FP32-градиенты весов; нужны поддерживаемое F16/BF16-оборудование и допустимая сетка запуска. TensorCore сообщает об отсутствии поддержки ошибкой; Auto переключается только при отсутствии возможностей, но не при ошибках компиляции или выполнения. Возможны копии для непрерывного размещения входов. Разный порядок редукции не гарантирует побитового совпадения.

Безопасный путь обучения экспертов: rudnn::moe::SwiGluExperts::forward_dispatched_training, затем ExpertTrainingCache::backward_with_strategy.