torch.nn.functional.scaled_addmm

torch.nn.functional.scaled_addmm(input, mat1, mat2, scale_a, scale_recipe_a, scale_b, scale_recipe_b, swizzle_a=None, swizzle_b=None, contraction_dim=(), use_fast_accum=False, *, beta=1.0, alpha=1.0)

Compute a scaled matrix product and add it to input.

The result is

$$\mathrm{out} = \beta\,\mathrm{input} + \alpha\,\mathrm{scaled\_mm}(\mathrm{mat1}, \mathrm{mat2}).$$

The scaling recipes and swizzles have the same meaning as in scaled_mm(). input must be a canonically contiguous, 16-byte-aligned matrix with shape (mat1.size(0), mat2.size(1)) and dtype float16, bfloat16, or float32. The result has the dtype of input; there is no separate output dtype. CUDA recipes are supported when their selected implementation uses cuBLASLt; non-cuBLAS fallbacks and ROCm are not supported.

Parameters:

• input (Tensor) – Matrix accumulated into the scaled matrix product.

• mat1 (Tensor) – Left matrix operand.

• mat2 (Tensor) – Right matrix operand.

• scale_a (Tensor | list[Tensor]) – Tensor containing decoding scaling factors for mat1.

• scale_recipe_a (_ScalingType | list[_ScalingType]) – Scaling recipe for mat1.

• scale_b (Tensor | list[Tensor]) – Tensor containing decoding scaling factors for mat2.

• scale_recipe_b (_ScalingType | list[_ScalingType]) – Scaling recipe for mat2.

• swizzle_a (_SwizzleType | list[_SwizzleType] | None) – Swizzling pattern, if any, for scale_a.

• swizzle_b (_SwizzleType | list[_SwizzleType] | None) – Swizzling pattern, if any, for scale_b.

• contraction_dim (list[int] | tuple[int, ...]) – Must be empty or (1, 0) (equivalent negative dimensions are also accepted).

• use_fast_accum (bool) – Whether to enable tensor-core fast accumulation.

• beta (float) – Multiplier for input.

• alpha (float) – Multiplier for the scaled matrix product.

Return type:

Tensor

Note

Fusing the addition removes an intermediate output rounding step, so the result need not be bitwise equal to a separate scaled matrix multiply followed by an addition.