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

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

Вычислительные библиотеки · ruBLAS · Тензоры и платформы · 中文

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

ruDNN обеспечивает операции нейронной сети. Пакет Cargo — ruDNN, а крейт Rust — rudnn.

Внимание и свертка имеют соответствующие функции автонастройки. См. Cargo.toml и экспорт модуля.

2. Механизм внимания и свёртка

Механизм внимания

rudnn::attention::tensor::attention принимает запрос, ключ, значение, необязательную маску, необязательные attn_bias, AttentionModuleOptions и AttentionStrategy. Он возвращает тензор устройства или AttentionSetupError.

Стратегии включают FlashBlackboxAccelerated, FlashUnit, Fallback и Autotune при включённом соответствующем feature. Без автонастройки по умолчанию используется Fallback, с ней — Autotune. Fallback использует несколько ядер на том же устройстве, а не бэкенд CPU.

Сопоставьте макет, маску и точность с выбранной стратегией. См. интерфейс внимания.

Свертка

rudnn::convolution::tensor::conv_forward принимает входные данные, вес, необязательные bias, ConvOptions и ConvStrategy. Он возвращает тензор устройства или ConvSetupError. Он преобразует входные данные в макет последнего канала для выполнения и преобразует выходные данные обратно. conv_forward_nhwc напрямую использует макет последнего канала.

Стратегии включают Direct, ImplicitGemm и дополнительный Autotune. Точка входа использует Direct для трехмерной свертки F32. Групповая свертка также использует Direct, если выбран ImplicitGemm. Таким образом, параметр стратегии не всегда является строгим требованием сохранить один алгоритм.

Один и тот же модуль предоставляет conv_data_backward и conv_weight_backward. Проверьте форму, варианты и требования к исполнению для каждого направления. См. интерфейс свертки.

3. Рабочий процесс MoE

rudnn::moe предоставляет следующую последовательность локальных вычислений:

Входные logits → softmax/top-k → компактное распределение → эксперты SwiGLU → взвешенное объединение.

Эти интерфейсы используют RudaTensor<R>. Точки входа вычислений возвращают Result с MoeError для недопустимых фигур, dtypes или других аргументов.

4. Контракт маршрутизации

Logits имеют форму [T, E] и используют неквантованные F32, F16 или BF16. RoutingOptions содержит top_k: usize и renormalize: bool, причём 1 ≤ top_k ≤ E.

Softmax вычисляет всех экспертов в FP32 перед выбором топ-k. Связи благоприятствуют более низким идентификаторам экспертов. При включенной перенормировке выбранные веса перенормируются, а затем приводятся к логитам dtype. Строки, содержащие NaN, положительную бесконечность или только отрицательную бесконечность, сохраняют веса NaN, а не используют равномерное распределение.

Выбранные экспертные индексы и веса имеют форму [T, top_k]. Индексы используют U32.

5. Веса экспертов и распределение

Входные токены имеют форму [T, H]; gate и up — [E, I, H]; down — [E, H, I]. Веса должны иметь общий dtype с плавающей точкой и находиться на одном устройстве. Число экспертов и размеры должны соответствовать маршрутизации и входу.

Распределение не отбрасывает токены ради соблюдения ограничения вместимости. Порядок атомарного назначения внутри эксперта не фиксирован; сохранённое отображение восстанавливает порядок токенов. Точка входа прямого прохода принимает заранее вычисленные logits, а не выполняет проекцию gate модели или загрузку файлов весов.

Источник: маршрутизация, отправка и объединение, эксперты и тесты.

6. Вызов MoE

Включите tensor-moe в вашей зависимости ruDNN. Подготовьте тензоры устройств с помощью приведенных выше макетов, затем создайте экспертные веса и запустите прямую операцию:

use rudnn::moe::{RoutingOptions, SwiGluExperts};

let experts = SwiGluExperts::new(gate, up, down)?;
let options = RoutingOptions {
    top_k: 2,
    renormalize: true,
};
let output = experts.forward(input, logits, options)?;

При этом для каждого токена выбирается два эксперта, поэтому E должно быть не менее 2. input имеет форму [T, H], logits имеет форму [T, E], а output имеет форму [T, H]. Повторно используйте experts в последующих входных пакетах без повторного создания объекта веса.

Чтобы проверить маршрутизацию или вставить собственную обработку между этапами, вызовите их отдельно:

use rudnn::moe::route;

let routing = route(logits, options)?;
let selected_experts = routing.expert_indices();
let selected_weights = routing.weights();
let dispatched = routing.dispatch(input)?;
let expert_output = experts.forward_dispatched(&dispatched)?;
let output = dispatched.combine(expert_output)?;

Это альтернативные формы; второй начинается со свежей партии input и logits. combine использует сопоставление диспетчеризации для восстановления порядка токенов и объединяет экспертные выходные данные с использованием весов маршрутизации.

Форма, dtype или несоответствие устройств возвращают MoeError. Чтобы прочитать результаты на хосте, используйте ruda_kernel::tensor::readback::into_data_sync(output), который ожидает результатов устройства и возвращает TensorData.

7. LayerNorm, RMSNorm и Softmax

Включите tensor-normalization и импортируйте эти функции из rudnn::normalization. Все они работают с конечной осью ввода и сохраняют форму:

Входные данные должны быть неквантованными с непустой конечной осью. Для длины конечной оси H gamma и beta должны иметь форму [H], не быть квантованными и иметь общее устройство ввода. epsilon должен быть конечным и положительным. Аффинные параметры остаются F32 даже для входов BF16/F16. В статистике и аффинной арифметике используется FP32, при этом выходные данные преобразуются во входные dtype.

use ruda_kernel::{dsl::Runtime, tensor::RudaTensor};
use rudnn::normalization::{NormalizationError, layer_norm, rms_norm, softmax_last_axis};

fn normalize<R: Runtime>(
    input: RudaTensor<R>,
    gamma: RudaTensor<R>,
    beta: Option<RudaTensor<R>>,
    epsilon: f32,
) -> Result<RudaTensor<R>, NormalizationError> {
    layer_norm(input, gamma, beta, epsilon)
}

fn normalize_rms<R: Runtime>(
    input: RudaTensor<R>,
    gamma: RudaTensor<R>,
    epsilon: f32,
) -> Result<RudaTensor<R>, NormalizationError> {
    rms_norm(input, gamma, epsilon)
}

fn probabilities<R: Runtime>(
    logits: RudaTensor<R>,
) -> Result<RudaTensor<R>, NormalizationError> {
    softmax_last_axis(logits)
}

8. Gated-delta: предварительное заполнение и рекуррентное вычисление

Включите tensor-gated-delta. Используйте chunk_gated_delta_rule(input, chunk_size) для предварительного заполнения фрагментированной последовательности и gated_delta_rule(input) для повторения каждого токена. Оба принимают GatedDeltaInput<R>:

Все тензоры должны быть неквантованными и находиться на одном устройстве. Поставка Q/K после нормализации для конкретной модели; точка входа не выполняет это за вас. Эта функция повторно использует предыдущий импорт Runtime и RudaTensor:

use rudnn::gated_delta::{
    GatedDeltaError, GatedDeltaInput, GatedDeltaOutput, chunk_gated_delta_rule,
};

fn delta_prefill<R: Runtime>(
    query: RudaTensor<R>,
    key: RudaTensor<R>,
    value: RudaTensor<R>,
    beta: RudaTensor<R>,
    log_decay: RudaTensor<R>,
    initial_state: RudaTensor<R>,
    query_scale: f32,
    chunk_size: usize,
) -> Result<GatedDeltaOutput<R>, GatedDeltaError> {
    chunk_gated_delta_rule(
        GatedDeltaInput {
            query, key, value, beta, log_decay, initial_state, query_scale,
        },
        chunk_size,
    )
}

Возвращаемый output имеет форму [B, H, T, V] и тот же dtype, что у query. final_state имеет тип F32 и форму [B, H, K, V]. Для следующего сегмента той же последовательности передайте этот final_state как initial_state. Для разных последовательностей храните отдельные состояния. Начальное состояние не перезаписывается на месте.

chunk_size должен быть положительным, а треугольная рабочая область размером 4 × (chunk_size² + chunk_size) байт не должна превышать лимит общей памяти устройства на рабочую группу. Точка входа дополняет последний блок; заполнение исключается из выхода. Недопустимые аргументы возвращают GatedDeltaError.

Для вызовов текста и изображений на уровне модели см. Руководство по выводу ruLLM.

9. Постраничное внимание и MLA

Включите tensor-paged-attention и используйте rudnn::paged_attention. HostPlan::new(page_size, pages, tables, lengths, sequence_ids, positions) проверяет метаданные планирования на хосте; positions содержит абсолютные позиции с отсчётом от нуля внутри каждой последовательности. DevicePlan::upload(host, &q) загружает метаданные на устройство и в очередь исполнения запроса. План можно использовать повторно только при неизменном расписании.

DevicePlan::attention(q, k, v, scale, causal) читает физические страницы кэша напрямую для упакованных prefill/decode переменной длины. Форма Q — [queries, Hq, D], K — [pages, page_size, Hkv, D], V — [pages, page_size, Hkv, Dv], результата — [queries, Hq, Dv]; Hq должно делиться на Hkv. Входы должны быть непрерывными неквантованными тензорами F32/F16/BF16 с одинаковыми dtype, устройством и очередью. Значения Q/K/V должны быть конечными, scale — конечным и положительным. D и Dv находятся в 1..=1024; устройство должно поддерживать plane из 32 или 64 линий и требуемую сетку запуска. Доступны прямой проход и обратный проход первого порядка; произвольные внешние маски и квантованные KV-кэши не поддерживаются.

DevicePlan::mla(q, qp, latent, kp, scale, causal) принимает запросы с поглощённой проекцией [queries, H, R], позиционные запросы [queries, H, P], общий латентный кэш [pages, page_size, 1, R] и позиционный кэш [pages, page_size, 1, P]. Результат — сжатый контекст [queries, H, R]; P находится в 1..=256. Позиционное кодирование применяется до вызова, проекции value/output — после. Используйте исходный масштаб QK модели, а не масштаб, вычисленный по сжатому рангу.

Эти методы используют путь без разбиения. Для разбиения истории создайте SplitWorkspace::new(&q, queries, heads, value_dim, splits) с числом splits в 2..=32 и вызовите attention_with_workspace или mla_with_workspace. Частичные статистики и их объединение используют FP32. Рабочая область ограничена 64 MiB и может использоваться повторно только при совпадающей форме в той же упорядоченной очереди исполнения.

DevicePlan::append(k, v, key_cache, value_cache) возвращает тензоры кэша, которые нужно сохранить для следующих вызовов. Разделяемые выделения памяти кэша копируются перед изменением; общие физические страницы префикса дополнительно требуют copy-on-write со стороны планировщика. Повторные записи в одну физическую позицию отклоняются.

10. Групповая sigmoid-маршрутизация MoE

При включённом tensor-moe функция route_sigmoid_grouped(logits, bias, options) возвращает RoutingPlan. Logits имеют форму [tokens, experts] и тип F32/F16/BF16; необязательный корректирующий bias — FP32 [experts] на том же устройстве и в той же очереди. Bias влияет только на выбор. Возвращаемые веса используют исходные sigmoid-оценки, необязательную повторную нормализацию и scale; при равенстве предпочтительны меньшие ID.

GroupRoutingOptions содержит top_k, groups, selected_groups, group_top_two, renormalize и scale. Число экспертов находится в 1..=1024, групп — в 1..=128, top-k — в 1..=64. Число экспертов должно делиться на число групп, число выбранных групп должно быть допустимым, а top-k не должно превышать суммарное число экспертов в них. group_top_two=true суммирует две наибольшие скорректированные оценки в каждой группе и требует не менее двух экспертов на группу; иначе используется максимум. Scale должен быть конечным и положительным.

SwiGluExperts::forward_sigmoid_grouped(input, logits, bias, options, strategy) выполняет маршрутизацию, распределение, вычисления экспертов и объединение. forward_dispatched_with_strategy выбирает GroupedStrategy::Scalar, Auto или TensorCore; существующие forward и forward_dispatched сохраняют Scalar. Для Tensor Core требуется соответствующая поддержка F16/BF16 устройством. Auto переходит на запасной путь только при неподдерживаемой конфигурации, но не при ошибках компиляции или исполнения. Проекции модели, общие эксперты и остаточные ветви остаются ответственностью вызывающего кода.

11. Обратный проход paged Attention и упорядоченные градиенты истории

DevicePlan::attention_backward(q, k, v, grad_out, scale, causal) возвращает AttentionBackward { dq, dk, dv }, а mla_backward(q, qp, latent, kp, grad_out, scale, causal) — MlaBackward { dq, dqp, dlatent, dkp }. dlatent включает вклады key и value. Вероятности пересчитываются при обратном проходе без сохранения полной матрицы оценок.

Небезопасные API attention_backward_selected_into и mla_backward_selected_into принимают независимо необязательные буферы градиентов. None исключает соответствующий выход и временную память истории этой ветви; другие производные всё ещё могут требовать исходное значение входа. Выходы не должны перекрываться со входами или друг с другом. Доступ выполняется в упорядоченной очереди плана. Градиенты истории этого пути используют FP32-атомики и не гарантируют побитовую детерминированность.

Для редукции истории без атомиков создайте OrderedBackwardWorkspace::new(&plan, &q)? и передайте изменяемую ссылку в attention_backward_ordered_into или mla_backward_ordered_into. Запрет перекрытия распространяется и на память workspace. Повторное использование требует совместимых неизменяемых метаданных расписания, числа запросов/голов, устройства и очереди. Статистика и обратный индекс вместе ограничены 64 MiB. Фиксированный порядок суммирования не гарантирует побитового совпадения между устройствами или с атомарным путём.

RUDA_PAGED_ORDERED_CACHE_ROWS=1 и RUDA_PAGED_ORDERED_COMPACT_HISTORY=1 включают последние две настройки при создании workspace. Отсутствие значения или 0 отключает их, остальные значения ошибочны. Изменение окружения не меняет существующий workspace.

Активность определяется достижимостью страниц в пределах эффективной длины KV последовательностей с запросами, а не ёмкостью таблицы или числом ненулевых градиентов. Первое включение загружает индекс размером physical_pages * 4 байт в рамках того же бюджета. Последующие переключения используют его повторно; отключение не освобождает индекс. history_compaction_pages() возвращает Option<(active_pages, inactive_pages)>, bytes() учитывает сохранённую память. Ошибка бюджета сохраняет прежний режим. Атомарный режим PyTorch по умолчанию не меняется; ускорение не гарантируется.

12. Обучение MoE первого порядка

selected_router_weights(&logits, &indices, options) возвращает FP32-веса [T, top_k]; selected_router_backward(&logits, &indices, &grad_weights, options) — градиенты [T, E] в dtype logits. Logits — непрерывные F32/F16/BF16, индексы — непрерывные U32/I32/I64; требуется 1 <= top_k <= min(E, 64), одинаковые устройство и очередь. RouterWeightOptions задаёт RouterScoring::Softmax или Sigmoid, необязательную перенормировку выбранных весов и последующее применение конечного положительного scale. grad_weights имеет тип FP32.

Повторяющиеся индексы имеют семантику gather. Недопустимые индексы дают NaN во всей строке без выхода за границы и синхронизации с хостом. Дифференцируются непрерывные веса при фиксированном выборе, но не решения top-k/группировки или корректирующий bias. RoutingPlan::into_training(logits, options) сохраняет выбор и пересчитывает веса. Используйте RouterTrainingPlan::routing() для распределения и backward(&grad_weights) для градиентов logits. Сохранённые входы нельзя изменять до обратного прохода.

Сохраните output и cache из SwiGluExperts::forward_dispatched_training(&dispatched, strategy). Цепочка обратного прохода:

1. dispatched.combine_backward(&expert_output, grad_output) возвращает dexpert и FP32-dweights.
2. cache.backward(dexpert) возвращает dinput распределённых строк и FP32-dgate, dup, ddown.
3. dispatched.dispatch_backward(dinput) суммирует выбранные строки к исходному token без повторного умножения на веса маршрутизации.
4. Передайте dweights в обратный проход роутера для получения градиентов logits.

combine_backward по умолчанию использует CombineGradientStrategy::Serial; combine_backward_with_strategy допускает Plane. Экспертный backward по умолчанию скалярный независимо от прямого прохода; backward_with_strategy принимает GroupedStrategy::Scalar, Auto или TensorCore. Tensor Core требует поддерживаемого F16/BF16-оборудования. Auto переключается только при отсутствии возможностей, а не при ошибках компиляции или выполнения.