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、および tensor-matmul-autotune が有効な場合の Autotune があります。この feature が無効なら既定は Ruda、有効なら Autotune です。量子化入力では Naive が最初の経路の失敗後に逆量子化して計算する場合がありますが、これはネイティブ INT4 計算ではありません。

CmmaResidueFirst は、実体化されたパディング行列を必要とせずに、部分的な K32 タイルを最初に処理する Tensor コア パスを明示的に選択します。 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] を生成します。無効なセットアップまたは欠落しているデバイス機能がある場合は、Naive に切り替えるのではなく、MatmulSetupError を返します。

行列乗算エントリ ポイント を参照してください。

ベクトルの外積

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] でインデックス付けされた重み行列を選択し、入力行とその行列の各行の間のドット積を計算します。最後の 2 つの重み次元は、具体化された転置を必要とせずに、転置された形式で参加します。

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] です。出力は入力と同じ dtype の [M, N] です。オペランドは非量子化で、同じデバイスと実行キュー上にある必要があります。

これは unsafe Rust API です。offsets は 0 から始まり M で終わる単調非減少の排他的プレフィックスであり、[offsets[e], offsets[e + 1]) の全行がエキスパート e に属する必要があります。row_experts は同一の不変なディスパッチを表す必要があります。形状チェックはこれらのデバイス上の値を検証しません。

GroupedStrategy::Scalar は既存のスカラーカーネルを使用します。TensorCore は一致する F16/BF16 入力、対応する 16×16×16 協調行列演算、32 レーンの plane、十分な共有メモリ、有効な起動グリッドを明示的に要求し、非対応の設定では GroupedMatmulError を返します。カーネルは FP32 累積を使用し、端の不完全なタイルも処理します。Auto は対応していればこの経路、それ以外はスカラーカーネルを選択します。コンパイル、起動、数値計算の失敗はフォールバック条件ではありません。

offsets を内部で構築する安全な MoE エントリポイントには、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] で、F32/F16/BF16 の dtype、デバイス、キューが一致する必要があります。dinput は入力 dtype、dweights は FP32 の [E, N, K] です。空のエキスパート区間の重み勾配はゼロです。行 ID と offsets は順伝播と同じ U32 レイアウト・不変の prefix 条件に従います。デバイスメタデータをホストへ読み戻さない unsafe API です。

既定の入口は GroupedStrategy::Scalar です。grouped_matmul_nt_backward_segmented_with_strategy(..., strategy) は順伝播と独立に Scalar、Auto、TensorCore を選択します。協調経路は 16×16×16 タイル、FP32 累積・重み勾配を使い、対応 F16/BF16 ハードウェアと起動制約を満たす必要があります。TensorCore は非対応時にエラー、Auto は能力不足のみでフォールバックします。コンパイル・実行エラーは隠しません。入力の連続化でコピーが必要な場合があります。異なる帰約順序のビット一致は保証しません。

安全なエキスパート学習入口には rudnn::moe::SwiGluExperts::forward_dispatched_training と ExpertTrainingCache::backward_with_strategy を使用します。