ttnn.sparse_matmul

ttnn.sparse_matmul = FastOperation(python_fully_qualified_name='ttnn.sparse_matmul', function=<ttnn._ttnn.operations.matmul.sparse_matmul_t object>, preprocess_golden_function_inputs=None, golden_function=None, postprocess_golden_function_outputs=None, is_cpp_operation=True, is_experimental=False)

Returns the matrix product of two tensors. Based on is_input_a_sparse, is_input_b_sparse and the sparsity tensor, some parts of the output computation is skipped.

The two input tensors must be be tiled and each have a rank of 4. The sparsity tensor must be a rank 4 tensor in row major layout.

Based on the input tensor shapes and is_input_a_sparse and is_input_b_sparse values, the output tensor shape is computed. See the supported modes table below.

Parameters:

  • input_tensor_a (ttnn.Tensor) – the first tensor to be multiplied. Needs to be on the device.

  • input_tensor_b (ttnn.Tensor) – the second tensor to be multiplied. Needs to be on the device.

Keyword Arguments:

  • sparsity (ttnn.Tensor) – the sparsity tensor containing the mask values. Needs to be on the device. The data type must be bfloat16.

  • program_config (ttnn.MatmulMultiCoreReuseMultiCast1DProgramConfig) – the program configuration for the matmul operation. Only this config type is supported. mcast_in0 must be set to True.

  • nnz (int, optional) – the number of non-zero values in the sparsity tensor. If not provided, it will be inferred from the sparsity tensor at runtime. If provided, it must be exactly equal to the number of non-zero entries in sparsity (i.e. nnz == count_nonzero(sparsity)). It is the user’s responsibility to set this value correctly. If the equality condition is not met, the kernel will hang (deadlock) on device, since the receiver and compute kernels loop exactly nnz times while the sender only multicasts once per non-zero sparsity entry. Because count_nonzero(sparsity) is data-dependent and only known at runtime, this cannot be validated on the host; the mismatch is asserted on-device (failing loudly under watcher) but otherwise results in a hang.

  • is_input_a_sparse (bool, optional) – boolean indicating whether input_tensor_a is sparse. Defaults to False. Together with is_input_b_sparse, it determines how the sparsity tensor is interpreted. See the supported modes table below.

  • is_input_b_sparse (bool, optional) – boolean indicating whether input_tensor_b is sparse. Defaults to True. Together with is_input_a_sparse, it determines how the sparsity tensor is interpreted. See the supported modes table below.

  • memory_config (ttnn.MemoryConfig, optional) – the memory configuration of the output tensor. Defaults to None, which will result in using ttnn.DRAM_MEMORY_CONFIG.

  • dtype (ttnn.DataType, optional) – the data type of the output tensor. Defaults to None.

  • compute_kernel_config (ttnn.DeviceComputeKernelConfig, optional) – the compute kernel configuration for the matmul operation. Defaults to None.

  • core_grid (ttnn.CoreGrid, optional) – the grid on which to distribute the sharded tensor on (writes to the cores L1s). Defaults to None.

  • output_tile (List of [int], optional) – Specifies the output tile configuration. Defaults to None.

  • optional_output_tensor (ttnn.Tensor, optional) – User provided on-device output tensor where the result of matmul is to be written. Defaults to None.

Returns:

ttnn.Tensor – the output tensor with sparse results.

Supported Modes

is_input_a_sparse

is_input_b_sparse

input_tensor_a shape

input_tensor_b shape

sparsity shape

nnz

output shape

True

True

[1, E, M, K]

[1, E, K, N]

[1, 1, 1, E]

None or 0 ≤ nnz ≤ E

[1, E, M, N]

False

True

[A, B, M, K]

[1, E, K, N]

[A, B, 1, E]

None or 0 ≤ nnz ≤ A * B * E

[A, B, 1, E, M, N]

True

False

[A, E, M, K]

[1, E, K, N]

[1, 1, A, E]

None or 0 ≤ nnz ≤ A * E

[A, E, M, N]

False

False

Invalid

–

–

–

–

Warning

When nnz is provided, it must equal count_nonzero(sparsity) exactly. Passing an nnz that does not match the actual number of non-zero entries in sparsity will cause the kernel to hang (deadlock) on device. Setting nnz correctly is the user’s responsibility. This condition cannot be checked on the host because count_nonzero(sparsity) is data-dependent and only known at runtime; it is validated on-device and will fail loudly (assert) when watcher is enabled, but otherwise manifests as a hang. If you cannot guarantee the exact count, omit nnz so it is inferred from the sparsity tensor at runtime.

Note

The input tensors support the following data types and layouts:

input_tensor_a

dtype

layout

BFLOAT4_B, BFLOAT8_B, BFLOAT16, FLOAT32

TILE
input_tensor_b

dtype

layout

BFLOAT4_B, BFLOAT8_B, BFLOAT16, FLOAT32

TILE
sparsity

dtype

layout

BFLOAT16

ROW_MAJOR
Memory Support:

The supported memory configurations for the two input tensors are program config dependent, as described below:

Supported Memory Configurations

Config

Input A

Input B

ttnn.MatmulMultiCoreReuseMultiCast1DProgramConfig with (mcast_in0=True)

Interleaved (L1/DRAM)

Interleaved (L1/DRAM)