ttnn.scale_mask_softmax_in_place

ttnn.scale_mask_softmax_in_place = FastOperation(python_fully_qualified_name='ttnn.scale_mask_softmax_in_place', function=<ttnn._ttnn.operations.normalization.scale_mask_softmax_in_place_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 in-place.

This operation modifies the input tensor directly and performs the following sequence:

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

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

  3. Computes softmax: input_tensor = softmax(input_tensor)

This in-place fused operation is commonly used in attention mechanisms and is memory-efficient as it reuses the input tensor for output, avoiding additional memory allocation.

Parameters:

  • input_tensor (ttnn.Tensor) – The input tensor to process. This tensor is modified in-place.

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

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

Keyword Arguments:

  • program_config (SoftmaxProgramConfig, optional) – Program configuration for the operation. Defaults to SoftmaxDefaultProgramConfig().

  • 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 False.

Returns:

ttnn.Tensor – The same tensor as input with the fused scale-mask-softmax operation applied in-place.

Note

The tensors support the following data types and layouts:

Input Tensor

Dtypes

Layouts

BFLOAT16, FLOAT32, BFLOAT8_B

TILE
Mask Tensor (optional)

Dtypes

Layouts

Ranks

BFLOAT16, BFLOAT8_B

TILE, ROW_MAJOR

2, 3, 4

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 unsharded ROW_MAJOR masks: intermediate dimensions (except last two) must be 1; last dimension must equal TILE_WIDTH; width must align to input tensor.

  • For sharded inputs: mask must be TILE layout with identical padded shape to input.

  • Internal block size constraints may restrict in-place operation for very large width tensors.