ttnn.scale_mask_softmax

ttnn.scale_mask_softmax = FastOperation(python_fully_qualified_name='ttnn.scale_mask_softmax', function=<ttnn._ttnn.operations.normalization.scale_mask_softmax_t object>, preprocess_golden_function_inputs=<function default_preprocess_golden_function_inputs>, golden_function=<function _golden_function>, postprocess_golden_function_outputs=<function default_postprocess_golden_function_outputs>, is_cpp_operation=True, is_experimental=False)

Computes a fused scale-mask-softmax operation along the last dimension of the input tensor.

This operation performs the following sequence:

  1. Optionally scales the input: scaled = input_tensor * scale (if scale is provided)

  2. Optionally applies mask: masked = scaled + mask (if mask is provided, with broadcasting)

  3. Computes softmax: output = softmax(masked)

This fused operation is commonly used in attention mechanisms where scaling and masking are applied before the softmax operation for efficiency.

Parameters:

  • input_tensor (ttnn.Tensor) – The input tensor to process.

  • scale (float, optional) – Scaling factor to multiply with input tensor.

  • mask (ttnn.Tensor, optional) – Attention mask tensor to add to scaled input.

Keyword Arguments:

  • memory_config (ttnn.MemoryConfig, optional) – Memory configuration for the output tensor. If not provided, inherits from input tensor.

  • is_causal_mask (bool, optional) – Whether the mask is a causal mask. Defaults to False.

  • compute_kernel_config (DeviceComputeKernelConfig, optional) – Compute kernel configuration for the operation.

  • numeric_stable (bool, optional) – Whether to use numerically stable softmax computation. Defaults to True.

Returns:

ttnn.Tensor – Output tensor with the fused scale-mask-softmax operation applied.

Note

The tensors support the following data types and layouts:

Input Tensor

Dtypes

Layouts

BFLOAT16, FLOAT32, BFLOAT8_B

TILE
Mask Tensor (optional)

Dtypes

Layouts

BFLOAT16, BFLOAT8_B

TILE, ROW_MAJOR

The output tensor will be in TILE layout and have the same dtype as the input_tensor

Limitations:

  • All tensors must be on-device.

  • For ROW_MAJOR masks: intermediate dimensions (except last two) must be 1; last dimension must equal TILE_WIDTH; width must align to input tensor’s tile width.