flashinfer.gemm.mm_bf16

flashinfer.gemm.mm_bf16(a: Tensor, b: Tensor, bias: Tensor | None = None, pdl: bool = False, out: Tensor | None = None, out_dtype: dtype = torch.bfloat16, backend: Literal['cutlass', 'tgv', 'auto'] = 'tgv') → Tensor

BF16 矩阵乘法

参数:

• a (torch.Tensor) – 输入张量，形状为 (m, k)，行主序的 bf16。

• b (torch.Tensor) – 权重张量，形状为 (k, n)，列主序的 bf16。

• bias (Optional[torch.Tensor]) – 可选的偏置张量，形状为 (n,)。如果提供，只能与 TGV 后端一起使用。默认为 None。

• pdl (bool) – 是否使用持久数据加载器模式。只能与 TGV 后端一起使用。默认为 False。

• out (Optional[torch.Tensor]) – 输出张量，形状为 (m, n)，bf16 或 fp16。如果提供，只能与 CUTLASS 后端一起使用。默认为 None。

• out_dtype (torch.dtype) – 输出 dtype，bf16 或 fp16。如果提供，只能与 CUTLASS 后端一起使用。默认为 torch.bfloat16。

• backend (Literal["cutlass", "tgv", "auto"]) – 用于该操作的后端。默认为 "tgv"。 "auto" 允许在启用自动调优时从所有可用后端中选择最佳策略。

返回值:

输出张量，形状为 (m, n)，行主序的 bf16 或 fp16。

返回值类型:

torch.Tensor

示例

>>> import torch
>>> import flashinfer
>>> # Using the TGV backend
>>> a = torch.randn([48, 64], device="cuda", dtype=torch.bfloat16)
>>> b = torch.randn([80, 64], device="cuda", dtype=torch.bfloat16).transpose(-2, -1)
>>> bias = torch.randn([80], device="cuda", dtype=torch.bfloat16)
>>> out = flashinfer.mm_bf16(a, b, bias=bias, pdl=True, backend="tgv")
>>> out.shape
torch.Size([48, 80])
>>> out.dtype
torch.bfloat16
>>> # Using the CUTLASS backend
>>> fp16_out = torch.empty([48, 80], device="cuda", dtype=torch.float16)
>>> out = flashinfer.mm_bf16(a, b, out=fp16_out, out_dtype=torch.float16, backend="cutlass")
>>> out.shape
torch.Size([48, 80])
>>> out.dtype
torch.float16