asc2.matmul

asc2.matmul(input: LocalTensor, other: LocalTensor, bias: LocalTensor | None = None, *, hf32: bool = False) → LocalTensor

Computes the matrix multiplication of input and other with optional bias.

Parameters:

  • input – The left operand (2D tensor in L0A)

  • other – The right operand (2D tensor in L0B)

  • bias – Optional bias tensor (1D tensor in BT)

  • hf32 – Enable the rounding to HF32 for input tensors with float32 dtype

Returns:

The result of the matrix multiplication (2D tensor in L0C)

Return type:

LocalTensor

Raises:

  • TypeError – If input or other is not a LocalTensor

  • RuntimeError – If input tensors are not 2D, have incompatible shapes, unsupported dtype, or bias has wrong shape/dtype

Note

Input tensors must have either float16, bfloat16, or float32 data type and compatible shapes. Result tensor type is always float32. Bias must be a 1D tensor of float16, bfloat16, or float32 with shape matching the last dimension of the output. Bias with float16 or bfloat16 dtype is automatically promoted to float32 to match the result type.

Examples

Basic matrix multiplication using operator syntax:

a = asc2.copy_in(a_gm, [0, 0], [64, 128], asc2.TensorLocation.L0A)
b = asc2.copy_in(b_gm, [0, 0], [128, 256], asc2.TensorLocation.L0B)
c = a @ b  # result shape: [64, 256], location: L0C

Matrix multiplication with bias:

a = asc2.copy_in(a_gm, [0, 0], [64, 128], asc2.TensorLocation.L0A)
b = asc2.copy_in(b_gm, [0, 0], [128, 256], asc2.TensorLocation.L0B)
bias = asc2.copy_in(bias_gm, [0], [256], asc2.TensorLocation.BT)
c = asc2.matmul(a, b, bias)

Matrix multiplication with HF32 mode (for float32 inputs):

a = asc2.copy_in(a_gm, [0, 0], [32, 64], asc2.TensorLocation.L0A)
b = asc2.copy_in(b_gm, [0, 0], [64, 64], asc2.TensorLocation.L0B)
c = asc2.matmul(a, b, hf32=True)

Store result to global memory:

c = a @ b
asc2.copy_out(c, c_gm, [0, 0])

This function can also be called via a binary operator on LocalTensor, as input @ other instead of matmul(input, other).