flashinfer.gemm.tgv_gemm_sm100

flashinfer.gemm.tgv_gemm_sm100(a: Tensor, b: Tensor, bias: Tensor, pdl: bool = False, out: Tensor | None = None, a_descale: Tensor | None = None, b_descale: Tensor | None = None) → Tensor

Perform TGV GEMM on SM100 architecture with automatic dtype detection.

Computes out = a @ b + bias. Dense inputs must share the same floating-point dtype (torch.bfloat16 or torch.float16). Block-scaled FP4 and FP8 inputs are also supported when scale factors are provided.

Parameters:

  • a (torch.Tensor) – First input tensor of shape (M, K) in row-major layout.

  • b (torch.Tensor) – Second input tensor of shape (K, N) in column-major layout (transposed from the typical PyTorch row-major convention).

  • bias (torch.Tensor) – Bias tensor of shape (N,) to add to each row of a @ b.

  • pdl (bool) – Whether to use PDL (Programmatic Dependent Launch). Defaults to False.

  • out (Optional[torch.Tensor]) – Pre-allocated output tensor of shape (M, N). If None, a new tensor is allocated.

  • a_descale (Optional[torch.Tensor]) – Scale factors for a. Required for block-scaled FP4/FP8 inputs.

  • b_descale (Optional[torch.Tensor]) – Scale factors for b. Required for block-scaled FP4/FP8 inputs.

Returns:

  • torch.Tensor – Output tensor of shape (M, N) in row-major layout.

  • Supported operand dtypes –

    • torch.bfloat16

    • torch.float16

    • torch.float4_e2m1fn_x2

    • torch.float8_e4m3fn

    • torch.float8_e5m2

  • Note –

    • Requires SM100 or SM103 architecture.

    • Dense inputs must have the same dtype and do not use scale factors.

    • Tensor b is expected to be in column-major layout (transposed from typical PyTorch row-major).

    • Block-scaled inputs require M <= 8 and flattened 128x4 scale-factor layouts.

    • NVFP4 requires two FP4 operands, FP8 E4M3 scales, and K divisible by 64.

    • MX scaling accepts FP4 or FP8 operands, including mixed operands, with FP8 E8M0 scales and K divisible by 128.

    • FP4 packs two values per byte, so its physical K dimension is K // 2.

    • For block-scaled inputs, bias and the output must be BF16 or FP16 and share the same dtype.