ruBLAS 用户指南

计算库 · Runtime API · English | 日本語 | Deutsch | Русский

1. 概述与功能入口

ruBLAS 组织线性代数计算及其设备 Tensor 入口。

Cargo package 为 rublas。通用路径选择所需 feature，并关闭不需要的默认功能。定义见 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 Core 路径，无需调用者物化补齐后的矩阵。lhs 和 rhs 必须是相同的非量化 BF16 或 F16，且矩阵维度匹配；设备还需支持该路径使用的矩阵指令。下面的函数返回 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，返回设备张量。非末维计算涉及维度置换与连续化。定义见叉积实现。

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>。它接收 input、weights、row_experts 三个 RudaTensor<R>，返回 Result<RudaTensor<R>, GroupedMatmulError>。

第 m 行选择编号为 row_experts[m] 的权重矩阵，计算输入行与该矩阵各行的内积。权重的最后两维按转置方式参与计算，不要求调用者先物化转置。

4. 分组矩阵乘的数值与边界

• K、N、E 必须大于零；M 可以为零。
• 超出专家范围的行编号表示 padding，输出零行。
• 当前内核采用 FP32 累计，最终转换为输入 dtype。
• 输入和权重需要的连续化由入口处理，可能引入数据复制。
• 相关元素数量不能超过 U32 索引范围。
• 参数不满足约定时返回 GroupedMatmulError；异步执行错误还需在回读或同步阶段处理。

源码：分组矩阵乘入口、内核。

5. 与其他库协作

ruDNN MoE 使用分组矩阵乘计算专家投影。量化权重路径位于独立的 tensor_int4 模块；不能将普通浮点分组矩阵乘当作 INT4 专家计算。

grouped_matmul_nt 使用标量累计。矩阵乘策略应按实际 dtype、shape 和后端分别测量。

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 接口：offsets 必须是从 0 开始、以 M 结束的非递减排他前缀，每个区间 [offsets[e], offsets[e + 1]) 内的行都属于专家 e。row_experts 必须描述同一份不可变分发结果。形状检查不会验证这些设备端数值。

GroupedStrategy::Scalar 使用现有标量内核。TensorCore 显式要求匹配的 F16/BF16 输入、受支持的 16×16×16 协作矩阵运算、32 lane 的 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]。空专家段的权重梯度为零。行编号和 offsets 沿用分段前向的 U32 布局及不可变前缀约束；该接口仍为 unsafe，不回读设备元数据进行校验。

默认入口选择 GroupedStrategy::Scalar。grouped_matmul_nt_backward_segmented_with_strategy(..., strategy) 可显式选择 Scalar、Auto 或 TensorCore，独立于前向策略。协作路径采用 16×16×16 tile、FP32 累加及 FP32 权重梯度，要求受支持的 F16/BF16 硬件和启动规模。TensorCore 在不支持时返回错误；Auto 仅针对能力不支持回退，不吞掉编译或执行错误。输入可能转换为连续布局。不同归约顺序不保证逐位相同。

需要安全的专家级训练入口时，使用 rudnn::moe::SwiGluExperts::forward_dispatched_training，再调用 ExpertTrainingCache::backward_with_strategy。