flashinfer.gemm

此模块提供一组 GEMM 操作。

BF16 GEMM

mm_bf16(a, b[, bias, pdl, out, out_dtype, ...])

MM BF16

bmm_bf16(A, B[, out, out_dtype, backend])

BMM BF16

FP4 GEMM

mm_fp4(a, b, a_descale, b_descale[, alpha, ...])

MM FP4

FP8 GEMM

bmm_fp8(A, B, A_scale, B_scale, dtype[, ...])

BMM FP8

gemm_fp8_nt_groupwise(a, b, a_scale, b_scale)

使用组间缩放执行 FP8 数据类型的矩阵乘法。

group_gemm_fp8_nt_groupwise(a, b, a_scale, ...)

执行使用组间缩放的 FP8 数据类型的组 GEMM。

group_deepgemm_fp8_nt_groupwise(a, b, ...[, ...])

执行使用 DeepGEMM 后端的 FP8 数据类型的分组矩阵乘法。

batch_deepgemm_fp8_nt_groupwise(a, b, ...[, ...])

执行使用 DeepGEMM 后端的 FP8 数据类型的批量矩阵乘法。

混合精度 GEMM (fp8 x fp4)

group_gemm_mxfp4_nt_groupwise(a, b, a_scale, ...)

执行使用组间缩放的 MXFP4 数据类型的组 GEMM。

分组 GEMM (Ampere/Hopper)

class flashinfer.gemm.SegmentGEMMWrapper(float_workspace_buffer: Tensor, backend: str = 'auto')

分段 GEMM 内核的包装器。

示例

>>> import torch
>>> from flashinfer import SegmentGEMMWrapper
>>> # create a 1MB workspace buffer
>>> workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device="cuda")
>>> segment_gemm = SegmentGEMMWrapper(workspace_buffer)
>>> seq_lens = torch.tensor([1, 2, 3, 4], dtype=torch.int64, device="cuda")
>>> # create packed input tensor (10 = 1 + 2 + 3 + 4)
>>> x = torch.randn(10, 128, device="cuda", dtype=torch.float16)
>>> # create weight tensor with 4 weights, each with 128 input and 256 output channels, column major
>>> weights = torch.randn(4, 256, 128, device="cuda", dtype=torch.float16)
>>> # compute the segment GEMM
>>> y = segment_gemm.run(x, weights, 4, True, seg_lens=seq_lens)
>>> y.shape
torch.Size([10, 256])
>>> y_ref_0 = torch.matmul(x[:1], weights[0].t())
>>> torch.allclose(y[:1], y_ref_0)
True
>>> y_ref_1 = torch.matmul(x[1:3], weights[1].t())
>>> torch.allclose(y[1:3], y_ref_1)
True
>>> y_ref_2 = torch.matmul(x[3:6], weights[2].t())
>>> torch.allclose(y[3:6], y_ref_2)
True
>>> y_ref_3 = torch.matmul(x[6:], weights[3].t())
>>> torch.allclose(y[6:], y_ref_3)
True
>>>
>>> # another example with weight indices
>>> weight_indices = torch.tensor([0, 1, 0, 1], dtype=torch.int64, device="cuda")
>>> y = segment_gemm.run(x, weights, 4, True, seg_lens=seq_lens, weight_indices=weight_indices)
>>> y.shape
torch.Size([10, 256])
>>> y_ref_0 = torch.matmul(x[:1], weights[0].t())
>>> torch.allclose(y[:1], y_ref_0)
True
>>> y_ref_1 = torch.matmul(x[1:3], weights[1].t())
>>> torch.allclose(y[1:3], y_ref_1)
True
>>> y_ref_2 = torch.matmul(x[3:6], weights[0].t())
>>> torch.allclose(y[3:6], y_ref_2)
True
>>> y_ref_3 = torch.matmul(x[6:], weights[1].t())
>>> torch.allclose(y[6:], y_ref_3)
True
__init__(float_workspace_buffer: Tensor, backend: str = 'auto') → None

初始化包装器。

参数:

float_workspace_buffer (torch.Tensor) – 内核的工作区缓冲区，我们用它来存储 cutlass 分段 GEMM 内核中的中间结果。建议大小为 128MB。

reset_workspace_buffer(float_workspace_buffer: Tensor, int_workspace_buffer: Tensor) → None

重置工作区缓冲区。

参数:

  • float_workspace_buffer (torch.Tensor) – 内核的新浮点工作区缓冲区。

  • int_workspace_buffer (torch.Tensor) – 内核的新整数工作区缓冲区。

run(x: Tensor, weights: Tensor, batch_size: int, weight_column_major: bool, out: Tensor | None = None, seg_lens: Tensor | None = None, seg_indptr: Tensor | None = None, weight_indices: Tensor | None = None) → Tensor

运行分段 GEMM 内核。

计算输入张量批次（具有可变数量的行，但具有固定数量的列）与具有固定数量的行和列的权重张量批次之间的矩阵乘法

\[y[i] = x[i] \times W[i]\]

如果提供了 weight_indices，我们将根据 weight_indices 张量中的索引选择权重张量

\[y[i] = x[i] \times W[\text{weight_indices}[i]]\]

我们使用不规则张量来表示输入张量 x 和输出张量 y，并且每个 x[i] 是连接张量的一个片段。请参阅 不规则张量教程 以获取更多详细信息。我们使用 seg_len 或 seg_indptr 张量（两者都可以）来指示每个片段的开始和结束位置，其中 seg_indptr 是 seg_lens 张量的累积和（在开头添加一个额外的 0）

\[\text{seg_indptr}[i] = \sum_{j=0}^{i-1} \text{seg_lens}[j], \quad \text{seg_indptr}[0] = 0\]
  • 如果提供了 seg_lens，则 x 的形状为 (sum(seg_lens), d_in)，y 的形状为

    (sum(seg_lens), d_out)，其中 d_in 是输入张量的列数，d_out 是输出张量的列数。

  • 如果提供了 seg_indptr，则 x 的形状为 (seg_indptr[-1], d_in)，y 的形状为

    (seg_indptr[-1], d_out).

参数:

  • x (torch.Tensor) – 形状为 (sum(seg_lens), d_in) 的输入张量。

  • weights (torch.Tensor) – 如果 weight_column_major 为 False，则形状为 (num_weights, d_in, d_out) 的 3D 权重张量，或者如果 weight_column_major 为 True，则形状为 (num_weights, d_out, d_in)。

  • batch_size (int) – 分段数。

  • weight_column_major (bool) – 权重张量是否为列主序。

  • out (Optional[torch.Tensor]) – 输出张量，形状为 (sum(seg_lens), d_out)。如果未提供，则将在内部创建一个新的张量。

  • seg_lens (Optional[torch.Tensor]) – 每个分段的长度，形状为 (batch_size,)，期望为 dtype 为 torch.int64 的 1D 张量。

  • seg_indptr (Optional[torch.Tensor]) – 分段的 indptr，形状为 (batch_size + 1,)，期望为 dtype 为 torch.int64 的 1D 张量。如果提供了此项，则将忽略 seg_lens，否则将从 seg_lens 内部计算 seg_indptr。

  • weight_indices (Optional[torch.Tensor]) – 要为每个分段选择的权重张量的索引，形状为 (batch_size,)。期望为 dtype 为 torch.int64 的 1D 张量。如果提供了此项，则权重张量将根据此张量中的索引进行选择。

返回值:

输出张量的形状为 (sum(seg_lens), d_out)。

返回值类型:

torch.Tensor