flashinfer.gemm.bmm_fp8

flashinfer.gemm.bmm_fp8(A: Tensor, B: Tensor, A_scale: Tensor, B_scale: Tensor, dtype: dtype, out: Tensor | None = None, backend: Literal['cudnn', 'cublas', 'cutlass', 'auto'] = 'cublas') → Tensor

BMM FP8

参数:

• A (torch.Tensor) – 输入张量，形状为 (b, m, k)，fp8 e4m3 或 fp8 e5m2。

• B (torch.Tensor) – Mat2 张量，形状为 (b, k, n)，应为列主序，fp8 e4m3 或 fp8 e5m2。

• A_scale (torch.Tensor) – A 的缩放张量，浮点数。

• B_scale (torch.Tensor) – B 的缩放张量，浮点数。

• dtype (torch.dtype) – 输出 dtype，bf16 或 fp16。

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

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

返回值:

out – 输出张量，形状为 (b, m, n)，bf16 或 fp16。

返回值类型:

torch.Tensor

示例

>>> import torch
>>> import torch.nn.functional as F
>>> import flashinfer
>>> def to_float8(x, dtype=torch.float8_e4m3fn):
...     finfo = torch.finfo(dtype)
...     min_val, max_val = x.aminmax()
...     amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12)
...     scale = finfo.max / amax
...     x_scl_sat = (x * scale).clamp(min=finfo.min, max=finfo.max)
...     return x_scl_sat.to(dtype), scale.float().reciprocal()
>>>
>>> input = torch.randn([16, 48, 64], device="cuda", dtype=torch.bfloat16)
>>> input_fp8, input_inv_s = to_float8(input, dtype=torch.float8_e4m3fn)
>>> # column major weight
>>> weight = torch.randn([16, 80, 64], device="cuda", dtype=torch.bfloat16).transpose(-2, -1)
>>> weight_fp8, weight_inv_s = to_float8(weight, dtype=torch.float8_e4m3fn)
>>> out = flashinfer.bmm_fp8(input_fp8, weight_fp8, input_inv_s, weight_inv_s, torch.bfloat16)
>>> out.shape
torch.Size([16, 48, 80])
>>> out.dtype
torch.bfloat16