ruDNN ユーザーガイド

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

計算ライブラリ · ruBLAS · テンソルとフレームワーク · 中文

1. 概要と特徴

ruDNN はニューラル ネットワーク操作を提供します。 Cargo パッケージは ruDNN、Rust クレートは rudnn です。

アテンションとコンボリューションには、それぞれ対応する自動調整機能があります。 Cargo.toml および モジュール エクスポート を参照してください。

2. Attention と畳み込み

Attention

rudnn::attention::tensor::attention は、クエリ、キー、値、オプションのマスク、オプションの attn_bias、AttentionModuleOptions、および AttentionStrategy を受け取ります。デバイス テンソルまたは AttentionSetupError を返します。

戦略には FlashBlackboxAccelerated、FlashUnit、Fallback、および対応 feature が有効な場合の Autotune があります。自動チューニングが無効なら既定は Fallback、有効なら Autotune です。Fallback は同じデバイス上で複数のカーネルを使い、CPU バックエンドは使いません。

レイアウト、マスク、精度を選択した戦略に合わせます。 アテンションインターフェイスを参照してください。

畳み込み

rudnn::convolution::tensor::conv_forward は、入力、重み、オプションの bias、ConvOptions、および ConvStrategy を受け取ります。デバイス テンソルまたは ConvSetupError を返します。実行のために入力をチャネル最後のレイアウトに変換し、出力を逆変換します。 conv_forward_nhwc は、チャネル最後のレイアウトを直接使用します。

戦略には、直接、ImplicitGemm、およびオプションの Autotune が含まれます。エントリ ポイントは、3 次元 F32 畳み込みに Direct を使用します。 ImplicitGemm が選択されている場合、グループ化コンボリューションでもダイレクトが使用されます。したがって、戦略パラメーターは、常に 1 つのアルゴリズムを保持するという厳密な要求ではありません。

同じモジュールで conv_data_backward と conv_weight_backward が提供されます。各方向の形状、オプション、施工条件をご確認ください。 コンボリューションインターフェイスを参照してください。

3. MoE ワークフロー

rudnn::moe は、次のローカル計算シーケンスを提供します。

入力 logits → softmax/top-k → コンパクトな振り分け → SwiGLU エキスパート → 重み付き結合。

これらのインターフェイスは RudaTensor<R> を使用します。計算エントリ ポイントは、無効なシェイプ、dtype、またはその他の引数に対して Result と MoeError を返します。

4. ルーティング契約

Logits の形状は [T, E] で、非量子化の F32、F16、BF16 を使います。RoutingOptions は top_k: usize と renormalize: bool を含み、1 ≤ top_k ≤ E が必要です。

Softmax は、top-k を選択する前に、FP32 のすべてのエキスパートを計算します。同点の場合、エキスパート ID が低いほど有利になります。再正規化を有効にすると、選択した重みが再正規化され、ロジット dtype にキャストされます。 NaN、正の無限大、または負の無限大のみを含む行は、一様分布を使用するのではなく、NaN の重みを保持します。

選択されたエキスパート インデックスとウェイトは両方とも形状 [T、top_k] です。インデックスはU32を使用します。

5. エキスパートの重みと振り分け

入力トークンの形状は [T, H]、gate と up は [E, I, H]、down は [E, H, I] です。重みは同じ浮動小数点 dtype とデバイスを使う必要があります。エキスパート数と各次元はルーティングおよび入力に一致していなければなりません。

振り分け処理は容量制限を守るためにトークンを破棄しません。エキスパート内のアトミックな割り当て順序は固定されず、保存した対応表でトークン順序を復元します。順伝播のエントリポイントは計算済み logits を受け取り、モデルのゲート射影や重みファイルの読み込みは行いません。

ソース: ルーティング、ディスパッチと結合、専門家、および テスト。

6. MoE の呼び出し

ruDNN 依存関係で tensor-moe を有効にします。上記のレイアウトでデバイス テンソルを準備し、エキスパート ウェイトを作成して、順方向操作を実行します。

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)?;

これにより、トークンごとに 2 人のエキスパートが選択されるため、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)?;

これらは代替形式です。 2 番目は、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 は有限かつ正でなければなりません。 BF16/F16 入力の場合でも、アフィン パラメーターは F32 のままです。統計とアフィン演算は 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) は物理キャッシュページを直接読み、パックされた可変長のプリフィルとデコードを処理します。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 です。デバイスは 32 または 64 レーンの plane と必要な起動グリッドをサポートする必要があります。順伝播と一階の逆伝播に対応します。任意の外部マスクや量子化 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 射影は呼び出し後に適用します。scale は圧縮ランクから導出せず、元のモデルの QK スケールを使用します。

上記メソッドは非分割経路を使用します。履歴を分割する場合は、splits を 2..=32 として SplitWorkspace::new(&q, queries, heads, value_dim, splits) を作成し、attention_with_workspace または mla_with_workspace を呼び出します。部分統計とマージは FP32 で処理します。ワークスペースの上限は 64 MiB で、形状が一致し、同じ順序付き実行キューを使用する場合にのみ再利用できます。

DevicePlan::append(k, v, key_cache, value_cache) は、以後の呼び出しで保持するキャッシュテンソルを返します。共有されたキャッシュ割り当ては変更前にコピーされます。共有プレフィックスの物理ページには、さらにスケジューラー側のコピーオンライトが必要です。同じ物理位置への重複書き込みは拒否されます。

10. グループ制限付き sigmoid MoE ルーティング

tensor-moe を有効にすると、route_sigmoid_grouped(logits, bias, options) が RoutingPlan を返します。logits は F32/F16/BF16 の [tokens, experts]、オプションの補正 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. ページ化 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 両方の寄与を含みます。完全なスコア行列を保存せず、逆伝播時に確率を再計算します。

unsafe API の attention_backward_selected_into と mla_backward_selected_into は、各勾配バッファを個別に省略できます。None の出力とその分岐専用の履歴領域は確保しませんが、他の導関数には元の入力値が必要な場合があります。出力は入力や他の出力と重複してはならず、アクセスは計画と同じ順序付きキューで行います。この経路の履歴勾配は FP32 アトミック加算を使用し、ビット単位の決定性は保証しません。

アトミック演算を使わない履歴帰約には OrderedBackwardWorkspace::new(&plan, &q)? を作り、可変参照で attention_backward_ordered_into または mla_backward_ordered_into に渡します。バッファ非重複条件はワークスペースにも適用されます。不変のスケジュール、クエリ数・ヘッド数、デバイス、キューが互換の場合のみ再利用します。統計量と逆引き情報の合計上限は 64 MiB です。固定加算順序でも、異なるデバイス間やアトミック経路とのビット一致は保証しません。

RUDA_PAGED_ORDERED_CACHE_ROWS=1 と RUDA_PAGED_ORDERED_COMPACT_HISTORY=1 は構築時に後者二つを有効にします。未設定または 0 は無効、それ以外はエラーです。環境変更は既存ワークスペースには反映されません。

ページ分類は、クエリのあるシーケンスの有効 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) は logits と同じ dtype の [T, E] 勾配を返します。logits は連続 F32/F16/BF16、indices は連続 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 勾配を得ます。保存入力を逆伝播前に変更しないでください。

SwiGluExperts::forward_dispatched_training(&dispatched, strategy) の output と cache を保持し、次の順に逆伝播します。

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 は順伝播と独立に scalar が既定で、backward_with_strategy は GroupedStrategy::Scalar、Auto、TensorCore を受け付けます。Tensor Core には対応 F16/BF16 ハードウェアが必要です。Auto は能力不足のみでフォールバックし、コンパイル・実行エラーを隠しません。