flashinfer.gemm.bmm_bf16

flashinfer.gemm.bmm_bf16(A: Tensor, B: Tensor, out: Tensor | None = None, out_dtype: dtype = torch.bfloat16, backend: Literal['cutlass'] = 'cutlass') → Tensor

BF16 批量矩阵乘法

参数:

• A (torch.Tensor) – 输入张量，形状为 (b, m, k)，bf16 格式，行优先布局。

• B (torch.Tensor) – 权重张量，形状为 (b, k, n)，bf16 格式，列优先布局。

• out (Optional[torch.Tensor]) – 输出张量，形状为 (b, m, n)，bf16 或 fp16 格式，默认为 None。

• out_dtype (torch.dtype) – 输出数据类型，bf16 (默认) 或 fp16。

• backend (Literal["cutlass"]) – 要使用的后端，默认为 “cutlass”。

返回值:

输出张量，形状为 (b, m, n)，bf16 或 fp16 格式，行优先布局。

返回值类型:

torch.Tensor

示例

>>> import torch
>>> import flashinfer
>>> input = torch.randn([16, 48, 64], device="cuda", dtype=torch.bfloat16)
>>> weight = torch.randn([16, 80, 64], device="cuda", dtype=torch.bfloat16).transpose(-2, -1)
>>> out = flashinfer.bmm_bf16(input, weight)
>>> out.shape
torch.Size([16, 48, 80])
>>> out.dtype
torch.bfloat16